Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve startup time by optimizing IMDS access #4649

Merged
merged 5 commits into from
Oct 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .cspell/custom-dictionary.txt
Original file line number Diff line number Diff line change
Expand Up @@ -435,3 +435,4 @@ firstbacalhauimage
RYUK
buildvcs
Nilf
IMDS
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
80 changes: 71 additions & 9 deletions pkg/s3/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading