From 142439e9d5797e32db06c8b65d2c99bd3d3968b6 Mon Sep 17 00:00:00 2001
From: Nicholas Hulston <nicholashulston@gmail.com>
Date: Thu, 30 Jan 2025 11:16:40 -0500
Subject: [PATCH] Move span pointer to inferred span; add env var to toggle
 span pointers (#607)

---
 src/index.ts               |  7 +++++++
 src/trace/listener.spec.ts |  1 +
 src/trace/listener.ts      | 21 ++++++++++++++++++---
 src/utils/span-pointers.ts |  2 +-
 4 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 80250755..1d401e4f 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -51,6 +51,7 @@ export const coldStartTracingEnvVar = "DD_COLD_START_TRACING";
 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";
 
 interface GlobalConfig {
   /**
@@ -95,6 +96,7 @@ export const defaultConfig: Config = {
   minColdStartTraceDuration: 3,
   coldStartTraceSkipLib: "",
   localTesting: false,
+  addSpanPointers: true,
 } as const;
 
 export const _metricsQueue: MetricsQueue = new MetricsQueue();
@@ -412,6 +414,11 @@ function getConfig(userConfig?: Partial<Config>): Config {
     config.localTesting = result === "true" || result === "1";
   }
 
+  if (userConfig === undefined || userConfig.addSpanPointers === undefined) {
+    const result = getEnvValue(addSpanPointersEnvVar, "true").toLowerCase();
+    config.addSpanPointers = result === "true";
+  }
+
   return config;
 }
 
diff --git a/src/trace/listener.spec.ts b/src/trace/listener.spec.ts
index 3fb7d115..fa192fe6 100644
--- a/src/trace/listener.spec.ts
+++ b/src/trace/listener.spec.ts
@@ -78,6 +78,7 @@ describe("TraceListener", () => {
     injectLogContext: false,
     minColdStartTraceDuration: 3,
     coldStartTraceSkipLib: "",
+    addSpanPointers: true,
   };
   const context = {
     invokedFunctionArn: "arn:aws:lambda:us-east-1:123456789101:function:my-lambda",
diff --git a/src/trace/listener.ts b/src/trace/listener.ts
index c03fb50c..a1531fde 100644
--- a/src/trace/listener.ts
+++ b/src/trace/listener.ts
@@ -69,6 +69,11 @@ export interface TraceConfig {
    * Libraries to ignore from cold start traces
    */
   coldStartTraceSkipLib: string;
+  /**
+   * Whether to enable span pointers
+   * @default true
+   */
+  addSpanPointers: boolean;
 }
 
 export class TraceListener {
@@ -137,7 +142,9 @@ export class TraceListener {
     this.triggerTags = extractTriggerTags(event, context, eventSource);
     this.stepFunctionContext = StepFunctionContextService.instance().context;
 
-    this.spanPointerAttributesList = getSpanPointerAttributes(eventSource, event);
+    if (this.config.addSpanPointers) {
+      this.spanPointerAttributesList = getSpanPointerAttributes(eventSource, event);
+    }
   }
 
   /**
@@ -201,9 +208,17 @@ export class TraceListener {
       }
     }
 
-    if (this.wrappedCurrentSpan && this.spanPointerAttributesList) {
+    let rootSpan = this.inferredSpan;
+    if (!rootSpan) {
+      rootSpan = this.wrappedCurrentSpan;
+    }
+    if (this.spanPointerAttributesList) {
       for (const attributes of this.spanPointerAttributesList) {
-        this.wrappedCurrentSpan.span.addSpanPointer(attributes.kind, attributes.direction, attributes.hash);
+        try {
+          rootSpan.span.addSpanPointer(attributes.kind, attributes.direction, attributes.hash);
+        } catch (e) {
+          logDebug("Failed to add span pointer");
+        }
       }
     }
     return false;
diff --git a/src/utils/span-pointers.ts b/src/utils/span-pointers.ts
index b3ac29e1..5d3a78a4 100644
--- a/src/utils/span-pointers.ts
+++ b/src/utils/span-pointers.ts
@@ -114,7 +114,7 @@ function processDynamoDbEvent(event: any): SpanPointerAttributes[] {
 
     const keys = record.dynamodb?.Keys;
     const eventSourceARN = record.eventSourceARN;
-    const tableName = record.eventSourceARN ? getTableNameFromARN(eventSourceARN) : undefined;
+    const tableName = eventSourceARN ? getTableNameFromARN(eventSourceARN) : undefined;
 
     if (!tableName || !keys) {
       logDebug("Unable to calculate hash because of missing parameters.");