Add RestartCount to tmde v4 response

aws · May 6, 2024 · bc83af0 · bc83af0
1 parent cd69dcd
commit bc83af0
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/agent/handlers/task_server_setup_test.go b/agent/handlers/task_server_setup_test.go
@@ -1801,8 +1801,9 @@ func TestV4ContainerMetadata(t *testing.T) {
 			setStateExpectations: func(state *mock_dockerstate.MockTaskEngineState) {
 				gomock.InOrder(
 					state.EXPECT().DockerIDByV3EndpointID(v3EndpointID).Return(containerID, true),
-					state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true),
+					state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true).AnyTimes(),
 					state.EXPECT().TaskByID(containerID).Return(task, true).Times(2),
+					state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true).AnyTimes(),
 				)
 			},
 			expectedStatusCode:   http.StatusOK,
@@ -1817,7 +1818,7 @@ func TestV4ContainerMetadata(t *testing.T) {
 					state.EXPECT().DockerIDByV3EndpointID(v3EndpointID).Return(containerID, true),
 					state.EXPECT().ContainerByID(containerID).Return(bridgeContainer, true),
 					state.EXPECT().TaskByID(containerID).Return(bridgeTask, true),
-					state.EXPECT().ContainerByID(containerID).Return(nil, false),
+					state.EXPECT().ContainerByID(containerID).Return(nil, false).AnyTimes(),
 				)
 			},
 			expectedStatusCode:   http.StatusInternalServerError,
@@ -2340,6 +2341,7 @@ func TestV4TaskMetadataWithTags(t *testing.T) {
 		gomock.InOrder(
 			state.EXPECT().TaskARNByV3EndpointID(v3EndpointID).Return(taskARN, true),
 			state.EXPECT().TaskByArn(taskARN).Return(task, true).AnyTimes(),
+			state.EXPECT().ContainerByID(containerID).Return(container1, true).AnyTimes(),
 			state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToDockerContainer, true),
 			state.EXPECT().TaskByArn(taskARN).Return(task, true).AnyTimes(),
 			state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
@@ -3002,6 +3004,7 @@ func TestGetTaskProtection(t *testing.T) {
 			state.EXPECT().TaskByArn(taskARN).Return(task, true).Times(2),
 			state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToDockerContainer, true),
 			state.EXPECT().TaskByArn(taskARN).Return(task, true),
+			state.EXPECT().ContainerByID(containerID).Return(container1, true).AnyTimes(),
 			state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
 		)
 	}

diff --git a/agent/handlers/v4/response.go b/agent/handlers/v4/response.go
@@ -55,9 +55,19 @@ func NewTaskResponse(
 		if err != nil {
 			return nil, err
 		}
+		dockerContainer, ok := state.ContainerByID(container.ID)
+		if !ok {
+			return nil, errors.Errorf(
+				"v4 container response: unable to find container '%s'", container.ID)
+		}
+		var restartCount int
+		if dockerContainer.Container.RestartPolicyEnabled() {
+			restartCount = dockerContainer.Container.RestartTracker.GetRestartCount()
+		}
 		containers = append(containers, tmdsv4.ContainerResponse{
 			ContainerResponse: &v2Resp.Containers[i],
 			Networks:          networks,
+			RestartCount:      restartCount,
 		})
 	}
 
@@ -87,9 +97,19 @@ func NewContainerResponse(
 	if err != nil {
 		return nil, err
 	}
+	dockerContainer, ok := state.ContainerByID(containerID)
+	if !ok {
+		return nil, errors.Errorf(
+			"v4 container response: unable to find container '%s'", containerID)
+	}
+	var restartCount int
+	if dockerContainer.Container.RestartPolicyEnabled() {
+		restartCount = dockerContainer.Container.RestartTracker.GetRestartCount()
+	}
 	return &tmdsv4.ContainerResponse{
 		ContainerResponse: container,
 		Networks:          networks,
+		RestartCount:      restartCount,
 	}, nil
 }