feat: support for DSM (#672)

* set checkpoint for all supported dsm queue types

* removed not needed checkpoint set

* clean

* clearer return type

* add tests

* add the manual checkpoint parameter

* Correct approach to ensure public api is used

* remove unused variable

* match logging patterns of repo

* lint

* lint

* add False manual_checkpoint false

* fix formatting

* optimization

* clean

* add prettier ignore

* add optional chaining

* add log debug for no arn case

* merge commit

* add checkpoints for headers

* Update SQS to work with batch processing

* Update Kinesis to work with batch processing

* Update SNS-SQS to work with batch processing

* Prettier changes

* Fix test for set sqs-sns

* Batch processing for sns

* Add a more extensive DD_DATA_STREAMS_ENABLED function

* Early continue when there is context for batch processing for all queues, rather than just kinesis

* Break out of the loop of records in batch processing when there is no DSM for speed

* Move the env var check on data streams enabled to config and break out of the loop of records earlier when data streams is not enabled

* test because i think the integration test failed for flake reasons

* Move from getConfig() to passed in config for dataStreamsEnabled

* Update the extraction functions to separate the setting of dsm checkpoints with the context stuff completely

* Move the config of dataStreamsEnabled out of the trace-wrapper but in the extractors

* Optimize such that you only parse the first messages headers once in the case where DSM is instrumented

* Update the dd-trace version

* Fix test description to be accurate for DSM is disabled

* empty

---------

Co-authored-by: Eric Firth <eric.firth@datadoghq.com>

Author: michael-zhao459

package.json

@@ -29,7 +29,7 @@
     "@types/node": "^20.12.10",
     "@types/promise-retry": "^1.1.3",
     "@types/shimmer": "^1.0.1",
-    "dd-trace": "^5.58.0",
+    "dd-trace": "^5.62.0",
     "jest": "^27.0.1",
     "mock-fs": "4.14.0",
     "nock": "13.5.4",

src/index.ts

@@ -52,6 +52,7 @@ export const minColdStartTraceDurationEnvVar = "DD_MIN_COLD_START_DURATION";
 export const coldStartTraceSkipLibEnvVar = "DD_COLD_START_TRACE_SKIP_LIB";
 export const localTestingEnvVar = "DD_LOCAL_TESTING";
 export const addSpanPointersEnvVar = "DD_TRACE_AWS_ADD_SPAN_POINTERS";
+export const dataStreamsEnabledEnvVar = "DD_DATA_STREAMS_ENABLED";
 
 interface GlobalConfig {
   /**
@@ -97,6 +98,7 @@ export const defaultConfig: Config = {
   coldStartTraceSkipLib: "",
   localTesting: false,
   addSpanPointers: true,
+  dataStreamsEnabled: false,
 } as const;
 
 export const _metricsQueue: MetricsQueue = new MetricsQueue();
@@ -419,6 +421,12 @@ function getConfig(userConfig?: Partial<Config>): Config {
     config.addSpanPointers = result === "true";
   }
 
+  if (userConfig === undefined || userConfig.dataStreamsEnabled === undefined) {
+    const result = getEnvValue(dataStreamsEnabledEnvVar, "false").toLowerCase();
+    const validEnabledValues = new Set(["true", "1", "yes", "y", "on"]);
+    config.dataStreamsEnabled = validEnabledValues.has(result);
+  }
+
   return config;
 }
 

src/trace/context/extractor.ts

@@ -86,12 +86,13 @@ export class TraceContextExtractor {
       return new HTTPEventTraceExtractor(this.tracerWrapper, this.config.decodeAuthorizerContext);
     }
 
-    if (EventValidator.isSNSEvent(event)) return new SNSEventTraceExtractor(this.tracerWrapper);
-    if (EventValidator.isSNSSQSEvent(event)) return new SNSSQSEventTraceExtractor(this.tracerWrapper);
+    if (EventValidator.isSNSEvent(event)) return new SNSEventTraceExtractor(this.tracerWrapper, this.config);
+    if (EventValidator.isSNSSQSEvent(event)) return new SNSSQSEventTraceExtractor(this.tracerWrapper, this.config);
     if (EventValidator.isEventBridgeSQSEvent(event)) return new EventBridgeSQSEventTraceExtractor(this.tracerWrapper);
     if (EventValidator.isAppSyncResolverEvent(event)) return new AppSyncEventTraceExtractor(this.tracerWrapper);
-    if (EventValidator.isSQSEvent(event)) return new SQSEventTraceExtractor(this.tracerWrapper);
-    if (EventValidator.isKinesisStreamEvent(event)) return new KinesisEventTraceExtractor(this.tracerWrapper);
+    if (EventValidator.isSQSEvent(event)) return new SQSEventTraceExtractor(this.tracerWrapper, this.config);
+    if (EventValidator.isKinesisStreamEvent(event))
+      return new KinesisEventTraceExtractor(this.tracerWrapper, this.config);
     if (EventValidator.isEventBridgeEvent(event)) return new EventBridgeEventTraceExtractor(this.tracerWrapper);
 
     return;

src/trace/context/extractors/kinesis.spec.ts

@@ -2,6 +2,15 @@ import { TracerWrapper } from "../../tracer-wrapper";
 import { KinesisEventTraceExtractor } from "./kinesis";
 
 let mockSpanContext: any = null;
+let mockDataStreamsCheckpointer: any = {
+  setConsumeCheckpoint: jest.fn(),
+};
+
+jest.mock("dd-trace/packages/dd-trace/src/datastreams/checkpointer", () => {
+  return {
+    DataStreamsCheckpointer: jest.fn().mockImplementation(() => mockDataStreamsCheckpointer),
+  };
+});
 
 // Mocking extract is needed, due to dd-trace being a No-op
 // if the detected environment is testing. This is expected, since
@@ -12,18 +21,31 @@ jest.mock("dd-trace", () => {
     ...ddTrace,
     _tracer: { _service: {} },
     extract: (_carrier: any, _headers: any) => mockSpanContext,
+    dataStreamsCheckpointer: mockDataStreamsCheckpointer,
   };
 });
 const spyTracerWrapper = jest.spyOn(TracerWrapper.prototype, "extract");
 
 describe("KinesisEventTraceExtractor", () => {
+  const mockConfig = {
+    autoPatchHTTP: true,
+    captureLambdaPayload: false,
+    captureLambdaPayloadMaxDepth: 10,
+    createInferredSpan: true,
+    encodeAuthorizerContext: true,
+    decodeAuthorizerContext: true,
+    mergeDatadogXrayTraces: false,
+    injectLogContext: false,
+    minColdStartTraceDuration: 3,
+    coldStartTraceSkipLib: "",
+    addSpanPointers: true,
+    dataStreamsEnabled: true,
+  };
+
   describe("extract", () => {
     beforeEach(() => {
       mockSpanContext = null;
-    });
-
-    afterEach(() => {
-      jest.resetModules();
+      mockDataStreamsCheckpointer.setConsumeCheckpoint.mockClear();
     });
 
     it("extracts trace context with valid payload", () => {
@@ -43,7 +65,7 @@ describe("KinesisEventTraceExtractor", () => {
               kinesisSchemaVersion: "1.0",
               partitionKey: "cdbfd750-cec0-4f0f-a4b0-82ae6152c7fb",
               sequenceNumber: "49625698045709644136382874226371117765484751339579768834",
-              data: "eyJJJ20gbWFkZSBvZiB3YXgsIExhcnJ5IjoiV2hhdCBhcmUgeW91IG1hZGUgb2Y/IiwiX2RhdGFkb2ciOnsieC1kYXRhZG9nLXRyYWNlLWlkIjoiNjY3MzA5NTE0MjIxMDM1NTM4IiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjEzNTA3MzUwMzU0OTc4MTE4MjgiLCJ4LWRhdGFkb2ctc2FtcGxlZCI6IjEiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIn19",
+              data: "eyJJJ20gbWFkZSBvZiB3YXgsIExhcnJ5IjoiV2hhdCBhcmUgeW91IG1hZGUgb2Y/IiwiX2RhdGFkb2ciOnsieC1kYXRhZG9nLXRyYWNlLWlkIjoiNjY3MzA5NTE0MjIxMDM1NTM4IiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjEzNTA3MzUwMzU0OTc4MTE4MjgiLCJ4LWRhdGFkb2ctc2FtcGxlZCI6IjEiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIiwiZGQtcGF0aHdheS1jdHgtYmFzZTY0Ijoic29tZS1iYXNlNjQtZW5jb2RlZC1jb250ZXh0In19Cg==",
               approximateArrivalTimestamp: 1642518727.248,
             },
             eventSource: "aws:kinesis",
@@ -57,7 +79,7 @@ describe("KinesisEventTraceExtractor", () => {
         ],
       };
 
-      const extractor = new KinesisEventTraceExtractor(tracerWrapper);
+      const extractor = new KinesisEventTraceExtractor(tracerWrapper, mockConfig);
 
       const traceContext = extractor.extract(payload);
       expect(traceContext).not.toBeNull();
@@ -67,25 +89,151 @@ describe("KinesisEventTraceExtractor", () => {
         "x-datadog-sampled": "1",
         "x-datadog-sampling-priority": "1",
         "x-datadog-trace-id": "667309514221035538",
+        "dd-pathway-ctx-base64": "some-base64-encoded-context",
       });
 
       expect(traceContext?.toTraceId()).toBe("667309514221035538");
       expect(traceContext?.toSpanId()).toBe("1350735035497811828");
       expect(traceContext?.sampleMode()).toBe("1");
       expect(traceContext?.source).toBe("event");
+
+      expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith(
+        "kinesis",
+        "arn:aws:kinesis:EXAMPLE",
+        {
+          "x-datadog-parent-id": "1350735035497811828",
+          "x-datadog-sampled": "1",
+          "x-datadog-sampling-priority": "1",
+          "x-datadog-trace-id": "667309514221035538",
+          "dd-pathway-ctx-base64": "some-base64-encoded-context",
+        },
+        false,
+      );
     });
 
     it.each([
-      ["Records", {}],
-      ["Records first entry", { Records: [] }],
-      ["valid data in kinesis", { Records: [{ kinesis: { data: "{" } }] }], // JSON.parse should fail
-      ["_datadog in data", { Records: [{ kinesis: { data: "e30=" } }] }],
-    ])("returns null and skips extracting when payload is missing '%s'", (_, payload) => {
+      ["Records", {}, 0],
+      ["Records first entry", { Records: [] }, 0],
+      ["valid data in kinesis", { Records: [{ kinesis: { data: "{" }, eventSourceARN: "arn:aws:kinesis:test" }] }, 1], // JSON.parse should fail
+      ["_datadog in data", { Records: [{ kinesis: { data: "e30=" }, eventSourceARN: "arn:aws:kinesis:test" }] }, 1],
+    ])("returns null and skips extracting when payload is missing '%s'", (_, payload, dsmCalls) => {
       const tracerWrapper = new TracerWrapper();
-      const extractor = new KinesisEventTraceExtractor(tracerWrapper);
+      const extractor = new KinesisEventTraceExtractor(tracerWrapper, mockConfig);
 
       const traceContext = extractor.extract(payload as any);
       expect(traceContext).toBeNull();
+
+      expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledTimes(dsmCalls);
+
+      if (dsmCalls > 0) {
+        expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith(
+          "kinesis",
+          "arn:aws:kinesis:test",
+          null,
+          false,
+        );
+      }
+    });
+
+    it("calls setConsumeCheckpoint for each record", () => {
+      mockSpanContext = {
+        toTraceId: () => "667309514221035538",
+        toSpanId: () => "1350735035497811828",
+        _sampling: {
+          priority: "1",
+        },
+      };
+      const tracerWrapper = new TracerWrapper();
+
+      const makeKinesisRecord = (eventSourceARN: string, hasDatadogHeaders: boolean) => {
+        const baseData = { some: "data" };
+        const dataWithHeaders = hasDatadogHeaders
+          ? {
+              ...baseData,
+              _datadog: {
+                "x-datadog-trace-id": "667309514221035538",
+                "x-datadog-parent-id": "1350735035497811828",
+                "x-datadog-sampled": "1",
+                "x-datadog-sampling-priority": "1",
+                "dd-pathway-ctx-base64": "some-base64-encoded-context",
+              },
+            }
+          : baseData;
+
+        const encodedData = Buffer.from(JSON.stringify(dataWithHeaders)).toString("base64");
+
+        return {
+          kinesis: {
+            kinesisSchemaVersion: "1.0",
+            partitionKey: "test-partition",
+            sequenceNumber: "12345",
+            data: encodedData,
+            approximateArrivalTimestamp: 1642518727.248,
+          },
+          eventSource: "aws:kinesis",
+          eventID: "test-event-id",
+          invokeIdentityArn: "arn:aws:iam::EXAMPLE",
+          eventVersion: "1.0",
+          eventName: "aws:kinesis:record",
+          eventSourceARN: eventSourceARN,
+          awsRegion: "us-east-1",
+        };
+      };
+
+      const payload = {
+        Records: [
+          makeKinesisRecord("arn:aws:kinesis:us-east-1:123456789012:stream/stream-1", true),
+          makeKinesisRecord("arn:aws:kinesis:us-east-1:123456789012:stream/stream-2", false),
+          makeKinesisRecord("arn:aws:kinesis:us-east-1:123456789012:stream/stream-3", true),
+          makeKinesisRecord("arn:aws:kinesis:us-east-1:123456789012:stream/stream-4", false),
+        ],
+      };
+
+      const extractor = new KinesisEventTraceExtractor(tracerWrapper, mockConfig);
+      const traceContext = extractor.extract(payload);
+      expect(traceContext).not.toBeNull();
+
+      expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledTimes(4);
+      expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenNthCalledWith(
+        1,
+        "kinesis",
+        "arn:aws:kinesis:us-east-1:123456789012:stream/stream-1",
+        {
+          "x-datadog-trace-id": "667309514221035538",
+          "x-datadog-parent-id": "1350735035497811828",
+          "x-datadog-sampled": "1",
+          "x-datadog-sampling-priority": "1",
+          "dd-pathway-ctx-base64": "some-base64-encoded-context",
+        },
+        false,
+      );
+      expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenNthCalledWith(
+        2,
+        "kinesis",
+        "arn:aws:kinesis:us-east-1:123456789012:stream/stream-2",
+        null,
+        false,
+      );
+      expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenNthCalledWith(
+        3,
+        "kinesis",
+        "arn:aws:kinesis:us-east-1:123456789012:stream/stream-3",
+        {
+          "x-datadog-trace-id": "667309514221035538",
+          "x-datadog-parent-id": "1350735035497811828",
+          "x-datadog-sampled": "1",
+          "x-datadog-sampling-priority": "1",
+          "dd-pathway-ctx-base64": "some-base64-encoded-context",
+        },
+        false,
+      );
+      expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenNthCalledWith(
+        4,
+        "kinesis",
+        "arn:aws:kinesis:us-east-1:123456789012:stream/stream-4",
+        null,
+        false,
+      );
     });
 
     it("returns null when extracted span context by tracer is null", () => {
@@ -112,10 +260,78 @@ describe("KinesisEventTraceExtractor", () => {
         ],
       };
 
-      const extractor = new KinesisEventTraceExtractor(tracerWrapper);
+      const extractor = new KinesisEventTraceExtractor(tracerWrapper, mockConfig);
 
       const traceContext = extractor.extract(payload);
       expect(traceContext).toBeNull();
     });
+
+    it("extracts trace context when DSM is disabled and does not call setConsumeCheckpoint", () => {
+      mockSpanContext = {
+        toTraceId: () => "667309514221035538",
+        toSpanId: () => "1350735035497811828",
+        _sampling: {
+          priority: "1",
+        },
+      };
+      const tracerWrapper = new TracerWrapper();
+
+      const makeKinesisRecord = (eventSourceARN: string, hasDatadogHeaders: boolean) => {
+        const baseData = { some: "data" };
+        const dataWithHeaders = hasDatadogHeaders
+          ? {
+              ...baseData,
+              _datadog: {
+                "x-datadog-trace-id": "667309514221035538",
+                "x-datadog-parent-id": "1350735035497811828",
+                "x-datadog-sampled": "1",
+                "x-datadog-sampling-priority": "1",
+                "dd-pathway-ctx-base64": "some-base64-encoded-context",
+              },
+            }
+          : baseData;
+
+        const encodedData = Buffer.from(JSON.stringify(dataWithHeaders)).toString("base64");
+
+        return {
+          kinesis: {
+            kinesisSchemaVersion: "1.0",
+            partitionKey: "test-partition",
+            sequenceNumber: "12345",
+            data: encodedData,
+            approximateArrivalTimestamp: 1642518727.248,
+          },
+          eventSource: "aws:kinesis",
+          eventID: "test-event-id",
+          invokeIdentityArn: "arn:aws:iam::EXAMPLE",
+          eventVersion: "1.0",
+          eventName: "aws:kinesis:record",
+          eventSourceARN: eventSourceARN,
+          awsRegion: "us-east-1",
+        };
+      };
+
+      const payload = {
+        Records: [
+          makeKinesisRecord("arn:aws:kinesis:us-east-1:123456789012:stream/stream-1", true),
+          makeKinesisRecord("arn:aws:kinesis:us-east-1:123456789012:stream/stream-2", false),
+          makeKinesisRecord("arn:aws:kinesis:us-east-1:123456789012:stream/stream-3", true),
+        ],
+      };
+
+      const disabledConfig = { ...mockConfig, dataStreamsEnabled: false };
+      const extractor = new KinesisEventTraceExtractor(tracerWrapper, disabledConfig);
+      const traceContext = extractor.extract(payload);
+
+      // Should still extract trace context from first record
+      expect(traceContext).not.toBeNull();
+      expect(traceContext?.toTraceId()).toBe("667309514221035538");
+      expect(traceContext?.toSpanId()).toBe("1350735035497811828");
+      expect(traceContext?.sampleMode()).toBe("1");
+      expect(traceContext?.source).toBe("event");
+
+      // Should NOT call setConsumeCheckpoint when DSM is disabled
+      expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledTimes(0);
+    });
   });
 });