add changes from review

Author: haouc

cmd/aws-k8s-agent/main.go

@@ -144,7 +144,8 @@ func _main() int {
 	}
 	// Measure node initialization duration
 	IPAMDNodeInitStartTime := time.Now()
-	ipamContext, err := ipamd.New(k8sClient, withApiServer)
+	ctx := context.Background()
+	ipamContext, err := ipamd.New(ctx, k8sClient, withApiServer)
 	IPAMDNodeInitDuration := time.Since(IPAMDNodeInitStartTime).Seconds()
 
 	if err != nil {
@@ -164,7 +165,7 @@ func _main() int {
 	}
 
 	// Pool manager
-	go ipamContext.StartNodeIPPoolManager(context.Background())
+	go ipamContext.StartNodeIPPoolManager(ctx)
 
 	if !utils.GetBoolAsStringEnvVar(envDisableMetrics, false) {
 		// Prometheus metrics

pkg/awsutils/awsutils.go

@@ -408,10 +408,7 @@ func (i instrumentedIMDS) GetMetadataWithContext(ctx context.Context, p string)
 }
 
 // New creates an EC2InstanceMetadataCache
-func New(useSubnetDiscovery, useCustomNetworking, disableLeakedENICleanup, v4Enabled, v6Enabled bool) (*EC2InstanceMetadataCache, error) {
-	// ctx is passed to initWithEC2Metadata func to cancel spawned go-routines when tests are run
-	ctx := context.Background()
-
+func New(ctx context.Context, useSubnetDiscovery, useCustomNetworking, disableLeakedENICleanup, v4Enabled, v6Enabled bool) (*EC2InstanceMetadataCache, error) {
 	awsconfig, err := awssession.New(ctx)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create aws session")
@@ -449,7 +446,7 @@ func New(useSubnetDiscovery, useCustomNetworking, disableLeakedENICleanup, v4Ena
 
 	// Clean up leaked ENIs in the background
 	if !disableLeakedENICleanup {
-		go wait.Forever(func() { cache.cleanUpLeakedENIs(context.Background()) }, time.Hour)
+		go wait.Forever(func() { cache.cleanUpLeakedENIs(ctx) }, time.Hour)
 	}
 	return cache, nil
 }
@@ -602,10 +599,7 @@ func (cache *EC2InstanceMetadataCache) getFilteredENIs(ctx context.Context, stor
 			isSecondarySubnet := cache.isENIInSecondarySubnet(ctx, eniID)
 
 			// Filter based on subnet type
-			if onlySecondarySubnets && !isSecondarySubnet {
-				continue
-			} else if !onlySecondarySubnets && isSecondarySubnet {
-				// When we want primary subnet ENIs, skip secondary subnet ENIs
+			if onlySecondarySubnets != isSecondarySubnet {
 				continue
 			}
 
@@ -676,8 +670,7 @@ func (cache *EC2InstanceMetadataCache) RefreshCustomSGIDs(ctx context.Context, d
 	sgIDs, err := cache.discoverCustomSecurityGroups(ctx)
 	if err != nil {
 		awsAPIErrInc("DiscoverCustomSecurityGroups", err)
-		log.Warnf("Failed to discover custom security groups: %v", err)
-		log.Info("Falling back to using primary security groups for ENIs in secondary subnets")
+		log.Warnf("Failed to discover custom security groups: %v. Falling back to using primary security groups for ENIs in secondary subnets", err)
 
 		// Clear custom security groups cache to trigger fallback behavior
 		cache.customSecurityGroups.Set([]string{})

pkg/ipamd/ipamd.go

@@ -394,9 +394,8 @@ func (c *IPAMContext) inInsufficientCidrCoolingPeriod() bool {
 
 // New retrieves IP address usage information from Instance MetaData service and Kubelet
 // then initializes IP address pool data store
-func New(k8sClient client.Client, withApiServer bool) (*IPAMContext, error) {
+func New(ctx context.Context, k8sClient client.Client, withApiServer bool) (*IPAMContext, error) {
 	prometheusRegister()
-	ctx := context.Background()
 	c := &IPAMContext{}
 	c.k8sClient = k8sClient
 	c.networkClient = networkutils.New()
@@ -407,7 +406,7 @@ func New(k8sClient client.Client, withApiServer bool) (*IPAMContext, error) {
 	c.enableIPv4 = isIPv4Enabled()
 	c.enableIPv6 = isIPv6Enabled()
 	c.disableENIProvisioning = disableENIProvisioning()
-	client, err := awsutils.New(c.useSubnetDiscovery, c.useCustomNetworking, disableLeakedENICleanup(), c.enableIPv4, c.enableIPv6)
+	client, err := awsutils.New(ctx, c.useSubnetDiscovery, c.useCustomNetworking, disableLeakedENICleanup(), c.enableIPv4, c.enableIPv6)
 	if err != nil {
 		return nil, errors.Wrap(err, "ipamd: can not initialize with AWS SDK interface")
 	}
@@ -446,16 +445,15 @@ func New(k8sClient client.Client, withApiServer bool) (*IPAMContext, error) {
 	c.myNodeName = os.Getenv(envNodeName)
 	c.withApiServer = withApiServer
 
-	if err := c.nodeInit(); err != nil {
+	if err := c.nodeInit(ctx); err != nil {
 		return nil, err
 	}
 	return c, nil
 }
 
-func (c *IPAMContext) nodeInit() error {
+func (c *IPAMContext) nodeInit(ctx context.Context) error {
 	prometheusmetrics.IpamdActionsInprogress.WithLabelValues("nodeInit").Add(float64(1))
 	defer prometheusmetrics.IpamdActionsInprogress.WithLabelValues("nodeInit").Sub(float64(1))
-	ctx := context.TODO()
 	log.Debugf("Start node init")
 
 	if err := c.initENIAndIPLimits(); err != nil {
@@ -582,10 +580,10 @@ func (c *IPAMContext) nodeInit() error {
 		// Refresh security groups and VPC CIDR blocks in the background
 		// Ignoring errors since we will retry in 30s
 		go wait.Forever(func() {
-			c.awsClient.RefreshSGIDs(context.Background(), primaryENIMac, c.dataStoreAccess)
+			c.awsClient.RefreshSGIDs(ctx, primaryENIMac, c.dataStoreAccess)
 			// Also refresh custom security groups for secondary subnets
 			if c.useSubnetDiscovery {
-				c.awsClient.RefreshCustomSGIDs(context.Background(), c.dataStoreAccess)
+				c.awsClient.RefreshCustomSGIDs(ctx, c.dataStoreAccess)
 			}
 		}, 30*time.Second)
 	}
@@ -2519,8 +2517,6 @@ func (c *IPAMContext) cleanupExcludedENI(ctx context.Context, eniID string) {
 
 // checkAndHandleAllENIExclusion checks all existing ENIs (primary and secondary) across all network cards for subnet exclusion
 func (c *IPAMContext) checkAndHandleAllENIExclusion(ctx context.Context) {
-	log.Debugf("checkAndHandleAllENIExclusion: checking all existing ENIs for subnet exclusion")
-
 	// Iterate over all datastores (one per network card)
 	for _, ds := range c.dataStoreAccess.DataStores {
 		// Get ENI information using the public method

pkg/ipamd/ipamd_test.go

@@ -222,7 +222,7 @@ func TestNodeInit(t *testing.T) {
 	// Add IPs
 	m.awsutils.EXPECT().AllocIPAddresses(gomock.Any(), gomock.Any(), gomock.Any())
 	os.Setenv("MY_NODE_NAME", myNodeName)
-	err := mockContext.nodeInit()
+	err := mockContext.nodeInit(context.TODO())
 	assert.NoError(t, err)
 }
 
@@ -324,7 +324,7 @@ func TestNodeInitwithPDenabledIPv4Mode(t *testing.T) {
 	m.k8sClient.Create(ctx, &fakeNode)
 
 	os.Setenv("MY_NODE_NAME", myNodeName)
-	err := mockContext.nodeInit()
+	err := mockContext.nodeInit(context.TODO())
 	assert.NoError(t, err)
 }
 
@@ -413,7 +413,7 @@ func TestNodeInitwithPDenabledIPv6Mode(t *testing.T) {
 	m.k8sClient.Create(ctx, &fakeNode)
 	os.Setenv("MY_NODE_NAME", myNodeName)
 
-	err := mockContext.nodeInit()
+	err := mockContext.nodeInit(context.TODO())
 	assert.NoError(t, err)
 }
 
@@ -2968,7 +2968,7 @@ func TestNodeInitPrimarySubnetExclusionWithExistingPodIPs(t *testing.T) {
 	// Add IPs
 	m.awsutils.EXPECT().AllocIPAddresses(gomock.Any(), gomock.Any(), gomock.Any()).Return(&ec2.AssignPrivateIpAddressesOutput{}, nil)
 	os.Setenv("MY_NODE_NAME", myNodeName)
-	err := mockContext.nodeInit()
+	err := mockContext.nodeInit(context.TODO())
 	assert.NoError(t, err)
 
 	// Verify that primary ENI exclusion is now always respected
@@ -3111,7 +3111,7 @@ func TestNodeInitPrimarySubnetExclusionWithoutExistingPodIPs(t *testing.T) {
 	// Add IPs
 	m.awsutils.EXPECT().AllocIPAddresses(gomock.Any(), gomock.Any(), gomock.Any()).Return(&ec2.AssignPrivateIpAddressesOutput{}, nil)
 	os.Setenv("MY_NODE_NAME", myNodeName)
-	err := mockContext.nodeInit()
+	err := mockContext.nodeInit(context.TODO())
 	assert.NoError(t, err)
 
 	// Verify that primary ENI remains excluded due to no existing pod IPs