feat: update DSM context passing

Author: jeastham1993

src/pricing-service/src/observability/observability.ts

@@ -5,7 +5,8 @@ import * as os from "os";
 
 const textEncoder = new TextEncoder();
 const logger = new Logger({});
-const cpuCount = os.cpus().length;
+
+export const DSM_PROPAGATION_KEY_BASE_64 = "dd-pathway-ctx-base64";
 
 export enum MessagingType {
   PUBLIC,
@@ -23,13 +24,21 @@ export interface SemanticConventions {
 
 export function startProcessSpanWithSemanticConventions(
   evt: CloudEvent<any>,
-  conventions: SemanticConventions
+  conventions: SemanticConventions,
 ): Span {
   const messageProcessingSpan = tracer.startSpan(`process ${evt.type}`, {
     childOf: conventions.parentSpan ?? undefined,
   });
-  const headers = {};
-  tracer.dataStreamsCheckpointer.setConsumeCheckpoint("eventbridge", evt.type, headers);
+
+  // Extract the DSM-specific header from the event data to set the consume checkpoint, enabling Datadog's Data Streams Monitoring to track this message consumption
+  const dsmHeader = (evt.data as Record<string, unknown>)[
+    DSM_PROPAGATION_KEY_BASE_64
+  ];
+  tracer.dataStreamsCheckpointer.setConsumeCheckpoint(
+    "eventbridge",
+    evt.type,
+    dsmHeader,
+  );
 
   try {
     messageProcessingSpan.addTags({
@@ -49,7 +58,7 @@ export function startProcessSpanWithSemanticConventions(
       "messaging.message.envelope.size": textEncoder.encode(JSON.stringify(evt))
         .length,
       "messaging.message.body.size": textEncoder.encode(
-        JSON.stringify(evt.data)
+        JSON.stringify(evt.data),
       ).length,
     });
 
@@ -73,18 +82,19 @@ export function startProcessSpanWithSemanticConventions(
 
 export function startPublishSpanWithSemanticConventions(
   evt: CloudEvent<any>,
-  conventions: SemanticConventions
-): Span {
+  conventions: SemanticConventions,
+): { span: Span; carrier: Record<string, string> } {
   const messagingSpan = tracer.startSpan(`publish ${evt.type}`, {
     childOf: conventions.parentSpan ?? undefined,
   });
 
+  const carrier = {};
+
   try {
-    const headers = {};
     tracer.dataStreamsCheckpointer.setProduceCheckpoint(
       "eventbridge",
       evt.type,
-      headers,
+      carrier,
     );
 
     messagingSpan.addTags({
@@ -102,7 +112,7 @@ export function startPublishSpanWithSemanticConventions(
       "messaging.message.envelope.size": textEncoder.encode(JSON.stringify(evt))
         .length,
       "messaging.message.body.size": textEncoder.encode(
-        JSON.stringify(evt.data)
+        JSON.stringify(evt.data),
       ).length,
       "messaging.operation.name": "send",
       "messaging.message.conversation_id": conventions.conversationId ?? "",
@@ -116,7 +126,7 @@ export function startPublishSpanWithSemanticConventions(
     logger.error(JSON.stringify(e));
   }
 
-  return messagingSpan;
+  return { span: messagingSpan, carrier };
 }
 
 export function addDefaultServiceTagsTo(span: Span | undefined | null) {

src/pricing-service/src/pricing-api/adapters/eventBridgeEventPublisher.ts

@@ -14,6 +14,7 @@ import { Span, tracer } from "dd-trace";
 import { CloudEvent } from "cloudevents";
 import { Logger } from "@aws-lambda-powertools/logger";
 import {
+  DSM_PROPAGATION_KEY_BASE_64,
   MessagingType,
   startPublishSpanWithSemanticConventions,
 } from "../../observability/observability";
@@ -33,7 +34,7 @@ export class EventBridgeEventPublisher implements EventPublisher {
 
   async publishPriceCalculatedEvent(
     evt: PriceCalculatedEventV1,
-    linkedTraceparent?: string
+    linkedTraceparent?: string,
   ): Promise<boolean> {
     const parentSpan = tracer.scope().active();
 
@@ -48,21 +49,27 @@ export class EventBridgeEventPublisher implements EventPublisher {
         traceparent: parentSpan?.context().toTraceparent(),
       });
 
-      messagingSpan = startPublishSpanWithSemanticConventions(
+      const publishSpan = startPublishSpanWithSemanticConventions(
         cloudEventWrapper,
         {
           publicOrPrivate: MessagingType.PUBLIC,
           messagingSystem: "eventbridge",
           destinationName: process.env.EVENT_BUS_NAME ?? "",
           parentSpan: parentSpan,
           linkedTraceparent,
-        }
+        },
       );
+      messagingSpan = publishSpan.span;
 
+      // Ensure the DSM context is injected
       const evtEntries: PutEventsRequestEntry[] = [
         {
           EventBusName: process.env.EVENT_BUS_NAME,
-          Detail: JSON.stringify(cloudEventWrapper),
+          Detail: JSON.stringify({
+            ...cloudEventWrapper,
+            DSM_PROPAGATION_KEY_BASE_64:
+              publishSpan.carrier[DSM_PROPAGATION_KEY_BASE_64],
+          }),
           DetailType: "pricing.pricingCalculated.v1",
           Source: `${process.env.ENV}.pricing`,
         },
@@ -71,12 +78,14 @@ export class EventBridgeEventPublisher implements EventPublisher {
       await this.client.send(
         new PutEventsCommand({
           Entries: evtEntries,
-        })
+        }),
       );
 
       messagingSpan?.finish();
     } catch (error: unknown) {
-      this.logger.error(error instanceof Error ? error.message : JSON.stringify(error));
+      this.logger.error(
+        error instanceof Error ? error.message : JSON.stringify(error),
+      );
       if (error instanceof Error) {
         const e = error as Error;
         const stack = e.stack!.split("\n").slice(1, 4).join("\n");

src/product-management-service/src/observability/main.go

@@ -29,6 +29,7 @@ type CloudEvent[T any] struct {
 	Time        string            `json:"time"`
 	TraceParent string            `json:"traceparent"`
 	Datadog     map[string]string `json:"_datadog,omitempty"`
+	DsmContext  string            `json:"dd-pathway-ctx-base64,omitempty"` // for backward compatibility with older services using the _datadog field as a simple string carrier
 }
 
 func NewCloudEvent[T any](ctx context.Context, evtType string, data T) CloudEvent[T] {

src/product-management-service/src/product-acl/pricing-changed-handler/main.go

@@ -132,10 +132,10 @@ func processMessage(ctx context.Context, record events.SQSMessage) error {
 
 	// Extract DSM context using the incoming ctx so the checkpoint is linked to
 	// the current trace, not an orphaned background context.
-	_, _ = tracer.SetDataStreamsCheckpointWithParams(datastreams.ExtractFromBase64Carrier(ctx, &evt), options.CheckpointParams{
+	dsm_ctx, _ := tracer.SetDataStreamsCheckpointWithParams(datastreams.ExtractFromBase64Carrier(ctx, &evt), options.CheckpointParams{
 		ServiceOverride: "productservice-acl",
 	}, "direction:in", productcore.ExternalPubSubName, "topic:"+evt.Type, "manual_checkpoint:true")
-	processSpan, _ := tracer.StartSpanFromContext(ctx, fmt.Sprintf("process %s", evt.Type), tracer.WithSpanLinks(spanLinks))
+	processSpan, _ := tracer.StartSpanFromContext(dsm_ctx, fmt.Sprintf("process %s", evt.Type), tracer.WithSpanLinks(spanLinks))
 	defer processSpan.Finish()
 
 	processSpan.SetTag("product.id", evt.Data.ProductId)
@@ -146,7 +146,7 @@ func processMessage(ctx context.Context, record events.SQSMessage) error {
 	processSpan.SetTag("messaging.operation.type", "process")
 	processSpan.SetTag("messaging.system", "aws_sqs")
 
-	_, err := eventTranslator.HandleProductPricingChanged(ctx, evt.Data)
+	_, err := eventTranslator.HandleProductPricingChanged(dsm_ctx, evt.Data)
 
 	if err != nil {
 		processSpan.SetTag("error", true)