Load Managed Daemon images in background (#3984)

Author: amogh09

agent/app/agent.go

@@ -433,19 +433,15 @@ func (agent *ecsAgent) doStart(containerChangeEventStream *eventstream.EventStre
 		return exitcodes.ExitTerminal
 	}
 
+	// Load Managed Daemon images asynchronously
+	agent.loadManagedDaemonImagesAsync(imageManager)
+
 	scManager := agent.serviceconnectManager
 	scManager.SetECSClient(client, agent.containerInstanceARN)
 	if loaded, _ := scManager.IsLoaded(agent.dockerClient); loaded {
 		imageManager.AddImageToCleanUpExclusionList(agent.serviceconnectManager.GetLoadedImageName())
 	}
 
-	// exclude all daemon images from cleanup
-	for _, csiDM := range agent.daemonManagers {
-		if loaded, _ := csiDM.IsLoaded(agent.dockerClient); loaded {
-			imageManager.AddImageToCleanUpExclusionList(csiDM.GetManagedDaemon().GetLoadedDaemonImageRef())
-		}
-	}
-
 	// Add container instance ARN to metadata manager
 	if agent.cfg.ContainerMetadataEnabled.Enabled() {
 		agent.metadataManager.SetContainerInstanceARN(agent.containerInstanceARN)
@@ -485,7 +481,7 @@ func (agent *ecsAgent) doStart(containerChangeEventStream *eventstream.EventStre
 	agent.startAsyncRoutines(containerChangeEventStream, credentialsManager, imageManager,
 		taskEngine, deregisterInstanceEventStream, client, taskHandler, attachmentEventHandler, state, doctor)
 	// TODO add EBS watcher to async routines
-	agent.startEBSWatcher(state, taskEngine)
+	agent.startEBSWatcher(state, taskEngine, agent.dockerClient)
 	// Start the acs session, which should block doStart
 	return agent.startACSSession(credentialsManager, taskEngine,
 		deregisterInstanceEventStream, client, state, taskHandler, doctor)
@@ -751,6 +747,38 @@ func (agent *ecsAgent) constructVPCSubnetAttributes() []*ecs.Attribute {
 	}
 }
 
+// Loads Managed Daemon images for all Managed Daemons registered on the Agent.
+// The images are loaded in the background. Successfully loaded images are added to
+// imageManager's cleanup exclusion list.
+func (agent *ecsAgent) loadManagedDaemonImagesAsync(imageManager engine.ImageManager) {
+	daemonManagers := agent.getDaemonManagers()
+	logger.Debug(fmt.Sprintf("Will load images for %d Managed Daemons", len(daemonManagers)))
+	for _, daemonManager := range daemonManagers {
+		go agent.loadManagedDaemonImage(daemonManager, imageManager)
+	}
+}
+
+// Loads Managed Daemon image and adds it to image cleanup exclusion list upon success.
+func (agent *ecsAgent) loadManagedDaemonImage(dm dm.DaemonManager, imageManager engine.ImageManager) {
+	imageRef := dm.GetManagedDaemon().GetImageRef()
+	logger.Info("Starting to load Managed Daemon image", logger.Fields{
+		field.ImageRef: imageRef,
+	})
+	image, err := dm.LoadImage(agent.ctx, agent.dockerClient)
+	if err != nil {
+		logger.Error("Failed to load Managed Daemon image", logger.Fields{
+			field.ImageRef: imageRef,
+			field.Error:    err,
+		})
+		return
+	}
+	logger.Info("Successfully loaded Managed Daemon image", logger.Fields{
+		field.ImageRef: imageRef,
+		field.ImageID:  image.ID,
+	})
+	imageManager.AddImageToCleanUpExclusionList(imageRef)
+}
+
 // registerContainerInstance registers the container instance ID for the ECS Agent
 func (agent *ecsAgent) registerContainerInstance(
 	client api.ECSClient,
@@ -1140,6 +1168,11 @@ func (agent *ecsAgent) setDaemonManager(key string, val dm.DaemonManager) {
 	agent.daemonManagers[key] = val
 }
 
+// Returns daemon managers map. Not designed to be thread-safe.
+func (agent *ecsAgent) getDaemonManagers() map[string]dm.DaemonManager {
+	return agent.daemonManagers
+}
+
 // setVPCSubnet sets the vpc and subnet ids for the agent by querying the
 // instance metadata service
 func (agent *ecsAgent) setVPCSubnet() (error, bool) {

agent/app/agent_capability.go

@@ -522,12 +522,16 @@ func (agent *ecsAgent) appendEBSTaskAttachCapabilities(capabilities []*ecs.Attri
 		if daemonDef.GetImageName() == md.EbsCsiDriver {
 			csiDaemonManager := dm.NewDaemonManager(daemonDef)
 			agent.setDaemonManager(md.EbsCsiDriver, csiDaemonManager)
-			if _, err := csiDaemonManager.LoadImage(agent.ctx, agent.dockerClient); err != nil {
-				logger.Error("Failed to load the EBS CSI Driver. This container instance will not be able to support EBS Task Attach",
+			imageExists, err := csiDaemonManager.ImageExists()
+			if !imageExists {
+				logger.Error(
+					"Either EBS Daemon image does not exist or failed to check its existence."+
+						" This container instance will not advertise EBS Task Attach capability.",
 					logger.Fields{
-						field.Error: err,
-					},
-				)
+						field.ImageName:    csiDaemonManager.GetManagedDaemon().GetImageName(),
+						field.ImageTARPath: csiDaemonManager.GetManagedDaemon().GetImageTarPath(),
+						field.Error:        err,
+					})
 				return capabilities
 			}
 		}

agent/app/agent_test.go

@@ -59,6 +59,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	aws_credentials "github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/docker/docker/api/types"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -1859,3 +1860,49 @@ func TestWaitUntilInstanceInServicePolling(t *testing.T) {
 		})
 	}
 }
+
+func TestLoadManagedDaemonImage(t *testing.T) {
+	tcs := []struct {
+		name                         string
+		setDaemonManagerExpectations func(*mock_daemonmanager.MockDaemonManager)
+		setImageManagerExpectations  func(*mock_engine.MockImageManager)
+	}{
+		{
+			name: "no exclusion list update if image load fails",
+			setDaemonManagerExpectations: func(mdm *mock_daemonmanager.MockDaemonManager) {
+				mdm.EXPECT().GetManagedDaemon().Return(md.NewManagedDaemon("name", "tag")).Times(2)
+				mdm.EXPECT().LoadImage(gomock.Any(), gomock.Any()).Return(nil, errors.New("error"))
+			},
+		},
+		{
+			name: "exclusion list is updated if image load succeeds",
+			setDaemonManagerExpectations: func(mdm *mock_daemonmanager.MockDaemonManager) {
+				mdm.EXPECT().GetManagedDaemon().Return(md.NewManagedDaemon("name", "tag")).Times(3)
+				mdm.EXPECT().
+					LoadImage(gomock.Any(), gomock.Any()).
+					Return(&types.ImageInspect{ID: "image-id"}, nil)
+			},
+			setImageManagerExpectations: func(mim *mock_engine.MockImageManager) {
+				mim.EXPECT().AddImageToCleanUpExclusionList("name:tag")
+			},
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			ctrl.Finish()
+
+			dockerClient := mock_dockerapi.NewMockDockerClient(ctrl)
+			daemonManager := mock_daemonmanager.NewMockDaemonManager(ctrl)
+			imageManager := mock_engine.NewMockImageManager(ctrl)
+
+			tc.setDaemonManagerExpectations(daemonManager)
+			if tc.setImageManagerExpectations != nil {
+				tc.setImageManagerExpectations(imageManager)
+			}
+
+			agent := &ecsAgent{ctx: context.Background(), dockerClient: dockerClient}
+			agent.loadManagedDaemonImage(daemonManager, imageManager)
+		})
+	}
+}

agent/app/agent_unix.go

@@ -21,6 +21,7 @@ import (
 
 	asmfactory "github.com/aws/amazon-ecs-agent/agent/asm/factory"
 	"github.com/aws/amazon-ecs-agent/agent/config"
+	"github.com/aws/amazon-ecs-agent/agent/dockerclient/dockerapi"
 	ebs "github.com/aws/amazon-ecs-agent/agent/ebs"
 	"github.com/aws/amazon-ecs-agent/agent/ecscni"
 	"github.com/aws/amazon-ecs-agent/agent/engine"
@@ -151,10 +152,14 @@ func (agent *ecsAgent) startENIWatcher(state dockerstate.TaskEngineState, stateC
 	return nil
 }
 
-func (agent *ecsAgent) startEBSWatcher(state dockerstate.TaskEngineState, taskEngine engine.TaskEngine) {
+func (agent *ecsAgent) startEBSWatcher(
+	state dockerstate.TaskEngineState,
+	taskEngine engine.TaskEngine,
+	dockerClient dockerapi.DockerClient,
+) {
 	if agent.ebsWatcher == nil {
 		seelog.Debug("Creating new EBS watcher...")
-		agent.ebsWatcher = ebs.NewWatcher(agent.ctx, state, taskEngine)
+		agent.ebsWatcher = ebs.NewWatcher(agent.ctx, state, taskEngine, dockerClient)
 		go agent.ebsWatcher.Start()
 	}
 }

agent/app/agent_windows.go

@@ -24,6 +24,7 @@ import (
 
 	asmfactory "github.com/aws/amazon-ecs-agent/agent/asm/factory"
 	"github.com/aws/amazon-ecs-agent/agent/data"
+	"github.com/aws/amazon-ecs-agent/agent/dockerclient/dockerapi"
 	"github.com/aws/amazon-ecs-agent/agent/ecscni"
 	"github.com/aws/amazon-ecs-agent/agent/engine"
 	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
@@ -99,9 +100,12 @@ func (agent *ecsAgent) startWindowsService() int {
 	return 0
 }
 
-func (agent *ecsAgent) startEBSWatcher(state dockerstate.TaskEngineState, taskEngine engine.TaskEngine) error {
+func (agent *ecsAgent) startEBSWatcher(
+	state dockerstate.TaskEngineState,
+	taskEngine engine.TaskEngine,
+	dockerClient dockerapi.DockerClient,
+) {
 	seelog.Debug("Windows EBS Watcher not implemented: No Op")
-	return nil
 }
 
 // handler implements https://godoc.org/golang.org/x/sys/windows/svc#Handler