aj/fix lazy imports (#393)

* fix: fix lazily loaded imports

* feat: Fix lazy loaded imports

* fix: format

src/runtime/index.ts

@@ -1,2 +1,2 @@
 export { load, loadSync } from "./user-function";
-export { subscribeToDC, getTraceTree, RequireNode } from "./require-tracer"
+export { subscribeToDC, getTraceTree, clearTraceTree, RequireNode } from "./require-tracer"

src/runtime/require-tracer.ts

@@ -22,7 +22,7 @@ export class RequireNode {
 
 const moduleLoadStartChannel = dc.channel('dd-trace:moduleLoadStart')
 const moduleLoadEndChannel = dc.channel('dd-trace:moduleLoadEnd')
-const rootNodes: RequireNode[] = []
+let rootNodes: RequireNode[] = []
 
 const requireStack: RequireNode[] = []
 const pushNode = (data: any) => {
@@ -56,3 +56,7 @@ export const subscribeToDC = () => {
 export const getTraceTree = (): RequireNode[] => {
   return rootNodes
 }
+
+export const clearTraceTree = () => {
+  rootNodes = []
+}

src/trace/cold-start-tracer.spec.ts

@@ -72,6 +72,7 @@ describe("ColdStartTracer", () => {
       currentSpanStartTime: 500,
       minDuration: 1,
       ignoreLibs: "",
+      isColdStart: true,
     };
     const coldStartTracer = new ColdStartTracer(coldStartConfig);
     coldStartTracer.trace(requireNodes);
@@ -162,6 +163,7 @@ describe("ColdStartTracer", () => {
       currentSpanStartTime: 500,
       minDuration: 1,
       ignoreLibs: "myChildModule,myCoreModule",
+      isColdStart: true,
     };
     const coldStartTracer = new ColdStartTracer(coldStartConfig);
     coldStartTracer.trace(requireNodes);
@@ -194,4 +196,86 @@ describe("ColdStartTracer", () => {
       service: "aws.lambda",
     });
   });
+  it("traces lazy imports", () => {
+    const requireNodes: RequireNode[] = [
+      {
+        id: "handler",
+        filename: "/var/task/handler.js",
+        startTime: 1,
+        endTime: 6,
+        children: [
+          {
+            id: "myChildModule",
+            filename: "/opt/nodejs/node_modules/my-child-module.js",
+            startTime: 2,
+            endTime: 3,
+          },
+          {
+            id: "myCoreModule",
+            filename: "http",
+            startTime: 4,
+            endTime: 5,
+          },
+          {
+            id: "aws-sdk",
+            filename: "/var/runtime/aws-sdk",
+            startTime: 4,
+            endTime: 5,
+          },
+        ],
+      } as any as RequireNode,
+    ];
+    const coldStartConfig: ColdStartTracerConfig = {
+      tracerWrapper: new TracerWrapper(),
+      parentSpan: {
+        span: {},
+        name: "my-parent-span",
+      } as any as SpanWrapper,
+      lambdaFunctionName: "my-function-name",
+      currentSpanStartTime: 500,
+      minDuration: 1,
+      ignoreLibs: "",
+      isColdStart: false,
+    };
+    const coldStartTracer = new ColdStartTracer(coldStartConfig);
+    coldStartTracer.trace(requireNodes);
+    expect(mockStartSpan).toHaveBeenCalledTimes(4);
+    expect(mockFinishSpan).toHaveBeenCalledTimes(4);
+    const span1 = mockStartSpan.mock.calls[0];
+    expect(span1[0]).toEqual("aws.lambda.require");
+    expect(span1[1].tags).toEqual({
+      operation_name: "aws.lambda.require",
+      "resource.name": "handler",
+      resource_names: "handler",
+      service: "aws.lambda",
+      filename: "/var/task/handler.js",
+    });
+    const span2 = mockStartSpan.mock.calls[1];
+    expect(span2[0]).toEqual("aws.lambda.require_layer");
+    expect(span2[1].tags).toEqual({
+      filename: "/opt/nodejs/node_modules/my-child-module.js",
+      operation_name: "aws.lambda.require_layer",
+      "resource.name": "myChildModule",
+      resource_names: "myChildModule",
+      service: "aws.lambda",
+    });
+    const span3 = mockStartSpan.mock.calls[2];
+    expect(span3[0]).toEqual("aws.lambda.require_core_module");
+    expect(span3[1].tags).toEqual({
+      filename: "http",
+      operation_name: "aws.lambda.require_core_module",
+      "resource.name": "myCoreModule",
+      resource_names: "myCoreModule",
+      service: "aws.lambda",
+    });
+    const span4 = mockStartSpan.mock.calls[3];
+    expect(span4[0]).toEqual("aws.lambda.require_runtime");
+    expect(span4[1].tags).toEqual({
+      filename: "/var/runtime/aws-sdk",
+      operation_name: "aws.lambda.require_runtime",
+      "resource.name": "aws-sdk",
+      resource_names: "aws-sdk",
+      service: "aws.lambda",
+    });
+  });
 });

src/trace/cold-start-tracer.ts

@@ -9,6 +9,7 @@ export interface ColdStartTracerConfig {
   currentSpanStartTime: number;
   minDuration: number;
   ignoreLibs: string;
+  isColdStart: boolean;
 }
 
 export class ColdStartTracer {
@@ -18,6 +19,7 @@ export class ColdStartTracer {
   private currentSpanStartTime: number;
   private minDuration: number;
   private ignoreLibs: string[];
+  private isColdStart: boolean;
 
   constructor(coldStartTracerConfig: ColdStartTracerConfig) {
     this.tracerWrapper = coldStartTracerConfig.tracerWrapper;
@@ -26,14 +28,21 @@ export class ColdStartTracer {
     this.currentSpanStartTime = coldStartTracerConfig.currentSpanStartTime;
     this.minDuration = coldStartTracerConfig.minDuration;
     this.ignoreLibs = coldStartTracerConfig.ignoreLibs.split(",");
+    this.isColdStart = coldStartTracerConfig.isColdStart;
   }
 
   trace(rootNodes: RequireNode[]) {
     const coldStartSpanStartTime = rootNodes[0]?.startTime;
     const coldStartSpanEndTime = Math.min(rootNodes[rootNodes.length - 1]?.endTime, this.currentSpanStartTime);
-    const coldStartSpan = this.createColdStartSpan(coldStartSpanStartTime, coldStartSpanEndTime, this.parentSpan);
+    let targetParentSpan: SpanWrapper | undefined;
+    if (this.isColdStart) {
+      const coldStartSpan = this.createColdStartSpan(coldStartSpanStartTime, coldStartSpanEndTime, this.parentSpan);
+      targetParentSpan = coldStartSpan;
+    } else {
+      targetParentSpan = this.parentSpan;
+    }
     for (const coldStartNode of rootNodes) {
-      this.traceTree(coldStartNode, coldStartSpan);
+      this.traceTree(coldStartNode, targetParentSpan);
     }
   }
 

src/trace/listener.ts

@@ -19,7 +19,7 @@ import { patchConsole } from "./patch-console";
 import { SpanContext, TraceOptions, TracerWrapper } from "./tracer-wrapper";
 import { SpanInferrer } from "./span-inferrer";
 import { SpanWrapper } from "./span-wrapper";
-import { getTraceTree } from "../runtime/index";
+import { getTraceTree, clearTraceTree } from "../runtime/index";
 export type TraceExtractor = (event: any, context: Context) => Promise<TraceContext> | TraceContext;
 
 export interface TraceConfig {
@@ -157,17 +157,19 @@ export class TraceListener {
       tagObject(this.tracerWrapper.currentSpan, "function.response", result);
     }
     const coldStartNodes = getTraceTree();
-    if (coldStartNodes.length > 0 && didFunctionColdStart()) {
+    if (coldStartNodes.length > 0) {
       const coldStartConfig: ColdStartTracerConfig = {
         tracerWrapper: this.tracerWrapper,
-        parentSpan: this.inferredSpan || this.wrappedCurrentSpan,
+        parentSpan: didFunctionColdStart() ? this.inferredSpan || this.wrappedCurrentSpan : this.wrappedCurrentSpan,
         lambdaFunctionName: this.context?.functionName,
         currentSpanStartTime: this.wrappedCurrentSpan?.startTime(),
         minDuration: this.config.minColdStartTraceDuration,
         ignoreLibs: this.config.coldStartTraceSkipLib,
+        isColdStart: didFunctionColdStart(),
       };
       const coldStartTracer = new ColdStartTracer(coldStartConfig);
       coldStartTracer.trace(coldStartNodes);
+      clearTraceTree();
     }
     if (this.triggerTags) {
       const statusCode = extractHTTPStatusCodeTag(this.triggerTags, result, isResponseStreamFunction);