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 {
@@ -80,6 +81,7 @@ export class TraceListener {
   private wrappedCurrentSpan?: SpanWrapper;
   private triggerTags?: { [key: string]: string };
   private lambdaSpanParentContext?: SpanContext;
+  private spanPointerAttributesList: object[] = [];
 
   public get currentTraceHeaders() {
     return this.contextService.currentTraceHeaders;
@@ -131,8 +133,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 +203,12 @@ export class TraceListener {
         }
       }
     }
+
+    if (this.wrappedCurrentSpan) {
+      for (const attributes of this.spanPointerAttributesList) {
+        this.wrappedCurrentSpan.span.addSpanPointer(attributes);
+      }
+    }
     return false;
   }
 

src/trace/trigger.spec.ts

@@ -166,15 +166,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(event);
+      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(event);
+      const triggerTags = extractTriggerTags(eventData, mockContext, eventSource);
       const isResponseStreamingFunction = false;
       for (const response of bufferedResponses) {
         const statusCode = extractHTTPStatusCodeTag(triggerTags, response.responseBody, isResponseStreamingFunction);
@@ -198,7 +200,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(event);
+      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

@@ -323,9 +323,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,151 @@
+import { getSpanPointerAttributes } from "./span-pointers";
+import { eventTypes } from "../trace/trigger";
+import { SPAN_LINK_KIND, S3_PTR_KIND, SPAN_POINTER_DIRECTION } from "dd-trace/packages/dd-trace/src/span_pointers";
+import * as spanPointers from "dd-trace/packages/dd-trace/src/span_pointers";
+
+// Mock the external dependencies
+jest.mock("./log", () => ({
+  logDebug: jest.fn(),
+}));
+
+describe("span-pointers utils", () => {
+  const mockS3PointerHash = "mock-hash-123";
+
+  beforeEach(() => {
+    // Mock the generateS3PointerHash function
+    jest.spyOn(spanPointers, "generateS3PointerHash").mockReturnValue(mockS3PointerHash);
+  });
+
+  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",
+                },
+              },
+            },
+          ],
+        };
+
+        const expected = [
+          {
+            "ptr.kind": S3_PTR_KIND,
+            "ptr.dir": SPAN_POINTER_DIRECTION.UPSTREAM,
+            "ptr.hash": mockS3PointerHash,
+            "link.kind": SPAN_LINK_KIND,
+          },
+        ];
+
+        const result = getSpanPointerAttributes(eventTypes.s3, event);
+        expect(result).toEqual(expected);
+        expect(spanPointers.generateS3PointerHash).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",
+                },
+              },
+            },
+            {
+              s3: {
+                bucket: { name: "bucket2" },
+                object: {
+                  key: "key2",
+                  eTag: "etag2",
+                },
+              },
+            },
+          ],
+        };
+
+        const expected = [
+          {
+            "ptr.kind": S3_PTR_KIND,
+            "ptr.dir": SPAN_POINTER_DIRECTION.UPSTREAM,
+            "ptr.hash": mockS3PointerHash,
+            "link.kind": SPAN_LINK_KIND,
+          },
+          {
+            "ptr.kind": S3_PTR_KIND,
+            "ptr.dir": SPAN_POINTER_DIRECTION.UPSTREAM,
+            "ptr.hash": mockS3PointerHash,
+            "link.kind": SPAN_LINK_KIND,
+          },
+        ];
+
+        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",
+                },
+              },
+            },
+          ],
+        };
+
+        const expected = [
+          {
+            "ptr.kind": S3_PTR_KIND,
+            "ptr.dir": SPAN_POINTER_DIRECTION.UPSTREAM,
+            "ptr.hash": mockS3PointerHash,
+            "link.kind": SPAN_LINK_KIND,
+          },
+        ];
+
+        const result = getSpanPointerAttributes(eventTypes.s3, event);
+        expect(result).toEqual(expected);
+      });
+    });
+  });
+});

src/utils/span-pointers.ts

@@ -0,0 +1,74 @@
+import { eventTypes } from "../trace/trigger";
+import { logDebug } from "./log";
+import {
+  SPAN_LINK_KIND,
+  S3_PTR_KIND,
+  SPAN_POINTER_DIRECTION,
+  generateS3PointerHash,
+} from "dd-trace/packages/dd-trace/src/span_pointers";
+
+interface SpanPointerAttributes {
+  "ptr.kind": string;
+  "ptr.dir": string;
+  "ptr.hash": string;
+  "link.kind": 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 = [];
+  const linkKind = SPAN_LINK_KIND;
+
+  for (const record of records) {
+    const eventName = record.eventName;
+    if (!["ObjectCreated:Put", "ObjectCreated:Copy", "ObjectCreated:CompleteMultipartUpload"].includes(eventName)) {
+      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;
+    const eTag = s3Event?.object?.eTag;
+
+    if (!bucketName || !objectKey || !eTag) {
+      logDebug("Unable to calculate span pointer hash because of missing parameters.");
+      continue;
+    }
+
+    const pointerHash = generateS3PointerHash(bucketName, objectKey, eTag);
+    const spanPointerAttributes = {
+      "ptr.kind": S3_PTR_KIND,
+      "ptr.dir": SPAN_POINTER_DIRECTION.UPSTREAM,
+      "ptr.hash": pointerHash,
+      "link.kind": linkKind,
+    };
+    spanPointerAttributesList.push(spanPointerAttributes);
+  }
+
+  return spanPointerAttributesList;
+}