fix: introduce otellambda instrumentation (#288)

Wrap our existing mux instrumentation with the `otellambda` package. The
upstream package sets a reasonable span name, kind, and injects
attributes that are consistent with the `faas` semantic convention.

In order to attach the payload to every invocation we still need to
provide our own Invoke method. The payload is added as an attribute.

cmd/forwarder/main.go

@@ -161,10 +161,7 @@ func realInit() (err error) {
 		return fmt.Errorf("failed to register functions: %w", err)
 	}
 
-	entrypoint = &tracing.LambdaHandler{
-		Handler: mux,
-		Tracer:  tracer,
-	}
+	entrypoint = tracing.NewLambdaHandler(mux, tracerProvider)
 	return nil
 }
 

cmd/subscriber/main.go

@@ -133,10 +133,7 @@ func realInit() error {
 		return fmt.Errorf("failed to register functions: %w", err)
 	}
 
-	entrypoint = &tracing.LambdaHandler{
-		Handler: mux,
-		Tracer:  tracer,
-	}
+	entrypoint = tracing.NewLambdaHandler(mux, tracerProvider)
 	return nil
 }
 

go.mod

@@ -61,6 +61,7 @@ require (
 	github.com/prometheus/common v0.53.0 // indirect
 	github.com/prometheus/procfs v0.15.0 // indirect
 	go.opentelemetry.io/contrib/bridges/prometheus v0.52.0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda v0.52.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc v1.27.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp v1.27.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.27.0 // indirect

go.sum

@@ -123,6 +123,8 @@ go.opentelemetry.io/contrib/detectors/aws/lambda v0.52.0 h1:aZAZUtzXd7dklbWDqf0s
 go.opentelemetry.io/contrib/detectors/aws/lambda v0.52.0/go.mod h1:0oIsMo4DjOJNB/mn/bssl3ba7TEZhSgqe26vRMoxHv4=
 go.opentelemetry.io/contrib/exporters/autoexport v0.52.0 h1:G/AGl5O78ZKHs63Rl65P1HyZfDnTyxjv8r7dbdZ9fB0=
 go.opentelemetry.io/contrib/exporters/autoexport v0.52.0/go.mod h1:WoVWPZjJ7EB5Z9aROW1DZuRIoFEemxmhCdZJlcjY2AE=
+go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda v0.52.0 h1:djhGC0347ZyXb4HB7qqVoNRlyEmRshXl2lzD2cU4fSI=
+go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda v0.52.0/go.mod h1:+PgU2bOkTOL1XcmdCe/8EA5H5XY/RRq2QqQRyjP5RZ0=
 go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-sdk-go-v2/otelaws v0.52.0 h1:kAytSRJYoIy4eJtDOfSGf9LOCD4QdXFN37YJs0+bYrw=
 go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-sdk-go-v2/otelaws v0.52.0/go.mod h1:l6VnFEqDdeMSMfwULTDDY9ewlnlVLhmvBainVT+h/Zs=
 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.52.0 h1:9l89oX4ba9kHbBol3Xin3leYJ+252h0zszDtBwyKe2A=

tracing/mux.go

@@ -4,28 +4,27 @@ import (
 	"context"
 
 	"github.com/aws/aws-lambda-go/lambda"
+	"go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda"
 	"go.opentelemetry.io/otel/attribute"
-	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
 )
 
+var payloadKey = attribute.Key("payload")
+
+func NewLambdaHandler(handler lambda.Handler, tp trace.TracerProvider) lambda.Handler {
+	return otellambda.WrapHandler(
+		&LambdaHandler{handler},
+		otellambda.WithTracerProvider(tp),
+	)
+}
+
 type LambdaHandler struct {
-	Tracer trace.Tracer
 	lambda.Handler
 }
 
-func (h *LambdaHandler) Invoke(ctx context.Context, payload []byte) (response []byte, err error) {
-	cctx, span := h.Tracer.Start(ctx, "Invoke",
-		trace.WithAttributes(attribute.String("payload", string(payload))),
-		trace.WithSpanKind(trace.SpanKindServer),
-	)
-	defer func() {
-		if err != nil {
-			span.RecordError(err)
-			span.SetStatus(codes.Error, err.Error())
-		}
-		span.End()
-	}()
-	response, err = h.Handler.Invoke(cctx, payload)
-	return
+func (h *LambdaHandler) Invoke(ctx context.Context, payload []byte) ([]byte, error) {
+	span := trace.SpanFromContext(ctx)
+	span.SetAttributes(payloadKey.String(string(payload)))
+	// nolint: wrapcheck
+	return h.Handler.Invoke(ctx, payload)
 }