-
Notifications
You must be signed in to change notification settings - Fork 414
/
aws.go
135 lines (112 loc) · 4.57 KB
/
aws.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2016 Datadog, Inc.
package aws
import (
"context"
"fmt"
"math"
"time"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
"github.com/aws/aws-sdk-go-v2/aws"
awsmiddleware "github.com/aws/aws-sdk-go-v2/aws/middleware"
"github.com/aws/smithy-go/middleware"
smithyhttp "github.com/aws/smithy-go/transport/http"
)
const (
tagAWSAgent = "aws.agent"
tagAWSService = "aws.service"
tagAWSOperation = "aws.operation"
tagAWSRegion = "aws.region"
tagAWSRequestID = "aws.request_id"
)
type spanTimestampKey struct{}
// AppendMiddleware takes the aws.Config and adds the Datadog tracing middleware into the APIOptions middleware stack.
// See https://aws.github.io/aws-sdk-go-v2/docs/middleware for more information.
func AppendMiddleware(awsCfg *aws.Config, opts ...Option) {
cfg := &config{}
defaults(cfg)
for _, opt := range opts {
opt(cfg)
}
tm := traceMiddleware{cfg: cfg}
awsCfg.APIOptions = append(awsCfg.APIOptions, tm.initTraceMiddleware, tm.startTraceMiddleware, tm.deserializeTraceMiddleware)
}
type traceMiddleware struct {
cfg *config
}
func (mw *traceMiddleware) initTraceMiddleware(stack *middleware.Stack) error {
return stack.Initialize.Add(middleware.InitializeMiddlewareFunc("InitTraceMiddleware", func(
ctx context.Context, in middleware.InitializeInput, next middleware.InitializeHandler,
) (
out middleware.InitializeOutput, metadata middleware.Metadata, err error,
) {
// Bind the timestamp to the context so that we can use it when we have enough information to start the trace.
ctx = context.WithValue(ctx, spanTimestampKey{}, time.Now())
return next.HandleInitialize(ctx, in)
}), middleware.Before)
}
func (mw *traceMiddleware) startTraceMiddleware(stack *middleware.Stack) error {
return stack.Initialize.Add(middleware.InitializeMiddlewareFunc("StartTraceMiddleware", func(
ctx context.Context, in middleware.InitializeInput, next middleware.InitializeHandler,
) (
out middleware.InitializeOutput, metadata middleware.Metadata, err error,
) {
operation := awsmiddleware.GetOperationName(ctx)
serviceID := awsmiddleware.GetServiceID(ctx)
opts := []ddtrace.StartSpanOption{
tracer.SpanType(ext.SpanTypeHTTP),
tracer.ServiceName(serviceName(mw.cfg, serviceID)),
tracer.ResourceName(fmt.Sprintf("%s.%s", serviceID, operation)),
tracer.Tag(tagAWSRegion, awsmiddleware.GetRegion(ctx)),
tracer.Tag(tagAWSOperation, operation),
tracer.Tag(tagAWSService, serviceID),
tracer.StartTime(ctx.Value(spanTimestampKey{}).(time.Time)),
tracer.Tag(ext.Component, "aws/aws-sdk-go-v2/aws"),
tracer.Tag(ext.SpanKind, ext.SpanKindClient),
}
if !math.IsNaN(mw.cfg.analyticsRate) {
opts = append(opts, tracer.Tag(ext.EventSampleRate, mw.cfg.analyticsRate))
}
span, spanctx := tracer.StartSpanFromContext(ctx, fmt.Sprintf("%s.request", serviceID), opts...)
// Handle initialize and continue through the middleware chain.
out, metadata, err = next.HandleInitialize(spanctx, in)
span.Finish(tracer.WithError(err))
return out, metadata, err
}), middleware.After)
}
func (mw *traceMiddleware) deserializeTraceMiddleware(stack *middleware.Stack) error {
return stack.Deserialize.Add(middleware.DeserializeMiddlewareFunc("DeserializeTraceMiddleware", func(
ctx context.Context, in middleware.DeserializeInput, next middleware.DeserializeHandler,
) (
out middleware.DeserializeOutput, metadata middleware.Metadata, err error,
) {
span, _ := tracer.SpanFromContext(ctx)
// Get values out of the request.
if req, ok := in.Request.(*smithyhttp.Request); ok {
span.SetTag(ext.HTTPMethod, req.Method)
span.SetTag(ext.HTTPURL, req.URL.String())
span.SetTag(tagAWSAgent, req.Header.Get("User-Agent"))
}
// Continue through the middleware chain which eventually sends the request.
out, metadata, err = next.HandleDeserialize(ctx, in)
// Get values out of the response.
if res, ok := out.RawResponse.(*smithyhttp.Response); ok {
span.SetTag(ext.HTTPCode, res.StatusCode)
}
// Extract the request id.
if requestID, ok := awsmiddleware.GetRequestIDMetadata(metadata); ok {
span.SetTag(tagAWSRequestID, requestID)
}
return out, metadata, err
}), middleware.Before)
}
func serviceName(cfg *config, serviceID string) string {
if cfg.serviceName != "" {
return cfg.serviceName
}
return fmt.Sprintf("aws.%s", serviceID)
}