Merge pull request #2621 from ConnorJC3/cherrypick-gca-fix-147

[CHERRYPICK release-1.47] Call sts:GetCallerIdentity Just in Time instead of blocking on startup

Author: k8s-ci-robot

cmd/main.go

@@ -23,8 +23,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/aws/aws-sdk-go-v2/config"
-	"github.com/aws/aws-sdk-go-v2/service/sts"
 	"github.com/kubernetes-sigs/aws-ebs-csi-driver/cmd/hooks"
 	"github.com/kubernetes-sigs/aws-ebs-csi-driver/pkg/cloud"
 	"github.com/kubernetes-sigs/aws-ebs-csi-driver/pkg/cloud/metadata"
@@ -181,30 +179,7 @@ func main() {
 		region = md.GetRegion()
 	}
 
-	var accountID string
-	if options.Mode == driver.ControllerMode || options.Mode == driver.AllMode {
-		cfg, err := config.LoadDefaultConfig(context.Background(), config.WithRegion(region))
-		if err != nil {
-			klog.ErrorS(err, "Failed to create AWS config for account ID retrieval")
-			klog.FlushAndExit(klog.ExitFlushTimeout, 1)
-		}
-
-		stsClient := sts.NewFromConfig(cfg)
-		resp, err := stsClient.GetCallerIdentity(context.Background(), &sts.GetCallerIdentityInput{})
-		if err != nil {
-			klog.ErrorS(err, "Failed to get AWS account ID, HyperPod functionality may not work")
-			// Continue without account ID - existing functionality should still work
-		} else {
-			accountID = *resp.Account
-			klog.V(5).InfoS("Retrieved AWS account ID for HyperPod operations", "accountID", accountID)
-		}
-	}
-
-	cloud, err := cloud.NewCloud(region, accountID, options.AwsSdkDebugLog, options.UserAgentExtra, options.Batching, options.DeprecatedMetrics)
-	if err != nil {
-		klog.ErrorS(err, "failed to create cloud service")
-		klog.FlushAndExit(klog.ExitFlushTimeout, 1)
-	}
+	cloud := cloud.NewCloud(region, options.AwsSdkDebugLog, options.UserAgentExtra, options.Batching, options.DeprecatedMetrics)
 
 	m, err := mounter.NewNodeMounter(options.WindowsHostProcess)
 	if err != nil {

pkg/cloud/cloud.go

@@ -33,6 +33,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
 	"github.com/aws/aws-sdk-go-v2/service/sagemaker"
+	"github.com/aws/aws-sdk-go-v2/service/sts"
 	"github.com/aws/smithy-go"
 	"github.com/kubernetes-sigs/aws-ebs-csi-driver/pkg/batcher"
 	dm "github.com/kubernetes-sigs/aws-ebs-csi-driver/pkg/cloud/devicemanager"
@@ -88,8 +89,9 @@ var (
 )
 
 const (
-	cacheForgetDelay        = 1 * time.Hour
-	volInitCacheForgetDelay = 6 * time.Hour
+	cacheForgetDelay            = 1 * time.Hour
+	volInitCacheForgetDelay     = 6 * time.Hour
+	getCallerIdentityRetryDelay = 30 * time.Second
 )
 
 // VolumeStatusInitializingState is const reported by EC2 DescribeVolumeStatus which AWS SDK does not have type for.
@@ -320,6 +322,7 @@ type batcherManager struct {
 }
 
 type cloud struct {
+	awsConfig             aws.Config
 	region                string
 	ec2                   EC2API
 	sm                    SageMakerAPI
@@ -331,18 +334,14 @@ type cloud struct {
 	latestClientTokens    expiringcache.ExpiringCache[string, int]
 	volumeInitializations expiringcache.ExpiringCache[string, volumeInitialization]
 	accountID             string
+	accountIDOnce         sync.Once
 }
 
 var _ Cloud = &cloud{}
 
 // NewCloud returns a new instance of AWS cloud
 // It panics if session is invalid.
-func NewCloud(region string, accountID string, awsSdkDebugLog bool, userAgentExtra string, batching bool, deprecatedMetrics bool) (Cloud, error) {
-	c := newEC2Cloud(region, accountID, awsSdkDebugLog, userAgentExtra, batching, deprecatedMetrics)
-	return c, nil
-}
-
-func newEC2Cloud(region string, accountID string, awsSdkDebugLog bool, userAgentExtra string, batchingEnabled bool, deprecatedMetrics bool) Cloud {
+func NewCloud(region string, awsSdkDebugLog bool, userAgentExtra string, batchingEnabled bool, deprecatedMetrics bool) Cloud {
 	cfg, err := config.LoadDefaultConfig(context.Background(), config.WithRegion(region))
 	if err != nil {
 		panic(err)
@@ -386,11 +385,12 @@ func newEC2Cloud(region string, accountID string, awsSdkDebugLog bool, userAgent
 
 	var bm *batcherManager
 	if batchingEnabled {
-		klog.V(4).InfoS("newEC2Cloud: batching enabled")
+		klog.V(4).InfoS("NewCloud: batching enabled")
 		bm = newBatcherManager(svc)
 	}
 
 	return &cloud{
+		awsConfig:             cfg,
 		region:                region,
 		dm:                    dm.NewDeviceManager(),
 		ec2:                   svc,
@@ -400,7 +400,6 @@ func newEC2Cloud(region string, accountID string, awsSdkDebugLog bool, userAgent
 		vwp:                   vwp,
 		likelyBadDeviceNames:  expiringcache.New[string, sync.Map](cacheForgetDelay),
 		latestClientTokens:    expiringcache.New[string, int](cacheForgetDelay),
-		accountID:             accountID,
 		volumeInitializations: expiringcache.New[string, volumeInitialization](volInitCacheForgetDelay),
 	}
 }
@@ -997,7 +996,11 @@ func (c *cloud) attachDiskHyperPod(ctx context.Context, volumeID, nodeID string)
 	klog.V(2).InfoS("AttachDisk: HyperPod node detected", "volumeID", volumeID, "nodeID", nodeID)
 
 	instanceID := getInstanceIDFromHyperPodNode(nodeID)
-	clusterArn := c.buildHyperPodClusterArn(nodeID)
+	accountID, err := c.getAccountID(ctx)
+	if err != nil {
+		return "", fmt.Errorf("failed to get account ID: %w", err)
+	}
+	clusterArn := buildHyperPodClusterArn(nodeID, c.region, accountID)
 
 	klog.V(5).InfoS("HyperPod attachment details",
 		"volumeID", volumeID,
@@ -1025,7 +1028,7 @@ func (c *cloud) attachDiskHyperPod(ctx context.Context, volumeID, nodeID string)
 
 	// Wait for attachment completion
 	deviceName := aws.ToString(resp.DeviceName)
-	_, err := c.WaitForAttachmentState(
+	_, err = c.WaitForAttachmentState(
 		ctx,
 		types.VolumeAttachmentStateAttached,
 		volumeID,
@@ -1099,7 +1102,11 @@ func (c *cloud) detachDiskHyperPod(ctx context.Context, volumeID, nodeID string)
 	klog.V(2).InfoS("DetachDisk: HyperPod node detected", "volumeID", volumeID, "nodeID", nodeID)
 
 	instanceID := getInstanceIDFromHyperPodNode(nodeID)
-	clusterArn := c.buildHyperPodClusterArn(nodeID)
+	accountID, err := c.getAccountID(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get account ID: %w", err)
+	}
+	clusterArn := buildHyperPodClusterArn(nodeID, c.region, accountID)
 
 	klog.V(4).InfoS("HyperPod detachment details",
 		"volumeID", volumeID,
@@ -1114,7 +1121,7 @@ func (c *cloud) detachDiskHyperPod(ctx context.Context, volumeID, nodeID string)
 	}
 	klog.V(4).InfoS("Calling DetachClusterNodeVolumeInput", "input", input)
 
-	_, err := c.sm.DetachClusterNodeVolume(ctx, input)
+	_, err = c.sm.DetachClusterNodeVolume(ctx, input)
 	if err != nil {
 		if isAWSHyperPodErrorIncorrectState(err) ||
 			isAWSHyperPodErrorInvalidAttachmentNotFound(err) ||
@@ -1450,9 +1457,9 @@ func getInstanceIDFromHyperPodNode(nodeID string) string {
 }
 
 // Only for hyperpod node, buildHyperPodClusterArn: arn:aws:sagemaker:region:account:cluster/clusterID.
-func (c *cloud) buildHyperPodClusterArn(nodeID string) string {
+func buildHyperPodClusterArn(nodeID string, region string, accountID string) string {
 	parts := strings.Split(nodeID, "-")
-	return fmt.Sprintf("arn:aws:sagemaker:%s:%s:cluster/%s", c.region, c.accountID, parts[1])
+	return fmt.Sprintf("arn:aws:sagemaker:%s:%s:cluster/%s", region, accountID, parts[1])
 }
 
 // For hyperpod node, AssociatedResource is in arn:aws:sagemaker:region:account:cluster/clusterID-instanceId format.
@@ -1916,6 +1923,54 @@ func (c *cloud) waitForVolume(ctx context.Context, volumeID string) error {
 	return err
 }
 
+// getAccountID returns the account ID of the AWS Account for the IAM credentials in use.
+//
+// In the first call (or any calls made before the first call succeeds), getAccountID
+// will attempt to determine the Account ID via sts:GetCallerIdentity.
+// This attempt will retry indefinitely, however getAccountID will return when ctx is cancelled,
+// leaving the account ID thread to run in the background.
+//
+// In subsequent calls (after the first success), getAccountID will use a cached value.
+func (c *cloud) getAccountID(ctx context.Context) (string, error) {
+	accountIDRetrieved := make(chan struct{}, 1)
+
+	// Start background thread if it isn't already.
+	// Intentionally runs in the background until account ID is retrieved, so we don't pass the context.
+	//nolint:contextcheck
+	go func() {
+		c.accountIDOnce.Do(func() {
+			for c.accountID == "" {
+				cfg, err := config.LoadDefaultConfig(context.Background(), config.WithRegion(c.region))
+				if err != nil {
+					klog.ErrorS(err, "Failed to create AWS config for account ID retrieval")
+				}
+
+				stsClient := sts.NewFromConfig(cfg)
+				resp, err := stsClient.GetCallerIdentity(context.Background(), &sts.GetCallerIdentityInput{})
+				if err != nil {
+					klog.ErrorS(err, "Failed to get AWS account ID, required for HyperPod operations, will retry")
+					time.Sleep(getCallerIdentityRetryDelay)
+				} else {
+					c.accountID = *resp.Account
+					klog.V(5).InfoS("Retrieved AWS account ID for HyperPod operations", "accountID", c.accountID)
+				}
+			}
+		})
+
+		// Once.Do blocks until the function exits, even if we aren't the first caller.
+		// So the account ID must be available now.
+		accountIDRetrieved <- struct{}{}
+	}()
+
+	select {
+	case <-ctx.Done():
+		return "", ctx.Err()
+
+	case <-accountIDRetrieved:
+		return c.accountID, nil
+	}
+}
+
 // isAWSError returns a boolean indicating whether the error is AWS-related
 // and has the given code. More information on AWS error codes at:
 // https://docs.aws.amazon.com/AWSEC2/latest/APIReference/errors-overview.html

pkg/cloud/cloud_test.go

@@ -89,7 +89,6 @@ func TestNewCloud(t *testing.T) {
 	testCases := []struct {
 		name              string
 		region            string
-		accountID         string
 		awsSdkDebugLog    bool
 		userAgentExtra    string
 		batchingEnabled   bool
@@ -98,15 +97,13 @@ func TestNewCloud(t *testing.T) {
 		{
 			name:            "success: with awsSdkDebugLog, userAgentExtra, and batchingEnabled",
 			region:          "us-east-1",
-			accountID:       "123456789012",
 			awsSdkDebugLog:  true,
 			userAgentExtra:  "example_user_agent_extra",
 			batchingEnabled: true,
 		},
 		{
 			name:           "success: with only awsSdkDebugLog, and userAgentExtra",
 			region:         "us-east-1",
-			accountID:      "123456789012",
 			awsSdkDebugLog: true,
 			userAgentExtra: "example_user_agent_extra",
 		},
@@ -116,10 +113,7 @@ func TestNewCloud(t *testing.T) {
 		},
 	}
 	for _, tc := range testCases {
-		ec2Cloud, err := NewCloud(tc.region, tc.accountID, tc.awsSdkDebugLog, tc.userAgentExtra, tc.batchingEnabled, tc.deprecatedMetrics)
-		if err != nil {
-			t.Fatalf("error %v", err)
-		}
+		ec2Cloud := NewCloud(tc.region, tc.awsSdkDebugLog, tc.userAgentExtra, tc.batchingEnabled, tc.deprecatedMetrics)
 		ec2CloudAscloud, ok := ec2Cloud.(*cloud)
 		if !ok {
 			t.Fatalf("could not assert object ec2Cloud as cloud type, %v", ec2Cloud)
@@ -132,6 +126,7 @@ func TestNewCloud(t *testing.T) {
 		}
 	}
 }
+
 func TestBatchDescribeVolumes(t *testing.T) {
 	t.Parallel()
 	testCases := []struct {
@@ -2442,11 +2437,15 @@ func TestBuildHyperPodClusterArn(t *testing.T) {
 	testCases := []struct {
 		name        string
 		nodeID      string
+		region      string
+		accountID   string
 		expectedArn string
 	}{
 		{
 			name:        "success: valid HyperPod node",
 			nodeID:      "hyperpod-abc123-i-1234567890abcdef0",
+			region:      "test-region",
+			accountID:   "123456789012",
 			expectedArn: "arn:aws:sagemaker:test-region:123456789012:cluster/abc123",
 		},
 	}
@@ -2456,13 +2455,7 @@ func TestBuildHyperPodClusterArn(t *testing.T) {
 			mockCtrl := gomock.NewController(t)
 			defer mockCtrl.Finish()
 
-			mockEC2 := NewMockEC2API(mockCtrl)
-			c := newCloud(mockEC2)
-			cloudInstance, ok := c.(*cloud)
-			if !ok {
-				t.Fatalf("could not assert cloudInstance as type cloud, %v", cloudInstance)
-			}
-			result := cloudInstance.buildHyperPodClusterArn(tc.nodeID)
+			result := buildHyperPodClusterArn(tc.nodeID, tc.region, tc.accountID)
 			assert.Equal(t, tc.expectedArn, result)
 		})
 	}

tests/e2e/dynamic_provisioning.go

@@ -621,10 +621,7 @@ var _ = Describe("[ebs-csi-e2e] [single-az] Dynamic Provisioning", func() {
 		availabilityZones := strings.Split(os.Getenv(awsAvailabilityZonesEnv), ",")
 		availabilityZone := availabilityZones[rand.Intn(len(availabilityZones))]
 		region := availabilityZone[0 : len(availabilityZone)-1]
-		cloud, err := awscloud.NewCloud(region, "", false, "", true, false)
-		if err != nil {
-			Fail(fmt.Sprintf("could not get NewCloud: %v", err))
-		}
+		cloud := awscloud.NewCloud(region, false, "", true, false)
 
 		test := testsuites.DynamicallyProvisionedReclaimPolicyTest{
 			CSIDriver: ebsDriver,

tests/e2e/pre_provsioning.go

@@ -87,10 +87,7 @@ var _ = Describe("[ebs-csi-e2e] [single-az] Pre-Provisioned", func() {
 			Tags:             map[string]string{awscloud.VolumeNameTagKey: dummyVolumeName, awscloud.AwsEbsDriverTagKey: "true"},
 		}
 		var err error
-		cloud, err = awscloud.NewCloud(region, "", false, "", true, false)
-		if err != nil {
-			Fail(fmt.Sprintf("could not get NewCloud: %v", err))
-		}
+		cloud = awscloud.NewCloud(region, false, "", true, false)
 		r1 := rand.New(rand.NewSource(time.Now().UnixNano()))
 		disk, err := cloud.CreateDisk(context.Background(), fmt.Sprintf("pvc-%d", r1.Uint64()), diskOptions)
 		if err != nil {
@@ -260,10 +257,7 @@ var _ = Describe("[ebs-csi-e2e] [single-az] Pre-Provisioned with Multi-Attach",
 			Tags:               map[string]string{awscloud.VolumeNameTagKey: dummyVolumeName, awscloud.AwsEbsDriverTagKey: "true"},
 		}
 		var err error
-		cloud, err = awscloud.NewCloud(region, "", false, "", true, false)
-		if err != nil {
-			Fail(fmt.Sprintf("could not get NewCloud: %v", err))
-		}
+		cloud = awscloud.NewCloud(region, false, "", true, false)
 		r1 := rand.New(rand.NewSource(time.Now().UnixNano()))
 		disk, err := cloud.CreateDisk(context.Background(), fmt.Sprintf("pvc-%d", r1.Uint64()), diskOptions)
 		if err != nil {