fix: OB-35469 add subscription request parent span as child of disco…

pkg/handler/subscriber/instrument.go

@@ -8,6 +8,8 @@ import (
 	"github.com/aws/aws-lambda-go/events"
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
 	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
+	"github.com/observeinc/aws-sam-apps/pkg/tracing"
+
 	"github.com/go-logr/logr"
 	"go.opentelemetry.io/contrib/propagators/b3"
 	"go.opentelemetry.io/otel/codes"
@@ -25,8 +27,6 @@ func (h *InstrumentedHandler) HandleSQS(ctx context.Context, request events.SQSE
 	for _, record := range request.Records {
 		var req Request
 		var err error
-		logger.V(3).Info("Getting context from message attributes")
-		ctx = NewSQSCarrier().Extract(ctx, record.MessageAttributes)
 
 		if err = json.Unmarshal([]byte(record.Body), &req); err == nil {
 			_, err = h.HandleRequest(ctx, &req)
@@ -66,7 +66,7 @@ func (q *InstrumentedSQSClient) SendMessage(ctx context.Context, msg *sqs.SendMe
 	logger := logr.FromContextOrDiscard(ctx)
 	logger.V(3).Info("Injecting context into message attributes")
 	msg.MessageAttributes = make(map[string]types.MessageAttributeValue)
-	if err := NewSQSCarrier().Inject(ctx, msg.MessageAttributes); err != nil {
+	if err := tracing.NewSQSCarrier().Inject(ctx, msg.MessageAttributes); err != nil {
 		return nil, fmt.Errorf("failed to inject context into message attributes: %w", err)
 	}
 	logger.V(3).Info("sending message %s", msg.MessageBody)

pkg/lambda/subscriber/lambda.go

@@ -130,6 +130,6 @@ func New(ctx context.Context, cfg *Config) (*Lambda, error) {
 		return nil, fmt.Errorf("failed to register functions: %w", err)
 	}
 
-	l.Entrypoint = tracing.NewLambdaHandler(mux, tracerProvider)
+	l.Entrypoint = tracing.WrapHandlerSQSContext(tracing.NewLambdaHandler(mux, tracerProvider))
 	return l, nil
 }

pkg/tracing/sqs.go

@@ -0,0 +1,49 @@
+package tracing
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+
+	"github.com/aws/aws-lambda-go/events"
+	"github.com/aws/aws-lambda-go/lambda"
+	"github.com/go-logr/logr"
+)
+
+type SQSWithContextHandler struct {
+	handler lambda.Handler
+}
+
+// Compile time check our Handler implements lambda.Handler.
+var _ lambda.Handler = SQSWithContextHandler{}
+
+// Invoke checks if the incoming payload is from a SQS event and, if so,
+// extracts the context from the SQS message and injects it into the context.
+// This means that the spans created with this context will appear as children
+// of a span in the discover request that created the SQS event
+func (h SQSWithContextHandler) Invoke(ctx context.Context, payload []byte) ([]byte, error) {
+	logger := logr.FromContextOrDiscard(ctx)
+	logger.V(3).Info("Getting context from message attributes")
+	var event events.SQSEvent
+	dec := json.NewDecoder(bytes.NewReader(payload))
+	dec.DisallowUnknownFields()
+	if err := dec.Decode(&event); err == nil {
+		for _, record := range event.Records {
+			ctx = NewSQSCarrier().Extract(ctx, record.MessageAttributes)
+			break
+		}
+	}
+
+	response, err := h.handler.Invoke(ctx, payload)
+	if err != nil {
+		return nil, err
+	}
+
+	return response, nil
+}
+
+// WrapHandlerSQSContext Provides a Handler which wraps an existing Handler while
+// injecting the SQS context into the context if the payload is a SQS event.
+func WrapHandlerSQSContext(handler lambda.Handler) lambda.Handler {
+	return SQSWithContextHandler{handler: handler}
+}

pkg/handler/subscriber/sqspropagator.go → pkg/tracing/sqspropagator.go

@@ -1,4 +1,4 @@
-package subscriber
+package tracing
 
 import (
 	"context"