IPAMD optimizations and makefile changes (#1975)

* IPAMD optimizations and makefile changes

* Minor comments

* Removed IMDS dependency

* fix test

* fix test

* fix test-format

Makefile

@@ -99,6 +99,8 @@ MULTI_PLATFORM_BUILD_TARGETS = 	linux/amd64,linux/arm64
 # Default to building an executable using the host's Go toolchain.
 .DEFAULT_GOAL = build-linux
 
+##@ Building
+
 # Build both CNI and metrics helper container images.
 all: docker docker-init docker-metrics   ## Builds Init, CNI and metrics helper container images.
 
@@ -137,6 +139,23 @@ docker-func-test: docker     ## Run the built CNI container image to use in func
 	docker run $(DOCKER_RUN_FLAGS) \
 		"$(IMAGE_NAME)"
 
+multi-arch-cni-build-push:		## Build multi-arch VPC CNI container image.
+	docker buildx build $(DOCKER_BUILD_FLAGS) \
+    		-f scripts/dockerfiles/Dockerfile.release \
+    		--platform "$(MULTI_PLATFORM_BUILD_TARGETS)"\
+    		-t "$(IMAGE_NAME)" \
+    		--push \
+    		.
+
+multi-arch-cni-init-build-push:     ## Build VPC CNI plugin Init container image.
+	docker buildx build $(DOCKER_BUILD_FLAGS) \
+		-f scripts/dockerfiles/Dockerfile.init \
+		--platform "$(MULTI_PLATFORM_BUILD_TARGETS)"\
+		-t "$(INIT_IMAGE_NAME)" \
+		--push \
+		.
+
+##@ Run Unit Tests 
 # Run unit tests
 unit-test: export AWS_VPC_K8S_CNI_LOG_FILE=stdout
 unit-test:    ## Run unit tests
@@ -155,6 +174,7 @@ unit-test-race:     ## Run unit tests with race detection (can only be run nativ
 	go test -v -cover -race -timeout 10s  ./pkg/eniconfig/...
 	go test -v -cover -race -timeout 10s  ./pkg/ipamd/...
 
+##@ Build and Run Unit Tests 
 # Build the unit test driver container image.
 build-docker-test:     ## Build the unit test driver container image.
 	docker build $(DOCKER_BUILD_FLAGS) \
@@ -167,6 +187,8 @@ docker-unit-tests: build-docker-test     ## Run unit tests inside of the testing
 	docker run $(DOCKER_RUN_ARGS) \
 		$(TEST_IMAGE_NAME) \
 		make unit-test
+
+##@ Build metrics helper agent 
 
 # Build metrics helper agent.
 build-metrics:     ## Build metrics helper agent.
@@ -180,6 +202,8 @@ docker-metrics:    ## Build metrics helper agent Docker image.
 		.
 	@echo "Built Docker image \"$(METRICS_IMAGE_NAME)\""
 
+##@ Run metrics helper Unit Tests 
+
 # Run metrics helper unit test suite (must be run natively).
 metrics-unit-test: CGO_ENABLED=1
 metrics-unit-test: GOARCH=
@@ -195,31 +219,6 @@ docker-metrics-test:     ## Run metrics helper unit test suite in a container.
 		$(GOLANG_IMAGE) \
 		make metrics-unit-test
 
-generate:
-	PATH=$(CURDIR)/scripts:$(PATH) go generate -x ./...
-	$(MAKE) format
-
-# Generate limit file go code
-# Generate eni-max-pods.txt file for EKS AMI
-generate-limits: GOOS=
-generate-limits:    ## Generate limit file go code
-	go run $(VENDOR_OVERRIDE_FLAG) scripts/gen_vpc_ip_limits.go
-
-multi-arch-cni-build-push:		## Build multi-arch VPC CNI container image.
-	docker buildx build $(DOCKER_BUILD_FLAGS) \
-    		-f scripts/dockerfiles/Dockerfile.release \
-    		--platform "$(MULTI_PLATFORM_BUILD_TARGETS)"\
-    		-t "$(IMAGE_NAME)" \
-    		--push \
-    		.
-
-multi-arch-cni-init-build-push:     ## Build VPC CNI plugin Init container image.
-	docker buildx build $(DOCKER_BUILD_FLAGS) \
-		-f scripts/dockerfiles/Dockerfile.init \
-		--platform "$(MULTI_PLATFORM_BUILD_TARGETS)"\
-		-t "$(INIT_IMAGE_NAME)" \
-		--push \
-		.
 # Fetch the CNI plugins
 plugins: FETCH_VERSION=1.1.1
 plugins: FETCH_URL=https://github.com/containernetworking/plugins/releases/download/v$(FETCH_VERSION)/cni-plugins-$(GOOS)-$(GOARCH)-v$(FETCH_VERSION).tgz
@@ -232,6 +231,8 @@ plugins:   ## Fetch the CNI plugins
 	@echo
 	curl -L $(FETCH_URL) | tar -zx $(PLUGIN_BINS)
 
+##@ Debug script 
+
 debug-script: FETCH_URL=https://raw.githubusercontent.com/awslabs/amazon-eks-ami/master/log-collector-script/linux/eks-log-collector.sh
 debug-script: VISIT_URL=https://github.com/awslabs/amazon-eks-ami/tree/master/log-collector-script/linux
 debug-script:    ## Fetching debug script from awslabs/amazon-eks-ami
@@ -243,6 +244,8 @@ debug-script:    ## Fetching debug script from awslabs/amazon-eks-ami
 	curl -L $(FETCH_URL) -o ./aws-cni-support.sh
 	chmod +x ./aws-cni-support.sh
 
+##@ Formatting 
+
 # Run all source code checks.
 check: check-format lint vet   ## Run all source code checks.
 
@@ -289,6 +292,18 @@ check-format: format
 version:
 	@echo ${VERSION}
 
+##@ Generate ENI/IP limits 
+
+generate:
+	PATH=$(CURDIR)/scripts:$(PATH) go generate -x ./...
+	$(MAKE) format
+
+# Generate limit file go code
+# Generate eni-max-pods.txt file for EKS AMI
+generate-limits: GOOS=
+generate-limits:    ## Generate limit file go code
+	go run $(VENDOR_OVERRIDE_FLAG) scripts/gen_vpc_ip_limits.go
+
 ekscharts-sync:
 	${MAKEFILE_PATH}/scripts/sync-to-eks-charts.sh -b ${HELM_CHART_NAME} -r ${REPO_FULL_NAME}
 
@@ -315,13 +330,16 @@ cleanup-ec2-sdk-override:
 	@if [ "$(EC2_SDK_OVERRIDE)" = "y" ] ; then \
 	    ./scripts/ec2_model_override/cleanup.sh ; \
 	fi
+
+##@ Cleanup
 
 # Clean temporary files and build artifacts from the project.
 clean:    ## Clean temporary files and build artifacts from the project.
 	@rm -f -- $(BINS)
 	@rm -f -- $(PLUGIN_BINS)
 	@rm -f -- coverage.txt
 
-help:           ## Show this help.
-	@grep -F -h "##" $(MAKEFILE_LIST) | grep -F -v grep | sed -e 's/\\$$//' \
-		| awk -F'[:#]' '{print $$1 = sprintf("%-30s", $$1), $$4}'
+##@ Helpers
+
+help:   ## Display this help
+	@awk 'BEGIN {FS = ":.*##"; printf "\nUsage:\n  make \033[36m<target>\033[0m\n"} /^[a-zA-Z_-]+:.*?##/ { printf "  \033[36m%-15s\033[0m %s\n", $$1, $$2 } /^##@/ { printf "\n\033[1m%s\033[0m\n", substr($$0, 5) } ' $(MAKEFILE_LIST)

pkg/awsutils/awsutils.go

@@ -138,7 +138,7 @@ type APIs interface {
 	AllocIPAddress(eniID string) error
 
 	// AllocIPAddresses allocates numIPs IP addresses on a ENI
-	AllocIPAddresses(eniID string, numIPs int) error
+	AllocIPAddresses(eniID string, numIPs int) (*ec2.AssignPrivateIpAddressesOutput, error)
 
 	// DeallocIPAddresses deallocates the list of IP addresses from a ENI
 	DeallocIPAddresses(eniID string, ips []string) error
@@ -1439,7 +1439,7 @@ func (cache *EC2InstanceMetadataCache) IsPrefixDelegationSupported() bool {
 }
 
 // AllocIPAddresses allocates numIPs of IP address on an ENI
-func (cache *EC2InstanceMetadataCache) AllocIPAddresses(eniID string, numIPs int) error {
+func (cache *EC2InstanceMetadataCache) AllocIPAddresses(eniID string, numIPs int) (*ec2.AssignPrivateIpAddressesOutput, error) {
 	var needIPs = numIPs
 
 	ipLimit := cache.GetENIIPv4Limit()
@@ -1450,7 +1450,7 @@ func (cache *EC2InstanceMetadataCache) AllocIPAddresses(eniID string, numIPs int
 
 	// If we don't need any more IPs, exit
 	if needIPs < 1 {
-		return nil
+		return nil, nil
 	}
 
 	log.Infof("Trying to allocate %d IP addresses on ENI %s", needIPs, eniID)
@@ -1479,11 +1479,11 @@ func (cache *EC2InstanceMetadataCache) AllocIPAddresses(eniID string, numIPs int
 		if containsPrivateIPAddressLimitExceededError(err) {
 			log.Debug("AssignPrivateIpAddresses returned PrivateIpAddressLimitExceeded. This can happen if the data store is out of sync." +
 				"Returning without an error here since we will verify the actual state by calling EC2 to see what addresses have already assigned to this ENI.")
-			return nil
+			return nil, nil
 		}
 		log.Errorf("Failed to allocate a private IP/Prefix addresses on ENI %v: %v", eniID, err)
 		awsAPIErrInc("AssignPrivateIpAddresses", err)
-		return err
+		return nil, err
 	}
 	if output != nil {
 		if cache.enablePrefixDelegation {
@@ -1492,7 +1492,7 @@ func (cache *EC2InstanceMetadataCache) AllocIPAddresses(eniID string, numIPs int
 			log.Infof("Allocated %d private IP addresses", len(output.AssignedPrivateIpAddresses))
 		}
 	}
-	return nil
+	return output, nil
 }
 
 func (cache *EC2InstanceMetadataCache) AllocIPv6Prefixes(eniID string) ([]*string, error) {

pkg/awsutils/awsutils_test.go

@@ -581,7 +581,7 @@ func TestAllocIPAddresses(t *testing.T) {
 	mockEC2.EXPECT().AssignPrivateIpAddressesWithContext(gomock.Any(), input, gomock.Any()).Return(nil, nil)
 
 	ins := &EC2InstanceMetadataCache{ec2SVC: mockEC2, instanceType: "c5n.18xlarge"}
-	err := ins.AllocIPAddresses(eniID, 5)
+	_, err := ins.AllocIPAddresses(eniID, 5)
 	assert.NoError(t, err)
 
 	// when required IP numbers(50) is higher than ENI's limit(49)
@@ -597,11 +597,11 @@ func TestAllocIPAddresses(t *testing.T) {
 	mockEC2.EXPECT().AssignPrivateIpAddressesWithContext(gomock.Any(), input, gomock.Any()).Return(&output, nil)
 
 	ins = &EC2InstanceMetadataCache{ec2SVC: mockEC2, instanceType: "c5n.18xlarge"}
-	err = ins.AllocIPAddresses(eniID, 50)
+	_, err = ins.AllocIPAddresses(eniID, 50)
 	assert.NoError(t, err)
 
 	// Adding 0 should do nothing
-	err = ins.AllocIPAddresses(eniID, 0)
+	_, err = ins.AllocIPAddresses(eniID, 0)
 	assert.NoError(t, err)
 }
 
@@ -618,7 +618,7 @@ func TestAllocIPAddressesAlreadyFull(t *testing.T) {
 	retErr := awserr.New("PrivateIpAddressLimitExceeded", "Too many IPs already allocated", nil)
 	mockEC2.EXPECT().AssignPrivateIpAddressesWithContext(gomock.Any(), input, gomock.Any()).Return(nil, retErr)
 	// If EC2 says that all IPs are already attached, we do nothing
-	err := ins.AllocIPAddresses(eniID, 14)
+	_, err := ins.AllocIPAddresses(eniID, 14)
 	assert.NoError(t, err)
 }
 
@@ -634,11 +634,11 @@ func TestAllocPrefixAddresses(t *testing.T) {
 	mockEC2.EXPECT().AssignPrivateIpAddressesWithContext(gomock.Any(), input, gomock.Any()).Return(nil, nil)
 
 	ins := &EC2InstanceMetadataCache{ec2SVC: mockEC2, instanceType: "c5n.18xlarge", enablePrefixDelegation: true}
-	err := ins.AllocIPAddresses(eniID, 1)
+	_, err := ins.AllocIPAddresses(eniID, 1)
 	assert.NoError(t, err)
 
 	// Adding 0 should do nothing
-	err = ins.AllocIPAddresses(eniID, 0)
+	_, err = ins.AllocIPAddresses(eniID, 0)
 	assert.NoError(t, err)
 }
 
@@ -655,7 +655,7 @@ func TestAllocPrefixesAlreadyFull(t *testing.T) {
 	retErr := awserr.New("PrivateIpAddressLimitExceeded", "Too many IPs already allocated", nil)
 	mockEC2.EXPECT().AssignPrivateIpAddressesWithContext(gomock.Any(), input, gomock.Any()).Return(nil, retErr)
 	// If EC2 says that all IPs are already attached, we do nothing
-	err := ins.AllocIPAddresses(eniID, 1)
+	_, err := ins.AllocIPAddresses(eniID, 1)
 	assert.NoError(t, err)
 }
 

pkg/ipamd/ipamd.go

@@ -16,7 +16,6 @@ package ipamd
 import (
 	"context"
 	"fmt"
-	"math"
 	"net"
 	"os"
 	"strconv"
@@ -859,7 +858,7 @@ func (c *IPAMContext) tryAllocateENI(ctx context.Context) error {
 
 	resourcesToAllocate := c.GetENIResourcesToAllocate()
 
-	err = c.awsClient.AllocIPAddresses(eni, resourcesToAllocate)
+	_, err = c.awsClient.AllocIPAddresses(eni, resourcesToAllocate)
 	if err != nil {
 		log.Warnf("Failed to allocate %d IP addresses on an ENI: %v", resourcesToAllocate, err)
 		// Continue to process the allocated IP addresses
@@ -932,24 +931,29 @@ func (c *IPAMContext) tryAssignIPs() (increasedPool bool, err error) {
 	if eni != nil && len(eni.AvailableIPv4Cidrs) < c.maxIPsPerENI {
 		currentNumberOfAllocatedIPs := len(eni.AvailableIPv4Cidrs)
 		// Try to allocate all available IPs for this ENI
-		err = c.awsClient.AllocIPAddresses(eni.ID, int(math.Min(float64(c.maxIPsPerENI-currentNumberOfAllocatedIPs), float64(toAllocate))))
+		resourcesToAllocate := min((c.maxIPsPerENI - currentNumberOfAllocatedIPs), toAllocate)
+		output, err := c.awsClient.AllocIPAddresses(eni.ID, resourcesToAllocate)
 		if err != nil {
 			log.Warnf("failed to allocate all available IP addresses on ENI %s, err: %v", eni.ID, err)
 			// Try to just get one more IP
-			err = c.awsClient.AllocIPAddresses(eni.ID, 1)
+			output, err = c.awsClient.AllocIPAddresses(eni.ID, 1)
 			if err != nil {
 				ipamdErrInc("increaseIPPoolAllocIPAddressesFailed")
 				return false, errors.Wrap(err, fmt.Sprintf("failed to allocate one IP addresses on ENI %s, err ", eni.ID))
 			}
 		}
-		// This call to EC2 is needed to verify which IPs got attached to this ENI.
-		ec2Addrs, err := c.awsClient.GetIPv4sFromEC2(eni.ID)
-		if err != nil {
+
+		if output == nil {
 			ipamdErrInc("increaseIPPoolGetENIaddressesFailed")
 			return true, errors.Wrap(err, "failed to get ENI IP addresses during IP allocation")
 		}
 
-		c.addENIsecondaryIPsToDataStore(ec2Addrs, eni.ID)
+		var ec2ip4s []*ec2.NetworkInterfacePrivateIpAddress
+		ec2Addrs := output.AssignedPrivateIpAddresses
+		for _, ec2Addr := range ec2Addrs {
+			ec2ip4s = append(ec2ip4s, &ec2.NetworkInterfacePrivateIpAddress{PrivateIpAddress: aws.String(aws.StringValue(ec2Addr.PrivateIpAddress))})
+		}
+		c.addENIsecondaryIPsToDataStore(ec2ip4s, eni.ID)
 		return true, nil
 	}
 	return false, nil
@@ -1001,21 +1005,22 @@ func (c *IPAMContext) tryAssignPrefixes() (increasedPool bool, err error) {
 	eni := c.dataStore.GetENINeedsIP(c.maxPrefixesPerENI, c.useCustomNetworking)
 	if eni != nil {
 		currentNumberOfAllocatedPrefixes := len(eni.AvailableIPv4Cidrs)
-		err = c.awsClient.AllocIPAddresses(eni.ID, min((c.maxPrefixesPerENI-currentNumberOfAllocatedPrefixes), toAllocate))
+		resourcesToAllocate := min((c.maxPrefixesPerENI - currentNumberOfAllocatedPrefixes), toAllocate)
+		output, err := c.awsClient.AllocIPAddresses(eni.ID, resourcesToAllocate)
 		if err != nil {
 			log.Warnf("failed to allocate all available IPv4 Prefixes on ENI %s, err: %v", eni.ID, err)
 			// Try to just get one more prefix
-			err = c.awsClient.AllocIPAddresses(eni.ID, 1)
+			output, err = c.awsClient.AllocIPAddresses(eni.ID, 1)
 			if err != nil {
 				ipamdErrInc("increaseIPPoolAllocIPAddressesFailed")
 				return false, errors.Wrap(err, fmt.Sprintf("failed to allocate one IPv4 prefix on ENI %s, err: %v", eni.ID, err))
 			}
 		}
-		ec2Prefixes, err := c.awsClient.GetIPv4PrefixesFromEC2(eni.ID)
-		if err != nil {
+		if output == nil {
 			ipamdErrInc("increaseIPPoolGetENIprefixedFailed")
 			return true, errors.Wrap(err, "failed to get ENI Prefix addresses during IPv4 Prefix allocation")
 		}
+		ec2Prefixes := output.AssignedIpv4Prefixes
 		c.addENIv4prefixesToDataStore(ec2Prefixes, eni.ID)
 		return true, nil
 	}
@@ -2043,10 +2048,10 @@ func (c *IPAMContext) GetENIResourcesToAllocate() int {
 		resourcesToAllocate = c.maxIPsPerENI
 		short, _, warmTargetDefined := c.datastoreTargetState()
 		if warmTargetDefined {
-			resourcesToAllocate = short
+			resourcesToAllocate = min(short, c.maxIPsPerENI)
 		}
 	} else {
-		resourcesToAllocate = c.getPrefixesNeeded()
+		resourcesToAllocate = min(c.getPrefixesNeeded(), c.maxPrefixesPerENI)
 	}
 	return resourcesToAllocate
 }