From a9f7b63ab2c56c4cbb8b6b5f78ea5bd138b4a110 Mon Sep 17 00:00:00 2001
From: fenxiong <fenxiong@amazon.com>
Date: Wed, 2 Sep 2020 16:38:17 -0700
Subject: [PATCH] engine: save dockerID in createContainer.

Previously, the dockerID in the dockerContainer struct of a container is only saved in termination handler. This means that if the agent is forcefully stopped (e.g. via SIGKILL instead of SIGTERM), the dockerID is not saved, and as a result loses track of container if it is stopped by SIGKILL. This commit fixes the issue by saving the dockerID  after creating the container.
---
 agent/engine/docker_task_engine.go      |  6 +++--
 agent/engine/docker_task_engine_test.go | 30 +++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/agent/engine/docker_task_engine.go b/agent/engine/docker_task_engine.go
index 263727b1dbe..f3521afac20 100644
--- a/agent/engine/docker_task_engine.go
+++ b/agent/engine/docker_task_engine.go
@@ -1102,9 +1102,11 @@ func (engine *DockerTaskEngine) createContainer(task *apitask.Task, container *a
 	if metadata.DockerID != "" {
 		seelog.Infof("Task engine [%s]: created docker container for task: %s -> %s",
 			task.Arn, container.Name, metadata.DockerID)
-		engine.state.AddContainer(&apicontainer.DockerContainer{DockerID: metadata.DockerID,
+		dockerContainer := &apicontainer.DockerContainer{DockerID: metadata.DockerID,
 			DockerName: dockerContainerName,
-			Container:  container}, task)
+			Container:  container}
+		engine.state.AddContainer(dockerContainer, task)
+		engine.saveDockerContainerData(dockerContainer)
 	}
 	container.SetLabels(config.Labels)
 	seelog.Infof("Task engine [%s]: created docker container for task: %s -> %s, took %s",
diff --git a/agent/engine/docker_task_engine_test.go b/agent/engine/docker_task_engine_test.go
index bf424fe41bc..7bc3a0ce8f0 100644
--- a/agent/engine/docker_task_engine_test.go
+++ b/agent/engine/docker_task_engine_test.go
@@ -699,6 +699,36 @@ func TestStopWithPendingStops(t *testing.T) {
 	// gets the pull image lock
 }
 
+func TestCreateContainerSaveDockerIDAndName(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
+	ctrl, client, _, privateTaskEngine, _, _, _ := mocks(t, ctx, &defaultConfig)
+	defer ctrl.Finish()
+	dataClient, cleanup := newTestDataClient(t)
+	defer cleanup()
+
+	taskEngine, _ := privateTaskEngine.(*DockerTaskEngine)
+	taskEngine.SetDataClient(dataClient)
+
+	sleepTask := testdata.LoadTask("sleep5")
+	sleepTask.Arn = testTaskARN
+	sleepContainer, _ := sleepTask.ContainerByName("sleep5")
+	sleepContainer.TaskARNUnsafe = testTaskARN
+
+	client.EXPECT().APIVersion().Return(defaultDockerClientAPIVersion, nil).AnyTimes()
+	client.EXPECT().CreateContainer(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(dockerapi.DockerContainerMetadata{
+		DockerID: testDockerID,
+	})
+	metadata := taskEngine.createContainer(sleepTask, sleepContainer)
+	require.NoError(t, metadata.Error)
+
+	containers, err := dataClient.GetContainers()
+	require.NoError(t, err)
+	require.Len(t, containers, 1)
+	assert.Equal(t, testDockerID, containers[0].DockerID)
+	assert.Contains(t, containers[0].DockerName, sleepContainer.Name)
+}
+
 func TestCreateContainerMetadata(t *testing.T) {
 	testcases := []struct {
 		name  string