Merge pull request #33 from DataDog/darcy.rayner/fix-timeout-no-return

Darcy.rayner/fix timeout no return

Author: DarcyRaynerDD

package.json

@@ -1,6 +1,6 @@
 {
   "name": "datadog-lambda-js",
-  "version": "0.7.0",
+  "version": "0.8.0",
   "description": "Lambda client library that supports hybrid tracing in node js",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",

src/metrics/listener.spec.ts

@@ -2,6 +2,7 @@ import { AWSError, KMS, Request } from "aws-sdk";
 import nock from "nock";
 
 import { LogLevel, setLogLevel } from "../utils";
+
 import { MetricsListener } from "./listener";
 
 const siteURL = "example.com";
@@ -17,6 +18,10 @@ class MockKMS {
 setLogLevel(LogLevel.NONE);
 
 describe("MetricsListener", () => {
+  afterEach(() => {
+    jest.restoreAllMocks();
+  });
+
   it("uses unencrypted api key by default", async () => {
     nock("https://api.example.com")
       .post("/api/v1/distribution_points?api_key=api-key")
@@ -74,4 +79,25 @@ describe("MetricsListener", () => {
     listener.sendDistributionMetric("my-metric", 10, "tag:a", "tag:b");
     await expect(listener.onCompleteInvocation()).resolves.toEqual(undefined);
   });
+
+  it("logs metrics when logForwarding is enabled", async () => {
+    const spy = jest.spyOn(process.stdout, "write");
+    jest.spyOn(Date, "now").mockImplementation(() => 1487076708000);
+    const kms = new MockKMS("kms-api-key-decrypted");
+    const listener = new MetricsListener(kms as any, {
+      apiKey: "api-key",
+      apiKeyKMS: "kms-api-key-encrypted",
+      enhancedMetrics: false,
+      logForwarding: true,
+      shouldRetryMetrics: false,
+      siteURL,
+    });
+    jest.useFakeTimers();
+
+    listener.onStartInvocation({});
+    listener.sendDistributionMetric("my-metric", 10, "tag:a", "tag:b");
+    await listener.onCompleteInvocation();
+
+    expect(spy).toHaveBeenCalledWith(`{"e":1487076708000,"m":"my-metric","t":["tag:a","tag:b"],"v":10}\n`);
+  });
 });

src/metrics/listener.ts

@@ -1,6 +1,6 @@
 import { promisify } from "util";
 
-import { logError } from "../utils";
+import { logDebug, logError } from "../utils";
 import { APIClient } from "./api";
 import { KMSService } from "./kms-service";
 import { Distribution } from "./model";
@@ -55,6 +55,9 @@ export class MetricsListener {
   }
 
   public onStartInvocation(_: any) {
+    if (this.config.logForwarding) {
+      return;
+    }
     this.currentProcessor = this.createProcessor(this.config, this.apiKey);
   }
 
@@ -124,8 +127,13 @@ export class MetricsListener {
       } catch (error) {
         logError("couldn't decrypt kms api key", { innerError: error });
       }
-    } else {
-      logError("api key not configured");
+    } else if (!config.logForwarding) {
+      const errorMessage = "api key not configured, see https://dtdg.co/sls-node-metrics";
+      if (config.logForwarding) {
+        logDebug(errorMessage);
+      } else {
+        logError(errorMessage);
+      }
     }
     return "";
   }

src/utils/handler.spec.ts

@@ -187,4 +187,42 @@ describe("wrap", () => {
     expect(calledComplete).toBeTruthy();
     expect(calledOriginalHandler).toBeFalsy();
   });
+
+  it("returns the first result to complete between the callback and the handler promise", async () => {
+    const handler: Handler = async (event, context, callback) => {
+      callback(null, { statusCode: 204, body: "The callback response" });
+      return { statusCode: 200, body: "The promise response" };
+    };
+
+    let calledOriginalHandler = false;
+
+    const wrappedHandler = wrap(handler, () => {}, async () => {});
+
+    const result = await wrappedHandler({}, mockContext, () => {
+      calledOriginalHandler = true;
+    });
+
+    expect(result).toEqual({ statusCode: 204, body: "The callback response" });
+    expect(calledOriginalHandler).toBeFalsy();
+  });
+
+  it("doesn't complete using non-promise return values", async () => {
+    const handler: Handler = (event, context, callback) => {
+      setTimeout(() => {
+        callback(null, { statusCode: 204, body: "The callback response" });
+      }, 10);
+      return ({ statusCode: 200, body: "The promise response" } as unknown) as void;
+    };
+
+    let calledOriginalHandler = false;
+
+    const wrappedHandler = wrap(handler, () => {}, async () => {});
+
+    const result = await wrappedHandler({}, mockContext, () => {
+      calledOriginalHandler = true;
+    });
+
+    expect(result).toEqual({ statusCode: 204, body: "The callback response" });
+    expect(calledOriginalHandler).toBeFalsy();
+  });
 });

src/utils/handler.ts

@@ -23,10 +23,8 @@ export function wrap<TEvent, TResult>(
       logError("Pre-lambda hook threw error", { innerError: error });
     }
     let result: TResult;
-    // Need to disable linter rule to explicitly assign the variable, otherwise TS
-    // won't reccognize that the var may be assigned in the catch block
-    // tslint:disable-next-line: no-unnecessary-initializer
-    let handlerError: Error | undefined = undefined;
+
+    let handlerError: Error | undefined;
 
     try {
       const wrappedHandler = onWrap !== undefined ? onWrap(promHandler) : promHandler;
@@ -69,9 +67,11 @@ export function promisifiedHandler<TEvent, TResult>(handler: Handler<TEvent, TRe
       };
     });
 
-    let promise = handler(event, context, modifiedCallback) as Promise<TResult> | undefined;
-    if (promise === undefined) {
-      promise = callbackProm;
+    const asyncProm = handler(event, context, modifiedCallback) as Promise<TResult> | undefined;
+    let promise: Promise<TResult> = callbackProm;
+    if (asyncProm !== undefined && typeof asyncProm.then === "function") {
+      // Mimics behaviour of lambda runtime, the first method of returning a result always wins.
+      promise = Promise.race([callbackProm, asyncProm]);
     }
     return promise;
   };