otellambda: allow custom attributes

CHANGELOG.md

@@ -10,6 +10,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ### Added
 
+- Add the `WithTraceAttributeFn` option to `go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda`. (#7556)
 - Set `SeverityText` field to logrus hook in `go.opentelemetry.io/contrib/bridges/otellogrus`. (#7553)
 - Add the unit `ns` to deprecated runtime metrics `process.runtime.go.gc.pause_total_ns` and `process.runtime.go.gc.pause_ns` in `go.opentelemetry.io/contrib/instrumentation/runtime`. (#7490)
 - The `go.opentelemetry.io/contrib/detectors/autodetect` package is added to automatically compose user defined `resource.Detector`s at runtime. (#7522)

instrumentation/github.com/aws/aws-lambda-go/otellambda/README.md

@@ -61,6 +61,7 @@ func main() {
 | `WithFlusher` | `otellambda.Flusher`  | This instrumentation will call the `ForceFlush` method of its `Flusher` at the end of each invocation. Should you be using asynchronous logic (such as `sddktrace's BatchSpanProcessor`) it is very import for spans to be `ForceFlush`'ed before [Lambda freezes](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-context.html) to avoid data delays. | `Flusher` with noop `ForceFlush`
 | `WithEventToCarrier` | `func(eventJSON []byte) propagation.TextMapCarrier{}` | Function for providing custom logic to support retrieving trace header from different event types that are handled by AWS Lambda (e.g., SQS, CloudWatch, Kinesis, API Gateway) and returning them in a `propagation.TextMapCarrier` which a Propagator can use to extract the trace header into the context. | Function which returns an empty `TextMapCarrier` - new spans will be part of a new Trace and have no parent past Lambda instrumentation span
 | `WithPropagator` | `propagation.Propagator` | The `Propagator` the instrumentation will use to extract trace information into the context. | `otel.GetTextMapPropagator()` |
+| `WithTraceAttributeFn` | `func(eventJSON []byte) []attribute.KeyValue` | Function to extract custom attributes from different event types (e.g., SQS, CloudWatch, Kinesis, API Gateway, custom event) and return them as a slice of `attribute.KeyValue` to be added to the span. | Function which returns an empty `[]]attribute.KeyValue` (no custom attributes) |
 
 ### Usage With Options Example
 
@@ -78,6 +79,12 @@ func mockEventToCarrier(eventJSON []byte) propagation.TextMapCarrier{
 	return propagation.HeaderCarrier{someHeaderKey: []string{request.Headers[someHeaderKey]}}
 }
 
+func mockTraceAttributeFn(eventJSON []byte) []attribute.KeyValue {
+	var request mockHTTPRequest
+	_ = json.unmarshal(eventJSON, &request)
+	return []attribute.KeyValue{attribute.String("mock.request.type", reflect.TypeOf(request).String())}
+}
+
 type mockPropagator struct{}
 // Extract - read from `someHeaderKey`
 // Inject
@@ -96,6 +103,7 @@ func main() {
 	lambda.Start(otellambda.InstrumentHandler(HandleRequest,
 		                                    otellambda.WithTracerProvider(tp),
 		                                    otellambda.WithFlusher(tp),
+		                                    otellambda.WithTraceAttributeFn(mockTraceAttributeFn),
 		                                    otellambda.WithEventToCarrier(mockEventToCarrier),
 		                                    otellambda.WithPropagator(mockPropagator{})))
 }

instrumentation/github.com/aws/aws-lambda-go/otellambda/config.go

@@ -6,6 +6,7 @@ package otellambda // import "go.opentelemetry.io/contrib/instrumentation/github
 import (
 	"context"
 
+	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 )
@@ -35,10 +36,18 @@ var _ Flusher = &noopFlusher{}
 // trace information is instead stored in the Lambda environment.
 type EventToCarrier func(eventJSON []byte) propagation.TextMapCarrier
 
+// TraceAttributeFn defines a function that extracts attributes
+// from the event JSON to be added to the span created by the instrumentation.
+type TraceAttributeFn func(eventJSON []byte) []attribute.KeyValue
+
 func emptyEventToCarrier([]byte) propagation.TextMapCarrier {
 	return propagation.HeaderCarrier{}
 }
 
+func emptyTraceAttributeFn([]byte) []attribute.KeyValue {
+	return []attribute.KeyValue{}
+}
+
 // Compile time check our emptyEventToCarrier implements EventToCarrier.
 var _ EventToCarrier = emptyEventToCarrier
 
@@ -80,6 +89,12 @@ type config struct {
 	// The default value of Propagator the global otel Propagator
 	// returned by otel.GetTextMapPropagator()
 	Propagator propagation.TextMapPropagator
+
+	// TraceAttributeFn is a function that returns custom attributes
+	// to be added to the span created by the instrumentation.
+	// The default value of TraceAttributeFn is nil, which means no attributes
+	// will be added to the span.
+	TraceAttributeFn TraceAttributeFn
 }
 
 // WithTracerProvider configures the TracerProvider used by the
@@ -114,3 +129,10 @@ func WithPropagator(propagator propagation.TextMapPropagator) Option {
 		c.Propagator = propagator
 	})
 }
+
+// WithTraceAttributeFn configures a function that returns custom attributes.
+func WithTraceAttributeFn(fn TraceAttributeFn) Option {
+	return optionFunc(func(c *config) {
+		c.TraceAttributeFn = fn
+	})
+}

instrumentation/github.com/aws/aws-lambda-go/otellambda/lambda.go

@@ -32,10 +32,11 @@ type instrumentor struct {
 
 func newInstrumentor(opts ...Option) instrumentor {
 	cfg := config{
-		TracerProvider: otel.GetTracerProvider(),
-		Flusher:        &noopFlusher{},
-		EventToCarrier: emptyEventToCarrier,
-		Propagator:     otel.GetTextMapPropagator(),
+		TracerProvider:   otel.GetTracerProvider(),
+		Flusher:          &noopFlusher{},
+		EventToCarrier:   emptyEventToCarrier,
+		Propagator:       otel.GetTextMapPropagator(),
+		TraceAttributeFn: emptyTraceAttributeFn,
 	}
 	for _, opt := range opts {
 		opt.apply(&cfg)
@@ -58,6 +59,8 @@ func (i *instrumentor) tracingBegin(ctx context.Context, eventJSON []byte) (cont
 	spanName := os.Getenv("AWS_LAMBDA_FUNCTION_NAME")
 
 	var attributes []attribute.KeyValue
+	customAttrs := i.configuration.TraceAttributeFn(eventJSON)
+	attributes = append(attributes, customAttrs...)
 	lc, ok := lambdacontext.FromContext(ctx)
 	if !ok {
 		errorLogger.Println("failed to load lambda context from context, ensure tracing enabled in Lambda")

instrumentation/github.com/aws/aws-lambda-go/otellambda/lambdatest_test.go

@@ -355,6 +355,16 @@ func mockRequestCarrier(eventJSON []byte) propagation.TextMapCarrier {
 	return propagation.HeaderCarrier{mockPropagatorKey: []string{event.Headers[mockPropagatorKey]}}
 }
 
+func mockTraceAttributeFn(eventJSON []byte) []attribute.KeyValue {
+	var event mockRequest
+	err := json.Unmarshal(eventJSON, &event)
+	if err != nil {
+		fmt.Println("event type: ", reflect.TypeOf(event))
+		panic("mockRequestCarrier only supports events of type mockRequest")
+	}
+	return []attribute.KeyValue{attribute.String("mock.request.type", reflect.TypeOf(event).String())}
+}
+
 func TestInstrumentHandlerTracingWithMockPropagator(t *testing.T) {
 	setEnvVars(t)
 	tp, memExporter := initMockTracerProvider()
@@ -398,3 +408,23 @@ func TestWrapHandlerTracingWithMockPropagator(t *testing.T) {
 	stub := memExporter.GetSpans()[0]
 	assertStubEqualsIgnoreTime(t, mockPropagatorTestsExpectedSpanStub, stub)
 }
+
+func TestWrapHandlerTracingWithTraceAttributeFn(t *testing.T) {
+	setEnvVars(t)
+	tp, memExporter := initMockTracerProvider()
+
+	// No flusher needed as SimpleSpanProcessor is synchronous
+	wrapped := otellambda.WrapHandler(emptyHandler{},
+		otellambda.WithTracerProvider(tp),
+		otellambda.WithTraceAttributeFn(mockTraceAttributeFn),
+	)
+
+	payload, _ := json.Marshal(mockPropagatorTestsEvent)
+	_, err := wrapped.Invoke(mockPropagatorTestsContext, payload)
+	assert.NoError(t, err)
+
+	assert.Len(t, memExporter.GetSpans(), 1)
+	stub := memExporter.GetSpans()[0]
+	expectedAttr := attribute.KeyValue{Key: "mock.request.type", Value: attribute.StringValue(reflect.TypeOf(mockPropagatorTestsEvent).String())}
+	assert.Contains(t, stub.Attributes, expectedAttr, "custom attribute 'mock.request.type' with value 'otellambda_test.mockRequest' not found")
+}