Add S3 Downstream Span Pointers

src/trace/listener.ts

@@ -2,7 +2,7 @@ import { Context } from "aws-lambda";
 
 import { patchHttp, unpatchHttp } from "./patch-http";
 
-import { extractTriggerTags, extractHTTPStatusCodeTag } from "./trigger";
+import { extractTriggerTags, extractHTTPStatusCodeTag, parseEventSource } from "./trigger";
 import { ColdStartTracerConfig, ColdStartTracer } from "./cold-start-tracer";
 import { logDebug, tagObject } from "../utils";
 import { didFunctionColdStart, isProactiveInitialization } from "../utils/cold-start";
@@ -17,6 +17,7 @@ import { TraceContext, TraceContextService, TraceSource } from "./trace-context-
 import { StepFunctionContext, StepFunctionContextService } from "./step-function-service";
 import { XrayService } from "./xray-service";
 import { AUTHORIZING_REQUEST_ID_HEADER } from "./context/extractors/http";
+import { getSpanPointerAttributes } from "../utils/span-pointers";
 export type TraceExtractor = (event: any, context: Context) => Promise<TraceContext> | TraceContext;
 
 export interface TraceConfig {
@@ -70,6 +71,12 @@ export interface TraceConfig {
   coldStartTraceSkipLib: string;
 }
 
+interface SpanPointerAttributes {
+  pointerKind: string;
+  pointerDirection: string;
+  pointerHash: string;
+}
+
 export class TraceListener {
   private contextService: TraceContextService;
   private context?: Context;
@@ -80,6 +87,7 @@ export class TraceListener {
   private wrappedCurrentSpan?: SpanWrapper;
   private triggerTags?: { [key: string]: string };
   private lambdaSpanParentContext?: SpanContext;
+  private spanPointerAttributesList: SpanPointerAttributes[] = [];
 
   public get currentTraceHeaders() {
     return this.contextService.currentTraceHeaders;
@@ -131,8 +139,14 @@ export class TraceListener {
 
     this.lambdaSpanParentContext = this.inferredSpan?.span || parentSpanContext;
     this.context = context;
-    this.triggerTags = extractTriggerTags(event, context);
+    const eventSource = parseEventSource(event);
+    this.triggerTags = extractTriggerTags(event, context, eventSource);
     this.stepFunctionContext = StepFunctionContextService.instance().context;
+
+    const result = getSpanPointerAttributes(eventSource, event);
+    if (result) {
+      this.spanPointerAttributesList.push(...result);
+    }
   }
 
   /**
@@ -195,6 +209,16 @@ export class TraceListener {
         }
       }
     }
+
+    if (this.wrappedCurrentSpan) {
+      for (const attributes of this.spanPointerAttributesList) {
+        this.wrappedCurrentSpan.span.addSpanPointer(
+          attributes.pointerKind,
+          attributes.pointerDirection,
+          attributes.pointerHash,
+        );
+      }
+    }
     return false;
   }
 

src/trace/trigger.spec.ts

@@ -174,15 +174,17 @@ describe("parseEventSource", () => {
   it("extracts all trigger tags", () => {
     for (let event of events) {
       const eventData = JSON.parse(readFileSync(`./event_samples/${event.file}`, "utf8"));
-      const triggerTags = extractTriggerTags(eventData, mockContext);
+      const eventSource = parseEventSource(eventData);
+      const triggerTags = extractTriggerTags(eventData, mockContext, eventSource);
       expect(triggerTags).toEqual(event.result);
     }
   });
 
   it("extracts the status code if API Gateway, ALB, or Function URL, otherwise do nothing, for buffered functions", () => {
     for (const event of events) {
       const eventData = JSON.parse(readFileSync(`./event_samples/${event.file}`, "utf8"));
-      const triggerTags = extractTriggerTags(eventData, mockContext);
+      const eventSource = parseEventSource(eventData);
+      const triggerTags = extractTriggerTags(eventData, mockContext, eventSource);
       const isResponseStreamingFunction = false;
       for (const response of bufferedResponses) {
         const statusCode = extractHTTPStatusCodeTag(triggerTags, response.responseBody, isResponseStreamingFunction);
@@ -206,7 +208,8 @@ describe("parseEventSource", () => {
   it("extracts the status code if API Gateway, ALB, or Function URL, otherwise do nothing, for streaming functions", () => {
     for (let event of events) {
       const eventData = JSON.parse(readFileSync(`./event_samples/${event.file}`, "utf8"));
-      const triggerTags = extractTriggerTags(eventData, mockContext);
+      const eventSource = parseEventSource(eventData);
+      const triggerTags = extractTriggerTags(eventData, mockContext, eventSource);
       const isResponseStreamingFunction = true;
       for (const response of streamingResponses) {
         const statusCode = extractHTTPStatusCodeTag(triggerTags, response.responseBody, isResponseStreamingFunction);

src/trace/trigger.ts

@@ -344,9 +344,8 @@ function extractHTTPTags(event: APIGatewayEvent | APIGatewayProxyEventV2 | ALBEv
 /**
  * extractTriggerTags parses the trigger event object for tags to be added to the span metadata
  */
-export function extractTriggerTags(event: any, context: Context) {
+export function extractTriggerTags(event: any, context: Context, eventSource: eventTypes | undefined) {
   let triggerTags: { [key: string]: string } = {};
-  const eventSource = parseEventSource(event);
   if (eventSource) {
     triggerTags["function_trigger.event_source"] = eventSource;
 

src/utils/span-pointers.spec.ts

@@ -0,0 +1,159 @@
+import { getSpanPointerAttributes } from "./span-pointers";
+import { eventTypes } from "../trace/trigger";
+
+// tslint:disable-next-line:no-var-requires
+const { S3_PTR_KIND, SPAN_POINTER_DIRECTION } = require("dd-trace/packages/dd-trace/src/constants");
+// tslint:disable-next-line:no-var-requires
+const util = require("dd-trace/packages/dd-trace/src/util");
+
+// Mock the external dependencies
+jest.mock("./log", () => ({
+  logDebug: jest.fn(),
+}));
+
+interface SpanPointerAttributes {
+  pointerKind: string;
+  pointerDirection: string;
+  pointerHash: string;
+}
+
+describe("span-pointers utils", () => {
+  const mockPointerHash = "mock-hash-123";
+
+  beforeEach(() => {
+    jest.spyOn(util, "generatePointerHash").mockReturnValue(mockPointerHash);
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe("getSpanPointerAttributes", () => {
+    it("returns undefined when eventSource is undefined", () => {
+      const result = getSpanPointerAttributes(undefined, {});
+      expect(result).toBeUndefined();
+    });
+
+    it("returns undefined for unsupported event types", () => {
+      const result = getSpanPointerAttributes("unsupported" as eventTypes, {});
+      expect(result).toBeUndefined();
+    });
+
+    describe("S3 event processing", () => {
+      it("processes single S3 record correctly", () => {
+        const event = {
+          Records: [
+            {
+              s3: {
+                bucket: { name: "test-bucket" },
+                object: {
+                  key: "test-key",
+                  eTag: "test-etag",
+                },
+              },
+              eventName: "ObjectCreated:SomeEventName",
+            },
+          ],
+        };
+
+        const expected: SpanPointerAttributes[] = [
+          {
+            pointerKind: S3_PTR_KIND,
+            pointerDirection: SPAN_POINTER_DIRECTION.UPSTREAM,
+            pointerHash: mockPointerHash,
+          },
+        ];
+
+        const result = getSpanPointerAttributes(eventTypes.s3, event);
+        expect(result).toEqual(expected);
+        expect(util.generatePointerHash).toHaveBeenCalledWith(["test-bucket", "test-key", "test-etag"]);
+      });
+
+      it("processes multiple S3 records correctly", () => {
+        const event = {
+          Records: [
+            {
+              s3: {
+                bucket: { name: "bucket1" },
+                object: {
+                  key: "key1",
+                  eTag: "etag1",
+                },
+              },
+              eventName: "ObjectCreated:SomeEventName",
+            },
+            {
+              s3: {
+                bucket: { name: "bucket2" },
+                object: {
+                  key: "key2",
+                  eTag: "etag2",
+                },
+              },
+              eventName: "ObjectCreated:SomeEventName",
+            },
+          ],
+        };
+
+        const expected: SpanPointerAttributes[] = [
+          {
+            pointerKind: S3_PTR_KIND,
+            pointerDirection: SPAN_POINTER_DIRECTION.UPSTREAM,
+            pointerHash: mockPointerHash,
+          },
+          {
+            pointerKind: S3_PTR_KIND,
+            pointerDirection: SPAN_POINTER_DIRECTION.UPSTREAM,
+            pointerHash: mockPointerHash,
+          },
+        ];
+
+        const result = getSpanPointerAttributes(eventTypes.s3, event);
+        expect(result).toEqual(expected);
+      });
+
+      it("handles empty Records array", () => {
+        const event = { Records: [] };
+        const result = getSpanPointerAttributes(eventTypes.s3, event);
+        expect(result).toEqual([]);
+      });
+
+      it("handles missing Records property", () => {
+        const event = {};
+        const result = getSpanPointerAttributes(eventTypes.s3, event);
+        expect(result).toEqual([]);
+      });
+
+      it("skips invalid records but processes valid ones", () => {
+        const event = {
+          Records: [
+            {
+              // Invalid record missing s3 property
+            },
+            {
+              s3: {
+                bucket: { name: "valid-bucket" },
+                object: {
+                  key: "valid-key",
+                  eTag: "valid-etag",
+                },
+              },
+              eventName: "ObjectCreated:SomeEventName",
+            },
+          ],
+        };
+
+        const expected: SpanPointerAttributes[] = [
+          {
+            pointerKind: S3_PTR_KIND,
+            pointerDirection: SPAN_POINTER_DIRECTION.UPSTREAM,
+            pointerHash: mockPointerHash,
+          },
+        ];
+
+        const result = getSpanPointerAttributes(eventTypes.s3, event);
+        expect(result).toEqual(expected);
+      });
+    });
+  });
+});

src/utils/span-pointers.ts

@@ -0,0 +1,85 @@
+import { eventTypes } from "../trace/trigger";
+import { logDebug } from "./log";
+
+interface SpanPointerAttributes {
+  pointerKind: string;
+  pointerDirection: string;
+  pointerHash: string;
+}
+
+/**
+ * Computes span pointer attributes
+ *
+ * @param {eventTypes} eventSource - The type of event being processed (e.g., S3, DynamoDB).
+ * @param {any} event - The event object containing source-specific data.
+ * @returns {SpanPointerAttributes[] | undefined} An array of span pointer attribute objects, or undefined if none could be computed.
+ */
+export function getSpanPointerAttributes(
+  eventSource: eventTypes | undefined,
+  event: any,
+): SpanPointerAttributes[] | undefined {
+  if (!eventSource) {
+    return;
+  }
+
+  switch (eventSource) {
+    case eventTypes.s3:
+      return processS3Event(event);
+    default:
+      logDebug(`Event type ${eventSource} not supported by span pointers.`);
+      return;
+  }
+}
+
+function processS3Event(event: any): SpanPointerAttributes[] {
+  const records = event.Records || [];
+  const spanPointerAttributesList: SpanPointerAttributes[] = [];
+
+  // Get dependencies from tracer only when needed
+  let constants;
+  let util;
+  try {
+    constants = require("dd-trace/packages/dd-trace/src/constants");
+    util = require("dd-trace/packages/dd-trace/src/util");
+  } catch (err) {
+    if (err instanceof Error) {
+      logDebug("Failed to load dd-trace span pointer dependencies", err);
+    }
+    return spanPointerAttributesList;
+  }
+
+  const { S3_PTR_KIND, SPAN_POINTER_DIRECTION } = constants;
+  const { generatePointerHash } = util;
+
+  for (const record of records) {
+    const eventName = record.eventName;
+    if (!eventName || !eventName.startsWith("ObjectCreated")) {
+      continue;
+    }
+    // Values are stored in the same place, regardless of AWS SDK v2/v3 or the event type.
+    // https://docs.aws.amazon.com/AmazonS3/latest/userguide/notification-content-structure.html
+    const s3Event = record?.s3;
+    const bucketName = s3Event?.bucket?.name;
+    const objectKey = s3Event?.object?.key;
+    let eTag = s3Event?.object?.eTag;
+
+    if (!bucketName || !objectKey || !eTag) {
+      logDebug("Unable to calculate span pointer hash because of missing parameters.");
+      continue;
+    }
+
+    // https://github.com/DataDog/dd-span-pointer-rules/blob/main/AWS/S3/Object/README.md
+    if (eTag.startsWith('"') && eTag.endsWith('"')) {
+      eTag = eTag.slice(1, -1);
+    }
+    const pointerHash = generatePointerHash([bucketName, objectKey, eTag]);
+    const spanPointerAttributes: SpanPointerAttributes = {
+      pointerKind: S3_PTR_KIND,
+      pointerDirection: SPAN_POINTER_DIRECTION.UPSTREAM,
+      pointerHash,
+    };
+    spanPointerAttributesList.push(spanPointerAttributes);
+  }
+
+  return spanPointerAttributesList;
+}