Inferred spans for Lambda function URLs (#247)

* Savepoint. Add trigger and event type logic. Start working out how to infer the span correctly

* feat: Create and finish inferred span. Handle parenting. Ensure completed span on invocation finish

* feat: Fix tests and guard clauses

* tests: Trigger specs added and working

* feat: Fix http method on createInferredSpanForLambdaUrl. Add test for SpanInferrer

* feat: Rename inferred-spans to span-inferrer, as it exports a class & follows command pattern. Update imports

* feat: Use the magical resource.name and span.type to correctly set the attributes for JS

* Switch to truthy case statement

Co-authored-by: Harvinder Ghotra <[email protected]>

Co-authored-by: chris.agocs <[email protected]>
Co-authored-by: Harvinder Ghotra <[email protected]>

Author: 3 people

event_samples/lambda-function-urls.json

@@ -0,0 +1,46 @@
+{
+  "version": "2.0",
+  "routeKey": "$default",
+  "rawPath": "/",
+  "rawQueryString": "",
+  "headers": {
+    "sec-fetch-mode": "navigate",
+    "sec-fetch-site": "none",
+    "accept-language": "en-US,en;q=0.9",
+    "x-forwarded-proto": "https",
+    "x-forwarded-port": "443",
+    "x-forwarded-for": "71.195.30.42",
+    "sec-fetch-user": "?1",
+    "pragma": "no-cache",
+    "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9",
+    "sec-ch-ua": "\"Google Chrome\";v=\"95\", \"Chromium\";v=\"95\", \";Not A Brand\";v=\"99\"",
+    "sec-ch-ua-mobile": "?0",
+    "x-amzn-trace-id": "Root=1-61953929-1ec00c3011062a48477b169e",
+    "sec-ch-ua-platform": "\"macOS\"",
+    "host": "a8hyhsshac.lambda-url.eu-south-1.amazonaws.com",
+    "upgrade-insecure-requests": "1",
+    "cache-control": "no-cache",
+    "accept-encoding": "gzip, deflate, br",
+    "sec-fetch-dest": "document",
+    "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/95.0.4638.69 Safari/537.36"
+  },
+  "requestContext": {
+    "accountId": "601427279990",
+    "apiId": "a8hyhsshac",
+    "domainName": "a8hyhsshac.lambda-url.eu-south-1.amazonaws.com",
+    "domainPrefix": "a8hyhsshac",
+    "http": {
+      "method": "GET",
+      "path": "/",
+      "protocol": "HTTP/1.1",
+      "sourceIp": "71.195.30.42",
+      "userAgent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/95.0.4638.69 Safari/537.36"
+    },
+    "requestId": "ec4d58f8-2b8b-4ceb-a1d5-2be7bff58505",
+    "routeKey": "$default",
+    "stage": "$default",
+    "time": "17/Nov/2021:17:17:29 +0000",
+    "timeEpoch": 1637169449721
+  },
+  "isBase64Encoded": false
+}

src/index.spec.ts

@@ -331,7 +331,7 @@ describe("datadog", () => {
 
     expect(mockedIncrementInvocations).toBeCalledTimes(1);
     expect(mockedIncrementInvocations).toBeCalledWith(expect.anything(), mockContext);
-    expect(logger.debug).toHaveBeenCalledTimes(9);
+    expect(logger.debug).toHaveBeenCalledTimes(10);
     expect(logger.debug).toHaveBeenLastCalledWith('{"status":"debug","message":"datadog:Unpatching HTTP libraries"}');
   });
 });

src/index.ts

@@ -150,6 +150,9 @@ export function datadog<TEvent, TResult>(
               if (traceListener.currentSpan) {
                 traceListener.currentSpan.setTag("http.status_code", statusCode);
               }
+              if (traceListener.currentInferredSpan) {
+                traceListener.currentInferredSpan.setTag("http.status_code", statusCode);
+              }
             }
           }
         }

src/trace/listener.ts

@@ -17,6 +17,7 @@ import { datadogLambdaVersion } from "../constants";
 import { Source, ddtraceVersion } from "./constants";
 import { patchConsole } from "./patch-console";
 import { SpanContext, TraceOptions, TracerWrapper } from "./tracer-wrapper";
+import { SpanInferrer } from "./span-inferrer";
 
 export type TraceExtractor = (event: any, context: Context) => TraceContext;
 
@@ -50,6 +51,8 @@ export class TraceListener {
   private context?: Context;
   private stepFunctionContext?: StepFunctionContext;
   private tracerWrapper: TracerWrapper;
+  private inferrer: SpanInferrer;
+  private inferredSpan: any;
 
   public triggerTags?: { [key: string]: string };
   public get currentTraceHeaders() {
@@ -58,9 +61,13 @@ export class TraceListener {
   public get currentSpan() {
     return this.tracerWrapper.currentSpan;
   }
+  public get currentInferredSpan() {
+    return this.inferredSpan;
+  }
   constructor(private config: TraceConfig, private handlerName: string) {
     this.tracerWrapper = new TracerWrapper();
     this.contextService = new TraceContextService(this.tracerWrapper);
+    this.inferrer = new SpanInferrer(this.tracerWrapper);
   }
 
   public onStartInvocation(event: any, context: Context) {
@@ -79,7 +86,8 @@ export class TraceListener {
     } else {
       logDebug("Not patching HTTP libraries", { autoPatchHTTP: this.config.autoPatchHTTP, tracerInitialized });
     }
-
+    logDebug("Creating inferred span");
+    this.inferredSpan = this.inferrer.createInferredSpan(event, context);
     this.context = context;
     this.triggerTags = extractTriggerTags(event, context);
     this.contextService.rootTraceContext = extractTraceContext(event, context, this.config.traceExtractor);
@@ -98,6 +106,10 @@ export class TraceListener {
       logDebug("Unpatching HTTP libraries");
       unpatchHttp();
     }
+    if (this.inferredSpan) {
+      logDebug("Finishing inferred span");
+      this.inferredSpan.finish(Date.now());
+    }
   }
 
   public onWrap<T = (...args: any[]) => any>(func: T): T {
@@ -147,8 +159,12 @@ export class TraceListener {
         ...this.stepFunctionContext,
       };
     }
-
-    if (parentSpanContext !== null) {
+    if (this.inferredSpan) {
+      options.childOf = this.inferredSpan;
+      if (parentSpanContext !== null) {
+        this.inferredSpan.childOf = parentSpanContext;
+      }
+    } else if (parentSpanContext !== null) {
       options.childOf = parentSpanContext;
     }
     options.type = "serverless";

src/trace/span-inferrer.spec.ts

@@ -0,0 +1,31 @@
+import { SpanInferrer } from "./span-inferrer";
+import { TracerWrapper } from "./tracer-wrapper";
+const lambdaURLEvent = require("../../event_samples/lambda-function-urls.json");
+
+const mockWrapper = {
+  startSpan: jest.fn(),
+};
+
+describe("SpanInferrer", () => {
+  it("creates an inferred span for lambda function URLs", () => {
+    const inferrer = new SpanInferrer(mockWrapper as unknown as TracerWrapper);
+    inferrer.createInferredSpan(lambdaURLEvent, {
+      awsRequestId: "abcd-1234",
+    } as any);
+
+    expect(mockWrapper.startSpan).toBeCalledWith("aws.lambda.url", {
+      service: "aws.lambda",
+      startTime: 1637169449721,
+      tags: {
+        endpoint: "/",
+        "http.method": "GET",
+        "http.url": "a8hyhsshac.lambda-url.eu-south-1.amazonaws.com/",
+        operation_name: "aws.lambda.url",
+        request_id: "abcd-1234",
+        "resource.name": "a8hyhsshac.lambda-url.eu-south-1.amazonaws.com/",
+        resource_names: "a8hyhsshac.lambda-url.eu-south-1.amazonaws.com/",
+        "span.type": "http",
+      },
+    });
+  });
+});

src/trace/span-inferrer.ts

@@ -0,0 +1,37 @@
+import { Context } from "aws-lambda";
+import { SpanOptions, TracerWrapper } from "./tracer-wrapper";
+import { eventSources, parseEventSource } from "./trigger";
+
+export class SpanInferrer {
+  traceWrapper: TracerWrapper;
+  constructor(traceWrapper: TracerWrapper) {
+    this.traceWrapper = traceWrapper;
+  }
+
+  public createInferredSpan(event: any, context: Context | undefined): any {
+    const eventSource = parseEventSource(event);
+    if (eventSource === eventSources.lambdaUrl) {
+      return this.createInferredSpanForLambdaUrl(event, context);
+    }
+  }
+
+  createInferredSpanForLambdaUrl(event: any, context: Context | undefined): any {
+    const options: SpanOptions = {};
+    const domain = event.requestContext.domainName;
+    const path = event.rawPath;
+    options.tags = {
+      operation_name: "aws.lambda.url",
+      "http.url": domain + path,
+      endpoint: path,
+      "http.method": event.requestContext.http.method,
+      resource_names: domain + path,
+      request_id: context?.awsRequestId,
+      "span.type": "http",
+      "resource.name": domain + path,
+    };
+    options.service = "aws.lambda";
+    options.startTime = event.requestContext.timeEpoch;
+
+    return this.traceWrapper.startSpan("aws.lambda.url", options);
+  }
+}

src/trace/tracer-wrapper.ts

@@ -16,6 +16,14 @@ export interface TraceOptions {
   childOf?: SpanContext;
 }
 
+export interface SpanOptions {
+  childOf?: SpanContext;
+  tags?: { [key: string]: any };
+  startTime?: number;
+  service?: string;
+  type?: string;
+}
+
 // TraceWrapper is used to remove dd-trace as a hard dependency from the npm package.
 // This lets a customer bring their own version of the tracer.
 export class TracerWrapper {
@@ -61,6 +69,13 @@ export class TracerWrapper {
     return this.tracer.wrap(name, options, fn);
   }
 
+  public startSpan<T = (...args: any[]) => any>(name: string, options: TraceOptions): T | null {
+    if (!this.isTracerAvailable) {
+      return null;
+    }
+    return this.tracer.startSpan(name, options);
+  }
+
   public traceContext(): TraceContext | undefined {
     if (!this.isTracerAvailable) {
       return;

src/trace/trigger.spec.ts

@@ -30,6 +30,15 @@ describe("parseEventSource", () => {
       },
       file: "api-gateway-v2.json",
     },
+    {
+      result: {
+        "function_trigger.event_source": "lambda-function-url",
+        "http.url": "a8hyhsshac.lambda-url.eu-south-1.amazonaws.com",
+        "http.url_details.path": "/",
+        "http.method": "GET",
+      },
+      file: "lambda-function-urls.json",
+    },
     {
       result: {
         "function_trigger.event_source": "application-load-balancer",
@@ -146,7 +155,14 @@ describe("parseEventSource", () => {
       for (let response of responses) {
         const statusCode = extractHTTPStatusCodeTag(triggerTags, response.responseBody);
         // We should always return a status code for API Gateway and ALB
-        if (["api-gateway-v1.json", "api-gateway-v2.json", "application-load-balancer.json"].includes(event.file)) {
+        if (
+          [
+            "api-gateway-v1.json",
+            "api-gateway-v2.json",
+            "application-load-balancer.json",
+            "lambda-function-urls.json",
+          ].includes(event.file)
+        ) {
           expect(statusCode).toEqual(response.expectedStatusCode);
         } else {
           expect(statusCode).toBeUndefined();