Merge Sharanya's RCI idempotency changes

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 1.22.1
+* Bug - Fixed a bug where agent can register container instance back to back and gets
+  assigned two container instance ARNs [#1579](https://github.com/aws/amazon-ecs-agent/pull/1579)
+
 ## 1.22.0
 * Feature - Add support for ECS Secrets integrating with AWS Systems Manager Parameter Store
 * Feature - Support for `--pid`, `--ipc` Docker run flags. [#1584](https://github.com/aws/amazon-ecs-agent/pull/1584)

agent/api/ecsclient/client.go

@@ -109,7 +109,7 @@ func (client *APIECSClient) CreateCluster(clusterName string) (string, error) {
 // instance ARN allows a container instance to update its registered
 // resources.
 func (client *APIECSClient) RegisterContainerInstance(containerInstanceArn string,
-	attributes []*ecs.Attribute, tags []*ecs.Tag) (string, string, error) {
+	attributes []*ecs.Attribute, tags []*ecs.Tag, registrationToken string) (string, string, error) {
 	clusterRef := client.config.Cluster
 	// If our clusterRef is empty, we should try to create the default
 	if clusterRef == "" {
@@ -120,7 +120,7 @@ func (client *APIECSClient) RegisterContainerInstance(containerInstanceArn strin
 		}()
 		// Attempt to register without checking existence of the cluster so we don't require
 		// excess permissions in the case where the cluster already exists and is active
-		containerInstanceArn, availabilityzone, err := client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags)
+		containerInstanceArn, availabilityzone, err := client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags, registrationToken)
 		if err == nil {
 			return containerInstanceArn, availabilityzone, nil
 		}
@@ -131,11 +131,11 @@ func (client *APIECSClient) RegisterContainerInstance(containerInstanceArn strin
 			return "", "", err
 		}
 	}
-	return client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags)
+	return client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags, registrationToken)
 }
 
 func (client *APIECSClient) registerContainerInstance(clusterRef string, containerInstanceArn string,
-	attributes []*ecs.Attribute, tags []*ecs.Tag) (string, string, error) {
+	attributes []*ecs.Attribute, tags []*ecs.Tag, registrationToken string) (string, string, error) {
 	registerRequest := ecs.RegisterContainerInstanceInput{Cluster: &clusterRef}
 	var registrationAttributes []*ecs.Attribute
 	if containerInstanceArn != "" {
@@ -169,6 +169,8 @@ func (client *APIECSClient) registerContainerInstance(clusterRef string, contain
 	}
 
 	registerRequest.TotalResources = resources
+
+	registerRequest.ClientToken = &registrationToken
 	resp, err := client.standardClient.RegisterContainerInstance(&registerRequest)
 	if err != nil {
 		seelog.Errorf("Unable to register as a container instance with ECS: %v", err)

agent/api/ecsclient/client_test.go

@@ -47,6 +47,7 @@ const (
 	configuredCluster = "mycluster"
 	iid               = "instanceIdentityDocument"
 	iidSignature      = "signature"
+	registrationToken = "clientToken"
 )
 
 var (
@@ -343,6 +344,7 @@ func TestReRegisterContainerInstance(t *testing.T) {
 		mc.EXPECT().RegisterContainerInstance(gomock.Any()).Do(func(req *ecs.RegisterContainerInstanceInput) {
 			assert.Equal(t, "arn:test", *req.ContainerInstanceArn, "Wrong container instance ARN")
 			assert.Equal(t, configuredCluster, *req.Cluster, "Wrong cluster")
+			assert.Equal(t, registrationToken, *req.ClientToken, "Wrong client token")
 			assert.Equal(t, iid, *req.InstanceIdentityDocument, "Wrong IID")
 			assert.Equal(t, iidSignature, *req.InstanceIdentityDocumentSignature, "Wrong IID sig")
 			assert.Equal(t, 4, len(req.TotalResources), "Wrong length of TotalResources")
@@ -376,13 +378,10 @@ func TestReRegisterContainerInstance(t *testing.T) {
 			nil),
 	)
 
-	arn, availabilityzone, err := client.RegisterContainerInstance("arn:test", capabilities, containerInstanceTags)
-	if err != nil {
-		t.Errorf("Should not be an error: %v", err)
-	}
-	if arn != "registerArn" {
-		t.Errorf("Wrong arn: %v", arn)
-	}
+	arn, availabilityzone, err := client.RegisterContainerInstance("arn:test", capabilities, containerInstanceTags, registrationToken)
+
+	assert.NoError(t, err)
+	assert.Equal(t, "registerArn", arn)
 	assert.Equal(t, "us-west-2b", availabilityzone, "availabilityZone is incorrect")
 }
 
@@ -410,6 +409,7 @@ func TestRegisterContainerInstance(t *testing.T) {
 		mc.EXPECT().RegisterContainerInstance(gomock.Any()).Do(func(req *ecs.RegisterContainerInstanceInput) {
 			assert.Nil(t, req.ContainerInstanceArn)
 			assert.Equal(t, configuredCluster, *req.Cluster, "Wrong cluster")
+			assert.Equal(t, registrationToken, *req.ClientToken, "Wrong client token")
 			assert.Equal(t, iid, *req.InstanceIdentityDocument, "Wrong IID")
 			assert.Equal(t, iidSignature, *req.InstanceIdentityDocumentSignature, "Wrong IID sig")
 			assert.Equal(t, 4, len(req.TotalResources), "Wrong length of TotalResources")
@@ -438,7 +438,7 @@ func TestRegisterContainerInstance(t *testing.T) {
 			nil),
 	)
 
-	arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags)
+	arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags, registrationToken)
 	assert.NoError(t, err)
 	assert.Equal(t, "registerArn", arn)
 	assert.Equal(t, "us-west-2b", availabilityzone)
@@ -472,6 +472,7 @@ func TestRegisterContainerInstanceNoIID(t *testing.T) {
 		mc.EXPECT().RegisterContainerInstance(gomock.Any()).Do(func(req *ecs.RegisterContainerInstanceInput) {
 			assert.Nil(t, req.ContainerInstanceArn)
 			assert.Equal(t, configuredCluster, *req.Cluster, "Wrong cluster")
+			assert.Equal(t, registrationToken, *req.ClientToken, "Wrong client token")
 			assert.Equal(t, "", *req.InstanceIdentityDocument, "Wrong IID")
 			assert.Equal(t, "", *req.InstanceIdentityDocumentSignature, "Wrong IID sig")
 			assert.Equal(t, 4, len(req.TotalResources), "Wrong length of TotalResources")
@@ -500,7 +501,7 @@ func TestRegisterContainerInstanceNoIID(t *testing.T) {
 			nil),
 	)
 
-	arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags)
+	arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags, registrationToken)
 	assert.NoError(t, err)
 	assert.Equal(t, "registerArn", arn)
 	assert.Equal(t, "us-west-2b", availabilityzone)
@@ -527,7 +528,7 @@ func TestRegisterContainerInstanceWithNegativeResource(t *testing.T) {
 		mockEC2Metadata.EXPECT().GetDynamicData(ec2.InstanceIdentityDocumentResource).Return("instanceIdentityDocument", nil),
 		mockEC2Metadata.EXPECT().GetDynamicData(ec2.InstanceIdentityDocumentSignatureResource).Return("signature", nil),
 	)
-	_, _, err := client.RegisterContainerInstance("", nil, nil)
+	_, _, err := client.RegisterContainerInstance("", nil, nil, "")
 	assert.Error(t, err, "Register resource with negative value should cause registration fail")
 }
 
@@ -557,7 +558,7 @@ func TestRegisterContainerInstanceWithEmptyTags(t *testing.T) {
 			nil),
 	)
 
-	_, _, err := client.RegisterContainerInstance("", nil, make([]*ecs.Tag, 0))
+	_, _, err := client.RegisterContainerInstance("", nil, make([]*ecs.Tag, 0), "")
 	assert.NoError(t, err)
 }
 
@@ -631,7 +632,7 @@ func TestRegisterBlankCluster(t *testing.T) {
 			nil),
 	)
 
-	arn, availabilityzone, err := client.RegisterContainerInstance("", nil, nil)
+	arn, availabilityzone, err := client.RegisterContainerInstance("", nil, nil, "")
 	if err != nil {
 		t.Errorf("Should not be an error: %v", err)
 	}

agent/api/interface.go

@@ -27,7 +27,7 @@ type ECSClient interface {
 	// instance ARN allows a container instance to update its registered
 	// resources.
 	RegisterContainerInstance(existingContainerInstanceArn string,
-		attributes []*ecs.Attribute, tags []*ecs.Tag) (string, string, error)
+		attributes []*ecs.Attribute, tags []*ecs.Tag, registrationToken string) (string, string, error)
 	// SubmitTaskStateChange sends a state change and returns an error
 	// indicating if it was submitted
 	SubmitTaskStateChange(change TaskStateChange) error

agent/app/agent.go

@@ -52,6 +52,7 @@ import (
 	aws_credentials "github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/defaults"
 	"github.com/cihub/seelog"
+	"github.com/pborman/uuid"
 )
 
 const (
@@ -107,6 +108,7 @@ type ecsAgent struct {
 	terminationHandler    sighandlers.TerminationHandler
 	mobyPlugins           mobypkgwrapper.Plugins
 	resourceFields        *taskresource.ResourceFields
+	registrationToken     string
 	availabilityZone      string
 }
 
@@ -241,7 +243,7 @@ func (agent *ecsAgent) doStart(containerChangeEventStream *eventstream.EventStre
 
 	// Initialize the state manager
 	stateManager, err := agent.newStateManager(taskEngine,
-		&agent.cfg.Cluster, &agent.containerInstanceARN, &currentEC2InstanceID, &agent.availabilityZone)
+		&agent.cfg.Cluster, &agent.containerInstanceARN, &currentEC2InstanceID, &agent.availabilityZone, &agent.registrationToken)
 	if err != nil {
 		seelog.Criticalf("Error creating state manager: %v", err)
 		return exitcodes.ExitTerminal
@@ -323,15 +325,15 @@ func (agent *ecsAgent) newTaskEngine(containerChangeEventStream *eventstream.Eve
 	}
 
 	// We try to set these values by loading the existing state file first
-	var previousCluster, previousEC2InstanceID, previousContainerInstanceArn, previousAZ string
+	var previousCluster, previousEC2InstanceID, previousContainerInstanceArn, previousAZ, previousRegistrationToken string
 	previousTaskEngine := engine.NewTaskEngine(agent.cfg, agent.dockerClient,
 		credentialsManager, containerChangeEventStream, imageManager, state,
 		agent.metadataManager, agent.resourceFields)
 
 	// previousStateManager is used to verify that our current runtime configuration is
 	// compatible with our past configuration as reflected by our state-file
 	previousStateManager, err := agent.newStateManager(previousTaskEngine, &previousCluster,
-		&previousContainerInstanceArn, &previousEC2InstanceID, &previousAZ)
+		&previousContainerInstanceArn, &previousEC2InstanceID, &previousAZ, &previousRegistrationToken)
 	if err != nil {
 		seelog.Criticalf("Error creating state manager: %v", err)
 		return nil, "", err
@@ -370,6 +372,7 @@ func (agent *ecsAgent) newTaskEngine(containerChangeEventStream *eventstream.Eve
 
 	// Use the values we loaded if there's no issue
 	agent.containerInstanceARN = previousContainerInstanceArn
+	agent.registrationToken = previousRegistrationToken
 
 	return previousTaskEngine, currentEC2InstanceID, nil
 }
@@ -416,7 +419,7 @@ func (agent *ecsAgent) newStateManager(
 	cluster *string,
 	containerInstanceArn *string,
 	savedInstanceID *string,
-	availabilityZone *string) (statemanager.StateManager, error) {
+	availabilityZone *string, registrationToken *string) (statemanager.StateManager, error) {
 
 	if !agent.cfg.Checkpoint {
 		return statemanager.NewNoopStateManager(), nil
@@ -431,6 +434,7 @@ func (agent *ecsAgent) newStateManager(
 		// This is for making testing easier as we can mock this
 		agent.saveableOptionFactory.AddSaveable("EC2InstanceID", savedInstanceID),
 		agent.saveableOptionFactory.AddSaveable("availabilityZone", availabilityZone),
+		agent.saveableOptionFactory.AddSaveable("RegistrationToken", registrationToken),
 	)
 }
 
@@ -479,13 +483,18 @@ func (agent *ecsAgent) registerContainerInstance(
 		tags = mergeTags(tags, ec2Tags)
 	}
 
+	if agent.registrationToken == "" {
+		agent.registrationToken = uuid.New()
+		stateManager.Save()
+	}
+
 	if agent.containerInstanceARN != "" {
 		seelog.Infof("Restored from checkpoint file. I am running as '%s' in cluster '%s'", agent.containerInstanceARN, agent.cfg.Cluster)
-		return agent.reregisterContainerInstance(client, capabilities, tags)
+		return agent.reregisterContainerInstance(client, capabilities, tags, agent.registrationToken)
 	}
 
 	seelog.Info("Registering Instance with ECS")
-	containerInstanceArn, availabilityZone, err := client.RegisterContainerInstance("", capabilities, tags)
+	containerInstanceArn, availabilityZone, err := client.RegisterContainerInstance("", capabilities, tags, agent.registrationToken)
 	if err != nil {
 		seelog.Errorf("Error registering: %v", err)
 		if retriable, ok := err.(apierrors.Retriable); ok && !retriable.Retry() {
@@ -513,8 +522,8 @@ func (agent *ecsAgent) registerContainerInstance(
 // registered with ECS. This is for cases where the ECS Agent is being restored
 // from a check point.
 func (agent *ecsAgent) reregisterContainerInstance(client api.ECSClient,
-	capabilities []*ecs.Attribute, tags []*ecs.Tag) error {
-	_, _, err := client.RegisterContainerInstance(agent.containerInstanceARN, capabilities, tags)
+	capabilities []*ecs.Attribute, tags []*ecs.Tag, registrationToken string) error {
+	_, _, err := client.RegisterContainerInstance(agent.containerInstanceARN, capabilities, tags, registrationToken)
 
 	if err == nil {
 		return nil

agent/app/agent_compatibility_linux_test.go

@@ -47,7 +47,7 @@ func TestCompatibilityEnabledSuccess(t *testing.T) {
 
 	gomock.InOrder(
 		saveableOptionFactory.EXPECT().AddSaveable(gomock.Any(), gomock.Any()).AnyTimes(),
-		stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
+		stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
 		stateManager.EXPECT().Load().AnyTimes(),
 		state.EXPECT().AllTasks().Return([]*apitask.Task{}),
 	)
@@ -79,7 +79,7 @@ func TestCompatibilityDefaultEnabledFail(t *testing.T) {
 	}
 	gomock.InOrder(
 		saveableOptionFactory.EXPECT().AddSaveable(gomock.Any(), gomock.Any()).AnyTimes(),
-		stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
+		stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
 		stateManager.EXPECT().Load().AnyTimes(),
 		state.EXPECT().AllTasks().Return(getTaskListWithOneBadTask()),
 	)
@@ -110,7 +110,7 @@ func TestCompatibilityExplicitlyEnabledFail(t *testing.T) {
 	}
 	gomock.InOrder(
 		saveableOptionFactory.EXPECT().AddSaveable(gomock.Any(), gomock.Any()).AnyTimes(),
-		stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
+		stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
 		stateManager.EXPECT().Load().AnyTimes(),
 		state.EXPECT().AllTasks().Return(getTaskListWithOneBadTask()),
 	)