From b82d0a774cb93a4c52cd2a33f8b01e3d68dde9c9 Mon Sep 17 00:00:00 2001
From: vela <wwvela@amazon.com>
Date: Thu, 18 Jul 2024 20:11:32 -0700
Subject: [PATCH] add retry to avoid ECR get auth token failed for some
 transient issues

---
 nodeadm/internal/aws/ecr/ecr.go | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/nodeadm/internal/aws/ecr/ecr.go b/nodeadm/internal/aws/ecr/ecr.go
index fc7ec4d55..4a9d444f0 100644
--- a/nodeadm/internal/aws/ecr/ecr.go
+++ b/nodeadm/internal/aws/ecr/ecr.go
@@ -5,11 +5,13 @@ import (
 	"fmt"
 	"net"
 	"strings"
+	"time"
 
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/ecr"
 	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/aws/imds"
 	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/system"
+	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/util"
 )
 
 // Returns the base64 encoded authorization token string for ECR of the format "AWS:XXXXX"
@@ -19,15 +21,19 @@ func GetAuthorizationToken(awsRegion string) (string, error) {
 		return "", err
 	}
 	ecrClient := ecr.NewFromConfig(awsConfig)
-
-	for i := 0; i < 3; i++ {
-		token, err := ecrClient.GetAuthorizationToken(context.Background(), &ecr.GetAuthorizationTokenInput{})
-		if err == nil {
-			authData := token.AuthorizationData[0].AuthorizationToken
-			return *authData, nil
+	var token *ecr.GetAuthorizationTokenOutput
+	err = util.RetryExponentialBackoff(3, 2*time.Second, func() error {
+		token, err = ecrClient.GetAuthorizationToken(context.Background(), &ecr.GetAuthorizationTokenInput{})
+		if err != nil {
+			return err
 		}
+		return nil
+	})
+	if err != nil {
+		return "", err
 	}
-	return "", err
+	authData := token.AuthorizationData[0].AuthorizationToken
+	return *authData, nil
 }
 
 func (r *ECRRegistry) GetSandboxImage() string {