Add serviceName to TMDE v4

aws · Aug 26, 2022 · bf161e1 · bf161e1
1 parent f4ca59d
commit bf161e1
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 7 deletions.
diff --git a/agent/handlers/task_server_setup_test.go b/agent/handlers/task_server_setup_test.go
@@ -1467,8 +1467,8 @@ func TestV4BridgeTaskMetadata(t *testing.T) {
 	gomock.InOrder(
 		state.EXPECT().TaskARNByV3EndpointID(v3EndpointID).Return(taskARN, true),
 		state.EXPECT().TaskByArn(taskARN).Return(bridgeTask, true),
-		state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToBridgeContainer, true),
 		state.EXPECT().TaskByArn(taskARN).Return(bridgeTask, true),
+		state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToBridgeContainer, true),
 		state.EXPECT().ContainerByID(containerID).Return(bridgeContainer, true),
 		state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
 	)
@@ -1502,8 +1502,8 @@ func TestV4BridgeTaskMetadataAllowMissingContainerNetwork(t *testing.T) {
 	gomock.InOrder(
 		state.EXPECT().TaskARNByV3EndpointID(v3EndpointID).Return(taskARN, true),
 		state.EXPECT().TaskByArn(taskARN).Return(bridgeTask, true),
-		state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToBridgeContainer, true),
 		state.EXPECT().TaskByArn(taskARN).Return(bridgeTask, true),
+		state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToBridgeContainer, true),
 		state.EXPECT().ContainerByID(containerID).Return(nil, false),
 		state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
 	)

diff --git a/agent/handlers/v4/response.go b/agent/handlers/v4/response.go
@@ -29,7 +29,8 @@ import (
 // with the v2 task response object.
 type TaskResponse struct {
 	*v2.TaskResponse
-	Containers []ContainerResponse `json:"Containers,omitempty"`
+	Containers  []ContainerResponse `json:"Containers,omitempty"`
+	ServiceName string              `json:"ServiceName,omitempty"`
 }
 
 // ContainerResponse is the v4 Container response. It augments the v4 Network response
@@ -83,6 +84,7 @@ func NewTaskResponse(
 	cluster string,
 	az string,
 	containerInstanceARN string,
+	serviceName string,
 	propagateTags bool,
 ) (*TaskResponse, error) {
 	// Construct the v2 response first.
@@ -109,6 +111,7 @@ func NewTaskResponse(
 	return &TaskResponse{
 		TaskResponse: v2Resp,
 		Containers:   containers,
+		ServiceName:  serviceName,
 	}, nil
 }
 

diff --git a/agent/handlers/v4/response_test.go b/agent/handlers/v4/response_test.go
@@ -39,6 +39,7 @@ const (
 	cluster                  = "default"
 	family                   = "sleep"
 	version                  = "1"
+	serviceName              = "someService"
 	containerID              = "cid"
 	containerName            = "sleepy"
 	imageName                = "busybox"
@@ -68,6 +69,7 @@ func TestNewTaskContainerResponses(t *testing.T) {
 		Arn:                 taskARN,
 		Family:              family,
 		Version:             version,
+		ServiceName:         serviceName,
 		DesiredStatusUnsafe: apitaskstatus.TaskRunning,
 		KnownStatusUnsafe:   apitaskstatus.TaskRunning,
 		ENIs: []*apieni.ENI{
@@ -134,7 +136,8 @@ func TestNewTaskContainerResponses(t *testing.T) {
 		state.EXPECT().TaskByArn(taskARN).Return(task, true),
 	)
 
-	taskResponse, err := NewTaskResponse(taskARN, state, ecsClient, cluster, availabilityZone, containerInstanceArn, false)
+	taskResponse, err := NewTaskResponse(taskARN, state, ecsClient, cluster,
+		availabilityZone, containerInstanceArn, task.ServiceName, false)
 	require.NoError(t, err)
 	_, err = json.Marshal(taskResponse)
 	require.NoError(t, err)
@@ -143,6 +146,7 @@ func TestNewTaskContainerResponses(t *testing.T) {
 	assert.Equal(t, eniIPv6Address, taskResponse.Containers[0].Networks[0].IPv6Addresses[0])
 	assert.Equal(t, ipv6SubnetCIDRBlock, taskResponse.Containers[0].Networks[0].IPv6SubnetCIDRBlock)
 	assert.Equal(t, subnetGatewayIPV4Address, taskResponse.Containers[0].Networks[0].SubnetGatewayIPV4Address)
+	assert.Equal(t, serviceName, taskResponse.ServiceName)
 
 	gomock.InOrder(
 		state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true),

diff --git a/agent/handlers/v4/task_metadata_handler.go b/agent/handlers/v4/task_metadata_handler.go
@@ -45,9 +45,12 @@ func TaskMetadataHandler(state dockerstate.TaskEngineState, ecsClient api.ECSCli
 			return
 		}
 
+		task, _ := state.TaskByArn(taskArn)
+
 		seelog.Infof("V4 taskMetadata handler: Writing response for task '%s'", taskArn)
 
-		taskResponse, err := NewTaskResponse(taskArn, state, ecsClient, cluster, az, containerInstanceArn, propagateTags)
+		taskResponse, err := NewTaskResponse(taskArn, state, ecsClient, cluster,
+			az, containerInstanceArn, task.ServiceName, propagateTags)
 		if err != nil {
 			errResponseJson, err := json.Marshal("Unable to generate metadata for v4 task: '" + taskArn + "'")
 			if e := utils.WriteResponseIfMarshalError(w, err); e != nil {
@@ -56,8 +59,6 @@ func TaskMetadataHandler(state dockerstate.TaskEngineState, ecsClient api.ECSCli
 			utils.WriteJSONToResponse(w, http.StatusInternalServerError, errResponseJson, utils.RequestTypeTaskMetadata)
 			return
 		}
-
-		task, _ := state.TaskByArn(taskArn)
 		// for non-awsvpc task mode
 		if !task.IsNetworkModeAWSVPC() {
 			// fill in non-awsvpc network details for container responses here