Refactor ACS heartbeat message handling

agent/acs/handler/acs_handler.go

@@ -34,6 +34,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/eventhandler"
 	"github.com/aws/amazon-ecs-agent/agent/eventstream"
 	"github.com/aws/amazon-ecs-agent/agent/version"
+	acssession "github.com/aws/amazon-ecs-agent/ecs-agent/acs/session"
 	rolecredentials "github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/doctor"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/utils/retry"
@@ -321,7 +322,12 @@ func (acsSession *session) startACSSession(client wsclient.ClientServer) error {
 
 	client.AddRequestHandler(payloadHandler.handlerFunc())
 
-	client.AddRequestHandler(HeartbeatHandlerFunc(client, acsSession.doctor))
+	responseSender := func(response interface{}) error {
+		return client.MakeRequest(response)
+	}
+
+	heartbeatResponder := acssession.NewHeartbeatResponder(acsSession.doctor, responseSender)
+	client.AddRequestHandler(heartbeatResponder.HandlerFunc())
 
 	updater.AddAgentUpdateHandlers(client, cfg, acsSession.state, acsSession.dataClient, acsSession.taskEngine)
 

ecs-agent/acs/session/heartbeat_responder.go

@@ -0,0 +1,71 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package session
+
+import (
+	"github.com/aws/amazon-ecs-agent/ecs-agent/acs/model/ecsacs"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/doctor"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/logger"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/logger/field"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/wsclient"
+
+	"github.com/aws/aws-sdk-go/aws"
+)
+
+// heartbeatResponder implements the wsclient.RequestResponder interface for responding
+// to ecsacs.HeartbeatMessage type.
+type heartbeatResponder struct {
+	doctor  *doctor.Doctor
+	respond wsclient.RespondFunc
+}
+
+// NewHeartbeatResponder returns an instance of the heartbeatResponder struct.
+func NewHeartbeatResponder(doctor *doctor.Doctor, responseSender wsclient.RespondFunc) wsclient.RequestResponder {
+	r := &heartbeatResponder{
+		doctor: doctor,
+	}
+	r.respond = ResponseToACSSender(r.Name(), responseSender)
+	return r
+}
+
+func (*heartbeatResponder) Name() string {
+	return "heartbeat message responder"
+}
+
+func (r *heartbeatResponder) HandlerFunc() wsclient.RequestHandler {
+	return r.processHeartbeatMessage
+}
+
+// processHeartbeatMessage processes an ACS heartbeat message.
+// This function is meant to be called from the ACS dispatcher and as such
+// should not block in any way to prevent starvation of the message handler.
+func (r *heartbeatResponder) processHeartbeatMessage(message *ecsacs.HeartbeatMessage) {
+	// Agent will run container instance healthchecks. They are triggered by ACS heartbeat.
+	// Results of healthchecks will be sent on to TACS.
+	go r.doctor.RunHealthchecks()
+
+	// Agent will send simple ack
+	ack := &ecsacs.HeartbeatAckRequest{
+		MessageId: message.MessageId,
+	}
+	go func() {
+		err := r.respond(ack)
+		if err != nil {
+			logger.Warn("Error acknowledging server heartbeat", logger.Fields{
+				field.MessageID: aws.StringValue(ack.MessageId),
+				field.Error:     err,
+			})
+		}
+	}()
+}

agent/acs/handler/heartbeat_handler_test.go → ecs-agent/acs/session/heartbeat_responder_test.go

@@ -14,18 +14,17 @@
 // express or implied. See the License for the specific language governing
 // permissions and limitations under the License.
 
-package handler
+package session
 
 import (
 	"testing"
 
-	"github.com/aws/amazon-ecs-agent/ecs-agent/acs/model/ecsacs"
-	"github.com/aws/amazon-ecs-agent/ecs-agent/doctor"
-	mock_wsclient "github.com/aws/amazon-ecs-agent/ecs-agent/wsclient/mock"
-
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
+
+	"github.com/aws/amazon-ecs-agent/ecs-agent/acs/model/ecsacs"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/doctor"
 )
 
 const (
@@ -87,16 +86,16 @@ func validateHeartbeatAck(t *testing.T, heartbeatReceived *ecsacs.HeartbeatMessa
 
 	ackSent := make(chan *ecsacs.HeartbeatAckRequest)
 
-	mockWsClient := mock_wsclient.NewMockClientServer(ctrl)
-	mockWsClient.EXPECT().MakeRequest(gomock.Any()).Do(func(message *ecsacs.HeartbeatAckRequest) {
-		ackSent <- message
-		close(ackSent)
-	}).Times(1)
-
 	emptyHealthchecksList := []doctor.Healthcheck{}
 	emptyDoctor, _ := doctor.NewDoctor(emptyHealthchecksList, "testCluster", "this:is:an:instance:arn")
 
-	handleSingleHeartbeatMessage(mockWsClient, emptyDoctor, heartbeatReceived)
+	testResponseSender := func(response interface{}) error {
+		resp := response.(*ecsacs.HeartbeatAckRequest)
+		ackSent <- resp
+		return nil
+	}
+	testHeartbeatResponder := NewHeartbeatResponder(emptyDoctor, testResponseSender)
+	testHeartbeatResponder.(*heartbeatResponder).processHeartbeatMessage(heartbeatReceived)
 
 	// wait till we send an
 	heartbeatAckSent := <-ackSent

ecs-agent/acs/session/session.go

@@ -0,0 +1,20 @@
+package session
+
+import (
+	"github.com/aws/amazon-ecs-agent/ecs-agent/logger"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/wsclient"
+)
+
+// ResponseToACSSender returns a wsclient.RespondFunc that a responder can invoke in response to receiving and
+// processing specific websocket request messages from ACS. The returned wsclient.RespondFunc:
+//  1. logs the response to be sent, as well as the name of the invoking responder
+//  2. sends the response request to ACS
+func ResponseToACSSender(responderName string, responseSender wsclient.RespondFunc) wsclient.RespondFunc {
+	return func(response interface{}) error {
+		logger.Debug("Sending response to ACS", logger.Fields{
+			"Name":     responderName,
+			"Response": response,
+		})
+		return responseSender(response)
+	}
+}

ecs-agent/logger/field/constants.go

@@ -49,4 +49,5 @@ const (
 	ContainerImage          = "containerImage"
 	ContainerExitCode       = "containerExitCode"
 	TMDSEndpointContainerID = "tmdsEndpointContainerID"
+	MessageID               = "messageID"
 )

ecs-agent/wsclient/client.go

@@ -99,10 +99,6 @@ type RequestHandler interface{}
 //	    respond  func(interface{}) error
 //	    dispatcher actor.Dispatcher
 //	}
-//	func(d *payloadmessagedispatcher) RegisterResponder(respond func(interface{}) error) error {
-//	    d.respond = respond
-//	    return nil
-//	}
 //	func(d *payloadmessagedispatcher) HandlerFunc() RequestHandler {
 //	    return func(payload *ecsacs.PayloadMessage) {
 //	        message := &actor.DispatcherMessage{
@@ -118,9 +114,6 @@ type RequestHandler interface{}
 type RequestResponder interface {
 	// Name returns the name of the responder. This is used mostly for logging.
 	Name() string
-	// RegisterResponder registers a function that can be invoked in response
-	// to receiving and processing a websocket request message.
-	RegisterResponder(RespondFunc)
 	// HandlerFunc returns the RequestHandler callback for a particular
 	// websocket request message type.
 	HandlerFunc() RequestHandler