Improve startup time by optimizing IMDS access (#4649)

.cspell/custom-dictionary.txt

@@ -435,3 +435,4 @@ firstbacalhauimage
 RYUK
 buildvcs
 Nilf
+IMDS

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/Masterminds/semver v1.5.0
 	github.com/aws/aws-sdk-go-v2 v1.30.4
 	github.com/aws/aws-sdk-go-v2/config v1.27.3
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.15.1
 	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.16.5
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.61.0
 	github.com/aws/smithy-go v1.20.4
@@ -89,7 +90,6 @@ require (
 	github.com/agnivade/levenshtein v1.1.1 // indirect
 	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.4 // indirect
 	github.com/aws/aws-sdk-go-v2/credentials v1.17.3 // indirect
-	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.15.1 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.16 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.16 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect

pkg/s3/utils.go

@@ -3,31 +3,93 @@ package s3
 import (
 	"context"
 	"os"
+	"strconv"
 	"strings"
+	"sync"
+	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
 	"github.com/rs/zerolog/log"
 )
 
-func DefaultAWSConfig() (aws.Config, error) {
-	// Set a default IMDC TTL of 1 hour if not set to avoid hitting the metadata service too often, which can slow down the
-	// node's startup time.
-	if _, ok := os.LookupEnv("AWS_EC2_METADATA_TTL"); !ok {
-		err := os.Setenv("AWS_EC2_METADATA_TTL", "3600")
-		if err != nil {
-			return aws.Config{}, err
+// IMDS (Instance Metadata Service) availability check results are cached globally
+// to avoid repeated timeout delays when IMDS is not available, which is the case
+// when running outside of AWS or when IMDS is disabled.
+var (
+	// imdsAvailable indicates whether IMDS is accessible
+	// This is set once during the first check and never modified afterwards
+	imdsAvailable bool
+
+	// imdsCheckOnce ensures the IMDS check is performed exactly once
+	imdsCheckOnce sync.Once
+)
+
+// checkIMDSAvailability determines if the AWS Instance Metadata Service (IMDS) is
+// accessible from the current environment. This function is safe to call multiple
+// times - the actual check will only be performed once, with subsequent calls
+// returning the cached result.
+//
+// Returns:
+//   - true if IMDS is available and responding
+//   - false if IMDS is disabled, unavailable, or times out
+func checkIMDSAvailability() bool {
+	imdsCheckOnce.Do(func() {
+		// Check if IMDS is explicitly disabled
+		// such as if `AWS_EC2_METADATA_DISABLED` is set to "true"
+		imdsOptions := imds.Options{}
+		if imdsOptions.ClientEnableState == imds.ClientDisabled {
+			imdsAvailable = false
+			return
 		}
-	}
+
+		// Attempt to access IMDS with configured timeout or default to 1 second
+		timeout := 1 * time.Second
+		if timeoutStr := os.Getenv("AWS_METADATA_SERVICE_TIMEOUT"); timeoutStr != "" {
+			if timeoutInt, err := strconv.Atoi(timeoutStr); err == nil && timeoutInt > 0 {
+				timeout = time.Duration(timeoutInt) * time.Second
+			}
+		}
+		ctx, cancel := context.WithTimeout(context.Background(), timeout)
+		defer cancel()
+
+		client := imds.New(imdsOptions)
+		_, err := client.GetMetadata(ctx, &imds.GetMetadataInput{
+			Path: "instance-id",
+		})
+
+		imdsAvailable = err == nil
+		if !imdsAvailable {
+			log.Debug().Msg("IMDS not available, will skip in future credential checks")
+		}
+	})
+
+	return imdsAvailable
+}
+
+// DefaultAWSConfig returns an AWS configuration with IMDS disabled if it's
+// determined to be unavailable or inaccessible. This prevents delays from
+// failed IMDS calls during credential retrieval.
+func DefaultAWSConfig() (aws.Config, error) {
 	var optFns []func(*config.LoadOptions) error
+
+	// If IMDS is not available, disable it in the config
+	if !checkIMDSAvailability() {
+		optFns = append(optFns, config.WithEC2IMDSClientEnableState(imds.ClientDisabled))
+	}
+
 	return config.LoadDefaultConfig(context.Background(), optFns...)
 }
 
 // HasValidCredentials returns true if the AWS config has valid credentials.
 func HasValidCredentials(config aws.Config) bool {
 	credentials, err := config.Credentials.Retrieve(context.Background())
 	if err != nil {
-		log.Debug().Err(err).Msg("Failed to check if we have valid AWS credentials")
+		// Only log if it's not an expected IMDS disabled error
+		if !strings.Contains(err.Error(), "EC2 IMDS") {
+			log.Debug().Err(err).Msg("Failed to check if we have valid AWS credentials")
+		}
 		return false
 	}
 	return credentials.HasKeys()