fix stale skew and delayed skew healing (#3359)

* fix stale skew

* fix comment

* fix test

* delete old test which is now covered in kitchensink

* regen after rebase

* add no-retry test

* regenerate

Author: lucix-aws

.changelog/9f49dccb1c42499f9a39521bea5bc73b.json

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	internalcontext "github.com/aws/aws-sdk-go-v2/internal/context"
@@ -43,6 +44,10 @@ type Attempt struct {
 	// A Meter instance for recording retry-related metrics.
 	OperationMeter metrics.Meter
 
+	// Initial clock skew that would have been saved from a previous operation
+	// call.
+	ClientSkew *atomic.Int64
+
 	retryer       aws.RetryerV2
 	requestCloner RequestCloner
 }
@@ -82,8 +87,12 @@ func (r Attempt) logf(logger logging.Logger, classification logging.Classificati
 func (r *Attempt) HandleFinalize(ctx context.Context, in smithymiddle.FinalizeInput, next smithymiddle.FinalizeHandler) (
 	out smithymiddle.FinalizeOutput, metadata smithymiddle.Metadata, err error,
 ) {
-	var attemptNum int
 	var attemptClockSkew time.Duration
+	if r.ClientSkew != nil {
+		attemptClockSkew = time.Duration(r.ClientSkew.Load())
+	}
+
+	var attemptNum int
 	var attemptResults AttemptResults
 
 	maxAttempts := r.retryer.MaxAttempts()
@@ -99,6 +108,8 @@ func (r *Attempt) HandleFinalize(ctx context.Context, in smithymiddle.FinalizeIn
 		attemptInput := in
 		attemptInput.Request = r.requestCloner(attemptInput.Request)
 
+		ctx = internalcontext.SetAttemptSkewContext(ctx, attemptClockSkew)
+
 		// Record the metadata for the for attempt being started.
 		attemptCtx := setRetryMetadata(ctx, retryMetadata{
 			AttemptNum:       attemptNum,
@@ -107,9 +118,6 @@ func (r *Attempt) HandleFinalize(ctx context.Context, in smithymiddle.FinalizeIn
 			AttemptClockSkew: attemptClockSkew,
 		})
 
-		// Setting clock skew to be used on other context (like signing)
-		ctx = internalcontext.SetAttemptSkewContext(ctx, attemptClockSkew)
-
 		var attemptResult AttemptResult
 
 		attemptCtx, span := tracing.StartSpan(attemptCtx, "Attempt", func(o *tracing.SpanOptions) {
@@ -149,6 +157,14 @@ func (r *Attempt) HandleFinalize(ctx context.Context, in smithymiddle.FinalizeIn
 		}
 	}
 
+	// this guarantees we are staying on top of the persistent skew value
+	// (either to apply it or to heal it back if the clocks realign)
+	if r.ClientSkew != nil {
+		if resultSkew, ok := awsmiddle.GetAttemptSkew(metadata); ok {
+			r.ClientSkew.Store(resultSkew.Nanoseconds())
+		}
+	}
+
 	addAttemptResults(&metadata, attemptResults)
 	return out, metadata, err
 }

aws/retry/middleware.go

@@ -8,11 +8,10 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
-	internalcontext "github.com/aws/aws-sdk-go-v2/internal/context"
-
 	"github.com/aws/aws-sdk-go-v2/aws/ratelimit"
 	"github.com/aws/aws-sdk-go-v2/internal/sdk"
 	"github.com/aws/smithy-go/middleware"
@@ -531,9 +530,13 @@ func TestClockSkew(t *testing.T) {
 	}
 	for name, tt := range cases {
 		t.Run(name, func(t *testing.T) {
+			skew := &atomic.Int64{}
+			skew.Store(tt.skew.Nanoseconds())
 			am := NewAttemptMiddleware(NewStandard(func(s *StandardOptions) {
-			}), func(i any) any { return i })
-			ctx := internalcontext.SetAttemptSkewContext(context.Background(), tt.skew)
+			}), func(i any) any { return i }, func(m *Attempt) {
+				m.ClientSkew = skew
+			})
+			ctx := context.Background()
 			_, metadata, err := am.HandleFinalize(ctx, middleware.FinalizeInput{}, middleware.FinalizeHandlerFunc(
 				func(ctx context.Context, in middleware.FinalizeInput) (
 					out middleware.FinalizeOutput, metadata middleware.Metadata, err error,

aws/retry/middleware_test.go

@@ -154,7 +154,10 @@ public List<RuntimeClientPlugin> getClientPlugins() {
                                 .resolvedFunction(SymbolUtils.createValueSymbolBuilder(
                                                 AwsRetryMiddlewareHelper.ADD_RETRY_MIDDLEWARES_HELPER)
                                         .build())
-                                .useClientOptions()
+                                .functionArguments(List.of(
+                                        buildPackageSymbol("options"),
+                                        buildPackageSymbol("c")
+                                ))
                                 .build())
                         .build(),
 

codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AssembleMiddlewareStack.java

@@ -45,10 +45,11 @@ private void generateRetryMiddlewareHelpers(GoWriter writer, String moduleName)
                 .addUseImports(SmithyGoDependency.SMITHY_HTTP_TRANSPORT)
                 .addUseImports(AwsGoDependency.AWS_RETRY)
                 .write(goTemplate("""
-                        func addRetry(stack *middleware.Stack, o Options) error {
+                        func addRetry(stack *middleware.Stack, o Options, c *Client) error {
                         attempt := retry.NewAttemptMiddleware(o.Retryer, smithyhttp.RequestCloner, func(m *retry.Attempt) {
                             m.LogAttempts = o.ClientLogMode.IsRetries()
                             m.OperationMeter = o.MeterProvider.Meter($S)
+                            m.ClientSkew = c.timeOffset
                         })
                         if err := stack.Finalize.Insert(attempt, "ResolveAuthScheme", middleware.Before); err != nil {
                             return err

codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AwsRetryMiddlewareHelper.java

@@ -21,31 +21,10 @@
  */
 public class ClockSkewGenerator implements GoIntegration {
     private static final String TIME_OFFSET = "timeOffset";
-    private static final String ADD_CLOCK_SKEW_BUILD = "addTimeOffsetBuild";
-    private static final String ADD_CLOCK_SKEW_BUILD_MIDDLEWARE = "AddTimeOffsetMiddleware";
 
     private static final Symbol TIME_OFFSET_RESOLVER = SymbolUtils.createValueSymbolBuilder(
             "initializeTimeOffsetResolver").build();
 
-    private static final Writable CLOCK_SKEW_INSERT_TEMPLATE = goTemplate("""
-                    $dep:D
-                    func $fn:L(stack $stack:P, c *Client) error {
-                        mw := $depalias:L.$middleware:L{Offset: c.$off:L}
-                        if err := stack.Build.Add(&mw, middleware.After); err != nil {
-                            return err
-                        }
-                        return stack.Deserialize.Insert(&mw, "$after:L", middleware.Before)
-                    }
-                    """,
-            Map.of(
-                    "fn", ADD_CLOCK_SKEW_BUILD,
-                    "stack", SmithyGoDependency.SMITHY_MIDDLEWARE.struct("Stack"),
-                    "depalias", INTERNAL_MIDDLEWARE.getAlias(),
-                    "middleware", ADD_CLOCK_SKEW_BUILD_MIDDLEWARE,
-                    "after", "RecordResponseTiming",
-                    "off", TIME_OFFSET,
-                    "dep", INTERNAL_MIDDLEWARE
-            ));
     private static final Writable TIME_OFFSET_RESOLVER_TEMPLATE = goTemplate(
             """
                     $import:D
@@ -68,16 +47,10 @@ public class ClockSkewGenerator implements GoIntegration {
     private static final ClientMemberResolver TIME_OFFSET_MEMBER_RESOLVER = ClientMemberResolver.builder()
             .resolver(TIME_OFFSET_RESOLVER)
             .build();
-    private static final MiddlewareRegistrar MIDDLEWARE = MiddlewareRegistrar.builder()
-            .resolvedFunction(SymbolUtils.createValueSymbolBuilder(ADD_CLOCK_SKEW_BUILD).build())
-            .functionArguments(ListUtils.of(
-                    SymbolUtils.createValueSymbolBuilder("c").build()
-            )).build();
     private static final List<RuntimeClientPlugin> CLIENT_PLUGINS = List.of(
             RuntimeClientPlugin.builder()
                     .addClientMember(TIME_OFFSET_MEMBER)
                     .addClientMemberResolver(TIME_OFFSET_MEMBER_RESOLVER)
-                    .registerMiddleware(MIDDLEWARE)
                     .build()
     );
 
@@ -93,7 +66,6 @@ public void writeAdditionalFiles(
 
         // generate code specific to service client
         goDelegator.useShapeWriter(service, writer -> {
-            writer.write(CLOCK_SKEW_INSERT_TEMPLATE);
             writer.write(TIME_OFFSET_RESOLVER_TEMPLATE);
         });
     }
@@ -102,4 +74,4 @@ public void writeAdditionalFiles(
     public List<RuntimeClientPlugin> getClientPlugins() {
         return CLIENT_PLUGINS;
     }
-}
+}