forked from open-telemetry/opentelemetry-collector-contrib
-
Notifications
You must be signed in to change notification settings - Fork 14
/
wrapper.go
105 lines (86 loc) · 3.23 KB
/
wrapper.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
// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0
package awsmiddleware // import "github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware"
import (
"context"
sdkmiddleware "github.com/aws/aws-sdk-go-v2/aws/middleware"
"github.com/aws/aws-sdk-go/aws/request"
"github.com/aws/smithy-go/middleware"
"github.com/aws/smithy-go/transport/http"
"github.com/google/uuid"
)
type (
requestIDKey struct{}
operationNameKey struct{}
)
func namedRequestHandler(handler RequestHandler) request.NamedHandler {
return request.NamedHandler{Name: handler.ID(), Fn: func(r *request.Request) {
ctx := mustRequestID(r.Context())
ctx = setOperationName(ctx, r.Operation.Name)
r.SetContext(ctx)
handler.HandleRequest(ctx, r.HTTPRequest)
}}
}
func namedResponseHandler(handler ResponseHandler) request.NamedHandler {
return request.NamedHandler{Name: handler.ID(), Fn: func(r *request.Request) {
handler.HandleResponse(r.Context(), r.HTTPResponse)
}}
}
type requestMiddleware struct {
RequestHandler
}
var _ middleware.BuildMiddleware = (*requestMiddleware)(nil)
func (r requestMiddleware) HandleBuild(ctx context.Context, in middleware.BuildInput, next middleware.BuildHandler) (out middleware.BuildOutput, metadata middleware.Metadata, err error) {
req, ok := in.Request.(*http.Request)
if ok {
ctx = mustRequestID(ctx)
ctx = setOperationName(ctx, sdkmiddleware.GetOperationName(ctx))
r.HandleRequest(ctx, req.Request)
}
return next.HandleBuild(ctx, in)
}
func withBuildOption(rmw *requestMiddleware, position middleware.RelativePosition) func(stack *middleware.Stack) error {
return func(stack *middleware.Stack) error {
return stack.Build.Add(rmw, position)
}
}
type responseMiddleware struct {
ResponseHandler
}
var _ middleware.DeserializeMiddleware = (*responseMiddleware)(nil)
func (r responseMiddleware) HandleDeserialize(ctx context.Context, in middleware.DeserializeInput, next middleware.DeserializeHandler) (out middleware.DeserializeOutput, metadata middleware.Metadata, err error) {
out, metadata, err = next.HandleDeserialize(ctx, in)
res, ok := out.RawResponse.(*http.Response)
if ok {
r.HandleResponse(ctx, res.Response)
}
return
}
func withDeserializeOption(rmw *responseMiddleware, position middleware.RelativePosition) func(stack *middleware.Stack) error {
return func(stack *middleware.Stack) error {
return stack.Deserialize.Add(rmw, position)
}
}
func mustRequestID(ctx context.Context) context.Context {
requestID := GetRequestID(ctx)
if requestID != "" {
return ctx
}
return setRequestID(ctx, uuid.NewString())
}
func setRequestID(ctx context.Context, id string) context.Context {
return context.WithValue(ctx, requestIDKey{}, id)
}
func setOperationName(ctx context.Context, name string) context.Context {
return context.WithValue(ctx, operationNameKey{}, name)
}
// GetRequestID retrieves the generated request ID from the context.
func GetRequestID(ctx context.Context) string {
requestID, _ := ctx.Value(requestIDKey{}).(string)
return requestID
}
// GetOperationName retrieves the service operation metadata from the context.
func GetOperationName(ctx context.Context) string {
operationName, _ := ctx.Value(operationNameKey{}).(string)
return operationName
}