Fixing task volumes from payload for EBS-backed tasks

agent/acs/session/payload_responder.go

@@ -22,6 +22,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/engine"
 	"github.com/aws/amazon-ecs-agent/agent/eventhandler"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/acs/model/ecsacs"
+	apiresource "github.com/aws/amazon-ecs-agent/ecs-agent/api/resource"
 	apitaskstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/task/status"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/logger"
@@ -106,6 +107,16 @@ func (pmHandler *payloadMessageHandler) addPayloadTasks(payload *ecsacs.PayloadM
 			allTasksOK = false
 			continue
 		}
+
+		// Note: If we receive an EBS-backed task, we'll also received an incomplete volume configuration in the list of Volumes
+		// To accomodate this, we'll first check if the task IS EBS-backed then we'll mark the corresponding Volume object to be
+		// of type "attachment". This volume object will be replaced by the newly created EBS volume configuration when we parse
+		// through the task attachments.
+		volName, ok := hasEBSAttachment(task)
+		if ok {
+			initializeAttachmentTypeVolume(task, volName)
+		}
+
 		apiTask, err := apitask.TaskFromACS(task, payload)
 		if err != nil {
 			pmHandler.handleInvalidTask(task, err, payload)
@@ -306,3 +317,26 @@ func isTaskStatusStopped(status apitaskstatus.TaskStatus) bool {
 func isTaskStatusNotStopped(status apitaskstatus.TaskStatus) bool {
 	return status != apitaskstatus.TaskStopped
 }
+
+func hasEBSAttachment(acsTask *ecsacs.Task) (string, bool) {
+	// TODO: This will only work if there's one EBS volume per task. If we there is a case where we have multi-attach for a task, this needs to be modified
+	for _, attachment := range acsTask.Attachments {
+		if *attachment.AttachmentType == apiresource.EBSTaskAttach {
+			for _, property := range attachment.AttachmentProperties {
+				if *property.Name == apiresource.VolumeNameKey {
+					return *property.Value, true
+				}
+			}
+		}
+	}
+	return "", false
+}
+
+func initializeAttachmentTypeVolume(acsTask *ecsacs.Task, volName string) {
+	for _, volume := range acsTask.Volumes {
+		if *volume.Name == volName && volume.Type == nil {
+			newType := "attachment"
+			volume.Type = &newType
+		}
+	}
+}

agent/acs/session/payload_responder_test.go

@@ -725,6 +725,12 @@ func TestHandlePayloadMessageAddedEBSToTask(t *testing.T) {
 					AttachmentType: aws.String(apiresource.EBSTaskAttach),
 				},
 			},
+			Volumes: []*ecsacs.Volume{
+				{
+					Name: aws.String(taskresourcevolume.TestVolumeName),
+					Type: aws.String(apitask.AttachmentType),
+				},
+			},
 		},
 	}
 

agent/api/task/task.go

@@ -332,7 +332,6 @@ func TaskFromACS(acsTask *ecsacs.Task, envelope *ecsacs.PayloadMessage) (*Task,
 	return task, nil
 }
 
-// TODO: Add unit test
 func (task *Task) RemoveVolume(index int) {
 	task.lock.Lock()
 	defer task.lock.Unlock()
@@ -343,7 +342,6 @@ func (task *Task) removeVolumeUnsafe(index int) {
 	if index < 0 || index >= len(task.Volumes) {
 		return
 	}
-	// temp := task.Volumes[:1]
 	out := make([]TaskVolume, 0)
 	out = append(out, task.Volumes[:index]...)
 	out = append(out, task.Volumes[index+1:]...)
@@ -352,17 +350,6 @@ func (task *Task) removeVolumeUnsafe(index int) {
 
 func (task *Task) initializeVolumes(cfg *config.Config, dockerClient dockerapi.DockerClient, ctx context.Context) error {
 	// TODO: Have EBS volumes use the DockerVolumeConfig to create the mountpoint
-	if task.IsEBSTaskAttachEnabled() {
-		ebsVolumes := task.GetEBSVolumeNames()
-		for index, tv := range task.Volumes {
-			volumeName := tv.Name
-			volumeType := tv.Type
-			if ebsVolumes[volumeName] && volumeType != apiresource.EBSTaskAttach {
-				task.RemoveVolume(index)
-			}
-		}
-	}
-
 	err := task.initializeDockerLocalVolumes(dockerClient, ctx)
 	if err != nil {
 		return apierrors.NewResourceInitError(task.Arn, err)
@@ -3485,28 +3472,6 @@ func (task *Task) isEBSTaskAttachEnabledUnsafe() bool {
 	return false
 }
 
-// TODO: Add unit tests
-func (task *Task) GetEBSVolumeNames() map[string]bool {
-	task.lock.RLock()
-	defer task.lock.RUnlock()
-	return task.getEBSVolumeNamesUnsafe()
-}
-
-func (task *Task) getEBSVolumeNamesUnsafe() map[string]bool {
-	volNames := map[string]bool{}
-	for _, tv := range task.Volumes {
-		switch tv.Volume.(type) {
-		case *taskresourcevolume.EBSTaskVolumeConfig:
-			logger.Debug("found ebs volume config")
-			ebsCfg := tv.Volume.(*taskresourcevolume.EBSTaskVolumeConfig)
-			volNames[ebsCfg.VolumeName] = true
-		default:
-			continue
-		}
-	}
-	return volNames
-}
-
 func (task *Task) IsServiceConnectBridgeModeApplicationContainer(container *apicontainer.Container) bool {
 	return container.GetNetworkModeFromHostConfig() == "container" && task.IsServiceConnectEnabled()
 }

agent/api/task/task_attachment_handler.go

@@ -110,6 +110,7 @@ func handleTaskAttachments(acsTask *ecsacs.Task, task *Task) error {
 			task.ServiceConnectConfig = scHandler.(*ServiceConnectAttachmentHandler).scConfig
 		}
 		if len(ebsVolumeAttachments) > 0 {
+			ebsVolumes := make(map[string]bool)
 			for _, attachment := range ebsVolumeAttachments {
 				ebs, err := taskresourcevolume.ParseEBSTaskVolumeAttachment(attachment)
 				if err != nil {
@@ -120,8 +121,17 @@ func handleTaskAttachments(acsTask *ecsacs.Task, task *Task) error {
 					Type:   apiresource.EBSTaskAttach,
 					Volume: ebs,
 				}
+				ebsVolumes[ebs.VolumeName] = true
 				task.Volumes = append(task.Volumes, taskVolume)
 			}
+			// We're removing all incorrect volume configuration that were intially passed over from ACS
+			for index, tv := range task.Volumes {
+				volumeName := tv.Name
+				volumeType := tv.Type
+				if ebsVolumes[volumeName] && volumeType != apiresource.EBSTaskAttach {
+					task.RemoveVolume(index)
+				}
+			}
 		}
 	}
 	return nil

agent/api/task/task_attachment_handler_test.go

@@ -274,6 +274,12 @@ func TestHandleTaskAttachmentWithEBSVolumeAttachment(t *testing.T) {
 						AttachmentType: stringToPointer(apiresource.EBSTaskAttach),
 					},
 				},
+				Volumes: []*ecsacs.Volume{
+					{
+						Name: strptr("test-volume"),
+						Type: strptr(AttachmentType),
+					},
+				},
 			}
 			testTask := &Task{}
 			err := handleTaskAttachments(testAcsTask, testTask)

agent/api/task/task_test.go

@@ -4691,6 +4691,15 @@ func TestTaskWithEBSVolumeAttachment(t *testing.T) {
 				AttachmentType: strptr(apiresource.EBSTaskAttach),
 			},
 		},
+		Volumes: []*ecsacs.Volume{
+			{
+				Name: strptr("test-volume"),
+				Type: strptr(AttachmentType),
+				Host: &ecsacs.HostVolumeProperties{
+					SourcePath: strptr("/host/path"),
+				},
+			},
+		},
 	}
 
 	testExpectedEBSCfg := &taskresourcevolume.EBSTaskVolumeConfig{
@@ -5234,3 +5243,38 @@ func TestToHostResources(t *testing.T) {
 		assert.Equal(t, len(tc.expectedResources["PORTS_UDP"].StringSetValue), len(calcResources["PORTS_UDP"].StringSetValue), "Error converting task UDP port resources")
 	}
 }
+
+func TestRemoveVolumes(t *testing.T) {
+	task := &Task{
+		Volumes: []TaskVolume{
+			{
+				Name: "volName",
+				Type: "host",
+				Volume: &taskresourcevolume.FSHostVolume{
+					FSSourcePath: "/host/path",
+				},
+			},
+		},
+	}
+	task.RemoveVolume(0)
+	assert.Equal(t, len(task.Volumes), 0)
+}
+
+func TestRemoveVolumeIndexOutOfBounds(t *testing.T) {
+	task := &Task{
+		Volumes: []TaskVolume{
+			{
+				Name: "volName",
+				Type: "host",
+				Volume: &taskresourcevolume.FSHostVolume{
+					FSSourcePath: "/host/path",
+				},
+			},
+		},
+	}
+	task.RemoveVolume(1)
+	assert.Equal(t, len(task.Volumes), 1)
+
+	task.RemoveVolume(-1)
+	assert.Equal(t, len(task.Volumes), 1)
+}

agent/api/task/taskvolume.go

@@ -32,6 +32,7 @@ const (
 	DockerVolumeType               = "docker"
 	EFSVolumeType                  = "efs"
 	FSxWindowsFileServerVolumeType = "fsxWindowsFileServer"
+	AttachmentType                 = "attachment"
 )
 
 // TaskVolume is a definition of all the volumes available for containers to
@@ -78,6 +79,9 @@ func (tv *TaskVolume) UnmarshalJSON(b []byte) error {
 		return tv.unmarshalFSxWindowsFileServerVolume(intermediate["fsxWindowsFileServerVolumeConfiguration"])
 	case apiresource.EBSTaskAttach:
 		return tv.unmarshalEBSVolume(intermediate["ebsVolumeConfiguration"])
+	case AttachmentType:
+		seelog.Warn("Obtaining the volume configuration from task attachments.")
+		return nil
 	default:
 		return errors.Errorf("unrecognized volume type: %q", tv.Type)
 	}