fix: omit creating cold start on LMI (#717)

* omit creating cold start on LMI

* fmt

Author: duncanista

src/index.ts

@@ -21,6 +21,7 @@ import {
   setSandboxInit,
   setLogger,
   setLogLevel,
+  isManagedInstancesMode,
 } from "./utils";
 import { getEnhancedMetricTags } from "./metrics/enhanced-metrics";
 import { DatadogTraceHeaders } from "./trace/context/extractor";
@@ -106,7 +107,9 @@ export const _metricsQueue: MetricsQueue = new MetricsQueue();
 let currentMetricsListener: MetricsListener | undefined;
 let currentTraceListener: TraceListener | undefined;
 
-if (getEnvValue(coldStartTracingEnvVar, "true").toLowerCase() === "true") {
+// Skip cold start tracing subscription in managed instances mode
+// In managed instances, the tracer library handles cold start independently
+if (getEnvValue(coldStartTracingEnvVar, "true").toLowerCase() === "true" && !isManagedInstancesMode()) {
   subscribeToDC();
 }
 

src/trace/listener.ts

@@ -5,7 +5,7 @@ import { patchHttp, unpatchHttp } from "./patch-http";
 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";
+import { didFunctionColdStart, isProactiveInitialization, isManagedInstancesMode } from "../utils/cold-start";
 import { datadogLambdaVersion } from "../constants";
 import { ddtraceVersion, parentSpanFinishTimeHeader, DD_SERVICE_ENV_VAR } from "./constants";
 import { patchConsole } from "./patch-console";
@@ -179,20 +179,27 @@ export class TraceListener {
     }
     const coldStartNodes = getTraceTree();
     if (coldStartNodes.length > 0) {
-      const coldStartConfig: ColdStartTracerConfig = {
-        tracerWrapper: this.tracerWrapper,
-        parentSpan:
-          didFunctionColdStart() || isProactiveInitialization()
-            ? this.inferredSpan || this.wrappedCurrentSpan
-            : this.wrappedCurrentSpan,
-        lambdaFunctionName: this.context?.functionName,
-        currentSpanStartTime: this.wrappedCurrentSpan?.startTime(),
-        minDuration: this.config.minColdStartTraceDuration,
-        ignoreLibs: this.config.coldStartTraceSkipLib,
-        isColdStart: didFunctionColdStart() || isProactiveInitialization(),
-      };
-      const coldStartTracer = new ColdStartTracer(coldStartConfig);
-      coldStartTracer.trace(coldStartNodes);
+      // Skip creating cold start spans in managed instances mode
+      // since the gap between the sandbox init and the function
+      // invocation might be too large to provide a useful trace and
+      // experience
+      if (!isManagedInstancesMode()) {
+        const coldStartConfig: ColdStartTracerConfig = {
+          tracerWrapper: this.tracerWrapper,
+          parentSpan:
+            didFunctionColdStart() || isProactiveInitialization()
+              ? this.inferredSpan || this.wrappedCurrentSpan
+              : this.wrappedCurrentSpan,
+          lambdaFunctionName: this.context?.functionName,
+          currentSpanStartTime: this.wrappedCurrentSpan?.startTime(),
+          minDuration: this.config.minColdStartTraceDuration,
+          ignoreLibs: this.config.coldStartTraceSkipLib,
+          isColdStart: didFunctionColdStart() || isProactiveInitialization(),
+        };
+        const coldStartTracer = new ColdStartTracer(coldStartConfig);
+        coldStartTracer.trace(coldStartNodes);
+      }
+      // Always clear the tree to prevent memory leaks, even if we skip span creation
       clearTraceTree();
     }
     if (this.triggerTags) {

src/utils/cold-start.spec.ts

@@ -1,4 +1,10 @@
-import { _resetColdStart, didFunctionColdStart, setSandboxInit, isProactiveInitialization } from "./cold-start";
+import {
+  _resetColdStart,
+  didFunctionColdStart,
+  setSandboxInit,
+  isProactiveInitialization,
+  isManagedInstancesMode,
+} from "./cold-start";
 
 beforeEach(_resetColdStart);
 afterAll(_resetColdStart);
@@ -45,4 +51,31 @@ describe("cold-start", () => {
     expect(didFunctionColdStart()).toEqual(false);
     expect(isProactiveInitialization()).toEqual(false);
   });
+
+  it("identifies managed instances mode when AWS_LAMBDA_INITIALIZATION_TYPE is set", () => {
+    const originalValue = process.env.AWS_LAMBDA_INITIALIZATION_TYPE;
+
+    process.env.AWS_LAMBDA_INITIALIZATION_TYPE = "lambda-managed-instances";
+    expect(isManagedInstancesMode()).toEqual(true);
+
+    process.env.AWS_LAMBDA_INITIALIZATION_TYPE = originalValue;
+  });
+
+  it("identifies non-managed instances mode when AWS_LAMBDA_INITIALIZATION_TYPE is not set", () => {
+    const originalValue = process.env.AWS_LAMBDA_INITIALIZATION_TYPE;
+
+    delete process.env.AWS_LAMBDA_INITIALIZATION_TYPE;
+    expect(isManagedInstancesMode()).toEqual(false);
+
+    process.env.AWS_LAMBDA_INITIALIZATION_TYPE = originalValue;
+  });
+
+  it("identifies non-managed instances mode when AWS_LAMBDA_INITIALIZATION_TYPE has different value", () => {
+    const originalValue = process.env.AWS_LAMBDA_INITIALIZATION_TYPE;
+
+    process.env.AWS_LAMBDA_INITIALIZATION_TYPE = "on-demand";
+    expect(isManagedInstancesMode()).toEqual(false);
+
+    process.env.AWS_LAMBDA_INITIALIZATION_TYPE = originalValue;
+  });
 });

src/utils/cold-start.ts

@@ -37,6 +37,16 @@ export function getSandboxInitTags(): string[] {
   return tags;
 }
 
+/**
+ * Checks if the Lambda function is running in managed instances mode.
+ * In managed instances mode, we should not create cold start tracing spans
+ * as the tracer library handles this independently.
+ * @returns true if running in managed instances mode, false otherwise
+ */
+export function isManagedInstancesMode(): boolean {
+  return process.env.AWS_LAMBDA_INITIALIZATION_TYPE === "lambda-managed-instances";
+}
+
 // For testing, reset the globals to their original values
 export function _resetColdStart() {
   functionDidColdStart = true;

src/utils/index.ts

@@ -1,4 +1,10 @@
-export { didFunctionColdStart, getSandboxInitTags, setSandboxInit, isProactiveInitialization } from "./cold-start";
+export {
+  didFunctionColdStart,
+  getSandboxInitTags,
+  setSandboxInit,
+  isProactiveInitialization,
+  isManagedInstancesMode,
+} from "./cold-start";
 export { promisifiedHandler } from "./handler";
 export { Timer } from "./timer";
 export { logWarning, logError, logDebug, Logger, setLogLevel, setLogger, LogLevel } from "./log";