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

feat: add disable switch for imds default timeout #2572

Merged
merged 2 commits into from Mar 21, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
8 changes: 8 additions & 0 deletions .changelog/bd8974c81b064061a388acea73c74956.json
@@ -0,0 +1,8 @@
{
"id": "bd8974c8-1b06-4061-a388-acea73c74956",
"type": "feature",
"description": "Add config switch `DisableDefaultTimeout` that allows you to disable the default operation timeout (5 seconds) for IMDS calls.",
"modules": [
"feature/ec2/imds"
]
}
4 changes: 4 additions & 0 deletions feature/ec2/imds/api_client.go
Expand Up @@ -185,6 +185,10 @@ type Options struct {
// [configuring IMDS]: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html
EnableFallback aws.Ternary

// By default, all IMDS client operations enforce a 5-second timeout. You
// can disable that behavior with this setting.
DisableDefaultTimeout bool

// provides the caching of API tokens used for operation calls. If unset,
// the API token will not be retrieved for the operation.
tokenProvider *tokenProvider
Expand Down
5 changes: 3 additions & 2 deletions feature/ec2/imds/doc.go
Expand Up @@ -3,8 +3,9 @@
//
// All Client operation calls have a default timeout. If the operation is not
// completed before this timeout expires, the operation will be canceled. This
// timeout can be overridden by providing Context with a timeout or deadline
// with calling the client's operations.
// timeout can be overridden through the following:
// - Set the options flag DisableDefaultTimeout
// - Provide a Context with a timeout or deadlin with calling the client's operations.
lucix-aws marked this conversation as resolved.
Show resolved Hide resolved
//
// See the EC2 IMDS user guide for more information on using the API.
// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html
Expand Down
6 changes: 6 additions & 0 deletions feature/ec2/imds/request_middleware.go
Expand Up @@ -56,6 +56,7 @@ func addRequestMiddleware(stack *middleware.Stack,

// Operation timeout
err = stack.Initialize.Add(&operationTimeout{
Disabled: options.DisableDefaultTimeout,
DefaultTimeout: defaultOperationTimeout,
}, middleware.Before)
if err != nil {
Expand Down Expand Up @@ -260,6 +261,7 @@ const (
// Otherwise the timeout cleanup will race the resource being consumed
// upstream.
type operationTimeout struct {
Disabled bool
DefaultTimeout time.Duration
}

Expand All @@ -270,6 +272,10 @@ func (m *operationTimeout) HandleInitialize(
) (
output middleware.InitializeOutput, metadata middleware.Metadata, err error,
) {
if m.Disabled {
return next.HandleInitialize(ctx, input)
}

if _, ok := ctx.Deadline(); !ok && m.DefaultTimeout != 0 {
var cancelFn func()
ctx, cancelFn = context.WithTimeout(ctx, m.DefaultTimeout)
Expand Down
23 changes: 23 additions & 0 deletions feature/ec2/imds/request_middleware_test.go
Expand Up @@ -220,6 +220,29 @@ func TestOperationTimeoutMiddleware_withCustomDeadline(t *testing.T) {
}
}

func TestOperationTimeoutMiddleware_Disabled(t *testing.T) {
m := &operationTimeout{
Disabled: true,
DefaultTimeout: time.Nanosecond,
}

_, _, err := m.HandleInitialize(context.Background(), middleware.InitializeInput{},
middleware.InitializeHandlerFunc(func(
ctx context.Context, input middleware.InitializeInput,
) (
out middleware.InitializeOutput, metadata middleware.Metadata, err error,
) {
if err := sdk.SleepWithContext(ctx, time.Second); err != nil {
return out, metadata, err
}

return out, metadata, nil
}))
if err != nil {
t.Fatalf("expect no error, got %v", err)
}
}

// Ensure that the response body is read in the deserialize middleware,
// ensuring that the timeoutOperation middleware won't race canceling the
// context with the upstream reading the response body.
Expand Down