[w3c] Support W3C trace context propagation (#429)

* refactor `trace-context-service.ts`

aiming to enhance its use, created for xray trace context handling, now we expect it to use it for all of the tracing context

* add `span-context-wrapper.ts`

simple wrapper to use instead of `TraceContext`, this allows us to type the object extracted by the tracer. I decided to wrap it into a class, to easily managea black-box class – yes, I know we are importing it in one of the constructors, but that is the easiest way to create one for trace context with Xray, StepFunction event, and for Custom Extractors, so we can tell the tracer what is a child of what.

* update `tracer-wrapper.ts`

make it complaint to use the refactored classes, also updated its unit tests to mock the new methods being used

* add `step-function-service.ts`

the main point of this class is to create a singleton state manager for the `StepFunctionContext` interface, it will run per invocation, and should not interfere with new invocations

* add `xray-service.ts`

this new class handles all logic related to Xray, previously, it was scattered in `context.ts`, but it actually didnt make sense, since we had multiple functions which should be in the same place, we could also keep a state since we are relying on a header and context which will always be deterministic, so no need to calculate it again in multiple functions

* add `extractor.ts`

this file will be in charge of extracting the trace context from the incoming context, event, step function context, or xray, separating the logic into an adapter pattern for a much simpler and easier way of handling this extraction

* add `http.ts` extractor

refactoring code to make sure we correctly handling an HTTP event for trace context

* add `custom.ts` extractor

allow custom extractor to be in a separate class, which then uses a helper method to generate an `SpanContextWrapper` to be returned, had to deprecate some fields, so to avoid breaking changes, I have marked some fields as nullable

* add `app-sync.ts` extractor

basically relies on `http.ts` to work

* add `sns.ts` extractor

same as other extractors, separation of concerns into its own file

* add `sqs.ts` extractor

same as other extractors, separation of concerns into its own file

* add `kinesis.ts` extractor

same as other extractors, separation of concerns into its own file

* add `sns-sqs.ts` extractor

same as other extractors, separation of concerns into its own file, I think we could do some more refactoring, but for now, I will leave sns-sqs and eventbridge-sqs as separate files

* add `event-bridge.ts` extractor

same as other extractors, separation of concerns into its own file

* add `event-bridge-sqs.ts` extractor

same as other extractors, separation of concerns into its own file, related to sns-sqs comment, maybe in the future we could explore refactoring this extraction, but for now, will stay as is

* add `lambda-context.ts` extractor

same as other extractors, separation of concerns into its own file, also made sure legacy payloads still work

* add `step-function.ts` extractor

same as other extractors, separation of concerns into its own file, uses the singleton class to get context

* add `index.ts` for extractors

basically to export them from the root

* refactor main `index.ts`

updated unit tests to make sure everything works as before with the new methods and mockings

* refactor `patch-http.ts`

w3c will not be supported here, since we would have to do manual extraction and injection, this will not be supported for customers not using `dd-trace`, so the refactoring is just to use the new methods and mockings

* refactor `patch-console.ts`

w3c will not be supported here, since we would have to do manual extraction and injection, this will not be supported for customers not using `dd-trace`, so the refactoring is just to use the new methods and mockings, if a new user wants to correlate traces with logs without `dd-trace`, they will have to use the datadog headers, this might be up for removal in the future

* refactor `listener.ts`

make use of the new services, also refactored unit tests to work as expected

* add `reset` method to `step-function-service.ts`

might need to re-explore this method, but for safety, I am adding it just to make sure we are indeed creating a new instance of the service on every invocation

* update `index.ts` in trace

export from the right file

* add `event-validator.ts`

this class aims to replace `event-type-guards.ts`, I think the file name is confusing, and we should use a class as a namespace here, nonetheless, I wonder if this change is necessary at all, since it would mostly use the same exact logic, but allowing it to be exported as a single class with static methods, this is exactly the same as what we do when importing the whole file.

* remove `context.ts`

decomissioning a large file and its unit tests, made sure to refactor every single method into multiple files and classes

* lint

* format

* when tracer is not present, mock an object which allows us to keep using the `SpanContextWrapper` class as usual

* remove state for `XrayService`, sadly, for some reason we have to generate the context everytime, I tried to remove this by having the `header` and `context` as properties, but somehow, it ends up being `undefined`, therefore the switch back

* get `spanContext` for parenting only when tracer is available

* minor changes to unit test so they can pass

* update function signature for `toString` in `SpanContextWrapper`

* lint

* fix unit tests imports

* rename `TraceHeaders` to `DatadogTraceHeaders`

* feat: Fix export of TraceHeaders, add note about deprecation. Fix relative import

* fix: s/breaking/major/

* feat: lint

---------

Co-authored-by: AJ Stuyvenberg <[email protected]>

src/index.spec.ts

@@ -2,17 +2,14 @@ import http from "http";
 import nock from "nock";
 
 import { Context, Handler } from "aws-lambda";
-import {
-  datadog,
-  getTraceHeaders,
-  sendDistributionMetric,
-  TraceHeaders,
-  sendDistributionMetricWithDate,
-} from "./index";
+import { datadog, getTraceHeaders, sendDistributionMetric, sendDistributionMetricWithDate } from "./index";
 import { incrementErrorsMetric, incrementInvocationsMetric } from "./metrics/enhanced-metrics";
 import { LogLevel, setLogLevel } from "./utils";
 import { HANDLER_STREAMING, STREAM_RESPONSE } from "./constants";
 import { PassThrough } from "stream";
+import { DatadogTraceHeaders } from "./trace/context/extractor";
+import { SpanContextWrapper } from "./trace/span-context-wrapper";
+import { TraceSource } from "./trace/trace-context-service";
 
 jest.mock("./metrics/enhanced-metrics");
 
@@ -28,6 +25,34 @@ const mockContext = {
 //   typeof OriginalListenerModule.MetricsListener
 // >;
 
+let mockSpanContextWrapper: any;
+let mockSpanContext: any;
+let mockTraceHeaders: Record<string, string> | undefined = undefined;
+let mockTraceSource: TraceSource | undefined = undefined;
+
+jest.mock("./trace/trace-context-service", () => {
+  class MockTraceContextService {
+    extract(event: any, context: Context): SpanContextWrapper {
+      return mockSpanContextWrapper;
+    }
+
+    get traceSource() {
+      return mockTraceSource;
+    }
+    get currentTraceContext() {
+      return mockSpanContextWrapper;
+    }
+
+    get currentTraceHeaders() {
+      return mockTraceHeaders;
+    }
+  }
+  return {
+    ...jest.requireActual("./trace/trace-context-service"),
+    TraceContextService: MockTraceContextService,
+  };
+});
+
 describe("datadog", () => {
   let traceId: string | undefined;
   let parentId: string | undefined;
@@ -43,6 +68,9 @@ describe("datadog", () => {
     callback(null, "Result");
   };
   beforeEach(() => {
+    mockTraceHeaders = undefined;
+    mockSpanContext = undefined;
+    mockSpanContextWrapper = undefined;
     traceId = undefined;
     parentId = undefined;
     sampled = undefined;
@@ -60,6 +88,12 @@ describe("datadog", () => {
 
   it("patches http request when autoPatch enabled", async () => {
     nock("http://www.example.com").get("/").reply(200, {});
+    mockTraceHeaders = {
+      "x-datadog-parent-id": "9101112",
+      "x-datadog-sampling-priority": "2",
+      "x-datadog-trace-id": "123456",
+    };
+
     const wrapped = datadog(handler, { forceWrap: true });
     await wrapped(
       {
@@ -189,7 +223,12 @@ describe("datadog", () => {
   });
 
   it("makes the current trace headers available", async () => {
-    let traceHeaders: Partial<TraceHeaders> = {};
+    mockTraceHeaders = {
+      "x-datadog-parent-id": "9101112",
+      "x-datadog-sampling-priority": "2",
+      "x-datadog-trace-id": "123456",
+    };
+    let traceHeaders: Partial<DatadogTraceHeaders> = {};
     const event = {
       headers: {
         "x-datadog-parent-id": "9101112",
@@ -214,6 +253,19 @@ describe("datadog", () => {
   });
 
   it("injects context into console.log messages", async () => {
+    mockSpanContext = {
+      toTraceId: () => "123456",
+      toSpanId: () => "9101112",
+      _sampling: {
+        priority: "2",
+      },
+    };
+    mockSpanContextWrapper = {
+      spanContext: mockSpanContext,
+      toTraceId: () => mockSpanContext.toTraceId(),
+      toSpanId: () => mockSpanContext.toSpanId(),
+    };
+
     const event = {
       headers: {
         "x-datadog-parent-id": "9101112",
@@ -238,6 +290,19 @@ describe("datadog", () => {
   it("injects context into console.log messages with env var", async () => {
     process.env.DD_LOGS_INJECTION = "true";
 
+    mockSpanContext = {
+      toTraceId: () => "123456",
+      toSpanId: () => "9101112",
+      _sampling: {
+        priority: "2",
+      },
+    };
+    mockSpanContextWrapper = {
+      spanContext: mockSpanContext,
+      toTraceId: () => mockSpanContext.toTraceId(),
+      toSpanId: () => mockSpanContext.toSpanId(),
+    };
+
     const event = {
       headers: {
         "x-datadog-parent-id": "9101112",
@@ -333,7 +398,7 @@ describe("datadog", () => {
 
     expect(mockedIncrementInvocations).toBeCalledTimes(1);
     expect(mockedIncrementInvocations).toBeCalledWith(expect.anything(), mockContext);
-    expect(logger.debug).toHaveBeenCalledTimes(11);
+    expect(logger.debug).toHaveBeenCalledTimes(8);
     expect(logger.debug).toHaveBeenLastCalledWith('{"status":"debug","message":"datadog:Unpatching HTTP libraries"}');
   });
 

src/index.ts

@@ -7,7 +7,7 @@ import {
   MetricsConfig,
   MetricsListener,
 } from "./metrics";
-import { TraceConfig, TraceHeaders, TraceListener } from "./trace";
+import { TraceConfig, TraceListener } from "./trace";
 import { subscribeToDC } from "./runtime";
 import {
   logDebug,
@@ -20,9 +20,10 @@ import {
   setLogLevel,
 } from "./utils";
 import { getEnhancedMetricTags } from "./metrics/enhanced-metrics";
+import { DatadogTraceHeaders } from "./trace/context/extractor";
 
-export { TraceHeaders } from "./trace";
-
+// Backwards-compatible export, TODO deprecate in next major
+export { DatadogTraceHeaders as TraceHeaders } from "./trace/context/extractor";
 export const apiKeyEnvVar = "DD_API_KEY";
 export const apiKeyKMSEnvVar = "DD_KMS_API_KEY";
 export const captureLambdaPayloadEnvVar = "DD_CAPTURE_LAMBDA_PAYLOAD";
@@ -278,7 +279,7 @@ export function sendDistributionMetric(name: string, value: number, ...tags: str
 /**
  * Retrieves the Datadog headers for the current trace.
  */
-export function getTraceHeaders(): Partial<TraceHeaders> {
+export function getTraceHeaders(): Partial<DatadogTraceHeaders> {
   if (currentTraceListener === undefined) {
     return {};
   }