Skip to content

Commit 44c6b76

Browse files
cyastellacyastella
committed
AvailabilityZone show on taskmetadata endpoint
1 parent 8d7d0d2 commit 44c6b76

21 files changed

+450
-137
lines changed

‎agent/api/ecsclient/client.go

+20-8
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ const (
4141
pollEndpointCacheSize = 1
4242
pollEndpointCacheTTL = 20 * time.Minute
4343
roundtripTimeout = 5 * time.Second
44+
azAttrName = "ecs.availability-zone"
4445
)
4546

4647
// APIECSClient implements ECSClient
@@ -108,7 +109,7 @@ func (client *APIECSClient) CreateCluster(clusterName string) (string, error) {
108109
// instance ARN allows a container instance to update its registered
109110
// resources.
110111
func (client *APIECSClient) RegisterContainerInstance(containerInstanceArn string,
111-
attributes []*ecs.Attribute, tags []*ecs.Tag) (string, error) {
112+
attributes []*ecs.Attribute, tags []*ecs.Tag) (string, string, error) {
112113
clusterRef := client.config.Cluster
113114
// If our clusterRef is empty, we should try to create the default
114115
if clusterRef == "" {
@@ -119,22 +120,22 @@ func (client *APIECSClient) RegisterContainerInstance(containerInstanceArn strin
119120
}()
120121
// Attempt to register without checking existence of the cluster so we don't require
121122
// excess permissions in the case where the cluster already exists and is active
122-
containerInstanceArn, err := client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags)
123+
containerInstanceArn, availabilityzone, err := client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags)
123124
if err == nil {
124-
return containerInstanceArn, nil
125+
return containerInstanceArn, availabilityzone, nil
125126
}
126127
// If trying to register fails, try to create the cluster before calling
127128
// register again
128129
clusterRef, err = client.CreateCluster(clusterRef)
129130
if err != nil {
130-
return "", err
131+
return "", "", err
131132
}
132133
}
133134
return client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags)
134135
}
135136

136137
func (client *APIECSClient) registerContainerInstance(clusterRef string, containerInstanceArn string,
137-
attributes []*ecs.Attribute, tags []*ecs.Tag) (string, error) {
138+
attributes []*ecs.Attribute, tags []*ecs.Tag) (string, string, error) {
138139
registerRequest := ecs.RegisterContainerInstanceInput{Cluster: &clusterRef}
139140
var registrationAttributes []*ecs.Attribute
140141
if containerInstanceArn != "" {
@@ -164,18 +165,29 @@ func (client *APIECSClient) registerContainerInstance(clusterRef string, contain
164165

165166
resources, err := client.getResources()
166167
if err != nil {
167-
return "", err
168+
return "", "", err
168169
}
169170

170171
registerRequest.TotalResources = resources
171172
resp, err := client.standardClient.RegisterContainerInstance(&registerRequest)
172173
if err != nil {
173174
seelog.Errorf("Unable to register as a container instance with ECS: %v", err)
174-
return "", err
175+
return "", "", err
176+
}
177+
178+
var availabilityzone = ""
179+
if resp != nil {
180+
for _, attr := range resp.ContainerInstance.Attributes {
181+
if aws.StringValue(attr.Name) == azAttrName {
182+
availabilityzone = aws.StringValue(attr.Value)
183+
break
184+
}
185+
}
175186
}
187+
176188
seelog.Info("Registered container instance with cluster!")
177189
err = validateRegisteredAttributes(registerRequest.Attributes, resp.ContainerInstance.Attributes)
178-
return aws.StringValue(resp.ContainerInstance.ContainerInstanceArn), err
190+
return aws.StringValue(resp.ContainerInstance.ContainerInstanceArn), availabilityzone, err
179191
}
180192

181193
func (client *APIECSClient) setInstanceIdentity(registerRequest ecs.RegisterContainerInstanceInput) ecs.RegisterContainerInstanceInput {

‎agent/api/ecsclient/client_test.go

+16-7
Original file line numberDiff line numberDiff line change
@@ -329,7 +329,8 @@ func TestReRegisterContainerInstance(t *testing.T) {
329329

330330
fakeCapabilities := []string{"capability1", "capability2"}
331331
expectedAttributes := map[string]string{
332-
"ecs.os-type": config.OSType,
332+
"ecs.os-type": config.OSType,
333+
"ecs.availability-zone": "us-west-2b",
333334
}
334335
for i := range fakeCapabilities {
335336
expectedAttributes[fakeCapabilities[i]] = ""
@@ -375,13 +376,14 @@ func TestReRegisterContainerInstance(t *testing.T) {
375376
nil),
376377
)
377378

378-
arn, err := client.RegisterContainerInstance("arn:test", capabilities, containerInstanceTags)
379+
arn, availabilityzone, err := client.RegisterContainerInstance("arn:test", capabilities, containerInstanceTags)
379380
if err != nil {
380381
t.Errorf("Should not be an error: %v", err)
381382
}
382383
if arn != "registerArn" {
383384
t.Errorf("Wrong arn: %v", arn)
384385
}
386+
assert.Equal(t, "us-west-2b", availabilityzone, "availabilityZone is incorrect")
385387
}
386388

387389
func TestRegisterContainerInstance(t *testing.T) {
@@ -398,6 +400,7 @@ func TestRegisterContainerInstance(t *testing.T) {
398400
"ecs.os-type": config.OSType,
399401
"my_custom_attribute": "Custom_Value1",
400402
"my_other_custom_attribute": "Custom_Value2",
403+
"ecs.availability-zone": "us-west-2b",
401404
}
402405
capabilities := buildAttributeList(fakeCapabilities, nil)
403406

@@ -435,9 +438,10 @@ func TestRegisterContainerInstance(t *testing.T) {
435438
nil),
436439
)
437440

438-
arn, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags)
441+
arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags)
439442
assert.NoError(t, err)
440443
assert.Equal(t, "registerArn", arn)
444+
assert.Equal(t, "us-west-2b", availabilityzone)
441445
}
442446

443447
func TestRegisterContainerInstanceNoIID(t *testing.T) {
@@ -460,6 +464,7 @@ func TestRegisterContainerInstanceNoIID(t *testing.T) {
460464
"ecs.os-type": config.OSType,
461465
"my_custom_attribute": "Custom_Value1",
462466
"my_other_custom_attribute": "Custom_Value2",
467+
"ecs.availability-zone": "us-west-2b",
463468
}
464469
capabilities := buildAttributeList(fakeCapabilities, nil)
465470

@@ -495,9 +500,10 @@ func TestRegisterContainerInstanceNoIID(t *testing.T) {
495500
nil),
496501
)
497502

498-
arn, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags)
503+
arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags)
499504
assert.NoError(t, err)
500505
assert.Equal(t, "registerArn", arn)
506+
assert.Equal(t, "us-west-2b", availabilityzone)
501507
}
502508

503509
// TestRegisterContainerInstanceWithNegativeResource tests the registeration should fail with negative resource
@@ -521,7 +527,7 @@ func TestRegisterContainerInstanceWithNegativeResource(t *testing.T) {
521527
mockEC2Metadata.EXPECT().GetDynamicData(ec2.InstanceIdentityDocumentResource).Return("instanceIdentityDocument", nil),
522528
mockEC2Metadata.EXPECT().GetDynamicData(ec2.InstanceIdentityDocumentSignatureResource).Return("signature", nil),
523529
)
524-
_, err := client.RegisterContainerInstance("", nil, nil)
530+
_, _, err := client.RegisterContainerInstance("", nil, nil)
525531
assert.Error(t, err, "Register resource with negative value should cause registration fail")
526532
}
527533

@@ -551,7 +557,7 @@ func TestRegisterContainerInstanceWithEmptyTags(t *testing.T) {
551557
nil),
552558
)
553559

554-
_, err := client.RegisterContainerInstance("", nil, make([]*ecs.Tag, 0))
560+
_, _, err := client.RegisterContainerInstance("", nil, make([]*ecs.Tag, 0))
555561
assert.NoError(t, err)
556562
}
557563

@@ -625,13 +631,16 @@ func TestRegisterBlankCluster(t *testing.T) {
625631
nil),
626632
)
627633

628-
arn, err := client.RegisterContainerInstance("", nil, nil)
634+
arn, availabilityzone, err := client.RegisterContainerInstance("", nil, nil)
629635
if err != nil {
630636
t.Errorf("Should not be an error: %v", err)
631637
}
632638
if arn != "registerArn" {
633639
t.Errorf("Wrong arn: %v", arn)
634640
}
641+
if availabilityzone != "" {
642+
t.Errorf("wrong availability zone: %v", availabilityzone)
643+
}
635644
}
636645

637646
func TestDiscoverTelemetryEndpoint(t *testing.T) {

‎agent/api/interface.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ type ECSClient interface {
2727
// instance ARN allows a container instance to update its registered
2828
// resources.
2929
RegisterContainerInstance(existingContainerInstanceArn string,
30-
attributes []*ecs.Attribute, tags []*ecs.Tag) (string, error)
30+
attributes []*ecs.Attribute, tags []*ecs.Tag) (string, string, error)
3131
// SubmitTaskStateChange sends a state change and returns an error
3232
// indicating if it was submitted
3333
SubmitTaskStateChange(change TaskStateChange) error

‎agent/api/mocks/api_mocks.go

+4-3
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎agent/app/agent.go

+11-7
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ type ecsAgent struct {
107107
terminationHandler sighandlers.TerminationHandler
108108
mobyPlugins mobypkgwrapper.Plugins
109109
resourceFields *taskresource.ResourceFields
110+
availabilityZone string
110111
}
111112

112113
// newAgent returns a new ecsAgent object, but does not start anything
@@ -240,7 +241,7 @@ func (agent *ecsAgent) doStart(containerChangeEventStream *eventstream.EventStre
240241

241242
// Initialize the state manager
242243
stateManager, err := agent.newStateManager(taskEngine,
243-
&agent.cfg.Cluster, &agent.containerInstanceARN, &currentEC2InstanceID)
244+
&agent.cfg.Cluster, &agent.containerInstanceARN, &currentEC2InstanceID, &agent.availabilityZone)
244245
if err != nil {
245246
seelog.Criticalf("Error creating state manager: %v", err)
246247
return exitcodes.ExitTerminal
@@ -321,15 +322,15 @@ func (agent *ecsAgent) newTaskEngine(containerChangeEventStream *eventstream.Eve
321322
}
322323

323324
// We try to set these values by loading the existing state file first
324-
var previousCluster, previousEC2InstanceID, previousContainerInstanceArn string
325+
var previousCluster, previousEC2InstanceID, previousContainerInstanceArn, previousAZ string
325326
previousTaskEngine := engine.NewTaskEngine(agent.cfg, agent.dockerClient,
326327
credentialsManager, containerChangeEventStream, imageManager, state,
327328
agent.metadataManager, agent.resourceFields)
328329

329330
// previousStateManager is used to verify that our current runtime configuration is
330331
// compatible with our past configuration as reflected by our state-file
331332
previousStateManager, err := agent.newStateManager(previousTaskEngine, &previousCluster,
332-
&previousContainerInstanceArn, &previousEC2InstanceID)
333+
&previousContainerInstanceArn, &previousEC2InstanceID, &previousAZ)
333334
if err != nil {
334335
seelog.Criticalf("Error creating state manager: %v", err)
335336
return nil, "", err
@@ -413,7 +414,8 @@ func (agent *ecsAgent) newStateManager(
413414
taskEngine engine.TaskEngine,
414415
cluster *string,
415416
containerInstanceArn *string,
416-
savedInstanceID *string) (statemanager.StateManager, error) {
417+
savedInstanceID *string,
418+
availabilityZone *string) (statemanager.StateManager, error) {
417419

418420
if !agent.cfg.Checkpoint {
419421
return statemanager.NewNoopStateManager(), nil
@@ -427,6 +429,7 @@ func (agent *ecsAgent) newStateManager(
427429
agent.saveableOptionFactory.AddSaveable("Cluster", cluster),
428430
// This is for making testing easier as we can mock this
429431
agent.saveableOptionFactory.AddSaveable("EC2InstanceID", savedInstanceID),
432+
agent.saveableOptionFactory.AddSaveable("availabilityZone", availabilityZone),
430433
)
431434
}
432435

@@ -481,7 +484,7 @@ func (agent *ecsAgent) registerContainerInstance(
481484
}
482485

483486
seelog.Info("Registering Instance with ECS")
484-
containerInstanceArn, err := client.RegisterContainerInstance("", capabilities, tags)
487+
containerInstanceArn, availabilityZone, err := client.RegisterContainerInstance("", capabilities, tags)
485488
if err != nil {
486489
seelog.Errorf("Error registering: %v", err)
487490
if retriable, ok := err.(apierrors.Retriable); ok && !retriable.Retry() {
@@ -499,6 +502,7 @@ func (agent *ecsAgent) registerContainerInstance(
499502
}
500503
seelog.Infof("Registration completed successfully. I am running as '%s' in cluster '%s'", containerInstanceArn, agent.cfg.Cluster)
501504
agent.containerInstanceARN = containerInstanceArn
505+
agent.availabilityZone = availabilityZone
502506
// Save our shiny new containerInstanceArn
503507
stateManager.Save()
504508
return nil
@@ -509,7 +513,7 @@ func (agent *ecsAgent) registerContainerInstance(
509513
// from a check point.
510514
func (agent *ecsAgent) reregisterContainerInstance(client api.ECSClient,
511515
capabilities []*ecs.Attribute, tags []*ecs.Tag) error {
512-
_, err := client.RegisterContainerInstance(agent.containerInstanceARN, capabilities, tags)
516+
_, _, err := client.RegisterContainerInstance(agent.containerInstanceARN, capabilities, tags)
513517

514518
if err == nil {
515519
return nil
@@ -551,7 +555,7 @@ func (agent *ecsAgent) startAsyncRoutines(
551555
statsEngine := stats.NewDockerStatsEngine(agent.cfg, agent.dockerClient, containerChangeEventStream)
552556

553557
// Start serving the endpoint to fetch IAM Role credentials and other task metadata
554-
go handlers.ServeTaskHTTPEndpoint(credentialsManager, state, agent.containerInstanceARN, agent.cfg, statsEngine)
558+
go handlers.ServeTaskHTTPEndpoint(credentialsManager, state, agent.containerInstanceARN, agent.cfg, statsEngine, agent.availabilityZone)
555559

556560
// Start sending events to the backend
557561
go eventhandler.HandleEngineEvents(taskEngine, client, taskHandler)

‎agent/app/agent_compatibility_linux_test.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ func TestCompatibilityEnabledSuccess(t *testing.T) {
4747

4848
gomock.InOrder(
4949
saveableOptionFactory.EXPECT().AddSaveable(gomock.Any(), gomock.Any()).AnyTimes(),
50-
stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
50+
stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
5151
stateManager.EXPECT().Load().AnyTimes(),
5252
state.EXPECT().AllTasks().Return([]*apitask.Task{}),
5353
)
@@ -79,7 +79,7 @@ func TestCompatibilityDefaultEnabledFail(t *testing.T) {
7979
}
8080
gomock.InOrder(
8181
saveableOptionFactory.EXPECT().AddSaveable(gomock.Any(), gomock.Any()).AnyTimes(),
82-
stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
82+
stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
8383
stateManager.EXPECT().Load().AnyTimes(),
8484
state.EXPECT().AllTasks().Return(getTaskListWithOneBadTask()),
8585
)
@@ -110,7 +110,7 @@ func TestCompatibilityExplicitlyEnabledFail(t *testing.T) {
110110
}
111111
gomock.InOrder(
112112
saveableOptionFactory.EXPECT().AddSaveable(gomock.Any(), gomock.Any()).AnyTimes(),
113-
stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
113+
stateManagerFactory.EXPECT().NewStateManager(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(stateManager, nil),
114114
stateManager.EXPECT().Load().AnyTimes(),
115115
state.EXPECT().AllTasks().Return(getTaskListWithOneBadTask()),
116116
)

0 commit comments

Comments
 (0)