feat: Handle provisioned concurrency and proactive initialization (#389)

* feat: Handle provisioned concurrency and proactive initialization

* feat: lint

* feat: lint

* feat: lint

---------

Co-authored-by: helen278 <helen.li@datadoghq.com>

Author: astuyve

src/index.ts

@@ -15,7 +15,7 @@ import {
   Logger,
   LogLevel,
   promisifiedHandler,
-  setColdStart,
+  setSandboxInit,
   setLogger,
   setLogLevel,
 } from "./utils";
@@ -92,6 +92,8 @@ if (getEnvValue(coldStartTracingEnvVar, "true").toLowerCase() === "true") {
   subscribeToDC();
 }
 
+const initTime = Date.now();
+
 /**
  * Wraps your AWS lambda handler functions to add tracing/metrics support
  * @param handler A lambda handler function.
@@ -131,8 +133,8 @@ export function datadog<TEvent, TResult>(
   let wrappedFunc: any;
   wrappedFunc = async (...args: any[]) => {
     const { event, context, responseStream } = extractArgs(isResponseStreamFunction, ...args);
-    setColdStart();
     const startTime = new Date();
+    setSandboxInit(initTime, startTime.getTime());
 
     currentMetricsListener = metricsListener;
     currentTraceListener = traceListener;

src/metrics/enhanced-metrics.ts

@@ -3,7 +3,7 @@ import { sendDistributionMetric } from "../index";
 
 import { Context } from "aws-lambda";
 import { parseTagsFromARN } from "../utils/arn";
-import { getColdStartTag } from "../utils/cold-start";
+import { getSandboxInitTags } from "../utils/cold-start";
 import { getProcessVersion } from "../utils/process-version";
 import { writeMetricToStdout } from "./metric-log";
 import { MetricsListener } from "./listener";
@@ -54,7 +54,7 @@ export function getEnhancedMetricTags(context: Context): string[] {
   if (context.invokedFunctionArn) {
     arnTags = parseTagsFromARN(context.invokedFunctionArn, context.functionVersion);
   }
-  const tags = [...arnTags, getColdStartTag(), `memorysize:${context.memoryLimitInMB}`, getVersionTag()];
+  const tags = [...arnTags, ...getSandboxInitTags(), `memorysize:${context.memoryLimitInMB}`, getVersionTag()];
 
   const runtimeTag = getRuntimeTag();
   if (runtimeTag) {

src/trace/listener.ts

@@ -12,7 +12,7 @@ import { extractTriggerTags, extractHTTPStatusCodeTag } from "./trigger";
 import { ColdStartTracerConfig, ColdStartTracer } from "./cold-start-tracer";
 
 import { logDebug, tagObject } from "../utils";
-import { didFunctionColdStart } from "../utils/cold-start";
+import { didFunctionColdStart, isProactiveInitialization } from "../utils/cold-start";
 import { datadogLambdaVersion } from "../constants";
 import { Source, ddtraceVersion, parentSpanFinishTimeHeader, authorizingRequestIdHeader } from "./constants";
 import { patchConsole } from "./patch-console";
@@ -160,12 +160,15 @@ export class TraceListener {
     if (coldStartNodes.length > 0) {
       const coldStartConfig: ColdStartTracerConfig = {
         tracerWrapper: this.tracerWrapper,
-        parentSpan: didFunctionColdStart() ? this.inferredSpan || this.wrappedCurrentSpan : this.wrappedCurrentSpan,
+        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(),
+        isColdStart: didFunctionColdStart() || isProactiveInitialization(),
       };
       const coldStartTracer = new ColdStartTracer(coldStartConfig);
       coldStartTracer.trace(coldStartNodes);
@@ -257,6 +260,9 @@ export class TraceListener {
         datadog_lambda: datadogLambdaVersion,
         dd_trace: ddtraceVersion,
       };
+      if (isProactiveInitialization()) {
+        options.tags.proactive_initialization = true;
+      }
       if (
         (this.contextService.traceSource === Source.Xray && this.config.mergeDatadogXrayTraces) ||
         this.contextService.traceSource === Source.Event

src/utils/cold-start.spec.ts

@@ -1,22 +1,48 @@
-import { _resetColdStart, didFunctionColdStart, setColdStart } from "./cold-start";
+import { _resetColdStart, didFunctionColdStart, setSandboxInit, isProactiveInitialization } from "./cold-start";
 
 beforeEach(_resetColdStart);
 afterAll(_resetColdStart);
 
 describe("cold-start", () => {
   it("identifies cold starts on the first execution", () => {
-    setColdStart();
+    setSandboxInit(0, 1);
     expect(didFunctionColdStart()).toEqual(true);
   });
 
   it("identifies non-cold starts on subsequent executions", () => {
-    setColdStart();
+    setSandboxInit(0, 1);
     expect(didFunctionColdStart()).toEqual(true);
 
-    setColdStart();
+    setSandboxInit(0, 1);
     expect(didFunctionColdStart()).toEqual(false);
 
-    setColdStart();
+    setSandboxInit(0, 1);
     expect(didFunctionColdStart()).toEqual(false);
   });
+
+  it("identifies proactive invocations on the first execution", () => {
+    setSandboxInit(0, 100000);
+    expect(didFunctionColdStart()).toEqual(false);
+    expect(isProactiveInitialization()).toEqual(true);
+
+    setSandboxInit(0, 1);
+    expect(didFunctionColdStart()).toEqual(false);
+
+    setSandboxInit(0, 1);
+    expect(didFunctionColdStart()).toEqual(false);
+  });
+
+  it("identifies non-proactive invocations on subsequent invocations", () => {
+    setSandboxInit(0, 100000);
+    expect(didFunctionColdStart()).toEqual(false);
+    expect(isProactiveInitialization()).toEqual(true);
+
+    setSandboxInit(0, 100000);
+    expect(didFunctionColdStart()).toEqual(false);
+    expect(isProactiveInitialization()).toEqual(false);
+
+    setSandboxInit(0, 100000);
+    expect(didFunctionColdStart()).toEqual(false);
+    expect(isProactiveInitialization()).toEqual(false);
+  });
 });

src/utils/cold-start.ts

@@ -1,27 +1,45 @@
 let functionDidColdStart = true;
+let proactiveInitialization = false;
 
 let isColdStartSet = false;
 
 /**
  * Use global variables to determine whether the container cold started
+ * and if the start was proactively initialized
  * On the first container run, isColdStartSet and functionDidColdStart are true
  * For subsequent executions isColdStartSet will be true and functionDidColdStart will be false
  */
-export function setColdStart() {
-  functionDidColdStart = !isColdStartSet;
+export function setSandboxInit(initTime: number, invocationStartTime: number) {
+  if (!isColdStartSet && invocationStartTime - initTime > 10_000) {
+    proactiveInitialization = true;
+    functionDidColdStart = false;
+  } else {
+    functionDidColdStart = !isColdStartSet;
+    proactiveInitialization = false;
+  }
   isColdStartSet = true;
 }
 
-export function didFunctionColdStart() {
+export function didFunctionColdStart(): boolean {
   return functionDidColdStart;
 }
 
-export function getColdStartTag() {
-  return `cold_start:${didFunctionColdStart()}`;
+export function isProactiveInitialization(): boolean {
+  return proactiveInitialization;
+}
+
+export function getSandboxInitTags(): string[] {
+  const tags = [`cold_start:${didFunctionColdStart()}`];
+  if (isProactiveInitialization()) {
+    tags.push("proactive_initialization:true");
+  }
+
+  return tags;
 }
 
 // For testing, reset the globals to their original values
 export function _resetColdStart() {
   functionDidColdStart = true;
+  proactiveInitialization = false;
   isColdStartSet = false;
 }

src/utils/index.ts

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