diff --git a/agent/engine/task_manager.go b/agent/engine/task_manager.go index 0c44f5d7fbe..bba3b40aca9 100644 --- a/agent/engine/task_manager.go +++ b/agent/engine/task_manager.go @@ -350,7 +350,8 @@ func (mtask *managedTask) steadyState() bool { } } -// cleanupCredentials removes credentials for a stopped task +// cleanupCredentials removes credentials for a stopped task (execution credentials are removed in cleanupTask +// due to its potential usage in the later phase of the task cleanup such as sending logs) func (mtask *managedTask) cleanupCredentials() { taskCredentialsID := mtask.GetCredentialsID() if taskCredentialsID != "" { @@ -1424,6 +1425,7 @@ func (mtask *managedTask) time() ttime.Time { } func (mtask *managedTask) cleanupTask(taskStoppedDuration time.Duration) { + taskExecutionCredentialsID := mtask.GetExecutionCredentialsID() cleanupTimeDuration := mtask.GetKnownStatusTime().Add(taskStoppedDuration).Sub(ttime.Now()) cleanupTime := make(<-chan time.Time) if cleanupTimeDuration < 0 { @@ -1464,6 +1466,14 @@ func (mtask *managedTask) cleanupTask(taskStoppedDuration time.Duration) { mtask.engine.sweepTask(mtask.Task) mtask.engine.deleteTask(mtask.Task) + // Remove TaskExecutionCredentials from credentialsManager + if taskExecutionCredentialsID != "" { + logger.Info("Cleaning up task's execution credentials", logger.Fields{ + field.TaskARN: mtask.Arn, + }) + mtask.credentialsManager.RemoveCredentials(taskExecutionCredentialsID) + } + // The last thing to do here is to cancel the context, which should cancel // all outstanding go routines associated with this managed task. mtask.cancel() diff --git a/agent/engine/task_manager_test.go b/agent/engine/task_manager_test.go index b22958fcf9c..6bb494e41b4 100644 --- a/agent/engine/task_manager_test.go +++ b/agent/engine/task_manager_test.go @@ -38,6 +38,7 @@ import ( apitaskstatus "github.com/aws/amazon-ecs-agent/agent/api/task/status" "github.com/aws/amazon-ecs-agent/agent/config" "github.com/aws/amazon-ecs-agent/agent/credentials" + mock_credentials "github.com/aws/amazon-ecs-agent/agent/credentials/mocks" "github.com/aws/amazon-ecs-agent/agent/data" mock_dockerapi "github.com/aws/amazon-ecs-agent/agent/dockerclient/dockerapi/mocks" "github.com/aws/amazon-ecs-agent/agent/engine/dependencygraph" @@ -1335,6 +1336,75 @@ func TestTaskWaitForExecutionCredentials(t *testing.T) { } } +func TestCleanupTaskWithExecutionCredentials(t *testing.T) { + cfg := getTestConfig() + ctrl := gomock.NewController(t) + mockTime := mock_ttime.NewMockTime(ctrl) + mockState := mock_dockerstate.NewMockTaskEngineState(ctrl) + mockClient := mock_dockerapi.NewMockDockerClient(ctrl) + mockImageManager := mock_engine.NewMockImageManager(ctrl) + mockCredentialsManager := mock_credentials.NewMockManager(ctrl) + mockResource := mock_taskresource.NewMockTaskResource(ctrl) + defer ctrl.Finish() + + ctx, cancel := context.WithCancel(context.TODO()) + defer cancel() + + taskEngine := &DockerTaskEngine{ + ctx: ctx, + cfg: &cfg, + dataClient: data.NewNoopClient(), + state: mockState, + client: mockClient, + imageManager: mockImageManager, + credentialsManager: mockCredentialsManager, + } + mTask := &managedTask{ + ctx: ctx, + cancel: cancel, + Task: testdata.LoadTask("sleep5"), + credentialsManager: mockCredentialsManager, + _time: mockTime, + engine: taskEngine, + acsMessages: make(chan acsTransition), + dockerMessages: make(chan dockerContainerChange), + resourceStateChangeEvent: make(chan resourceStateChange), + cfg: taskEngine.cfg, + } + + mTask.Task.ResourcesMapUnsafe = make(map[string][]taskresource.TaskResource) + mTask.Task.SetExecutionRoleCredentialsID("executionRoleCredentialsId") + mTask.AddResource("mockResource", mockResource) + mTask.SetKnownStatus(apitaskstatus.TaskStopped) + mTask.SetSentStatus(apitaskstatus.TaskStopped) + container := mTask.Containers[0] + dockerContainer := &apicontainer.DockerContainer{ + DockerName: "dockerContainer", + } + + // Expectations for triggering cleanup + now := mTask.GetKnownStatusTime() + taskStoppedDuration := 1 * time.Minute + mockTime.EXPECT().Now().Return(now).AnyTimes() + cleanupTimeTrigger := make(chan time.Time) + mockTime.EXPECT().After(gomock.Any()).Return(cleanupTimeTrigger) + go func() { + cleanupTimeTrigger <- now + }() + + // Expectations to verify the execution credentials get removed + mockCredentialsManager.EXPECT().RemoveCredentials("executionRoleCredentialsId") + + // Expectations to verify that the task gets removed + mockState.EXPECT().ContainerMapByArn(mTask.Arn).Return(map[string]*apicontainer.DockerContainer{container.Name: dockerContainer}, true) + mockClient.EXPECT().RemoveContainer(gomock.Any(), dockerContainer.DockerName, gomock.Any()).Return(nil) + mockImageManager.EXPECT().RemoveContainerReferenceFromImageState(container).Return(nil) + mockState.EXPECT().RemoveTask(mTask.Task) + mockResource.EXPECT().Cleanup() + mockResource.EXPECT().GetName() + mTask.cleanupTask(taskStoppedDuration) +} + func TestCleanupTaskWithInvalidInterval(t *testing.T) { ctrl := gomock.NewController(t) mockTime := mock_ttime.NewMockTime(ctrl)