Skip to content

Commit

Permalink
Add Outpost ARN to RCI call
Browse files Browse the repository at this point in the history
  • Loading branch information
@shubham2892
shubham2892 committed Oct 23, 2019
1 parent 18080d9 commit f223d8a
Show file tree
Hide file tree
Showing 10 changed files with 204 additions and 61 deletions.
40 changes: 31 additions & 9 deletions agent/api/ecsclient/client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,8 +110,10 @@ func (client *APIECSClient) CreateCluster(clusterName string) (string, error) {
// ContainerInstanceARN if successful. Supplying a non-empty container
// instance ARN allows a container instance to update its registered
// resources.
func (client *APIECSClient) RegisterContainerInstance(containerInstanceArn string,
attributes []*ecs.Attribute, tags []*ecs.Tag, registrationToken string, platformDevices []*ecs.PlatformDevice) (string, string, error) {
func (client *APIECSClient) RegisterContainerInstance(containerInstanceArn string, attributes []*ecs.Attribute,
tags []*ecs.Tag, registrationToken string, platformDevices []*ecs.PlatformDevice,
outpostARN string) (string, string, error) {

clusterRef := client.config.Cluster
// If our clusterRef is empty, we should try to create the default
if clusterRef == "" {
Expand All @@ -122,7 +124,8 @@ 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, registrationToken, platformDevices)
containerInstanceArn, availabilityzone, err := client.registerContainerInstance(clusterRef,
containerInstanceArn, attributes, tags, registrationToken, platformDevices, outpostARN)
if err == nil {
return containerInstanceArn, availabilityzone, nil
}
Expand All @@ -136,11 +139,14 @@ func (client *APIECSClient) RegisterContainerInstance(containerInstanceArn strin
}
}
}
return client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags, registrationToken, platformDevices)
return client.registerContainerInstance(clusterRef, containerInstanceArn, attributes, tags, registrationToken,
platformDevices, outpostARN)
}

func (client *APIECSClient) registerContainerInstance(clusterRef string, containerInstanceArn string,
attributes []*ecs.Attribute, tags []*ecs.Tag, registrationToken string, platformDevices []*ecs.PlatformDevice) (string, string, error) {
attributes []*ecs.Attribute, tags []*ecs.Tag, registrationToken string,
platformDevices []*ecs.PlatformDevice, outpostARN string) (string, string, error) {

registerRequest := ecs.RegisterContainerInstanceInput{Cluster: &clusterRef}
var registrationAttributes []*ecs.Attribute
if containerInstanceArn != "" {
Expand All @@ -162,6 +168,8 @@ func (client *APIECSClient) registerContainerInstance(clusterRef string, contain

// Add additional attributes such as the os type
registrationAttributes = append(registrationAttributes, client.getAdditionalAttributes()...)
registrationAttributes = append(registrationAttributes, client.getOutpostAttribute(outpostARN)...)

registerRequest.Attributes = registrationAttributes
if len(tags) > 0 {
registerRequest.Tags = tags
Expand Down Expand Up @@ -317,10 +325,24 @@ func validateRegisteredAttributes(expectedAttributes, actualAttributes []*ecs.At
}

func (client *APIECSClient) getAdditionalAttributes() []*ecs.Attribute {
return []*ecs.Attribute{{
Name: aws.String("ecs.os-type"),
Value: aws.String(config.OSType),
}}
return []*ecs.Attribute{
{
Name: aws.String("ecs.os-type"),
Value: aws.String(config.OSType),
},
}
}

func (client *APIECSClient) getOutpostAttribute(outpostARN string) []*ecs.Attribute {
if len(outpostARN) > 0 {
return []*ecs.Attribute{
{
Name: aws.String("ecs.outpost-arn"),
Value: aws.String(outpostARN),
},
}
}
return []*ecs.Attribute{}
}

func (client *APIECSClient) getCustomAttributes() []*ecs.Attribute {
Expand Down
37 changes: 24 additions & 13 deletions agent/api/ecsclient/client_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ var (
},
}
containerInstanceTagsMap = map[string]string{

"my_key1": "my_val1",
"my_key2": "my_val2",
}
Expand Down Expand Up @@ -336,6 +337,7 @@ func TestReRegisterContainerInstance(t *testing.T) {
expectedAttributes := map[string]string{
"ecs.os-type": config.OSType,
"ecs.availability-zone": "us-west-2b",
"ecs.outpost-arn": "test:arn:outpost",
}
for i := range fakeCapabilities {
expectedAttributes[fakeCapabilities[i]] = ""
Expand All @@ -355,8 +357,8 @@ func TestReRegisterContainerInstance(t *testing.T) {
resource, ok := findResource(req.TotalResources, "PORTS_UDP")
assert.True(t, ok, `Could not find resource "PORTS_UDP"`)
assert.Equal(t, "STRINGSET", *resource.Type, `Wrong type for resource "PORTS_UDP"`)
// "ecs.os-type" and the 2 that we specified as additionalAttributes
assert.Equal(t, 3, len(req.Attributes), "Wrong number of Attributes")
// "ecs.os-type", ecs.outpost-arn and the 2 that we specified as additionalAttributes
assert.Equal(t, 4, len(req.Attributes), "Wrong number of Attributes")
reqAttributes := func() map[string]string {
rv := make(map[string]string, len(req.Attributes))
for i := range req.Attributes {
Expand All @@ -382,7 +384,8 @@ func TestReRegisterContainerInstance(t *testing.T) {
nil),
)

arn, availabilityzone, err := client.RegisterContainerInstance("arn:test", capabilities, containerInstanceTags, registrationToken, nil)
arn, availabilityzone, err := client.RegisterContainerInstance("arn:test", capabilities,
containerInstanceTags, registrationToken, nil, "test:arn:outpost")

assert.NoError(t, err)
assert.Equal(t, "registerArn", arn)
Expand All @@ -404,6 +407,7 @@ func TestRegisterContainerInstance(t *testing.T) {
"my_custom_attribute": "Custom_Value1",
"my_other_custom_attribute": "Custom_Value2",
"ecs.availability-zone": "us-west-2b",
"ecs.outpost-arn": "test:arn:outpost",
}
capabilities := buildAttributeList(fakeCapabilities, nil)
platformDevices := []*ecs.PlatformDevice{
Expand Down Expand Up @@ -434,8 +438,8 @@ func TestRegisterContainerInstance(t *testing.T) {
resource, ok := findResource(req.TotalResources, "PORTS_UDP")
assert.True(t, ok, `Could not find resource "PORTS_UDP"`)
assert.Equal(t, "STRINGSET", *resource.Type, `Wrong type for resource "PORTS_UDP"`)
// 3 from expectedAttributes and 2 from additionalAttributes
assert.Equal(t, 5, len(req.Attributes), "Wrong number of Attributes")
// 3 from expectedAttributes and 3 from additionalAttributes
assert.Equal(t, 6, len(req.Attributes), "Wrong number of Attributes")
for i := range req.Attributes {
if strings.Contains(*req.Attributes[i].Name, "capability") {
assert.Contains(t, fakeCapabilities, *req.Attributes[i].Name)
Expand All @@ -457,7 +461,8 @@ func TestRegisterContainerInstance(t *testing.T) {
nil),
)

arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags, registrationToken, platformDevices)
arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities,
containerInstanceTags, registrationToken, platformDevices, "test:arn:outpost")
assert.NoError(t, err)
assert.Equal(t, "registerArn", arn)
assert.Equal(t, "us-west-2b", availabilityzone)
Expand All @@ -484,6 +489,7 @@ func TestRegisterContainerInstanceNoIID(t *testing.T) {
"my_custom_attribute": "Custom_Value1",
"my_other_custom_attribute": "Custom_Value2",
"ecs.availability-zone": "us-west-2b",
"ecs.outpost-arn": "test:arn:outpost",
}
capabilities := buildAttributeList(fakeCapabilities, nil)

Expand All @@ -498,8 +504,8 @@ func TestRegisterContainerInstanceNoIID(t *testing.T) {
resource, ok := findResource(req.TotalResources, "PORTS_UDP")
assert.True(t, ok, `Could not find resource "PORTS_UDP"`)
assert.Equal(t, "STRINGSET", *resource.Type, `Wrong type for resource "PORTS_UDP"`)
// 3 from expectedAttributes and 2 from additionalAttributes
assert.Equal(t, 5, len(req.Attributes), "Wrong number of Attributes")
// 3 from expectedAttributes and 3 from additionalAttributes
assert.Equal(t, 6, len(req.Attributes), "Wrong number of Attributes")
for i := range req.Attributes {
if strings.Contains(*req.Attributes[i].Name, "capability") {
assert.Contains(t, fakeCapabilities, *req.Attributes[i].Name)
Expand All @@ -520,7 +526,8 @@ func TestRegisterContainerInstanceNoIID(t *testing.T) {
nil),
)

arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities, containerInstanceTags, registrationToken, nil)
arn, availabilityzone, err := client.RegisterContainerInstance("", capabilities,
containerInstanceTags, registrationToken, nil, "test:arn:outpost")
assert.NoError(t, err)
assert.Equal(t, "registerArn", arn)
assert.Equal(t, "us-west-2b", availabilityzone)
Expand All @@ -547,7 +554,8 @@ 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, "", nil)
_, _, err := client.RegisterContainerInstance("", nil, nil,
"", nil, "")
assert.Error(t, err, "Register resource with negative value should cause registration fail")
}

Expand Down Expand Up @@ -577,7 +585,8 @@ func TestRegisterContainerInstanceWithEmptyTags(t *testing.T) {
nil),
)

_, _, err := client.RegisterContainerInstance("", nil, make([]*ecs.Tag, 0), "", nil)
_, _, err := client.RegisterContainerInstance("", nil, make([]*ecs.Tag, 0),
"", nil, "")
assert.NoError(t, err)
}

Expand Down Expand Up @@ -651,7 +660,8 @@ func TestRegisterBlankCluster(t *testing.T) {
nil),
)

arn, availabilityzone, err := client.RegisterContainerInstance("", nil, nil, "", nil)
arn, availabilityzone, err := client.RegisterContainerInstance("", nil, nil,
"", nil, "")
if err != nil {
t.Errorf("Should not be an error: %v", err)
}
Expand Down Expand Up @@ -705,7 +715,8 @@ func TestRegisterBlankClusterNotCreatingClusterWhenErrorNotClusterNotFound(t *te
nil),
)

arn, _, err := client.RegisterContainerInstance("", nil, nil, "", nil)
arn, _, err := client.RegisterContainerInstance("", nil, nil, "",
nil, "")
assert.NoError(t, err, "Should not return error")
assert.Equal(t, "registerArn", arn, "Wrong arn")
}
Expand Down
3 changes: 2 additions & 1 deletion agent/api/interface.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,8 @@ type ECSClient interface {
// instance ARN allows a container instance to update its registered
// resources.
RegisterContainerInstance(existingContainerInstanceArn string,
attributes []*ecs.Attribute, tags []*ecs.Tag, registrationToken string, platformDevices []*ecs.PlatformDevice) (string, string, error)
attributes []*ecs.Attribute, tags []*ecs.Tag, registrationToken string, platformDevices []*ecs.PlatformDevice,
outpostARN string) (string, string, error)
// SubmitTaskStateChange sends a state change and returns an error
// indicating if it was submitted
SubmitTaskStateChange(change TaskStateChange) error
Expand Down
8 changes: 4 additions & 4 deletions agent/api/mocks/api_mocks.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

26 changes: 21 additions & 5 deletions agent/app/agent.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -435,6 +435,17 @@ func (agent *ecsAgent) getEC2InstanceID() string {
return instanceID
}

// getoutpostARN gets the Outpost ARN from the metadata service
func (agent *ecsAgent) getoutpostARN() string {
outpostARN, err := agent.ec2MetadataClient.OutpostARN()
if err != nil {
seelog.Warnf(
"Unable to obtain Outpost ARN from EC2 Metadata: %v", err)
return ""
}
return outpostARN
}

// newStateManager creates a new state manager object for the task engine.
// Rest of the parameters are pointers and it's expected that all of these
// will be backfilled when state manager's Load() method is invoked
Expand Down Expand Up @@ -508,13 +519,16 @@ func (agent *ecsAgent) registerContainerInstance(

platformDevices := agent.getPlatformDevices()

outpostARN := agent.getoutpostARN()

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, uuid.New(), platformDevices)
return agent.reregisterContainerInstance(client, capabilities, tags, uuid.New(), platformDevices, outpostARN)
}

seelog.Info("Registering Instance with ECS")
containerInstanceArn, availabilityZone, err := client.RegisterContainerInstance("", capabilities, tags, uuid.New(), platformDevices)
containerInstanceArn, availabilityZone, err := client.RegisterContainerInstance("",
capabilities, tags, uuid.New(), platformDevices, outpostARN)
if err != nil {
seelog.Errorf("Error registering: %v", err)
if retriable, ok := err.(apierrors.Retriable); ok && !retriable.Retry() {
Expand All @@ -541,9 +555,11 @@ func (agent *ecsAgent) registerContainerInstance(
// reregisterContainerInstance registers a container instance that has already been
// 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, registrationToken string, platformDevices []*ecs.PlatformDevice) error {
_, availabilityZone, err := client.RegisterContainerInstance(agent.containerInstanceARN, capabilities, tags, registrationToken, platformDevices)
func (agent *ecsAgent) reregisterContainerInstance(client api.ECSClient, capabilities []*ecs.Attribute,
tags []*ecs.Tag, registrationToken string, platformDevices []*ecs.PlatformDevice, outpostARN string) error {
_, availabilityZone, err := client.RegisterContainerInstance(agent.containerInstanceARN, capabilities, tags,
registrationToken, platformDevices, outpostARN)

//set az to agent
agent.availabilityZone = availabilityZone

Expand Down
Loading

0 comments on commit f223d8a

Please sign in to comment.