extend new tracing instrumentation for the SDK (#2766)

Makefile

@@ -86,10 +86,16 @@ generate: smithy-generate update-requires gen-repo-mod-replace update-module-met
 gen-config-asserts gen-internal-codegen copy-attributevalue-feature gen-mod-dropreplace-smithy-. min-go-version-. \
 tidy-modules-. add-module-license-files gen-aws-ptrs format
 
-generate-tmpreplace-smithy: smithy-generate update-requires gen-repo-mod-replace update-module-metadata smithy-annotate-stable \
-gen-config-asserts gen-internal-codegen copy-attributevalue-feature gen-mod-replace-smithy-. min-go-version-. \
+generate-tmpreplace-smithy: smithy-generate update-requires gen-repo-mod-replace gen-mod-replace-smithy-. update-module-metadata smithy-annotate-stable \
+gen-config-asserts gen-internal-codegen copy-attributevalue-feature min-go-version-. \
 tidy-modules-. add-module-license-files gen-aws-ptrs format gen-mod-dropreplace-smithy-. reset-sum
 
+# stripped-down regenerate script that eliminates a lot of the cruft you don't
+# need in development (that takes time)
+# modify this with whatever service you're working on
+generate-dev: smithy-generate update-requires gen-repo-mod-replace gen-mod-replace-smithy-config gen-mod-replace-smithy-aws gen-mod-replace-smithy-service_s3 update-module-metadata smithy-annotate-stable \
+gen-config-asserts gen-internal-codegen tidy-modules-config tidy-modules-aws tidy-modules-service_s3 format-dev
+
 reset-sum:
 	find . -name go.sum -exec git checkout -- {} \;
 
@@ -123,6 +129,9 @@ smithy-go-publish-local:
 format:
 	gofmt -w -s .
 
+format-dev:
+	gofmt -w -s service/s3
+
 gen-config-asserts:
 	@echo "Generating SDK config package implementor assertions"
 	cd config \

aws/middleware/request_id_retriever.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/aws/smithy-go/middleware"
+	"github.com/aws/smithy-go/tracing"
 	smithyhttp "github.com/aws/smithy-go/transport/http"
 )
 
@@ -45,6 +46,9 @@ func (m *RequestIDRetriever) HandleDeserialize(ctx context.Context, in middlewar
 		if v := resp.Header.Get(h); len(v) != 0 {
 			// set reqID on metadata for successful responses.
 			SetRequestIDMetadata(&metadata, v)
+
+			span, _ := tracing.GetSpan(ctx)
+			span.SetProperty("aws.request_id", v)
 			break
 		}
 	}

aws/retry/middleware.go

@@ -16,6 +16,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/internal/sdk"
 	"github.com/aws/smithy-go/logging"
 	smithymiddle "github.com/aws/smithy-go/middleware"
+	"github.com/aws/smithy-go/tracing"
 	"github.com/aws/smithy-go/transport/http"
 )
 
@@ -97,7 +98,13 @@ func (r *Attempt) HandleFinalize(ctx context.Context, in smithymiddle.FinalizeIn
 		ctx = internalcontext.SetAttemptSkewContext(ctx, attemptClockSkew)
 
 		var attemptResult AttemptResult
+
+		attemptCtx, span := tracing.StartSpan(attemptCtx, "Attempt", func(o *tracing.SpanOptions) {
+			o.Properties.Set("operation.attempt", attemptNum)
+		})
 		out, attemptResult, releaseRetryToken, err = r.handleAttempt(attemptCtx, attemptInput, releaseRetryToken, next)
+		span.End()
+
 		attemptClockSkew, _ = awsmiddle.GetAttemptSkew(attemptResult.ResponseMetadata)
 
 		// AttemptResult Retried states that the attempt was not successful, and

aws/signer/v4/middleware.go

@@ -15,6 +15,7 @@ import (
 	internalauth "github.com/aws/aws-sdk-go-v2/internal/auth"
 	"github.com/aws/aws-sdk-go-v2/internal/sdk"
 	"github.com/aws/smithy-go/middleware"
+	"github.com/aws/smithy-go/tracing"
 	smithyhttp "github.com/aws/smithy-go/transport/http"
 )
 
@@ -161,6 +162,9 @@ func (m *ComputePayloadSHA256) HandleFinalize(
 		return next.HandleFinalize(ctx, in)
 	}
 
+	_, span := tracing.StartSpan(ctx, "ComputePayloadSHA256")
+	defer span.End()
+
 	req, ok := in.Request.(*smithyhttp.Request)
 	if !ok {
 		return out, metadata, &HashComputationError{
@@ -186,6 +190,7 @@ func (m *ComputePayloadSHA256) HandleFinalize(
 
 	ctx = SetPayloadHash(ctx, hex.EncodeToString(hash.Sum(nil)))
 
+	span.End()
 	return next.HandleFinalize(ctx, in)
 }
 

aws/transport/http/client.go

@@ -1,13 +1,16 @@
 package http
 
 import (
+	"context"
 	"crypto/tls"
-	"github.com/aws/aws-sdk-go-v2/aws"
 	"net"
 	"net/http"
 	"reflect"
 	"sync"
 	"time"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/smithy-go/tracing"
 )
 
 // Defaults for the HTTPTransportBuilder.
@@ -179,7 +182,7 @@ func defaultHTTPTransport() *http.Transport {
 
 	tr := &http.Transport{
 		Proxy:                 http.ProxyFromEnvironment,
-		DialContext:           dialer.DialContext,
+		DialContext:           traceDialContext(dialer.DialContext),
 		TLSHandshakeTimeout:   DefaultHTTPTransportTLSHandleshakeTimeout,
 		MaxIdleConns:          DefaultHTTPTransportMaxIdleConns,
 		MaxIdleConnsPerHost:   DefaultHTTPTransportMaxIdleConnsPerHost,
@@ -194,6 +197,35 @@ func defaultHTTPTransport() *http.Transport {
 	return tr
 }
 
+type dialContext func(ctx context.Context, network, addr string) (net.Conn, error)
+
+func traceDialContext(dc dialContext) dialContext {
+	return func(ctx context.Context, network, addr string) (net.Conn, error) {
+		span, _ := tracing.GetSpan(ctx)
+		span.SetProperty("net.peer.name", addr)
+
+		conn, err := dc(ctx, network, addr)
+		if err != nil {
+			return conn, err
+		}
+
+		raddr := conn.RemoteAddr()
+		if raddr == nil {
+			return conn, err
+		}
+
+		host, port, err := net.SplitHostPort(raddr.String())
+		if err != nil { // don't blow up just because we couldn't parse
+			span.SetProperty("net.peer.addr", raddr.String())
+		} else {
+			span.SetProperty("net.peer.host", host)
+			span.SetProperty("net.peer.port", port)
+		}
+
+		return conn, err
+	}
+}
+
 // shallowCopyStruct creates a shallow copy of the passed in source struct, and
 // returns that copy of the same struct type.
 func shallowCopyStruct(src interface{}) interface{} {

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

@@ -20,8 +20,11 @@
 import static software.amazon.smithy.go.codegen.SymbolUtils.buildPackageSymbol;
 
 import java.util.List;
+import java.util.Map;
+
 import software.amazon.smithy.aws.go.codegen.customization.AdjustAwsRestJsonContentType;
 import software.amazon.smithy.aws.traits.auth.UnsignedPayloadTrait;
+import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.go.codegen.GoDelegator;
 import software.amazon.smithy.go.codegen.GoSettings;
@@ -31,6 +34,7 @@
 import software.amazon.smithy.go.codegen.integration.GoIntegration;
 import software.amazon.smithy.go.codegen.integration.MiddlewareRegistrar;
 import software.amazon.smithy.go.codegen.integration.RuntimeClientPlugin;
+import software.amazon.smithy.go.codegen.middleware.FinalizeStepMiddleware;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.EventStreamIndex;
 import software.amazon.smithy.utils.ListUtils;
@@ -168,6 +172,15 @@ public List<RuntimeClientPlugin> getClientPlugins() {
                                         .build()
                         ).build(),
 
+                // wrap the retry loop in a span
+                RuntimeClientPlugin.builder()
+                        .registerMiddleware(
+                                MiddlewareRegistrar.builder()
+                                        .resolvedFunction(buildPackageSymbol("addSpanRetryLoop"))
+                                        .useClientOptions()
+                                        .build()
+                        ).build(),
+
                 // Add Client UserAgent
                 RuntimeClientPlugin.builder()
                         .registerMiddleware(MiddlewareRegistrar.builder()
@@ -205,12 +218,35 @@ public List<RuntimeClientPlugin> getClientPlugins() {
     public void writeAdditionalFiles(GoSettings settings, Model model, SymbolProvider symbolProvider, GoDelegator goDelegator) {
         goDelegator.useFileWriter("api_client.go", settings.getModuleName(), writer -> {
             writer.write(addMiddleware());
+            writer.write(spanRetryLoopMiddleware());
             if (hasSigV4X(model, settings.getService(model))) {
                 writer.write(addSigV4XMiddleware());
             }
         });
     }
 
+    private GoWriter.Writable spanRetryLoopMiddleware() {
+        return new FinalizeStepMiddleware() {
+            public String getStructName() {
+                return "spanRetryLoop";
+            }
+
+            public Map<String, Symbol> getFields() {
+                return Map.of("options", buildPackageSymbol("Options"));
+            }
+
+            public GoWriter.Writable getFuncBody() {
+                return goTemplate("""
+                        tracer := operationTracer(m.options.TracerProvider)
+                        ctx, span := tracer.StartSpan(ctx, "RetryLoop")
+                        defer span.End()
+
+                        return next.HandleFinalize(ctx, in)
+                        """);
+            }
+        };
+    }
+
     private GoWriter.Writable addMiddleware() {
         return goTemplate("""
                 $D $D $D
@@ -229,6 +265,10 @@ func addRawResponseToMetadata(stack *middleware.Stack) error {
                 func addRecordResponseTiming(stack *middleware.Stack) error {
                     return stack.Deserialize.Add(&awsmiddleware.RecordResponseTiming{}, middleware.After)
                 }
+
+                func addSpanRetryLoop(stack *middleware.Stack, options Options) error {
+                    return stack.Finalize.Insert(&spanRetryLoop{options: options}, "Retry", middleware.Before)
+                }
                 """, SmithyGoDependency.SMITHY_MIDDLEWARE, AwsGoDependency.AWS_MIDDLEWARE, SmithyGoDependency.SMITHY_HTTP_TRANSPORT);
     }
 

service/internal/s3shared/metadata_retriever.go

@@ -5,6 +5,7 @@ import (
 
 	awsmiddleware "github.com/aws/aws-sdk-go-v2/aws/middleware"
 	"github.com/aws/smithy-go/middleware"
+	"github.com/aws/smithy-go/tracing"
 	smithyhttp "github.com/aws/smithy-go/transport/http"
 )
 
@@ -30,6 +31,8 @@ func (m *metadataRetriever) HandleDeserialize(ctx context.Context, in middleware
 ) {
 	out, metadata, err = next.HandleDeserialize(ctx, in)
 
+	span, _ := tracing.GetSpan(ctx)
+
 	resp, ok := out.RawResponse.(*smithyhttp.Response)
 	if !ok {
 		// No raw response to wrap with.
@@ -40,12 +43,14 @@ func (m *metadataRetriever) HandleDeserialize(ctx context.Context, in middleware
 	if v := resp.Header.Get("X-Amz-Request-Id"); len(v) != 0 {
 		// set reqID on metadata for successful responses.
 		awsmiddleware.SetRequestIDMetadata(&metadata, v)
+		span.SetProperty("aws.request_id", v)
 	}
 
 	// look up host-id
 	if v := resp.Header.Get("X-Amz-Id-2"); len(v) != 0 {
 		// set reqID on metadata for successful responses.
 		SetHostIDMetadata(&metadata, v)
+		span.SetProperty("aws.extended_request_id", v)
 	}
 
 	return out, metadata, err