Lazy load metrics listener (#320)

* add lazy loading of MetricsListener

* remove console log

* add test and extensionEnabled method

* add back fileExists method

Co-authored-by: AJ Stuyvenberg <[email protected]>

Author: amymin00

src/index.spec.ts

@@ -11,6 +11,7 @@ import {
 } from "./index";
 import { incrementErrorsMetric, incrementInvocationsMetric } from "./metrics/enhanced-metrics";
 import { LogLevel, setLogLevel } from "./utils";
+import mock from "mock-fs";
 
 jest.mock("./metrics/enhanced-metrics");
 
@@ -319,6 +320,27 @@ describe("datadog", () => {
     expect(mockedIncrementErrors).toBeCalledTimes(0);
   });
 
+  it("doesn't increment enhanced metrics when using extension", async () => {
+    process.env.DD_ENHANCED_METRICS = "false";
+    mock({
+      "/opt/extensions/datadog-agent": Buffer.from([0]),
+    });
+
+    const handlerError: Handler = (event, context, callback) => {
+      throw Error("Some error");
+    };
+
+    const wrappedHandler = datadog(handlerError, { forceWrap: true });
+
+    const result = wrappedHandler({}, mockContext, () => {});
+    await expect(result).rejects.toEqual(Error("Some error"));
+
+    expect(mockedIncrementInvocations).toBeCalledTimes(0);
+    expect(mockedIncrementErrors).toBeCalledTimes(0);
+
+    mock.restore();
+  });
+
   it("use custom logger to log debug messages", async () => {
     const logger = {
       debug: jest.fn(),

src/index.ts

@@ -1,11 +1,7 @@
 import { Context, Handler } from "aws-lambda";
-import {
-  incrementErrorsMetric,
-  incrementInvocationsMetric,
-  KMSService,
-  MetricsConfig,
-  MetricsListener,
-} from "./metrics";
+import { MetricsConfig } from "./metrics";
+import type { MetricsListener as MetricsListenerType } from "./metrics";
+
 import { TraceConfig, TraceHeaders, TraceListener } from "./trace";
 import {
   logDebug,
@@ -17,6 +13,8 @@ import {
   setLogger,
   setLogLevel,
 } from "./utils";
+import { getExtensionPath } from "./utils/extension-path";
+import { existsSync } from "fs";
 
 export { TraceHeaders } from "./trace";
 
@@ -77,8 +75,9 @@ export const defaultConfig: Config = {
   siteURL: "",
 } as const;
 
-let currentMetricsListener: MetricsListener | undefined;
+let currentMetricsListener: MetricsListenerType | undefined;
 let currentTraceListener: TraceListener | undefined;
+let isExtension: boolean;
 
 /**
  * Wraps your AWS lambda handler functions to add tracing/metrics support
@@ -97,7 +96,20 @@ export function datadog<TEvent, TResult>(
   config?: Partial<Config>,
 ): Handler<TEvent, TResult> {
   const finalConfig = getConfig(config);
-  const metricsListener = new MetricsListener(new KMSService(), finalConfig);
+  let metricsListener: MetricsListenerType;
+  let incrementErrorsMetric: any;
+  let incrementInvocationsMetric: any;
+  if (!isExtensionEnabled()) {
+    const {
+      MetricsListener,
+      KMSService,
+      incrementErrorsMetric: errorFunc,
+      incrementInvocationsMetric: invoFunc,
+    } = require("./metrics");
+    incrementErrorsMetric = errorFunc;
+    incrementInvocationsMetric = invoFunc;
+    metricsListener = new MetricsListener(new KMSService(), finalConfig);
+  }
 
   const traceListener = new TraceListener(finalConfig);
 
@@ -122,7 +134,7 @@ export function datadog<TEvent, TResult>(
     try {
       await traceListener.onStartInvocation(event, context);
       await metricsListener.onStartInvocation(event);
-      if (finalConfig.enhancedMetrics) {
+      if (finalConfig.enhancedMetrics && !isExtensionEnabled()) {
         incrementInvocationsMetric(metricsListener, context);
       }
     } catch (err) {
@@ -145,7 +157,7 @@ export function datadog<TEvent, TResult>(
             localResult,
             finalConfig.captureLambdaPayload,
           );
-          if (responseIs5xxError) {
+          if (responseIs5xxError && !isExtensionEnabled()) {
             incrementErrorsMetric(metricsListener, context);
           }
         }
@@ -156,7 +168,7 @@ export function datadog<TEvent, TResult>(
       error = err;
     }
     try {
-      if (didThrow && finalConfig.enhancedMetrics) {
+      if (didThrow && finalConfig.enhancedMetrics && !isExtensionEnabled()) {
         incrementErrorsMetric(metricsListener, context);
       }
       await metricsListener.onCompleteInvocation();
@@ -299,3 +311,13 @@ function getRuntimeTag(): string {
   const version = process.version;
   return `dd_lambda_layer:datadog-node${version}`;
 }
+
+function isExtensionEnabled(): boolean {
+  if (isExtension !== undefined) {
+    return isExtension;
+  }
+
+  const extensionPath = getExtensionPath();
+  isExtension = existsSync(extensionPath);
+  return isExtension;
+}

src/metrics/extension.ts

@@ -1,15 +1,16 @@
 import { URL } from "url";
 import { get, post, logDebug, logError } from "../utils";
+import { getExtensionPath } from "../utils/extension-path";
 import fs from "fs";
 
 export const AGENT_URL = "http://127.0.0.1:8124";
 const HELLO_PATH = "/lambda/hello";
 const FLUSH_PATH = "/lambda/flush";
-const EXTENSION_PATH = "/opt/extensions/datadog-agent";
 const AGENT_TIMEOUT_MS = 100;
 
 export async function isAgentRunning() {
-  const extensionExists = await fileExists(EXTENSION_PATH);
+  const extensionPath = getExtensionPath();
+  const extensionExists = await fileExists(extensionPath);
   if (!extensionExists) {
     logDebug(`Agent isn't present in sandbox`);
     return false;

src/utils/extension-path.ts

@@ -0,0 +1,5 @@
+const EXTENSION_PATH = "/opt/extensions/datadog-agent";
+
+export function getExtensionPath() {
+  return EXTENSION_PATH;
+}