Fix concurrent writes during image cleanup

This patch inspects the cleanup and resolves the inherent concurrent map write
issue reported. An unit test has been added to increase confidence in the fix.

Fixes aws#707

Signed-off-by: Vinothkumar Siddharth <sidvin@amazon.com>

agent/engine/docker_image_manager.go

@@ -193,8 +193,6 @@ func (imageManager *dockerImageManager) getImageState(containerImageID string) (
 
 // removeImageState removes the imageState from the list of imageState objects in ImageManager
 func (imageManager *dockerImageManager) removeImageState(imageStateToBeRemoved *image.ImageState) {
-	imageManager.updateLock.Lock()
-	defer imageManager.updateLock.Unlock()
 	for i, imageState := range imageManager.imageStates {
 		if imageState.Image.ImageID == imageStateToBeRemoved.Image.ImageID {
 			// Image State found; hence remove it
@@ -277,6 +275,8 @@ func (imageManager *dockerImageManager) performPeriodicImageCleanup(ctx context.
 }
 
 func (imageManager *dockerImageManager) removeUnusedImages() {
+	imageManager.updateLock.Lock()
+	defer imageManager.updateLock.Unlock()
 	imageManager.imageStatesConsideredForDeletion = make(map[string]*image.ImageState)
 	for _, imageState := range imageManager.getAllImageStates() {
 		imageManager.imageStatesConsideredForDeletion[imageState.Image.ImageID] = imageState
@@ -305,8 +305,6 @@ func (imageManager *dockerImageManager) removeLeastRecentlyUsedImage() error {
 }
 
 func (imageManager *dockerImageManager) getUnusedImageForDeletion() *image.ImageState {
-	imageManager.updateLock.RLock()
-	defer imageManager.updateLock.RUnlock()
 	candidateImageStatesForDeletion := imageManager.getCandidateImagesForDeletion()
 	if len(candidateImageStatesForDeletion) < 1 {
 		seelog.Infof("No eligible images for deletion for this cleanup cycle")

agent/engine/docker_image_manager_test.go

@@ -17,6 +17,7 @@ package engine
 import (
 	"errors"
 	"reflect"
+	"sync"
 	"testing"
 	"time"
 
@@ -27,6 +28,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/statemanager"
 	docker "github.com/fsouza/go-dockerclient"
 	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/require"
 	"golang.org/x/net/context"
 )
 
@@ -951,3 +953,69 @@ func TestGetImageStateFromImageNameNoImageState(t *testing.T) {
 		t.Error("Incorrect image state retrieved by image name")
 	}
 }
+
+func TestConcurrentRemoveUnusedImages(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	client := NewMockDockerClient(ctrl)
+
+	imageManager := &dockerImageManager{
+		client: client,
+		state:  dockerstate.NewTaskEngineState(),
+		minimumAgeBeforeDeletion: config.DefaultImageDeletionAge,
+		numImagesToDelete:        config.DefaultNumImagesToDeletePerCycle,
+		imageCleanupTimeInterval: config.DefaultImageCleanupTimeInterval,
+	}
+
+	imageManager.SetSaver(statemanager.NewNoopStateManager())
+	container := &api.Container{
+		Name:  "testContainer",
+		Image: "testContainerImage",
+	}
+	sourceImage := &image.Image{
+		ImageID: "sha256:qwerty",
+	}
+	sourceImage.Names = append(sourceImage.Names, container.Image)
+	imageInspected := &docker.Image{
+		ID: "sha256:qwerty",
+	}
+	client.EXPECT().InspectImage(container.Image).Return(imageInspected, nil).AnyTimes()
+	err := imageManager.RecordContainerReference(container)
+	if err != nil {
+		t.Error("Error in adding container to an existing image state")
+	}
+
+	require.Equal(t, 1, len(imageManager.imageStates))
+
+	err = imageManager.RemoveContainerReferenceFromImageState(container)
+	if err != nil {
+		t.Error("Error removing container reference from image state")
+	}
+
+	imageState, _ := imageManager.getImageState(imageInspected.ID)
+	imageState.RemoveImageName(container.Image)
+	imageState.PulledAt = time.Now().AddDate(0, -2, 0)
+	imageState.LastUsedAt = time.Now().AddDate(0, -2, 0)
+
+	client.EXPECT().RemoveImage(sourceImage.ImageID, removeImageTimeout).Return(nil)
+	require.Equal(t, 1, len(imageManager.imageStates))
+
+	numRoutines := 1000
+	var waitGroup sync.WaitGroup
+	waitGroup.Add(numRoutines)
+
+	ok := make(chan bool)
+
+	for i := 0; i < numRoutines; i++ {
+		go func() {
+			<-ok
+			imageManager.removeUnusedImages()
+			waitGroup.Done()
+		}()
+
+	}
+
+	close(ok)
+	waitGroup.Wait()
+	require.Equal(t, 0, len(imageManager.imageStates))
+}