agent/app/agent.go: Reset agent state if the instance id changed

Author: richardpen

agent/app/agent.go

@@ -232,6 +232,8 @@ func (agent *ecsAgent) newTaskEngine(containerChangeEventStream *eventstream.Eve
 		log.Warnf(instanceIDMismatchErrorFormat,
 			previousEC2InstanceID, currentEC2InstanceID)
 
+		// Reset agent state as a new container instance
+		state.Reset()
 		// Reset taskEngine; all the other values are still default
 		return engine.NewTaskEngine(agent.cfg, agent.dockerClient, credentialsManager,
 			containerChangeEventStream, imageManager, state), currentEC2InstanceID, nil

agent/app/agent_test.go

@@ -287,6 +287,7 @@ func TestNewTaskEngineRestoreFromCheckpointPreviousEC2InstanceIDLoadedHappyPath(
 			gomock.Any(), gomock.Any(), gomock.Any()).Return(
 			statemanager.NewNoopStateManager(), nil),
 		ec2MetadataClient.EXPECT().InstanceIdentityDocument().Return(iid, nil),
+		state.EXPECT().Reset(),
 	)
 
 	ctx, cancel := context.WithCancel(context.TODO())

agent/engine/dockerstate/docker_task_engine_state.go

@@ -52,6 +52,8 @@ type TaskEngineState interface {
 	AddImageState(imageState *image.ImageState)
 	// RemoveTask removes a task from the state
 	RemoveTask(task *api.Task)
+	// Reset resets all the fileds in the state
+	Reset()
 	// RemoveImageState removes an image.ImageState
 	RemoveImageState(imageState *image.ImageState)
 	json.Marshaler
@@ -96,6 +98,18 @@ func newDockerTaskEngineState() *DockerTaskEngineState {
 	}
 }
 
+// Reset resets all the states
+func (state *DockerTaskEngineState) Reset() {
+	state.lock.Lock()
+	defer state.lock.Unlock()
+
+	state.tasks = make(map[string]*api.Task)
+	state.idToTask = make(map[string]string)
+	state.taskToID = make(map[string]map[string]*api.DockerContainer)
+	state.idToContainer = make(map[string]*api.DockerContainer)
+	state.imageStates = make(map[string]*image.ImageState)
+}
+
 // AllTasks returns all of the tasks
 func (state *DockerTaskEngineState) AllTasks() []*api.Task {
 	state.lock.RLock()

agent/engine/dockerstate/mocks/dockerstate_mocks.go

@@ -146,6 +146,14 @@ func (_mr *_MockTaskEngineStateRecorder) RemoveTask(arg0 interface{}) *gomock.Ca
 	return _mr.mock.ctrl.RecordCall(_mr.mock, "RemoveTask", arg0)
 }
 
+func (_m *MockTaskEngineState) Reset() {
+	_m.ctrl.Call(_m, "Reset")
+}
+
+func (_mr *_MockTaskEngineStateRecorder) Reset() *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "Reset")
+}
+
 func (_m *MockTaskEngineState) TaskByArn(_param0 string) (*api.Task, bool) {
 	ret := _m.ctrl.Call(_m, "TaskByArn", _param0)
 	ret0, _ := ret[0].(*api.Task)