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

extend new tracing instrumentation for the SDK #2766

Merged
merged 5 commits into from
Sep 3, 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
13 changes: 11 additions & 2 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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 \
Copy link
Contributor

Choose a reason for hiding this comment

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

nice

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 -- {} \;

Expand Down Expand Up @@ -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 \
Expand Down
4 changes: 4 additions & 0 deletions aws/middleware/request_id_retriever.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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
}
}
Expand Down
7 changes: 7 additions & 0 deletions aws/retry/middleware.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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
Expand Down
5 changes: 5 additions & 0 deletions aws/signer/v4/middleware.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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{
Expand All @@ -186,6 +190,7 @@ func (m *ComputePayloadSHA256) HandleFinalize(

ctx = SetPayloadHash(ctx, hex.EncodeToString(hash.Sum(nil)))

span.End()
return next.HandleFinalize(ctx, in)
}

Expand Down
36 changes: 34 additions & 2 deletions aws/transport/http/client.go
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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{} {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand All @@ -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);
}

Expand Down
5 changes: 5 additions & 0 deletions service/internal/s3shared/metadata_retriever.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand All @@ -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.
Expand All @@ -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
Expand Down
Loading
Loading