Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add restart count to TMDE v4 response #4166

Merged
merged 2 commits into from
May 22, 2024
Merged

Add restart count to TMDE v4 response #4166

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
69b05a2
Add restart count to TMDE v4 response
sparrc May 6, 2024
1e27c89
Add warning log message if container is not found in internal state
sparrc May 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 25 additions & 11 deletions agent/handlers/task_server_setup_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1381,7 +1381,7 @@ func testTMDSRequest[R TMDSResponse](t *testing.T, tc TMDSTestCase[R]) {
// Parse the response body
var actualResponseBody R
err = json.Unmarshal(recorder.Body.Bytes(), &actualResponseBody)
require.NoError(t, err)
require.NoError(t, err, recorder.Body.String())

// Assert status code and body
assert.Equal(t, tc.expectedStatusCode, recorder.Code)
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -1832,7 +1833,7 @@ func TestV4ContainerMetadata(t *testing.T) {
state.EXPECT().DockerIDByV3EndpointID(v3EndpointID).Return(containerID, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainerNoNetwork, true),
state.EXPECT().TaskByID(containerID).Return(bridgeTask, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainerNoNetwork, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainerNoNetwork, true).AnyTimes(),
)
},
expectedStatusCode: http.StatusInternalServerError,
Expand All @@ -1846,9 +1847,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(bridgeContainer, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainer, true).AnyTimes(),
state.EXPECT().TaskByID(containerID).Return(bridgeTask, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainer, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainer, true).AnyTimes(),
)
},
expectedStatusCode: http.StatusOK,
Expand Down Expand Up @@ -1939,8 +1940,10 @@ func TestV4TaskMetadata(t *testing.T) {
gomock.InOrder(
state.EXPECT().TaskARNByV3EndpointID(v3EndpointID).Return(taskARN, true),
state.EXPECT().TaskByArn(taskARN).Return(task, true).Times(2),
state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true).AnyTimes(),
state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToDockerContainer, true),
state.EXPECT().TaskByArn(taskARN).Return(task, true),
state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true).AnyTimes(),
state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
)
},
Expand All @@ -1955,8 +1958,10 @@ func TestV4TaskMetadata(t *testing.T) {
gomock.InOrder(
state.EXPECT().TaskARNByV3EndpointID(v3EndpointID).Return(taskARN, true),
state.EXPECT().TaskByArn(taskARN).Return(pulledTask, true).Times(2),
state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true).AnyTimes(),
state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToDockerContainer, true),
state.EXPECT().TaskByArn(taskARN).Return(pulledTask, true),
state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true).AnyTimes(),
state.EXPECT().PulledContainerMapByArn(taskARN).Return(pulledContainerNameToDockerContainer, true),
)
},
Expand All @@ -1972,8 +1977,9 @@ func TestV4TaskMetadata(t *testing.T) {
state.EXPECT().TaskARNByV3EndpointID(v3EndpointID).Return(taskARN, true),
state.EXPECT().TaskByArn(taskARN).Return(bridgeTask, true).Times(2),
state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToBridgeContainer, true),
state.EXPECT().ContainerByID(containerID).Return(nil, false),
state.EXPECT().ContainerByID(containerID).Return(nil, false).AnyTimes(),
state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
state.EXPECT().ContainerByID(containerID).Return(nil, false).AnyTimes(),
)
},
expectedStatusCode: http.StatusOK,
Expand All @@ -1988,8 +1994,9 @@ func TestV4TaskMetadata(t *testing.T) {
state.EXPECT().TaskARNByV3EndpointID(v3EndpointID).Return(taskARN, true),
state.EXPECT().TaskByArn(taskARN).Return(bridgeTask, true).Times(2),
state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToBridgeContainer, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainerNoNetwork, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainerNoNetwork, true).AnyTimes(),
state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainerNoNetwork, true).AnyTimes(),
)
},
expectedStatusCode: http.StatusOK,
Expand All @@ -2004,8 +2011,9 @@ func TestV4TaskMetadata(t *testing.T) {
state.EXPECT().TaskARNByV3EndpointID(v3EndpointID).Return(taskARN, true),
state.EXPECT().TaskByArn(taskARN).Return(bridgeTask, true).Times(2),
state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToBridgeContainer, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainer, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainer, true).AnyTimes(),
state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainer, true).AnyTimes(),
)
},
expectedStatusCode: http.StatusOK,
Expand Down Expand Up @@ -2340,8 +2348,10 @@ 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(bridgeContainer, true).AnyTimes(),
state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToDockerContainer, true),
state.EXPECT().TaskByArn(taskARN).Return(task, true).AnyTimes(),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainer, true).AnyTimes(),
state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
)
}
Expand Down Expand Up @@ -2495,8 +2505,9 @@ func TestV4TaskMetadataWithTags(t *testing.T) {
state.EXPECT().TaskARNByV3EndpointID(v3EndpointID).Return(taskARN, true),
state.EXPECT().TaskByArn(taskARN).Return(bridgeTask, true).Times(2),
state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToBridgeContainer, true),
state.EXPECT().ContainerByID(containerID).Return(nil, false),
state.EXPECT().ContainerByID(containerID).Return(nil, false).AnyTimes(),
state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
state.EXPECT().ContainerByID(containerID).Return(nil, false).AnyTimes(),
)
},
setECSClientExpectations: happyECSClientExpectations,
Expand All @@ -2515,8 +2526,9 @@ func TestV4TaskMetadataWithTags(t *testing.T) {
state.EXPECT().TaskARNByV3EndpointID(v3EndpointID).Return(taskARN, true),
state.EXPECT().TaskByArn(taskARN).Return(bridgeTask, true).Times(2),
state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToBridgeContainer, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainerNoNetwork, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainerNoNetwork, true).AnyTimes(),
state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
state.EXPECT().ContainerByID(containerID).Return(bridgeContainerNoNetwork, true).AnyTimes(),
)
},
setECSClientExpectations: happyECSClientExpectations,
Expand Down Expand Up @@ -3002,6 +3014,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(dockerContainer, true).AnyTimes(),
state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
)
}
Expand Down Expand Up @@ -3273,6 +3286,7 @@ func TestUpdateTaskProtection(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(dockerContainer, true).AnyTimes(),
state.EXPECT().PulledContainerMapByArn(taskARN).Return(nil, true),
)
}
Expand Down
43 changes: 36 additions & 7 deletions agent/handlers/v4/response.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@ import (
"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
v2 "github.com/aws/amazon-ecs-agent/agent/handlers/v2"
"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs"
"github.com/aws/amazon-ecs-agent/ecs-agent/logger"
"github.com/aws/amazon-ecs-agent/ecs-agent/logger/field"
ni "github.com/aws/amazon-ecs-agent/ecs-agent/netlib/model/networkinterface"
tmdsresponse "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/response"
"github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/utils"
Expand All @@ -28,7 +30,7 @@ import (
)

// NewTaskResponse creates a new v4 response object for the task. It augments v2 task response
// with additional network interface fields.
// with additional fields for the v4 response.
func NewTaskResponse(
taskARN string,
state dockerstate.TaskEngineState,
Expand All @@ -55,10 +57,12 @@ func NewTaskResponse(
if err != nil {
return nil, err
}
containers = append(containers, tmdsv4.ContainerResponse{
v4Response := &tmdsv4.ContainerResponse{
ContainerResponse: &v2Resp.Containers[i],
Networks: networks,
})
}
v4Response = augmentContainerResponse(container.ID, state, v4Response)
containers = append(containers, *v4Response)
}

return &tmdsv4.TaskResponse{
Expand All @@ -69,8 +73,8 @@ func NewTaskResponse(
}, nil
}

// NewContainerResponse creates a new v4 container response based on container id. It augments
// v4 container response with additional network interface fields.
// NewContainerResponse creates a new v4 container response based on container id. It augments
// v2 container response with additional fields for v4.
func NewContainerResponse(
containerID string,
state dockerstate.TaskEngineState,
Expand All @@ -87,10 +91,11 @@ func NewContainerResponse(
if err != nil {
return nil, err
}
return &tmdsv4.ContainerResponse{
v4Response := tmdsv4.ContainerResponse{
ContainerResponse: container,
Networks: networks,
}, nil
}
return augmentContainerResponse(containerID, state, &v4Response), nil
}

// toV4NetworkResponse converts v2 network response to v4. Additional fields are only
Expand Down Expand Up @@ -121,6 +126,30 @@ func toV4NetworkResponse(
return resp, nil
}

// augmentContainerResponse augments the container response with additional fields.
func augmentContainerResponse(
containerID string,
state dockerstate.TaskEngineState,
v4Response *tmdsv4.ContainerResponse,
) *tmdsv4.ContainerResponse {
dockerContainer, ok := state.ContainerByID(containerID)
if !ok {
// did not find container, continue on and try next container(s)
// we don't return error here to avoid disrupting all of a TMDS response
// on a single missing container.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we add a log at warning or debug level if the container cannot be found from the state?
Like https://github.com/aws/amazon-ecs-agent/blob/master/agent/handlers/v4/stats_response.go#L41-L46.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes I can add that

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

logger.Warn("V4 container response: unable to find container in internal state",
logger.Fields{
field.RuntimeID: containerID,
})
return v4Response
}
if dockerContainer.Container.RestartPolicyEnabled() {
restartCount := dockerContainer.Container.RestartTracker.GetRestartCount()
v4Response.RestartCount = &restartCount
}
return v4Response
}

// newNetworkInterfaceProperties creates the NetworkInterfaceProperties object for a given
// task.
func newNetworkInterfaceProperties(task *apitask.Task) (tmdsv4.NetworkInterfaceProperties, error) {
Expand Down
15 changes: 6 additions & 9 deletions agent/handlers/v4/response_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,11 +132,10 @@ func TestNewTaskContainerResponses(t *testing.T) {
containerNameToDockerContainer := map[string]*apicontainer.DockerContainer{
taskARN: dockerContainer,
}
gomock.InOrder(
state.EXPECT().TaskByArn(taskARN).Return(task, true),
state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToDockerContainer, true),
state.EXPECT().TaskByArn(taskARN).Return(task, true),
)
state.EXPECT().TaskByArn(taskARN).Return(task, true)
state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true).AnyTimes()
state.EXPECT().ContainerMapByArn(taskARN).Return(containerNameToDockerContainer, true)
state.EXPECT().TaskByArn(taskARN).Return(task, true)

taskResponse, err := NewTaskResponse(taskARN, state, ecsClient, cluster,
availabilityZone, vpcID, containerInstanceArn, task.ServiceName, false)
Expand All @@ -150,10 +149,8 @@ func TestNewTaskContainerResponses(t *testing.T) {
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),
state.EXPECT().TaskByID(containerID).Return(task, true).Times(2),
)
state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true).AnyTimes()
state.EXPECT().TaskByID(containerID).Return(task, true).Times(2)
containerResponse, err := NewContainerResponse(containerID, state)
require.NoError(t, err)
_, err = json.Marshal(containerResponse)
Expand Down
Loading