[extension/awsmiddleware] Set request ID and operation name in contex…

…t. (#132)

cmd/configschema/go.mod

@@ -268,7 +268,7 @@ require (
 	github.com/alecthomas/participle/v2 v2.0.0 // indirect
 	github.com/alecthomas/units v0.0.0-20211218093645-b94a6e3cc137 // indirect
 	github.com/aliyun/aliyun-log-go-sdk v0.1.54 // indirect
-	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231020160851-a8ff477e82af // indirect
+	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231023152757-c6e2437e6590 // indirect
 	github.com/amazon-contributing/opentelemetry-collector-contrib/override/aws v0.0.0-20230818193829-04a761abd409 // indirect
 	github.com/andybalholm/brotli v1.0.5 // indirect
 	github.com/antonmedv/expr v1.15.0 // indirect

cmd/otelcontribcol/go.mod

@@ -289,7 +289,7 @@ require (
 	github.com/alecthomas/participle/v2 v2.0.0 // indirect
 	github.com/alecthomas/units v0.0.0-20211218093645-b94a6e3cc137 // indirect
 	github.com/aliyun/aliyun-log-go-sdk v0.1.54 // indirect
-	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231020160851-a8ff477e82af // indirect
+	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231023152757-c6e2437e6590 // indirect
 	github.com/amazon-contributing/opentelemetry-collector-contrib/override/aws v0.0.0-20230818193829-04a761abd409 // indirect
 	github.com/andybalholm/brotli v1.0.5 // indirect
 	github.com/antonmedv/expr v1.15.0 // indirect

exporter/awscloudwatchlogsexporter/go.mod

@@ -3,7 +3,7 @@ module github.com/open-telemetry/opentelemetry-collector-contrib/exporter/awsclo
 go 1.20
 
 require (
-	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231020160851-a8ff477e82af
+	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231023152757-c6e2437e6590
 	github.com/aws/aws-sdk-go v1.45.2
 	github.com/cenkalti/backoff/v4 v4.2.1
 	github.com/google/uuid v1.3.1

exporter/awsemfexporter/go.mod

@@ -3,7 +3,7 @@ module github.com/open-telemetry/opentelemetry-collector-contrib/exporter/awsemf
 go 1.20
 
 require (
-	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231020160851-a8ff477e82af
+	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231023152757-c6e2437e6590
 	github.com/aws/aws-sdk-go v1.45.24
 	github.com/google/uuid v1.3.1
 	github.com/open-telemetry/opentelemetry-collector-contrib/internal/aws/awsutil v0.84.0

exporter/awsxrayexporter/go.mod

@@ -3,7 +3,7 @@ module github.com/open-telemetry/opentelemetry-collector-contrib/exporter/awsxra
 go 1.20
 
 require (
-	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231020160851-a8ff477e82af
+	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231023152757-c6e2437e6590
 	github.com/aws/aws-sdk-go v1.45.2
 	github.com/open-telemetry/opentelemetry-collector-contrib/internal/aws/awsutil v0.84.0
 	github.com/open-telemetry/opentelemetry-collector-contrib/internal/aws/xray v0.84.0

extension/awsmiddleware/README.md

@@ -13,18 +13,20 @@ The `awsmiddleware.RequestHandler` interface contains the following methods:
 ```
 ID() string
 Position() HandlerPosition
-HandleRequest(id string, r *http.Request)
+HandleRequest(ctx context.Context, r *http.Request)
 ```
 
 The `awsmiddleware.ResponseHandler` interface contains the following methods:
 ```
 ID() string
 Position() HandlerPosition
-HandleResponse(id string, r *http.Response)
+HandleResponse(ctx context.Context, r *http.Response)
 ```
 
 - `ID` uniquely identifies a handler. Middleware will fail if there is clashing 
 - `Position` determines whether the handler is appended to the front or back of the existing list. Insertion is done
 in the order of the handlers provided.
 - `HandleRequest/Response` provides a hook to handle the request/response before and after they've been sent along
-with an attached request ID.
+with the context.
+
+There are a functions available that can be used to extract metadata from the context.

extension/awsmiddleware/config.go

@@ -9,10 +9,8 @@ import (
 	"go.opentelemetry.io/collector/component"
 )
 
-type ID = component.ID
-
 // getMiddleware retrieves the extension implementing Middleware based on the middlewareID.
-func getMiddleware(extensions map[component.ID]component.Component, middlewareID ID) (Middleware, error) {
+func getMiddleware(extensions map[component.ID]component.Component, middlewareID component.ID) (Middleware, error) {
 	if extension, found := extensions[middlewareID]; found {
 		if middleware, ok := extension.(Middleware); ok {
 			return middleware, nil
@@ -24,7 +22,7 @@ func getMiddleware(extensions map[component.ID]component.Component, middlewareID
 
 // GetConfigurer retrieves the extension implementing Middleware based on the middlewareID and
 // wraps it in a Configurer.
-func GetConfigurer(extensions map[component.ID]component.Component, middlewareID ID) (*Configurer, error) {
+func GetConfigurer(extensions map[component.ID]component.Component, middlewareID component.ID) (*Configurer, error) {
 	middleware, err := getMiddleware(extensions, middlewareID)
 	if err != nil {
 		return nil, err

extension/awsmiddleware/middleware.go

@@ -4,6 +4,7 @@
 package awsmiddleware // import "github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware"
 
 import (
+	"context"
 	"encoding"
 	"errors"
 	"fmt"
@@ -86,13 +87,13 @@ type handlerConfig interface {
 // RequestHandler allows for custom processing of requests.
 type RequestHandler interface {
 	handlerConfig
-	HandleRequest(id string, r *http.Request)
+	HandleRequest(ctx context.Context, r *http.Request)
 }
 
 // ResponseHandler allows for custom processing of responses.
 type ResponseHandler interface {
 	handlerConfig
-	HandleResponse(id string, r *http.Response)
+	HandleResponse(ctx context.Context, r *http.Response)
 }
 
 // Middleware defines the request and response handlers to be configured

extension/awsmiddleware/middleware_test.go

@@ -29,12 +29,13 @@ const (
 type testHandler struct {
 	id             string
 	position       HandlerPosition
-	handleRequest  func(id string, r *http.Request)
-	handleResponse func(id string, r *http.Response)
+	handleRequest  func(ctx context.Context, r *http.Request)
+	handleResponse func(ctx context.Context, r *http.Response)
 	start          time.Time
 	end            time.Time
 	requestIDs     []string
 	responseIDs    []string
+	operations     []string
 }
 
 var _ RequestHandler = (*testHandler)(nil)
@@ -48,19 +49,21 @@ func (t *testHandler) Position() HandlerPosition {
 	return t.position
 }
 
-func (t *testHandler) HandleRequest(id string, r *http.Request) {
+func (t *testHandler) HandleRequest(ctx context.Context, r *http.Request) {
 	t.start = time.Now()
-	t.requestIDs = append(t.requestIDs, id)
+	t.requestIDs = append(t.requestIDs, GetRequestID(ctx))
+	t.operations = append(t.operations, GetOperationName(ctx))
 	if t.handleRequest != nil {
-		t.handleRequest(id, r)
+		t.handleRequest(ctx, r)
 	}
 }
 
-func (t *testHandler) HandleResponse(id string, r *http.Response) {
+func (t *testHandler) HandleResponse(ctx context.Context, r *http.Response) {
 	t.end = time.Now()
-	t.responseIDs = append(t.responseIDs, id)
+	t.responseIDs = append(t.responseIDs, GetRequestID(ctx))
+	t.operations = append(t.operations, GetOperationName(ctx))
 	if t.handleResponse != nil {
-		t.handleResponse(id, r)
+		t.handleResponse(ctx, r)
 	}
 }
 
@@ -72,8 +75,8 @@ type recordOrder struct {
 	order []string
 }
 
-func (ro *recordOrder) Handle(id string) func(string, *http.Request) {
-	return func(string, *http.Request) {
+func (ro *recordOrder) Handle(id string) func(context.Context, *http.Request) {
+	return func(context.Context, *http.Request) {
 		ro.order = append(ro.order, id)
 	}
 }
@@ -237,6 +240,9 @@ func TestRoundTripSDKv1(t *testing.T) {
 	assert.NotNil(t, output)
 	assert.GreaterOrEqual(t, recorder.Latency(), testLatency)
 	assert.Equal(t, recorder.requestIDs, recorder.responseIDs)
+	for _, operation := range recorder.operations {
+		assert.Equal(t, "ListBuckets", operation)
+	}
 }
 
 func TestRoundTripSDKv2(t *testing.T) {
@@ -252,13 +258,16 @@ func TestRoundTripSDKv2(t *testing.T) {
 	assert.NotNil(t, output)
 	assert.GreaterOrEqual(t, recorder.Latency(), testLatency)
 	assert.Equal(t, recorder.requestIDs, recorder.responseIDs)
+	for _, operation := range recorder.operations {
+		assert.Equal(t, "ListBuckets", operation)
+	}
 }
 
 func userAgentHandler() RequestHandler {
 	return &testHandler{
 		id:       "test.UserAgent",
 		position: Before,
-		handleRequest: func(_ string, r *http.Request) {
+		handleRequest: func(_ context.Context, r *http.Request) {
 			r.Header.Set("User-Agent", testUserAgent)
 		},
 	}

extension/awsmiddleware/mock.go

@@ -4,6 +4,7 @@
 package awsmiddleware // import "github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware"
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/stretchr/testify/mock"
@@ -53,12 +54,12 @@ func (m *MockHandler) Position() HandlerPosition {
 	return args.Get(0).(HandlerPosition)
 }
 
-func (m *MockHandler) HandleRequest(id string, r *http.Request) {
-	m.Called(id, r)
+func (m *MockHandler) HandleRequest(ctx context.Context, r *http.Request) {
+	m.Called(ctx, r)
 }
 
-func (m *MockHandler) HandleResponse(id string, r *http.Response) {
-	m.Called(id, r)
+func (m *MockHandler) HandleResponse(ctx context.Context, r *http.Response) {
+	m.Called(ctx, r)
 }
 
 // MockExtensionsHost only mocks the GetExtensions function.

extension/awsmiddleware/options.go

@@ -1,7 +1,7 @@
 // Copyright The OpenTelemetry Authors
 // SPDX-License-Identifier: Apache-2.0
 
-package awsmiddleware
+package awsmiddleware // import "github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware"
 
 import (
 	"github.com/aws/aws-sdk-go-v2/aws"

extension/awsmiddleware/wrapper.go

@@ -6,28 +6,30 @@ package awsmiddleware // import "github.com/amazon-contributing/opentelemetry-co
 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 key struct{}
-
-var requestID key
+type (
+	requestIDKey     struct{}
+	operationNameKey struct{}
+)
 
 func namedRequestHandler(handler RequestHandler) request.NamedHandler {
 	return request.NamedHandler{Name: handler.ID(), Fn: func(r *request.Request) {
-		ctx, id := setID(r.Context())
+		ctx := mustRequestID(r.Context())
+		ctx = setOperationName(ctx, r.Operation.Name)
 		r.SetContext(ctx)
-		handler.HandleRequest(id, r.HTTPRequest)
+		handler.HandleRequest(ctx, r.HTTPRequest)
 	}}
 }
 
 func namedResponseHandler(handler ResponseHandler) request.NamedHandler {
 	return request.NamedHandler{Name: handler.ID(), Fn: func(r *request.Request) {
-		id, _ := getID(r.Context())
-		handler.HandleResponse(id, r.HTTPResponse)
+		handler.HandleResponse(r.Context(), r.HTTPResponse)
 	}}
 }
 
@@ -40,9 +42,9 @@ 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 {
-		var id string
-		ctx, id = setID(ctx)
-		r.HandleRequest(id, req.Request)
+		ctx = mustRequestID(ctx)
+		ctx = setOperationName(ctx, sdkmiddleware.GetOperationName(ctx))
+		r.HandleRequest(ctx, req.Request)
 	}
 	return next.HandleBuild(ctx, in)
 }
@@ -63,8 +65,7 @@ func (r responseMiddleware) HandleDeserialize(ctx context.Context, in middleware
 	out, metadata, err = next.HandleDeserialize(ctx, in)
 	res, ok := out.RawResponse.(*http.Response)
 	if ok {
-		id, _ := getID(ctx)
-		r.HandleResponse(id, res.Response)
+		r.HandleResponse(ctx, res.Response)
 	}
 	return
 }
@@ -75,16 +76,30 @@ func withDeserializeOption(rmw *responseMiddleware, position middleware.Relative
 	}
 }
 
-func setID(ctx context.Context) (context.Context, string) {
-	id, ok := getID(ctx)
-	if !ok {
-		id = uuid.NewString()
-		return context.WithValue(ctx, requestID, id), id
+func mustRequestID(ctx context.Context) context.Context {
+	requestID := GetRequestID(ctx)
+	if requestID != "" {
+		return ctx
 	}
-	return ctx, id
+	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
 }
 
-func getID(ctx context.Context) (string, bool) {
-	id, ok := ctx.Value(requestID).(string)
-	return id, ok
+// GetOperationName retrieves the service operation metadata from the context.
+func GetOperationName(ctx context.Context) string {
+	operationName, _ := ctx.Value(operationNameKey{}).(string)
+	return operationName
 }

go.mod

@@ -265,7 +265,7 @@ require (
 	github.com/alecthomas/participle/v2 v2.0.0 // indirect
 	github.com/alecthomas/units v0.0.0-20211218093645-b94a6e3cc137 // indirect
 	github.com/aliyun/aliyun-log-go-sdk v0.1.54 // indirect
-	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231020160851-a8ff477e82af // indirect
+	github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware v0.0.0-20231023152757-c6e2437e6590 // indirect
 	github.com/amazon-contributing/opentelemetry-collector-contrib/override/aws v0.0.0-20230818193829-04a761abd409 // indirect
 	github.com/andybalholm/brotli v1.0.5 // indirect
 	github.com/antonmedv/expr v1.15.0 // indirect