Add receive timeout

Author: mostafa

cmd/run.go

@@ -407,6 +407,7 @@ var runCmd = &cobra.Command{
 			// Fill the missing and zero values with the default ones.
 			clients[name].TCPKeepAlivePeriod = clients[name].GetTCPKeepAlivePeriod()
 			clients[name].ReceiveDeadline = clients[name].GetReceiveDeadline()
+			clients[name].ReceiveTimeout = clients[name].GetReceiveTimeout()
 			clients[name].SendDeadline = clients[name].GetSendDeadline()
 			clients[name].ReceiveChunkSize = clients[name].GetReceiveChunkSize()
 

config/config.go

@@ -113,6 +113,7 @@ func (c *Config) LoadDefaults(ctx context.Context) {
 		TCPKeepAlivePeriod: DefaultTCPKeepAlivePeriod,
 		ReceiveChunkSize:   DefaultChunkSize,
 		ReceiveDeadline:    DefaultReceiveDeadline,
+		ReceiveTimeout:     DefaultReceiveTimeout,
 		SendDeadline:       DefaultSendDeadline,
 	}
 

config/constants.go

@@ -101,6 +101,7 @@ const (
 	DefaultSendDeadline       = 0
 	DefaultTCPKeepAlivePeriod = 30 * time.Second
 	DefaultTCPKeepAlive       = false
+	DefaultReceiveTimeout     = 0
 
 	// Pool constants.
 	EmptyPoolCapacity        = 0

config/getters.go

@@ -126,6 +126,14 @@ func (c Client) GetReceiveDeadline() time.Duration {
 	return c.ReceiveDeadline
 }
 
+// GetReceiveTimeout returns the receive timeout from config file or default value.
+func (c Client) GetReceiveTimeout() time.Duration {
+	if c.ReceiveTimeout <= 0 {
+		return DefaultReceiveTimeout
+	}
+	return c.ReceiveTimeout
+}
+
 // GetSendDeadline returns the send deadline from config file or default value.
 func (c Client) GetSendDeadline() time.Duration {
 	if c.SendDeadline <= 0 {

config/types.go

@@ -44,6 +44,7 @@ type Client struct {
 	TCPKeepAlivePeriod time.Duration `json:"tcpKeepAlivePeriod" jsonschema:"oneof_type=string;integer"`
 	ReceiveChunkSize   int           `json:"receiveChunkSize"`
 	ReceiveDeadline    time.Duration `json:"receiveDeadline" jsonschema:"oneof_type=string;integer"`
+	ReceiveTimeout     time.Duration `json:"receiveTimeout" jsonschema:"oneof_type=string;integer"`
 	SendDeadline       time.Duration `json:"sendDeadline" jsonschema:"oneof_type=string;integer"`
 }
 

gatewayd.yaml

@@ -33,6 +33,7 @@ clients:
     tcpKeepAlivePeriod: 30s # duration
     receiveChunkSize: 8192
     receiveDeadline: 0s # duration, 0ms/0s means no deadline
+    receiveTimeout: 0s # duration, 0ms/0s means no timeout
     sendDeadline: 0s # duration, 0ms/0s means no deadline
 
 pools:

network/client.go

@@ -34,6 +34,7 @@ type Client struct {
 	ReceiveChunkSize   int
 	ReceiveDeadline    time.Duration
 	SendDeadline       time.Duration
+	ReceiveTimeout     time.Duration
 	ID                 string
 	Network            string // tcp/udp/unix
 	Address            string
@@ -168,12 +169,21 @@ func (c *Client) Receive() (int, []byte, *gerr.GatewayDError) {
 	_, span := otel.Tracer(config.TracerName).Start(c.ctx, "Receive")
 	defer span.End()
 
+	var ctx context.Context
+	var cancel context.CancelFunc
+	if c.ReceiveTimeout > 0 {
+		ctx, cancel = context.WithTimeout(c.ctx, c.ReceiveTimeout)
+		defer cancel()
+	} else {
+		ctx = context.Background()
+	}
+
 	var received int
 	buffer := bytes.NewBuffer(nil)
 	// Read the data in chunks.
 	select { //nolint:gosimple
 	case <-time.After(time.Millisecond):
-		for {
+		for ctx.Err() == nil {
 			chunk := make([]byte, c.ReceiveChunkSize)
 			read, err := c.Conn.Read(chunk)
 			if err != nil {

network/client_test.go

@@ -30,6 +30,7 @@ func CreateNewClient(t *testing.T) *Client {
 			Address:            "localhost:5432",
 			ReceiveChunkSize:   config.DefaultChunkSize,
 			ReceiveDeadline:    config.DefaultReceiveDeadline,
+			ReceiveTimeout:     config.DefaultReceiveTimeout,
 			SendDeadline:       config.DefaultSendDeadline,
 			TCPKeepAlive:       false,
 			TCPKeepAlivePeriod: config.DefaultTCPKeepAlivePeriod,