From d88c4ba263d395edc1eb442010322f7ba8b9a20d Mon Sep 17 00:00:00 2001 From: Dane H Lim Date: Mon, 19 Aug 2024 15:02:57 -0700 Subject: [PATCH] Stop tasks via task engine AddTask for full task sync --- .../task_stop_verification_ack_responder.go | 5 +++-- ...sk_stop_verification_ack_responder_test.go | 19 +++++++++++++++++++ agent/engine/common_integ_testutil.go | 13 ------------- agent/engine/engine_windows_integ_test.go | 4 ---- 4 files changed, 22 insertions(+), 19 deletions(-) diff --git a/agent/acs/session/task_stop_verification_ack_responder.go b/agent/acs/session/task_stop_verification_ack_responder.go index 96420bc7a4e..00981f18982 100644 --- a/agent/acs/session/task_stop_verification_ack_responder.go +++ b/agent/acs/session/task_stop_verification_ack_responder.go @@ -41,8 +41,9 @@ func (ts *taskStopper) StopTask(taskARN string) { logger.Info("Stopping task from task stop verification ACK: %s", logger.Fields{ loggerfield.TaskARN: task.Arn, }) - task.SetDesiredStatus(apitaskstatus.TaskStopped) - task.UpdateDesiredStatus() + taskWithDesiredStatusStopped := *task + taskWithDesiredStatusStopped.SetDesiredStatus(apitaskstatus.TaskStopped) + ts.taskEngine.AddTask(&taskWithDesiredStatusStopped) if err := ts.dataClient.SaveTask(task); err != nil { logger.Error("Failed to save data for task", logger.Fields{ loggerfield.TaskARN: task.Arn, diff --git a/agent/acs/session/task_stop_verification_ack_responder_test.go b/agent/acs/session/task_stop_verification_ack_responder_test.go index f7f28826ada..dc4600656d4 100644 --- a/agent/acs/session/task_stop_verification_ack_responder_test.go +++ b/agent/acs/session/task_stop_verification_ack_responder_test.go @@ -20,6 +20,7 @@ import ( "github.com/aws/amazon-ecs-agent/agent/api/container" "github.com/aws/amazon-ecs-agent/agent/api/task" + apitask "github.com/aws/amazon-ecs-agent/agent/api/task" "github.com/aws/amazon-ecs-agent/agent/data" mock_engine "github.com/aws/amazon-ecs-agent/agent/engine/mocks" "github.com/aws/amazon-ecs-agent/ecs-agent/acs/model/ecsacs" @@ -120,7 +121,9 @@ func TestTaskStopVerificationAckResponderStopsMultipleTasks(t *testing.T) { } tester.taskEngine.EXPECT().GetTaskByArn(taskARN2).Return(tasksOnInstance[taskARN2], true) + mockTaskEngineExpectAddTaskWithArnAndDesiredStatus(t, tester, tasksOnInstance, taskARN2, apitaskstatus.TaskStopped) tester.taskEngine.EXPECT().GetTaskByArn(taskARN3).Return(tasksOnInstance[taskARN3], true) + mockTaskEngineExpectAddTaskWithArnAndDesiredStatus(t, tester, tasksOnInstance, taskARN3, apitaskstatus.TaskStopped) handleTaskStopVerificationAck := tester.taskStopVerificationAckResponder.HandlerFunc().(func(message *ecsacs.TaskStopVerificationAck)) @@ -172,8 +175,11 @@ func TestTaskStopVerificationAckResponderStopsAllTasks(t *testing.T) { } tester.taskEngine.EXPECT().GetTaskByArn(taskARN1).Return(tasksOnInstance[taskARN1], true) + mockTaskEngineExpectAddTaskWithArnAndDesiredStatus(t, tester, tasksOnInstance, taskARN1, apitaskstatus.TaskStopped) tester.taskEngine.EXPECT().GetTaskByArn(taskARN2).Return(tasksOnInstance[taskARN2], true) + mockTaskEngineExpectAddTaskWithArnAndDesiredStatus(t, tester, tasksOnInstance, taskARN2, apitaskstatus.TaskStopped) tester.taskEngine.EXPECT().GetTaskByArn(taskARN3).Return(tasksOnInstance[taskARN3], true) + mockTaskEngineExpectAddTaskWithArnAndDesiredStatus(t, tester, tasksOnInstance, taskARN3, apitaskstatus.TaskStopped) handleTaskStopVerificationAck := tester.taskStopVerificationAckResponder.HandlerFunc().(func(message *ecsacs.TaskStopVerificationAck)) @@ -186,3 +192,16 @@ func TestTaskStopVerificationAckResponderStopsAllTasks(t *testing.T) { assert.Equal(t, apicontainerstatus.ContainerStopped, task.Containers[0].GetDesiredStatus()) } } + +// mockTaskEngineExpectAddTaskWithArnAndDesiredStatus expects the mock task engine to call AddTask and asserts that the +// added task has task ARN `taskARN` and desired status `desiredStatus` before setting and updating the desired status +// of the task and its containers on the instance. +func mockTaskEngineExpectAddTaskWithArnAndDesiredStatus(t *testing.T, tester *taskStopVerificationAckTestHelper, + tasksOnInstance map[string]*task.Task, taskARN string, desiredStatus apitaskstatus.TaskStatus) { + tester.taskEngine.EXPECT().AddTask(gomock.Any()).Do(func(addedTask *apitask.Task) { + assert.Equal(t, taskARN, addedTask.Arn) + assert.Equal(t, desiredStatus, addedTask.GetDesiredStatus()) + tasksOnInstance[taskARN].SetDesiredStatus(desiredStatus) + tasksOnInstance[taskARN].UpdateDesiredStatus() + }) +} diff --git a/agent/engine/common_integ_testutil.go b/agent/engine/common_integ_testutil.go index 949f96ec75f..71e275f24f3 100644 --- a/agent/engine/common_integ_testutil.go +++ b/agent/engine/common_integ_testutil.go @@ -52,11 +52,6 @@ var ( sdkClientFactory sdkclientfactory.Factory ) -const ( - taskSteadyStatePollInterval = 100 * time.Millisecond - taskSteadyStatePollIntervalJitter = 10 * time.Millisecond -) - func init() { sdkClientFactory = sdkclientfactory.NewFactory(context.TODO(), dockerEndpoint) } @@ -124,10 +119,6 @@ func setupGMSALinux(cfg *config.Config, state dockerstate.TaskEngineState, t *te taskEngine := NewDockerTaskEngine(cfg, dockerClient, credentialsManager, eventstream.NewEventStream("ENGINEINTEGTEST", context.Background()), imageManager, &hostResourceManager, state, metadataManager, resourceFields, execcmd.NewManager(), engineserviceconnect.NewManager(), daemonManagers) - // Set the steady state poll interval to a low value so that tasks transition from their current state to their - // desired state faster. This prevents tests from appearing to hang while waiting for state change events. - taskEngine.taskSteadyStatePollInterval = taskSteadyStatePollInterval - taskEngine.taskSteadyStatePollIntervalJitter = taskSteadyStatePollIntervalJitter taskEngine.MustInit(context.TODO()) return taskEngine, func() { taskEngine.Shutdown() @@ -269,10 +260,6 @@ func SetupIntegTestTaskEngine(cfg *config.Config, state dockerstate.TaskEngineSt taskEngine := NewDockerTaskEngine(cfg, dockerClient, credentialsManager, eventstream.NewEventStream("ENGINEINTEGTEST", context.Background()), imageManager, &hostResourceManager, state, metadataManager, nil, execcmd.NewManager(), engineserviceconnect.NewManager(), daemonManagers) - // Set the steady state poll interval to a low value so that tasks transition from their current state to their - // desired state faster. This prevents tests from appearing to hang while waiting for state change events. - taskEngine.taskSteadyStatePollInterval = taskSteadyStatePollInterval - taskEngine.taskSteadyStatePollIntervalJitter = taskSteadyStatePollIntervalJitter taskEngine.MustInit(context.TODO()) return taskEngine, func() { taskEngine.Shutdown() diff --git a/agent/engine/engine_windows_integ_test.go b/agent/engine/engine_windows_integ_test.go index d69446e2398..1472a6e5582 100644 --- a/agent/engine/engine_windows_integ_test.go +++ b/agent/engine/engine_windows_integ_test.go @@ -560,10 +560,6 @@ func setupGMSA(cfg *config.Config, state dockerstate.TaskEngineState, t *testing taskEngine := NewDockerTaskEngine(cfg, dockerClient, credentialsManager, eventstream.NewEventStream("ENGINEINTEGTEST", context.Background()), imageManager, &hostResourceManager, state, metadataManager, resourceFields, execcmd.NewManager(), engineserviceconnect.NewManager()) - // Set the steady state poll interval to a low value so that tasks transition from their current state to their - // desired state faster. This prevents tests from appearing to hang while waiting for state change events. - taskEngine.taskSteadyStatePollInterval = taskSteadyStatePollInterval - taskEngine.taskSteadyStatePollIntervalJitter = taskSteadyStatePollIntervalJitter taskEngine.MustInit(context.TODO()) return taskEngine, func() { taskEngine.Shutdown()