AvailabilityZone show on taskmetadata endpoint

agent/api/ecsclient/client.go

@@ -41,6 +41,7 @@ const (
 	pollEndpointCacheSize = 1
 	pollEndpointCacheTTL  = 20 * time.Minute
 	roundtripTimeout      = 5 * time.Second
+	azAttrName            = "ecs.availability-zone"
 )
 
 // APIECSClient implements ECSClient
@@ -108,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, error) {
+	attributes []*ecs.Attribute, tags []*ecs.Tag) (string, string, error) {
 	clusterRef := client.config.Cluster
 	// If our clusterRef is empty, we should try to create the default
 	if clusterRef == "" {
@@ -119,22 +120,22 @@ 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, err := client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags)
+		containerInstanceArn, availabilityzone, err := client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags)
 		if err == nil {
-			return containerInstanceArn, nil
+			return containerInstanceArn, availabilityzone, nil
 		}
 		// If trying to register fails, try to create the cluster before calling
 		// register again
 		clusterRef, err = client.CreateCluster(clusterRef)
 		if err != nil {
-			return "", err
+			return "", "", err
 		}
 	}
 	return client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags)
 }
 
 func (client *APIECSClient) registerContainerInstance(clusterRef string, containerInstanceArn string,
-	attributes []*ecs.Attribute, tags []*ecs.Tag) (string, error) {
+	attributes []*ecs.Attribute, tags []*ecs.Tag) (string, string, error) {
 	registerRequest := ecs.RegisterContainerInstanceInput{Cluster: &clusterRef}
 	var registrationAttributes []*ecs.Attribute
 	if containerInstanceArn != "" {
@@ -164,18 +165,29 @@ func (client *APIECSClient) registerContainerInstance(clusterRef string, contain
 
 	resources, err := client.getResources()
 	if err != nil {
-		return "", err
+		return "", "", err
 	}
 
 	registerRequest.TotalResources = resources
 	resp, err := client.standardClient.RegisterContainerInstance(&registerRequest)
 	if err != nil {
 		seelog.Errorf("Unable to register as a container instance with ECS: %v", err)
-		return "", err
+		return "", "", err
+	}
+
+	var availabilityzone = ""
+	if resp != nil {
+		for _, attr := range resp.ContainerInstance.Attributes {
+			if aws.StringValue(attr.Name) == azAttrName {
+				availabilityzone = aws.StringValue(attr.Value)
+				break
+			}
+		}
 	}
+
 	seelog.Info("Registered container instance with cluster!")
 	err = validateRegisteredAttributes(registerRequest.Attributes, resp.ContainerInstance.Attributes)
-	return aws.StringValue(resp.ContainerInstance.ContainerInstanceArn), err
+	return aws.StringValue(resp.ContainerInstance.ContainerInstanceArn), availabilityzone, err
 }
 
 func (client *APIECSClient) setInstanceIdentity(registerRequest ecs.RegisterContainerInstanceInput) ecs.RegisterContainerInstanceInput {

agent/api/ecsclient/client_test.go

@@ -329,7 +329,8 @@ func TestReRegisterContainerInstance(t *testing.T) {
 
 	fakeCapabilities := []string{"capability1", "capability2"}
 	expectedAttributes := map[string]string{
-		"ecs.os-type": config.OSType,
+		"ecs.os-type":           config.OSType,
+		"ecs.availability-zone": "us-west-2b",
 	}
 	for i := range fakeCapabilities {
 		expectedAttributes[fakeCapabilities[i]] = ""
@@ -375,13 +376,14 @@ func TestReRegisterContainerInstance(t *testing.T) {
 			nil),
 	)
 
-	arn, err := client.RegisterContainerInstance("arn:test", capabilities, containerInstanceTags)
+	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)
 	}
+	assert.Equal(t, "us-west-2b", availabilityzone, "availabilityZone is incorrect")
 }
 
 func TestRegisterContainerInstance(t *testing.T) {
@@ -398,6 +400,7 @@ func TestRegisterContainerInstance(t *testing.T) {
 		"ecs.os-type":               config.OSType,
 		"my_custom_attribute":       "Custom_Value1",
 		"my_other_custom_attribute": "Custom_Value2",
+		"ecs.availability-zone":     "us-west-2b",
 	}
 	capabilities := buildAttributeList(fakeCapabilities, nil)
 
@@ -435,9 +438,10 @@ func TestRegisterContainerInstance(t *testing.T) {
 			nil),
 	)
 
-	arn, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags)
+	arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags)
 	assert.NoError(t, err)
 	assert.Equal(t, "registerArn", arn)
+	assert.Equal(t, "us-west-2b", availabilityzone)
 }
 
 func TestRegisterContainerInstanceNoIID(t *testing.T) {
@@ -460,6 +464,7 @@ func TestRegisterContainerInstanceNoIID(t *testing.T) {
 		"ecs.os-type":               config.OSType,
 		"my_custom_attribute":       "Custom_Value1",
 		"my_other_custom_attribute": "Custom_Value2",
+		"ecs.availability-zone":     "us-west-2b",
 	}
 	capabilities := buildAttributeList(fakeCapabilities, nil)
 
@@ -495,9 +500,10 @@ func TestRegisterContainerInstanceNoIID(t *testing.T) {
 			nil),
 	)
 
-	arn, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags)
+	arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags)
 	assert.NoError(t, err)
 	assert.Equal(t, "registerArn", arn)
+	assert.Equal(t, "us-west-2b", availabilityzone)
 }
 
 // TestRegisterContainerInstanceWithNegativeResource tests the registeration should fail with negative resource
@@ -521,7 +527,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")
 }
 
@@ -551,7 +557,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)
 }
 
@@ -625,13 +631,16 @@ func TestRegisterBlankCluster(t *testing.T) {
 			nil),
 	)
 
-	arn, err := client.RegisterContainerInstance("", nil, nil)
+	arn, availabilityzone, err := client.RegisterContainerInstance("", nil, nil)
 	if err != nil {
 		t.Errorf("Should not be an error: %v", err)
 	}
 	if arn != "registerArn" {
 		t.Errorf("Wrong arn: %v", arn)
 	}
+	if availabilityzone != "" {
+		t.Errorf("wrong availability zone: %v", availabilityzone)
+	}
 }
 
 func TestDiscoverTelemetryEndpoint(t *testing.T) {

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, error)
+		attributes []*ecs.Attribute, tags []*ecs.Tag) (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

@@ -107,6 +107,7 @@ type ecsAgent struct {
 	terminationHandler    sighandlers.TerminationHandler
 	mobyPlugins           mobypkgwrapper.Plugins
 	resourceFields        *taskresource.ResourceFields
+	availabilityZone      string
 }
 
 // newAgent returns a new ecsAgent object, but does not start anything
@@ -240,7 +241,7 @@ func (agent *ecsAgent) doStart(containerChangeEventStream *eventstream.EventStre
 
 	// Initialize the state manager
 	stateManager, err := agent.newStateManager(taskEngine,
-		&agent.cfg.Cluster, &agent.containerInstanceARN, &currentEC2InstanceID)
+		&agent.cfg.Cluster, &agent.containerInstanceARN, &currentEC2InstanceID, &agent.availabilityZone)
 	if err != nil {
 		seelog.Criticalf("Error creating state manager: %v", err)
 		return exitcodes.ExitTerminal
@@ -321,15 +322,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 string
+	var previousCluster, previousEC2InstanceID, previousContainerInstanceArn, previousAZ 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)
+		&previousContainerInstanceArn, &previousEC2InstanceID, &previousAZ)
 	if err != nil {
 		seelog.Criticalf("Error creating state manager: %v", err)
 		return nil, "", err
@@ -413,7 +414,8 @@ func (agent *ecsAgent) newStateManager(
 	taskEngine engine.TaskEngine,
 	cluster *string,
 	containerInstanceArn *string,
-	savedInstanceID *string) (statemanager.StateManager, error) {
+	savedInstanceID *string,
+	availabilityZone *string) (statemanager.StateManager, error) {
 
 	if !agent.cfg.Checkpoint {
 		return statemanager.NewNoopStateManager(), nil
@@ -427,6 +429,7 @@ func (agent *ecsAgent) newStateManager(
 		agent.saveableOptionFactory.AddSaveable("Cluster", cluster),
 		// This is for making testing easier as we can mock this
 		agent.saveableOptionFactory.AddSaveable("EC2InstanceID", savedInstanceID),
+		agent.saveableOptionFactory.AddSaveable("availabilityZone", availabilityZone),
 	)
 }
 
@@ -481,7 +484,7 @@ func (agent *ecsAgent) registerContainerInstance(
 	}
 
 	seelog.Info("Registering Instance with ECS")
-	containerInstanceArn, err := client.RegisterContainerInstance("", capabilities, tags)
+	containerInstanceArn, availabilityZone, err := client.RegisterContainerInstance("", capabilities, tags)
 	if err != nil {
 		seelog.Errorf("Error registering: %v", err)
 		if retriable, ok := err.(apierrors.Retriable); ok && !retriable.Retry() {
@@ -499,6 +502,7 @@ func (agent *ecsAgent) registerContainerInstance(
 	}
 	seelog.Infof("Registration completed successfully. I am running as '%s' in cluster '%s'", containerInstanceArn, agent.cfg.Cluster)
 	agent.containerInstanceARN = containerInstanceArn
+	agent.availabilityZone = availabilityZone
 	// Save our shiny new containerInstanceArn
 	stateManager.Save()
 	return nil
@@ -509,7 +513,7 @@ func (agent *ecsAgent) registerContainerInstance(
 // 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)
+	_, _, err := client.RegisterContainerInstance(agent.containerInstanceARN, capabilities, tags)
 
 	if err == nil {
 		return nil
@@ -551,7 +555,7 @@ func (agent *ecsAgent) startAsyncRoutines(
 	statsEngine := stats.NewDockerStatsEngine(agent.cfg, agent.dockerClient, containerChangeEventStream)
 
 	// Start serving the endpoint to fetch IAM Role credentials and other task metadata
-	go handlers.ServeTaskHTTPEndpoint(credentialsManager, state, agent.containerInstanceARN, agent.cfg, statsEngine)
+	go handlers.ServeTaskHTTPEndpoint(credentialsManager, state, agent.containerInstanceARN, agent.cfg, statsEngine, agent.availabilityZone)
 
 	// Start sending events to the backend
 	go eventhandler.HandleEngineEvents(taskEngine, client, taskHandler)

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()).AnyTimes().Return(stateManager, nil),
+		stateManagerFactory.EXPECT().NewStateManager(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()).AnyTimes().Return(stateManager, nil),
+		stateManagerFactory.EXPECT().NewStateManager(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()).AnyTimes().Return(stateManager, nil),
+		stateManagerFactory.EXPECT().NewStateManager(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()),
 	)