Add otelaws middleware to support trace propagation in the AWS SDK v2…

… module (#2856)

* Add middleware for trace propagation

* Add tests and make propagator configurable

* Fix func doc

* Update changelog

* Fix linter errors

* Fix changelog again

Co-authored-by: Anthony Mirabella <a9@aneurysm9.com>

CHANGELOG.md

@@ -8,6 +8,10 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased]
 
+### Added
+
+- Add trace context propagation support to `instrumentation/github.com/aws/aws-sdk-go-v2/otelaws` (#2856).
+
 ## [1.11.0/0.36.3/0.5.1]
 
 ### Changed

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/aws.go

@@ -25,6 +25,7 @@ import (
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/propagation"
 	semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
 	"go.opentelemetry.io/otel/trace"
 )
@@ -40,6 +41,7 @@ type AttributeSetter func(context.Context, middleware.InitializeInput) []attribu
 
 type otelMiddlewares struct {
 	tracer          trace.Tracer
+	propagator      propagation.TextMapPropagator
 	attributeSetter []AttributeSetter
 }
 
@@ -110,12 +112,29 @@ func (m otelMiddlewares) deserializeMiddleware(stack *middleware.Stack) error {
 		middleware.Before)
 }
 
+func (m otelMiddlewares) finalizeMiddleware(stack *middleware.Stack) error {
+	return stack.Finalize.Add(middleware.FinalizeMiddlewareFunc("OTelFinalizeMiddleware", func(
+		ctx context.Context, in middleware.FinalizeInput, next middleware.FinalizeHandler) (
+		out middleware.FinalizeOutput, metadata middleware.Metadata, err error) {
+		// Propagate the Trace information by injecting it into the HTTP request.
+		switch req := in.Request.(type) {
+		case *smithyhttp.Request:
+			m.propagator.Inject(ctx, propagation.HeaderCarrier(req.Header))
+		default:
+		}
+
+		return next.HandleFinalize(ctx, in)
+	}),
+		middleware.Before)
+}
+
 // AppendMiddlewares attaches OTel middlewares to the AWS Go SDK V2 for instrumentation.
 // OTel middlewares can be appended to either all aws clients or a specific operation.
 // Please see more details in https://aws.github.io/aws-sdk-go-v2/docs/middleware/
 func AppendMiddlewares(apiOptions *[]func(*middleware.Stack) error, opts ...Option) {
 	cfg := config{
-		TracerProvider: otel.GetTracerProvider(),
+		TracerProvider:    otel.GetTracerProvider(),
+		TextMapPropagator: otel.GetTextMapPropagator(),
 	}
 	for _, opt := range opts {
 		opt.apply(&cfg)
@@ -127,6 +146,7 @@ func AppendMiddlewares(apiOptions *[]func(*middleware.Stack) error, opts ...Opti
 
 	m := otelMiddlewares{tracer: cfg.TracerProvider.Tracer(tracerName,
 		trace.WithInstrumentationVersion(SemVersion())),
+		propagator:      cfg.TextMapPropagator,
 		attributeSetter: cfg.AttributeSetter}
-	*apiOptions = append(*apiOptions, m.initializeMiddlewareBefore, m.initializeMiddlewareAfter, m.deserializeMiddleware)
+	*apiOptions = append(*apiOptions, m.initializeMiddlewareBefore, m.initializeMiddlewareAfter, m.finalizeMiddleware, m.deserializeMiddleware)
 }

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/aws_test.go

@@ -0,0 +1,80 @@
+// Copyright The OpenTelemetry Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package otelaws
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/aws/smithy-go/middleware"
+	smithyhttp "github.com/aws/smithy-go/transport/http"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"go.opentelemetry.io/otel/propagation"
+)
+
+type mockPropagator struct {
+	injectKey   string
+	injectValue string
+}
+
+func (p mockPropagator) Inject(ctx context.Context, carrier propagation.TextMapCarrier) {
+	carrier.Set(p.injectKey, p.injectValue)
+}
+func (p mockPropagator) Extract(ctx context.Context, carrier propagation.TextMapCarrier) context.Context {
+	return context.TODO()
+}
+func (p mockPropagator) Fields() []string {
+	return []string{}
+}
+
+func Test_otelMiddlewares_finalizeMiddleware(t *testing.T) {
+	stack := middleware.Stack{
+		Finalize: middleware.NewFinalizeStep(),
+	}
+
+	propagator := mockPropagator{
+		injectKey:   "mock-key",
+		injectValue: "mock-value",
+	}
+
+	m := otelMiddlewares{
+		propagator: propagator,
+	}
+
+	err := m.finalizeMiddleware(&stack)
+	require.NoError(t, err)
+
+	input := &smithyhttp.Request{
+		Request: &http.Request{
+			Header: http.Header{},
+		},
+	}
+
+	next := middleware.HandlerFunc(func(ctx context.Context, input interface{}) (output interface{}, metadata middleware.Metadata, err error) {
+		return nil, middleware.Metadata{}, nil
+	})
+
+	_, _, _ = stack.Finalize.HandleMiddleware(context.Background(), input, next)
+
+	// Assert header has been updated with injected values
+	key := http.CanonicalHeaderKey(propagator.injectKey)
+	value := propagator.injectValue
+
+	assert.Contains(t, input.Header, key)
+	assert.Contains(t, input.Header[key], value)
+}

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/config.go

@@ -15,12 +15,14 @@
 package otelaws // import "go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-sdk-go-v2/otelaws"
 
 import (
+	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 )
 
 type config struct {
-	TracerProvider  trace.TracerProvider
-	AttributeSetter []AttributeSetter
+	TracerProvider    trace.TracerProvider
+	TextMapPropagator propagation.TextMapPropagator
+	AttributeSetter   []AttributeSetter
 }
 
 // Option applies an option value.
@@ -46,6 +48,16 @@ func WithTracerProvider(provider trace.TracerProvider) Option {
 	})
 }
 
+// WithTextMapPropagator specifies a Text Map Propagator to use when propagating context.
+// If none is specified, the global TextMapPropagator is used.
+func WithTextMapPropagator(propagator propagation.TextMapPropagator) Option {
+	return optionFunc(func(cfg *config) {
+		if propagator != nil {
+			cfg.TextMapPropagator = propagator
+		}
+	})
+}
+
 // WithAttributeSetter specifies an attribute setter function for setting service specific attributes.
 // If none is specified, the service will be determined by the DefaultAttributeSetter function and the corresponding attributes will be included.
 func WithAttributeSetter(attributesetters ...AttributeSetter) Option {

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/config_test.go

@@ -0,0 +1,33 @@
+// Copyright The OpenTelemetry Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package otelaws
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"go.opentelemetry.io/otel"
+)
+
+func TestWithTextMapPropagator(t *testing.T) {
+	cfg := config{}
+	propagator := otel.GetTextMapPropagator()
+
+	option := WithTextMapPropagator(propagator)
+	option.apply(&cfg)
+
+	assert.Equal(t, cfg.TextMapPropagator, propagator)
+}