Merge pull request #28 from DataDog/stephenf/add-enhanced-metrics

Add enhanced lambda metrics

Author: swfunc

README.md

@@ -79,6 +79,10 @@ How much logging datadog-lambda-layer-js should do. Set this to "debug" for exte
 
 If you have the Datadog Lambda Log forwarder enabled and are sending custom metrics, set this to true so your metrics will be sent via logs, (instead of being sent at the end of your lambda invocation).
 
+### DD_ENHANCED_METRICS
+
+If you set the value of this variable to "true" then the Lambda layer will increment a Lambda integration metric called `aws.lambda.enhanced.invocations` with each invocation and `aws.lambda.enhanced.errors` if the invocation results in an error. These metrics are tagged with the function name, region, and account, as well as `cold_start:true|false`.
+
 ## Usage
 
 Datadog needs to be able to read headers from the incoming Lambda event.

package.json

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

src/index.spec.ts

@@ -1,9 +1,25 @@
 import http from "http";
 import nock from "nock";
 
+import { Context, Handler } from "aws-lambda";
 import { datadog, getTraceHeaders, sendDistributionMetric, TraceHeaders } from "./index";
+import { incrementErrorsMetric, incrementInvocationsMetric } from "./metrics/enhanced-metrics";
+import { MetricsListener } from "./metrics/listener";
 import { LogLevel, setLogLevel } from "./utils";
-import tracer, { Span } from "dd-trace";
+
+jest.mock("./metrics/enhanced-metrics");
+
+const mockedIncrementErrors = incrementErrorsMetric as jest.Mock<typeof incrementErrorsMetric>;
+const mockedIncrementInvocations = incrementInvocationsMetric as jest.Mock<typeof incrementInvocationsMetric>;
+
+const mockARN = "arn:aws:lambda:us-east-1:123497598159:function:my-test-lambda";
+const mockContext = ({
+  invokedFunctionArn: mockARN,
+} as any) as Context;
+
+// const MockedListener = OriginalListenerModule.MetricsListener as jest.Mocked<
+//   typeof OriginalListenerModule.MetricsListener
+// >;
 
 describe("datadog", () => {
   let traceId: string | undefined;
@@ -27,6 +43,9 @@ describe("datadog", () => {
     oldEnv = process.env;
     process.env = { ...oldEnv };
     nock.cleanAll();
+
+    mockedIncrementErrors.mockClear();
+    mockedIncrementInvocations.mockClear();
   });
   afterEach(() => {
     process.env = oldEnv;
@@ -162,4 +181,70 @@ describe("datadog", () => {
       "x-datadog-trace-id": "123456",
     });
   });
+
+  it("increments invocations for each function call with env var", async () => {
+    process.env.DD_ENHANCED_METRICS = "true";
+    const wrapped = datadog(handler);
+
+    await wrapped({}, mockContext, () => {});
+
+    expect(mockedIncrementInvocations).toBeCalledTimes(1);
+    expect(mockedIncrementInvocations).toBeCalledWith(mockARN);
+
+    await wrapped({}, mockContext, () => {});
+    await wrapped({}, mockContext, () => {});
+    await wrapped({}, mockContext, () => {});
+
+    expect(mockedIncrementInvocations).toBeCalledTimes(4);
+  });
+
+  it("increments errors correctly with env var", async () => {
+    process.env.DD_ENHANCED_METRICS = "true";
+
+    const handlerError: Handler = (event, context, callback) => {
+      throw Error("Some error");
+    };
+
+    const wrappedHandler = datadog(handlerError);
+
+    const result = wrappedHandler({}, mockContext, () => {});
+    await expect(result).rejects.toEqual(Error("Some error"));
+
+    expect(mockedIncrementInvocations).toBeCalledTimes(1);
+    expect(mockedIncrementErrors).toBeCalledTimes(1);
+
+    expect(mockedIncrementInvocations).toBeCalledWith(mockARN);
+    expect(mockedIncrementErrors).toBeCalledWith(mockARN);
+  });
+
+  it("increments errors and invocations with config setting", async () => {
+    const handlerError: Handler = (event, context, callback) => {
+      throw Error("Some error");
+    };
+
+    const wrappedHandler = datadog(handlerError, { enhancedMetrics: true });
+
+    const result = wrappedHandler({}, mockContext, () => {});
+    await expect(result).rejects.toEqual(Error("Some error"));
+
+    expect(mockedIncrementInvocations).toBeCalledTimes(1);
+    expect(mockedIncrementErrors).toBeCalledTimes(1);
+
+    expect(mockedIncrementInvocations).toBeCalledWith(mockARN);
+    expect(mockedIncrementErrors).toBeCalledWith(mockARN);
+  });
+
+  it("doesn't increment enhanced metrics without env var or config", async () => {
+    const handlerError: Handler = (event, context, callback) => {
+      throw Error("Some error");
+    };
+
+    const wrappedHandler = datadog(handlerError);
+
+    const result = wrappedHandler({}, mockContext, () => {});
+    await expect(result).rejects.toEqual(Error("Some error"));
+
+    expect(mockedIncrementInvocations).toBeCalledTimes(0);
+    expect(mockedIncrementErrors).toBeCalledTimes(0);
+  });
 });

src/index.ts

@@ -1,6 +1,12 @@
 import { Handler } from "aws-lambda";
 
-import { KMSService, MetricsConfig, MetricsListener } from "./metrics";
+import {
+  incrementErrorsMetric,
+  incrementInvocationsMetric,
+  KMSService,
+  MetricsConfig,
+  MetricsListener,
+} from "./metrics";
 import { TraceConfig, TraceHeaders, TraceListener } from "./trace";
 import { logError, LogLevel, setLogLevel, wrap } from "./utils";
 
@@ -11,6 +17,7 @@ const apiKeyKMSEnvVar = "DD_KMS_API_KEY";
 const siteURLEnvVar = "DD_SITE";
 const logLevelEnvVar = "DD_LOG_LEVEL";
 const logForwardingEnvVar = "DD_FLUSH_TO_LOG";
+const enhancedMetricsEnvVar = "DD_ENHANCED_METRICS";
 
 const defaultSiteURL = "datadoghq.com";
 
@@ -32,6 +39,7 @@ export const defaultConfig: Config = {
   apiKeyKMS: "",
   autoPatchHTTP: true,
   debugLogging: false,
+  enhancedMetrics: false,
   logForwarding: false,
   shouldRetryMetrics: false,
   siteURL: "",
@@ -72,8 +80,14 @@ export function datadog<TEvent, TResult>(
       for (const listener of listeners) {
         listener.onStartInvocation(event, context);
       }
+      if (finalConfig.enhancedMetrics) {
+        incrementInvocationsMetric(context.invokedFunctionArn);
+      }
     },
-    async () => {
+    async (event, context, error?) => {
+      if (finalConfig.enhancedMetrics && error) {
+        incrementErrorsMetric(context.invokedFunctionArn);
+      }
       // Completion hook, (called once per handler invocation)
       for (const listener of listeners) {
         await listener.onCompleteInvocation();
@@ -142,11 +156,15 @@ function getConfig(userConfig?: Partial<Config>): Config {
     const result = getEnvValue(logForwardingEnvVar, "false").toLowerCase();
     config.logForwarding = result === "true";
   }
+  if (userConfig === undefined || userConfig.enhancedMetrics === undefined) {
+    const result = getEnvValue(enhancedMetricsEnvVar, "false").toLowerCase();
+    config.enhancedMetrics = result === "true";
+  }
 
   return config;
 }
 
-function getEnvValue(key: string, defaultValue: string): string {
+export function getEnvValue(key: string, defaultValue: string): string {
   const val = process.env[key];
   return val !== undefined ? val : defaultValue;
 }

src/metrics/enhanced-metrics.ts

@@ -0,0 +1,16 @@
+import { getEnvValue, sendDistributionMetric } from "../index";
+
+import { parseTagsFromARN } from "../utils/arn";
+import { getColdStartTag } from "../utils/cold-start";
+
+const ENHANCED_LAMBDA_METRICS_NAMESPACE = "aws.lambda.enhanced";
+
+export function incrementInvocationsMetric(functionARN: string): void {
+  const tags = [...parseTagsFromARN(functionARN), getColdStartTag()];
+  sendDistributionMetric(`${ENHANCED_LAMBDA_METRICS_NAMESPACE}.invocations`, 1, ...tags);
+}
+
+export function incrementErrorsMetric(functionARN: string): void {
+  const tags = [...parseTagsFromARN(functionARN), getColdStartTag()];
+  sendDistributionMetric(`${ENHANCED_LAMBDA_METRICS_NAMESPACE}.errors`, 1, ...tags);
+}

src/metrics/index.ts

@@ -1,2 +1,3 @@
 export { MetricsConfig, MetricsListener } from "./listener";
 export { KMSService } from "./kms-service";
+export { incrementErrorsMetric, incrementInvocationsMetric } from "./enhanced-metrics";

src/metrics/listener.spec.ts

@@ -26,6 +26,7 @@ describe("MetricsListener", () => {
     const listener = new MetricsListener(kms as any, {
       apiKey: "api-key",
       apiKeyKMS: "kms-api-key-encrypted",
+      enhancedMetrics: false,
       logForwarding: false,
       shouldRetryMetrics: false,
       siteURL,
@@ -46,6 +47,7 @@ describe("MetricsListener", () => {
     const listener = new MetricsListener(kms as any, {
       apiKey: "",
       apiKeyKMS: "kms-api-key-encrypted",
+      enhancedMetrics: false,
       logForwarding: false,
       shouldRetryMetrics: false,
       siteURL,
@@ -62,6 +64,7 @@ describe("MetricsListener", () => {
     const listener = new MetricsListener(kms as any, {
       apiKey: "",
       apiKeyKMS: "kms-api-key-encrypted",
+      enhancedMetrics: false,
       logForwarding: false,
       shouldRetryMetrics: false,
       siteURL,

src/metrics/listener.ts

@@ -38,6 +38,12 @@ export interface MetricsConfig {
    * @default false
    */
   logForwarding: boolean;
+
+  /**
+   * Whether to increment invocations and errors Lambda integration metrics from this layer.
+   * @default false
+   */
+  enhancedMetrics: boolean;
 }
 
 export class MetricsListener {

src/trace/listener.ts

@@ -5,6 +5,8 @@ import { extractTraceContext } from "./context";
 import { patchHttp, unpatchHttp } from "./patch-http";
 import { TraceContextService } from "./trace-context-service";
 
+import { didFunctionColdStart } from "../utils/cold-start";
+
 export interface TraceConfig {
   /**
    * Whether to automatically patch all outgoing http requests with Datadog's hybrid tracing headers.
@@ -16,7 +18,6 @@ export interface TraceConfig {
 export class TraceListener {
   private contextService = new TraceContextService();
   private context?: Context;
-  private coldstart = true;
 
   public get currentTraceHeaders() {
     return this.contextService.currentTraceHeaders;
@@ -37,7 +38,6 @@ export class TraceListener {
     if (this.config.autoPatchHTTP) {
       unpatchHttp();
     }
-    this.coldstart = false;
   }
 
   public onWrap<T = (...args: any[]) => any>(func: T): T {
@@ -46,7 +46,7 @@ export class TraceListener {
     const options: SpanOptions & TraceOptions = {};
     if (this.context) {
       options.tags = {
-        cold_start: this.coldstart,
+        cold_start: didFunctionColdStart(),
         function_arn: this.context.invokedFunctionArn,
         request_id: this.context.awsRequestId,
         resource_names: this.context.functionName,

src/utils/arn.spec.ts

@@ -0,0 +1,35 @@
+import { parseLambdaARN, parseTagsFromARN } from "./arn";
+
+describe("arn utils", () => {
+  it("parses arn properties", () => {
+    expect(parseLambdaARN("arn:aws:lambda:us-east-1:123497598159:function:my-test-lambda")).toEqual({
+      account_id: "123497598159",
+      functionname: "my-test-lambda",
+      region: "us-east-1",
+    });
+  });
+
+  it("parses arn properties with version alias", () => {
+    expect(parseLambdaARN("arn:aws:lambda:us-east-1:123497598159:function:my-test-lambda:my-version-alias")).toEqual({
+      account_id: "123497598159",
+      functionname: "my-test-lambda",
+      region: "us-east-1",
+    });
+  });
+
+  it("parses arn tags", () => {
+    const parsedTags = parseTagsFromARN("arn:aws:lambda:us-east-1:123497598159:function:my-test-lambda");
+    for (const tag of ["account_id:123497598159", "functionname:my-test-lambda", "region:us-east-1"]) {
+      expect(parsedTags).toContain(tag);
+    }
+  });
+
+  it("parses arn tags with version", () => {
+    const parsedTags = parseTagsFromARN(
+      "arn:aws:lambda:us-east-1:123497598159:function:my-test-lambda:my-version-alias",
+    );
+    for (const tag of ["account_id:123497598159", "functionname:my-test-lambda", "region:us-east-1"]) {
+      expect(parsedTags).toContain(tag);
+    }
+  });
+});