Refactor ACS attach task ENI message handling

aws · Jun 19, 2023 · 0de7229 · 0de7229
1 parent f14bee8
commit 0de7229
Show file tree

Hide file tree

Showing 15 changed files with 794 additions and 595 deletions.
diff --git a/agent/acs/handler/acs_handler.go b/agent/acs/handler/acs_handler.go
@@ -263,19 +263,6 @@ func (acsSession *session) startACSSession(client wsclient.ClientServer) error {
 		dataClient: acsSession.dataClient,
 	}
 
-	// Add handler to ack task ENI attach message
-	eniAttachHandler := newAttachTaskENIHandler(
-		acsSession.ctx,
-		cfg.Cluster,
-		acsSession.containerInstanceARN,
-		client,
-		eniHandler,
-	)
-	eniAttachHandler.start()
-	defer eniAttachHandler.stop()
-
-	client.AddRequestHandler(eniAttachHandler.handlerFunc())
-
 	// Add handler to ack instance ENI attach message
 	instanceENIAttachHandler := newAttachInstanceENIHandler(
 		acsSession.ctx,
@@ -325,9 +312,13 @@ func (acsSession *session) startACSSession(client wsclient.ClientServer) error {
 	responseSender := func(response interface{}) error {
 		return client.MakeRequest(response)
 	}
-
-	heartbeatResponder := acssession.NewHeartbeatResponder(acsSession.doctor, responseSender)
-	client.AddRequestHandler(heartbeatResponder.HandlerFunc())
+	responders := []wsclient.RequestResponder{
+		acssession.NewAttachTaskENIResponder(eniHandler, responseSender),
+		acssession.NewHeartbeatResponder(acsSession.doctor, responseSender),
+	}
+	for _, r := range responders {
+		client.AddRequestHandler(r.HandlerFunc())
+	}
 
 	updater.AddAgentUpdateHandlers(client, cfg, acsSession.state, acsSession.dataClient, acsSession.taskEngine)
 

diff --git a/agent/acs/handler/attach_eni_handler_common_test.go b/agent/acs/handler/attach_eni_handler_common_test.go
@@ -173,3 +173,43 @@ func testHandleENIAttachment(t *testing.T, attachmentType, taskArn string) {
 	assert.NoError(t, err)
 	assert.Len(t, res, 1)
 }
+
+// TestHandleExpiredENIAttachmentTaskENI tests handling an expired task eni
+func TestHandleExpiredENIAttachmentTaskENI(t *testing.T) {
+	testHandleExpiredENIAttachment(t, apieni.ENIAttachmentTypeTaskENI, taskArn)
+}
+
+// TestHandleExpiredENIAttachmentInstanceENI tests handling an expired instance eni
+func TestHandleExpiredENIAttachmentInstanceENI(t *testing.T) {
+	testHandleExpiredENIAttachment(t, apieni.ENIAttachmentTypeInstanceENI, "")
+}
+
+func testHandleExpiredENIAttachment(t *testing.T, attachmentType, taskArn string) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	// Set expiresAt to a value in the past.
+	expiresAt := time.Unix(time.Now().Unix()-1, 0)
+
+	taskEngineState := dockerstate.NewTaskEngineState()
+	dataClient := data.NewNoopClient()
+
+	eniAttachment := &apieni.ENIAttachment{
+		AttachmentInfo: attachmentinfo.AttachmentInfo{
+			TaskARN:       taskArn,
+			AttachmentARN: attachmentArn,
+			ExpiresAt:     expiresAt,
+		},
+		AttachmentType: attachmentType,
+		MACAddress:     randomMAC,
+	}
+	eniHandler := &eniHandler{
+		state:      taskEngineState,
+		dataClient: dataClient,
+	}
+
+	// Expect an error starting the timer because of <=0 duration.
+	err := eniHandler.HandleENIAttachment(eniAttachment)
+	assert.Error(t, err)
+	assert.Equal(t, true, eniAttachment.HasExpired())
+}
diff --git a/agent/acs/handler/attach_instance_eni_handler_test.go b/agent/acs/handler/attach_instance_eni_handler_test.go
@@ -26,6 +26,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
 	mock_dockerstate "github.com/aws/amazon-ecs-agent/agent/engine/dockerstate/mocks"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/acs/model/ecsacs"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/acs/session/testconst"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/api/attachmentinfo"
 	apieni "github.com/aws/amazon-ecs-agent/ecs-agent/api/eni"
 	mock_wsclient "github.com/aws/amazon-ecs-agent/ecs-agent/wsclient/mock"
@@ -42,8 +43,8 @@ func TestInvalidAttachInstanceENIMessage(t *testing.T) {
 	}{
 		{
 			message: &ecsacs.AttachInstanceNetworkInterfacesMessage{
-				ClusterArn:               aws.String(clusterName),
-				ContainerInstanceArn:     aws.String(containerInstanceArn),
+				ClusterArn:               aws.String(testconst.ClusterName),
+				ContainerInstanceArn:     aws.String(testconst.ContainerInstanceARN),
 				ElasticNetworkInterfaces: []*ecsacs.ElasticNetworkInterface{},
 				WaitTimeoutMs:            aws.Int64(waitTimeoutMillis),
 			},
@@ -52,7 +53,7 @@ func TestInvalidAttachInstanceENIMessage(t *testing.T) {
 		{
 			message: &ecsacs.AttachInstanceNetworkInterfacesMessage{
 				MessageId:                aws.String(eniMessageId),
-				ContainerInstanceArn:     aws.String(containerInstanceArn),
+				ContainerInstanceArn:     aws.String(testconst.ContainerInstanceARN),
 				ElasticNetworkInterfaces: []*ecsacs.ElasticNetworkInterface{},
 				WaitTimeoutMs:            aws.Int64(waitTimeoutMillis),
 			},
@@ -61,7 +62,7 @@ func TestInvalidAttachInstanceENIMessage(t *testing.T) {
 		{
 			message: &ecsacs.AttachInstanceNetworkInterfacesMessage{
 				MessageId:                aws.String(eniMessageId),
-				ClusterArn:               aws.String(clusterName),
+				ClusterArn:               aws.String(testconst.ClusterName),
 				ElasticNetworkInterfaces: []*ecsacs.ElasticNetworkInterface{},
 				WaitTimeoutMs:            aws.Int64(waitTimeoutMillis),
 			},
@@ -70,16 +71,16 @@ func TestInvalidAttachInstanceENIMessage(t *testing.T) {
 		{
 			message: &ecsacs.AttachInstanceNetworkInterfacesMessage{
 				MessageId:     aws.String(eniMessageId),
-				ClusterArn:    aws.String(clusterName),
+				ClusterArn:    aws.String(testconst.ClusterName),
 				WaitTimeoutMs: aws.Int64(waitTimeoutMillis),
 			},
 			description: "Message without network interfaces should be invalid",
 		},
 		{
 			message: &ecsacs.AttachInstanceNetworkInterfacesMessage{
 				MessageId:            aws.String(eniMessageId),
-				ClusterArn:           aws.String(clusterName),
-				ContainerInstanceArn: aws.String(containerInstanceArn),
+				ClusterArn:           aws.String(testconst.ClusterName),
+				ContainerInstanceArn: aws.String(testconst.ContainerInstanceARN),
 				ElasticNetworkInterfaces: []*ecsacs.ElasticNetworkInterface{
 					{
 						MacAddress: aws.String(randomMAC),
@@ -97,8 +98,8 @@ func TestInvalidAttachInstanceENIMessage(t *testing.T) {
 		{
 			message: &ecsacs.AttachInstanceNetworkInterfacesMessage{
 				MessageId:            aws.String(eniMessageId),
-				ClusterArn:           aws.String(clusterName),
-				ContainerInstanceArn: aws.String(containerInstanceArn),
+				ClusterArn:           aws.String(testconst.ClusterName),
+				ContainerInstanceArn: aws.String(testconst.ContainerInstanceARN),
 				ElasticNetworkInterfaces: []*ecsacs.ElasticNetworkInterface{
 					{},
 				},
@@ -109,8 +110,8 @@ func TestInvalidAttachInstanceENIMessage(t *testing.T) {
 		{
 			message: &ecsacs.AttachInstanceNetworkInterfacesMessage{
 				MessageId:            aws.String(eniMessageId),
-				ClusterArn:           aws.String(clusterName),
-				ContainerInstanceArn: aws.String(containerInstanceArn),
+				ClusterArn:           aws.String(testconst.ClusterName),
+				ContainerInstanceArn: aws.String(testconst.ContainerInstanceARN),
 				ElasticNetworkInterfaces: []*ecsacs.ElasticNetworkInterface{
 					{
 						Ec2Id: aws.String("1"),
@@ -123,8 +124,8 @@ func TestInvalidAttachInstanceENIMessage(t *testing.T) {
 		{
 			message: &ecsacs.AttachInstanceNetworkInterfacesMessage{
 				MessageId:            aws.String(eniMessageId),
-				ClusterArn:           aws.String(clusterName),
-				ContainerInstanceArn: aws.String(containerInstanceArn),
+				ClusterArn:           aws.String(testconst.ClusterName),
+				ContainerInstanceArn: aws.String(testconst.ContainerInstanceARN),
 				ElasticNetworkInterfaces: []*ecsacs.ElasticNetworkInterface{
 					{
 						MacAddress: aws.String(randomMAC),
@@ -153,7 +154,7 @@ func TestInstanceENIAckSingleMessage(t *testing.T) {
 
 	ctx := context.TODO()
 	mockWSClient := mock_wsclient.NewMockClientServer(ctrl)
-	handler := newAttachInstanceENIHandler(ctx, clusterName, containerInstanceArn, mockWSClient,
+	handler := newAttachInstanceENIHandler(ctx, testconst.ClusterName, testconst.ContainerInstanceARN, mockWSClient,
 		&eniHandler{
 			state:      taskEngineState,
 			dataClient: dataClient,
@@ -176,8 +177,8 @@ func TestInstanceENIAckSingleMessage(t *testing.T) {
 	}
 	message := &ecsacs.AttachInstanceNetworkInterfacesMessage{
 		MessageId:            aws.String(eniMessageId),
-		ClusterArn:           aws.String(clusterName),
-		ContainerInstanceArn: aws.String(containerInstanceArn),
+		ClusterArn:           aws.String(testconst.ClusterName),
+		ContainerInstanceArn: aws.String(testconst.ContainerInstanceARN),
 		ElasticNetworkInterfaces: []*ecsacs.ElasticNetworkInterface{
 			&mockNetInterface1,
 		},
@@ -201,7 +202,7 @@ func TestInstanceENIAckSingleMessageDuplicateENIAttachmentMessageStartsTimer(t *
 
 	ctx := context.TODO()
 	mockWSClient := mock_wsclient.NewMockClientServer(ctrl)
-	handler := newAttachInstanceENIHandler(ctx, clusterName, containerInstanceArn, mockWSClient,
+	handler := newAttachInstanceENIHandler(ctx, testconst.ClusterName, testconst.ContainerInstanceARN, mockWSClient,
 		&eniHandler{
 			state:      mockState,
 			dataClient: dataClient,
@@ -237,8 +238,8 @@ func TestInstanceENIAckSingleMessageDuplicateENIAttachmentMessageStartsTimer(t *
 	}
 	message := &ecsacs.AttachInstanceNetworkInterfacesMessage{
 		MessageId:            aws.String(eniMessageId),
-		ClusterArn:           aws.String(clusterName),
-		ContainerInstanceArn: aws.String(containerInstanceArn),
+		ClusterArn:           aws.String(testconst.ClusterName),
+		ContainerInstanceArn: aws.String(testconst.ContainerInstanceARN),
 		ElasticNetworkInterfaces: []*ecsacs.ElasticNetworkInterface{
 			&mockNetInterface1,
 		},
@@ -261,7 +262,7 @@ func TestInstanceENIAckHappyPath(t *testing.T) {
 	dataClient := data.NewNoopClient()
 
 	mockWSClient := mock_wsclient.NewMockClientServer(ctrl)
-	handler := newAttachInstanceENIHandler(ctx, clusterName, containerInstanceArn, mockWSClient,
+	handler := newAttachInstanceENIHandler(ctx, testconst.ClusterName, testconst.ContainerInstanceARN, mockWSClient,
 		&eniHandler{
 			state:      taskEngineState,
 			dataClient: dataClient,
@@ -284,8 +285,8 @@ func TestInstanceENIAckHappyPath(t *testing.T) {
 	}
 	message := &ecsacs.AttachInstanceNetworkInterfacesMessage{
 		MessageId:            aws.String(eniMessageId),
-		ClusterArn:           aws.String(clusterName),
-		ContainerInstanceArn: aws.String(containerInstanceArn),
+		ClusterArn:           aws.String(testconst.ClusterName),
+		ContainerInstanceArn: aws.String(testconst.ContainerInstanceARN),
 		ElasticNetworkInterfaces: []*ecsacs.ElasticNetworkInterface{
 			&mockNetInterface1,
 		},