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.

Sign up for GitHub

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

Jump to bottom

Fix unreleased capacity after retryToken is requested #1424

Closed
wants to merge 1 commit into from
Closed

Fix unreleased capacity after retryToken is requested #1424

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
7 changes: 6 additions & 1 deletion aws/retry/metadata.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
package retry

import (
awsmiddle "github.com/aws/aws-sdk-go-v2/aws/middleware"
"github.com/aws/smithy-go/middleware"

awsmiddle "github.com/aws/aws-sdk-go-v2/aws/middleware"
)

// attemptResultsKey is a metadata accessor key to retrieve metadata
Expand Down Expand Up @@ -39,6 +40,10 @@ type AttemptResult struct {

// ResponseMetadata is any existing metadata passed via the response middlewares.
ResponseMetadata middleware.Metadata

// ReleaseRetryTokenFn if defined, needs to be called on success to release possible
// spent capacity.
ReleaseRetryTokenFn func(error) error
}

// addAttemptResults adds attempt results to middleware metadata
Expand Down
28 changes: 20 additions & 8 deletions aws/retry/middleware.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,14 @@ import (
"strings"
"time"

"github.com/aws/aws-sdk-go-v2/aws"
awsmiddle "github.com/aws/aws-sdk-go-v2/aws/middleware"
"github.com/aws/aws-sdk-go-v2/internal/sdk"
"github.com/aws/smithy-go/logging"
"github.com/aws/smithy-go/middleware"
smithymiddle "github.com/aws/smithy-go/middleware"
"github.com/aws/smithy-go/transport/http"

"github.com/aws/aws-sdk-go-v2/aws"
awsmiddle "github.com/aws/aws-sdk-go-v2/aws/middleware"
"github.com/aws/aws-sdk-go-v2/internal/sdk"
)

// RequestCloner is a function that can take an input request type and clone the request
Expand All @@ -34,8 +35,9 @@ type Attempt struct {
// This will include logging retry attempts, unretryable errors, and when max attempts are reached.
LogAttempts bool

retryer aws.Retryer
requestCloner RequestCloner
retryer aws.Retryer
requestCloner RequestCloner
releaseRetryToken func(error) error
}

// NewAttemptMiddleware returns a new Attempt retry middleware.
Expand Down Expand Up @@ -84,7 +86,9 @@ func (r Attempt) HandleFinalize(ctx context.Context, in smithymiddle.FinalizeInp
var attemptResult AttemptResult

out, attemptResult, err = r.handleAttempt(attemptCtx, attemptInput, next)

if attemptResult.ReleaseRetryTokenFn != nil {
r.releaseRetryToken = attemptResult.ReleaseRetryTokenFn
}
var ok bool
attemptClockSkew, ok = awsmiddle.GetAttemptSkew(attemptResult.ResponseMetadata)
if !ok {
Expand Down Expand Up @@ -117,7 +121,7 @@ func (r Attempt) handleAttempt(ctx context.Context, in smithymiddle.FinalizeInpu
attemptResult.Err = err
}()

relRetryToken := r.retryer.GetInitialToken()
relInitialToken := r.retryer.GetInitialToken()
logger := smithymiddle.GetLogger(ctx)
service, operation := awsmiddle.GetServiceID(ctx), awsmiddle.GetOperationName(ctx)

Expand All @@ -140,11 +144,17 @@ func (r Attempt) handleAttempt(ctx context.Context, in smithymiddle.FinalizeInpu
out, metadata, err = next.HandleFinalize(ctx, in)
attemptResult.ResponseMetadata = metadata

if releaseError := relRetryToken(err); releaseError != nil && err != nil {
if releaseError := relInitialToken(err); releaseError != nil && err != nil {
err = fmt.Errorf("failed to release token after request error, %w", err)
return out, attemptResult, err
}

if r.releaseRetryToken != nil {
if err := r.releaseRetryToken(err); err != nil {
return out, attemptResult, err
}
}

if err == nil {
return out, attemptResult, err
}
Expand Down Expand Up @@ -172,6 +182,8 @@ func (r Attempt) handleAttempt(ctx context.Context, in smithymiddle.FinalizeInpu
return out, attemptResult, reqErr
}

attemptResult.ReleaseRetryTokenFn = relRetryToken

retryDelay, reqErr := r.retryer.RetryDelay(attemptNum, err)
if reqErr != nil {
return out, attemptResult, reqErr
Expand Down
61 changes: 56 additions & 5 deletions aws/retry/middleware_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package retry

import (
"context"
"errors"
"fmt"
"net/http"
"reflect"
Expand All @@ -10,11 +11,13 @@ import (
"testing"
"time"

"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/internal/sdk"
"github.com/aws/smithy-go/middleware"
smithyhttp "github.com/aws/smithy-go/transport/http"
"github.com/google/go-cmp/cmp"

"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/aws/ratelimit"
"github.com/aws/aws-sdk-go-v2/internal/sdk"
)

func TestMetricsHeaderMiddleware(t *testing.T) {
Expand Down Expand Up @@ -427,13 +430,27 @@ func TestAttemptMiddleware(t *testing.T) {
if diff := cmp.Diff(recorded, tt.Expect); len(diff) > 0 {
t.Error(diff)
}

attemptResults, ok := GetAttemptResults(metadata)
if !ok {
t.Fatalf("expected metadata to contain attempt results, got none")
}
if e, a := tt.ExpectResults, attemptResults; !reflect.DeepEqual(e, a) {
t.Fatalf("expected %v, got %v", e, a)
e, a := tt.ExpectResults, attemptResults
if len(e.Results) != len(a.Results) {
t.Fatalf("expected len %d got %d", len(e.Results), len(a.Results))
}
for i, eRes := range e.Results {
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a side effect of altering the return value with adding a function into it. Can no longer be deepEqualed. If anyone has a nice idea on how to fix that, I'm open for suggestions.

Sadly, it's late and I can't think of anything better ATM. :)

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can experiment with a mock retry provider in these cases.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should be able to change how these values are compared to use cmp.Diff. This is commonly used for other data types in the v2 SDK.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can try that, sure.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Though I'm not sure how that works for functions.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for clarifying that. Yes, you're correct cmp.Diff doens't work with functions, checking for not-nil is pretty much the only way to compare in this case.

if eRes.Retried != a.Results[i].Retried {
t.Fatalf("expected retried %v got %v", eRes.Retried, a.Results[i].Retried)
}
if eRes.Retryable != a.Results[i].Retryable {
t.Fatalf("expected retryable %v got %v", eRes.Retryable, a.Results[i].Retryable)
}
if eRes.Err != nil && eRes.Err.Error() != a.Results[i].Err.Error() {
t.Fatalf("expected err %v got %v", eRes.Err, a.Results[i].Err)
}
if !reflect.DeepEqual(eRes.ResponseMetadata, a.Results[i].ResponseMetadata) {
t.Fatalf("expected response metadata %v got %v", eRes.ResponseMetadata, a.Results[i].ResponseMetadata)
}
}

for i, attempt := range attemptResults.Results {
Expand All @@ -446,6 +463,40 @@ func TestAttemptMiddleware(t *testing.T) {
}
}

func TestReleaseRetryLock(t *testing.T) {
standard := NewStandard(func(s *StandardOptions) {
s.MaxAttempts = 3
s.RateLimiter = ratelimit.NewTokenRateLimit(10)
s.RetryCost = 10
})
am := NewAttemptMiddleware(standard, func(i interface{}) interface{} {
return i
})
f := func(retries *[]retryMetadata) middleware.FinalizeHandler {
num := 0
return middleware.FinalizeHandlerFunc(func(ctx context.Context, in middleware.FinalizeInput) (out middleware.FinalizeOutput, metadata middleware.Metadata, err error) {
m, ok := getRetryMetadata(ctx)
if ok {
*retries = append(*retries, m)
}
if num > 0 {
return out, metadata, err
}
num++
return out, metadata, mockRetryableError{b: true}
})
}
var recorded []retryMetadata
_, _, err := am.HandleFinalize(context.Background(), middleware.FinalizeInput{}, f(&recorded))
if err != nil {
t.Fatal(err)
}
_, err = standard.GetRetryToken(context.Background(), errors.New("retryme"))
if err != nil {
t.Fatal(err)
}
}

// mockRawResponseKey is used to test the behavior when response metadata is
// nested within the attempt request.
type mockRawResponseKey struct{}
Expand Down
2 changes: 1 addition & 1 deletion aws/retryer.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ type Retryer interface {
// Returning the token release function, or error.
GetRetryToken(ctx context.Context, opErr error) (releaseToken func(error) error, err error)

// GetInitalToken returns the initial request token that can increment the
// GetInitialToken returns the initial request token that can increment the
// retry token pool if the request is successful.
GetInitialToken() (releaseToken func(error) error)
}
Expand Down