Skip to content

Commit 1c22e4d

Browse files
jchen6585jchen6585
authored
Refactor AllocENI (#2640)
1 parent 1dd45b8 commit 1c22e4d

File tree

5 files changed

+263
-129
lines changed

5 files changed

+263
-129
lines changed

‎pkg/awsutils/awsutils.go

+29-9
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ var (
123123
// APIs defines interfaces calls for adding/getting/deleting ENIs/secondary IPs. The APIs are not thread-safe.
124124
type APIs interface {
125125
// AllocENI creates an ENI and attaches it to the instance
126-
AllocENI(useCustomCfg bool, sg []*string, subnet string) (eni string, err error)
126+
AllocENI(useCustomCfg bool, sg []*string, subnet string, numIPs int) (eni string, err error)
127127

128128
// FreeENI detaches ENI interface and deletes it
129129
FreeENI(eniName string) error
@@ -740,8 +740,8 @@ func (cache *EC2InstanceMetadataCache) awsGetFreeDeviceNumber() (int, error) {
740740

741741
// AllocENI creates an ENI and attaches it to the instance
742742
// returns: newly created ENI ID
743-
func (cache *EC2InstanceMetadataCache) AllocENI(useCustomCfg bool, sg []*string, subnet string) (string, error) {
744-
eniID, err := cache.createENI(useCustomCfg, sg, subnet)
743+
func (cache *EC2InstanceMetadataCache) AllocENI(useCustomCfg bool, sg []*string, subnet string, numIPs int) (string, error) {
744+
eniID, err := cache.createENI(useCustomCfg, sg, subnet, numIPs)
745745
if err != nil {
746746
return "", errors.Wrap(err, "AllocENI: failed to create ENI")
747747
}
@@ -812,7 +812,7 @@ func (cache *EC2InstanceMetadataCache) attachENI(eniID string) (string, error) {
812812
}
813813

814814
// return ENI id, error
815-
func (cache *EC2InstanceMetadataCache) createENI(useCustomCfg bool, sg []*string, subnet string) (string, error) {
815+
func (cache *EC2InstanceMetadataCache) createENI(useCustomCfg bool, sg []*string, subnet string, numIPs int) (string, error) {
816816
eniDescription := eniDescriptionPrefix + cache.instanceID
817817
tags := map[string]string{
818818
eniCreatedAtTagKey: time.Now().Format(time.RFC3339),
@@ -826,14 +826,34 @@ func (cache *EC2InstanceMetadataCache) createENI(useCustomCfg bool, sg []*string
826826
Tags: convertTagsToSDKTags(tags),
827827
},
828828
}
829+
var needIPs = numIPs
829830

830-
input := &ec2.CreateNetworkInterfaceInput{
831-
Description: aws.String(eniDescription),
832-
Groups: aws.StringSlice(cache.securityGroups.SortedList()),
833-
SubnetId: aws.String(cache.subnetID),
834-
TagSpecifications: tagSpec,
831+
ipLimit := cache.GetENIIPv4Limit()
832+
if ipLimit < needIPs {
833+
needIPs = ipLimit
835834
}
836835

836+
log.Infof("Trying to allocate %d IP addresses on new ENI", needIPs)
837+
log.Debugf("PD enabled - %t", cache.enablePrefixDelegation)
838+
input := &ec2.CreateNetworkInterfaceInput{}
839+
840+
if cache.enablePrefixDelegation {
841+
input = &ec2.CreateNetworkInterfaceInput{
842+
Description: aws.String(eniDescription),
843+
Groups: aws.StringSlice(cache.securityGroups.SortedList()),
844+
SubnetId: aws.String(cache.subnetID),
845+
TagSpecifications: tagSpec,
846+
Ipv4PrefixCount: aws.Int64(int64(needIPs)),
847+
}
848+
} else {
849+
input = &ec2.CreateNetworkInterfaceInput{
850+
Description: aws.String(eniDescription),
851+
Groups: aws.StringSlice(cache.securityGroups.SortedList()),
852+
SubnetId: aws.String(cache.subnetID),
853+
TagSpecifications: tagSpec,
854+
SecondaryPrivateIpAddressCount: aws.Int64(int64(needIPs)),
855+
}
856+
}
837857
if useCustomCfg {
838858
log.Info("Using a custom network config for the new ENI")
839859
if len(sg) != 0 {

‎pkg/awsutils/awsutils_test.go

+128-9
Original file line numberDiff line numberDiff line change
@@ -401,11 +401,12 @@ func TestAllocENI(t *testing.T) {
401401
mockEC2.EXPECT().ModifyNetworkInterfaceAttributeWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
402402

403403
cache := &EC2InstanceMetadataCache{
404-
ec2SVC: mockEC2,
405-
imds: TypedIMDS{mockMetadata},
404+
ec2SVC: mockEC2,
405+
imds: TypedIMDS{mockMetadata},
406+
instanceType: "c5n.18xlarge",
406407
}
407408

408-
_, err := cache.AllocENI(false, nil, "")
409+
_, err := cache.AllocENI(false, nil, "", 5)
409410
assert.NoError(t, err)
410411
}
411412

@@ -435,11 +436,12 @@ func TestAllocENINoFreeDevice(t *testing.T) {
435436
mockEC2.EXPECT().DeleteNetworkInterfaceWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
436437

437438
cache := &EC2InstanceMetadataCache{
438-
ec2SVC: mockEC2,
439-
imds: TypedIMDS{mockMetadata},
439+
ec2SVC: mockEC2,
440+
imds: TypedIMDS{mockMetadata},
441+
instanceType: "c5n.18xlarge",
440442
}
441443

442-
_, err := cache.AllocENI(false, nil, "")
444+
_, err := cache.AllocENI(false, nil, "", 5)
443445
assert.Error(t, err)
444446
}
445447

@@ -471,11 +473,128 @@ func TestAllocENIMaxReached(t *testing.T) {
471473
mockEC2.EXPECT().DeleteNetworkInterfaceWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
472474

473475
cache := &EC2InstanceMetadataCache{
474-
ec2SVC: mockEC2,
475-
imds: TypedIMDS{mockMetadata},
476+
ec2SVC: mockEC2,
477+
imds: TypedIMDS{mockMetadata},
478+
instanceType: "c5n.18xlarge",
476479
}
477480

478-
_, err := cache.AllocENI(false, nil, "")
481+
_, err := cache.AllocENI(false, nil, "", 5)
482+
assert.Error(t, err)
483+
}
484+
485+
func TestAllocENIWithIPAddresses(t *testing.T) {
486+
ctrl, mockEC2 := setup(t)
487+
defer ctrl.Finish()
488+
489+
// when required IP numbers(5) is below ENI's limit(30)
490+
currentEniID := eniID
491+
eni := ec2.CreateNetworkInterfaceOutput{NetworkInterface: &ec2.NetworkInterface{NetworkInterfaceId: &currentEniID}}
492+
mockEC2.EXPECT().CreateNetworkInterfaceWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(&eni, nil)
493+
494+
ec2ENIs := make([]*ec2.InstanceNetworkInterface, 0)
495+
deviceNum1 := int64(0)
496+
ec2ENI := &ec2.InstanceNetworkInterface{Attachment: &ec2.InstanceNetworkInterfaceAttachment{DeviceIndex: &deviceNum1}}
497+
ec2ENIs = append(ec2ENIs, ec2ENI)
498+
499+
deviceNum2 := int64(3)
500+
ec2ENI = &ec2.InstanceNetworkInterface{Attachment: &ec2.InstanceNetworkInterfaceAttachment{DeviceIndex: &deviceNum2}}
501+
ec2ENIs = append(ec2ENIs, ec2ENI)
502+
503+
result := &ec2.DescribeInstancesOutput{
504+
Reservations: []*ec2.Reservation{{Instances: []*ec2.Instance{{NetworkInterfaces: ec2ENIs}}}}}
505+
mockEC2.EXPECT().DescribeInstancesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(result, nil)
506+
attachmentID := "eni-attach-58ddda9d"
507+
attachResult := &ec2.AttachNetworkInterfaceOutput{
508+
AttachmentId: &attachmentID}
509+
mockEC2.EXPECT().AttachNetworkInterfaceWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(attachResult, nil)
510+
mockEC2.EXPECT().ModifyNetworkInterfaceAttributeWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
511+
512+
cache := &EC2InstanceMetadataCache{ec2SVC: mockEC2, instanceType: "c5n.18xlarge"}
513+
_, err := cache.AllocENI(false, nil, subnetID, 5)
514+
assert.NoError(t, err)
515+
516+
// when required IP numbers(50) is higher than ENI's limit(49)
517+
mockEC2.EXPECT().CreateNetworkInterfaceWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(&eni, nil)
518+
mockEC2.EXPECT().DescribeInstancesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(result, nil)
519+
mockEC2.EXPECT().AttachNetworkInterfaceWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(attachResult, nil)
520+
mockEC2.EXPECT().ModifyNetworkInterfaceAttributeWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
521+
cache = &EC2InstanceMetadataCache{ec2SVC: mockEC2, instanceType: "c5n.18xlarge"}
522+
_, err = cache.AllocENI(false, nil, subnetID, 49)
523+
assert.NoError(t, err)
524+
}
525+
526+
func TestAllocENIWithIPAddressesAlreadyFull(t *testing.T) {
527+
ctrl, mockEC2 := setup(t)
528+
defer ctrl.Finish()
529+
530+
mockMetadata := testMetadata(nil)
531+
532+
retErr := awserr.New("PrivateIpAddressLimitExceeded", "Too many IPs already allocated", nil)
533+
mockEC2.EXPECT().CreateNetworkInterfaceWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, retErr)
534+
535+
cache := &EC2InstanceMetadataCache{
536+
ec2SVC: mockEC2,
537+
imds: TypedIMDS{mockMetadata},
538+
instanceType: "t3.xlarge",
539+
}
540+
_, err := cache.AllocENI(true, nil, "", 14)
541+
assert.Error(t, err)
542+
}
543+
544+
func TestAllocENIWithPrefixAddresses(t *testing.T) {
545+
ctrl, mockEC2 := setup(t)
546+
defer ctrl.Finish()
547+
548+
mockMetadata := testMetadata(nil)
549+
550+
currentEniID := eniID
551+
eni := ec2.CreateNetworkInterfaceOutput{NetworkInterface: &ec2.NetworkInterface{NetworkInterfaceId: &currentEniID}}
552+
mockEC2.EXPECT().CreateNetworkInterfaceWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(&eni, nil)
553+
554+
ec2ENIs := make([]*ec2.InstanceNetworkInterface, 0)
555+
deviceNum1 := int64(0)
556+
ec2ENI := &ec2.InstanceNetworkInterface{Attachment: &ec2.InstanceNetworkInterfaceAttachment{DeviceIndex: &deviceNum1}}
557+
ec2ENIs = append(ec2ENIs, ec2ENI)
558+
559+
deviceNum2 := int64(3)
560+
ec2ENI = &ec2.InstanceNetworkInterface{Attachment: &ec2.InstanceNetworkInterfaceAttachment{DeviceIndex: &deviceNum2}}
561+
ec2ENIs = append(ec2ENIs, ec2ENI)
562+
563+
result := &ec2.DescribeInstancesOutput{
564+
Reservations: []*ec2.Reservation{{Instances: []*ec2.Instance{{NetworkInterfaces: ec2ENIs}}}}}
565+
mockEC2.EXPECT().DescribeInstancesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(result, nil)
566+
attachmentID := "eni-attach-58ddda9d"
567+
attachResult := &ec2.AttachNetworkInterfaceOutput{
568+
AttachmentId: &attachmentID}
569+
mockEC2.EXPECT().AttachNetworkInterfaceWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(attachResult, nil)
570+
mockEC2.EXPECT().ModifyNetworkInterfaceAttributeWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
571+
572+
cache := &EC2InstanceMetadataCache{
573+
ec2SVC: mockEC2,
574+
imds: TypedIMDS{mockMetadata},
575+
instanceType: "c5n.18xlarge",
576+
enablePrefixDelegation: true,
577+
}
578+
_, err := cache.AllocENI(false, nil, subnetID, 1)
579+
assert.NoError(t, err)
580+
}
581+
582+
func TestAllocENIWithPrefixesAlreadyFull(t *testing.T) {
583+
ctrl, mockEC2 := setup(t)
584+
defer ctrl.Finish()
585+
586+
mockMetadata := testMetadata(nil)
587+
588+
retErr := awserr.New("PrivateIpAddressLimitExceeded", "Too many IPs already allocated", nil)
589+
mockEC2.EXPECT().CreateNetworkInterfaceWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, retErr)
590+
591+
cache := &EC2InstanceMetadataCache{
592+
ec2SVC: mockEC2,
593+
imds: TypedIMDS{mockMetadata},
594+
instanceType: "c5n.18xlarge",
595+
enablePrefixDelegation: true,
596+
}
597+
_, err := cache.AllocENI(true, nil, "", 1)
479598
assert.Error(t, err)
480599
}
481600

0 commit comments

Comments
 (0)