acs,tcs handlers: increase websocket rw timeout

Increase the websocket read and write timeouts
as per review comments

Author: aaithal

agent/acs/handler/acs_handler.go

@@ -48,6 +48,9 @@ const (
 	// without disconnecting
 	heartbeatTimeout = 1 * time.Minute
 	heartbeatJitter  = 1 * time.Minute
+	// wsRWTimeout is the duration of read and write deadline for the
+	// websocket connection
+	wsRWTimeout = 2*heartbeatTimeout + heartbeatJitter
 
 	inactiveInstanceReconnectDelay = 1 * time.Hour
 
@@ -376,7 +379,7 @@ func (acsSession *session) heartbeatJitter() time.Duration {
 
 // createACSClient creates the ACS Client using the specified URL
 func (acsResources *acsSessionResources) createACSClient(url string, cfg *config.Config) wsclient.ClientServer {
-	return acsclient.New(url, cfg, acsResources.credentialsProvider, heartbeatTimeout+heartbeatJitter)
+	return acsclient.New(url, cfg, acsResources.credentialsProvider, wsRWTimeout)
 }
 
 // connectedToACS records a successful connection to ACS
@@ -437,11 +440,11 @@ func anyMessageHandler(timer ttime.Timer, client wsclient.ClientServer) func(int
 	return func(interface{}) {
 		seelog.Debug("ACS activity occurred")
 		// Reset read deadline as there's activity on the channel
-		if err := client.SetReadDeadline(time.Now().Add(heartbeatTimeout + heartbeatJitter)); err != nil {
-			seelog.Warn("Unable to extend read deadline for ACS connection: %v", err)
+		if err := client.SetReadDeadline(time.Now().Add(wsRWTimeout)); err != nil {
+			seelog.Warnf("Unable to extend read deadline for ACS connection: %v", err)
 		}
 
-		// Reset heearbeat timer
+		// Reset hearbeat timer
 		timer.Reset(utils.AddJitter(heartbeatTimeout, heartbeatJitter))
 	}
 }

agent/acs/handler/acs_handler_test.go

@@ -803,9 +803,7 @@ func TestConnectionIsClosedOnIdle(t *testing.T) {
 		_heartbeatTimeout:    20 * time.Millisecond,
 		_heartbeatJitter:     10 * time.Millisecond,
 	}
-	go func() {
-		acsSession.startACSSession(mockWsClient)
-	}()
+	go acsSession.startACSSession(mockWsClient)
 
 	// Wait for connection to be closed. If the connection is not closed
 	// due to inactivity, the test will time out

agent/tcs/handler/handler.go

@@ -36,8 +36,11 @@ const (
 	defaultPublishMetricsInterval = 20 * time.Second
 
 	// The maximum time to wait between heartbeats without disconnecting
-	defaultHeartbeatTimeout            = 1 * time.Minute
-	defaultHeartbeatJitter             = 1 * time.Minute
+	defaultHeartbeatTimeout = 1 * time.Minute
+	defaultHeartbeatJitter  = 1 * time.Minute
+	// wsRWTimeout is the duration of read and write deadline for the
+	// websocket connection
+	wsRWTimeout                        = 2*defaultHeartbeatTimeout + defaultHeartbeatJitter
 	deregisterContainerInstanceHandler = "TCSDeregisterContainerInstanceHandler"
 )
 
@@ -92,7 +95,9 @@ func startTelemetrySession(params TelemetrySessionParams, statsEngine stats.Engi
 		return err
 	}
 	url := formatURL(tcsEndpoint, params.Cfg.Cluster, params.ContainerInstanceArn)
-	return startSession(url, params.Cfg, params.CredentialProvider, statsEngine, defaultHeartbeatTimeout, defaultHeartbeatJitter, defaultPublishMetricsInterval, params.DeregisterInstanceEventStream)
+	return startSession(url, params.Cfg, params.CredentialProvider, statsEngine,
+		defaultHeartbeatTimeout, defaultHeartbeatJitter, defaultPublishMetricsInterval,
+		params.DeregisterInstanceEventStream)
 }
 
 func startSession(url string,
@@ -103,7 +108,7 @@ func startSession(url string,
 	publishMetricsInterval time.Duration,
 	deregisterInstanceEventStream *eventstream.EventStream) error {
 	client := tcsclient.New(url, cfg, credentialProvider, statsEngine,
-		publishMetricsInterval, defaultHeartbeatTimeout+defaultHeartbeatJitter)
+		publishMetricsInterval, wsRWTimeout)
 	defer client.Close()
 
 	err := deregisterInstanceEventStream.Subscribe(deregisterContainerInstanceHandler, client.Disconnect)
@@ -112,24 +117,24 @@ func startSession(url string,
 	}
 	defer deregisterInstanceEventStream.Unsubscribe(deregisterContainerInstanceHandler)
 
+	err = client.Connect()
+	if err != nil {
+		seelog.Errorf("Error connecting to TCS: %v", err.Error())
+		return err
+	}
+	seelog.Info("Connected to TCS endpoint")
 	// start a timer and listens for tcs heartbeats/acks. The timer is reset when
 	// we receive a heartbeat from the server or when a publish metrics message
 	// is acked.
 	timer := time.AfterFunc(utils.AddJitter(heartbeatTimeout, heartbeatJitter), func() {
 		// Close the connection if there haven't been any messages received from backend
 		// for a long time.
-		seelog.Info("TCS Connection hasn't had a heartbeat or an ack message in too long of a timeout; disconnecting")
+		seelog.Info("TCS Connection hasn't had any activity for too long; disconnecting")
 		client.Disconnect()
 	})
 	defer timer.Stop()
 	client.AddRequestHandler(heartbeatHandler(timer))
 	client.AddRequestHandler(ackPublishMetricHandler(timer))
-	err = client.Connect()
-	if err != nil {
-		seelog.Errorf("Error connecting to TCS: %v", err.Error())
-		return err
-	}
-	seelog.Info("Connected to TCS endpoint")
 	client.SetAnyRequestHandler(anyMessageHandler(client))
 	return client.Serve()
 }
@@ -155,9 +160,9 @@ func ackPublishMetricHandler(timer *time.Timer) func(*ecstcs.AckPublishMetric) {
 // connection is active
 func anyMessageHandler(client wsclient.ClientServer) func(interface{}) {
 	return func(interface{}) {
-		seelog.Trace("TCS activity occured")
+		seelog.Trace("TCS activity occurred")
 		// Reset read deadline as there's activity on the channel
-		if err := client.SetReadDeadline(time.Now().Add(defaultHeartbeatTimeout + defaultHeartbeatJitter)); err != nil {
+		if err := client.SetReadDeadline(time.Now().Add(wsRWTimeout)); err != nil {
 			seelog.Warnf("Unable to extend read deadline for TCS connection: %v", err)
 		}
 	}

agent/wsclient/client.go

@@ -305,7 +305,7 @@ func (cs *ClientServerImpl) WriteMessage(send []byte) error {
 
 	// This is just future proofing. Ignore the error as the gorilla websocket
 	// library returns 'nil' anyway for SetWriteDeadline
-	// https://github.com/gorilla/websocket/blob/master/conn.go#L761
+	// https://github.com/gorilla/websocket/blob/4201258b820c74ac8e6922fc9e6b52f71fe46f8d/conn.go#L761
 	if err := cs.conn.SetWriteDeadline(time.Now().Add(cs.RWTimeout)); err != nil {
 		seelog.Warnf("Unable to set write deadline for websocket connection: %v for %s", err, cs.URL)
 	}