feat(createLambdaExtensionClient): Allow metrics propagation via Data…

…dog Lambda Extension (#196) * feat(createLambdaExtensionClient): Allow metrics propagation via Datadog Lambda Extension * Explicitly remove `decrement` from `createLambdaExtensionClient` * Cleanup and provide apt description. * Close up test coverage * Add README entry * Apply suggestions from code review Co-authored-by: Ryan Ling <[email protected]> * Drop README as it's autogenerated by GitHub * Export `createLambdaExtensionClient` * Move `@types/aws-lambda` into deps to fix TS type checking * Add migration notice and add test for `createLambdaExtensionClient` export * Fix README * Update package.json Co-authored-by: Ryan Ling <[email protected]> * Fix up `yarn.lock` --------- Co-authored-by: Ryan Ling <[email protected]>
seek-oss · Mar 20, 2023 · 059b25e · 059b25e
1 parent 8f81026
commit 059b25e
Show file tree

Hide file tree

Showing 10 changed files with 355 additions and 16 deletions.
diff --git a/README.md b/README.md
@@ -12,16 +12,6 @@ Helpers for sending [Datadog custom metrics](https://docs.datadoghq.com/develope
 yarn add seek-datadog-custom-metrics
 ```
 
-## Table of contents
-
-- [Tagging convention](#tagging-convention)
-- [API reference](#api-reference)
-  - [createStatsDClient](#createstatsdclient)
-  - [createNoOpClient](#createnoopclient)
-  - [createTimedSpan](#createtimedspan)
-  - [httpTracingConfig](#httptracingconfig)
-- [Contributing](https://github.com/seek-oss/datadog-custom-metrics/blob/master/CONTRIBUTING.md)
-
 ## Tagging convention
 
 All custom metrics are prefixed by `AppConfig.name`.
@@ -56,6 +46,38 @@ const errorHandler = (err: Error) => {
 const metricsClient = createStatsDClient(StatsD, config, errorHandler);
 ```
 
+### `createLambdaExtensionClient`
+
+`createLambdaExtensionClient` creates a [Lambda extension](https://docs.datadoghq.com/serverless/libraries_integrations/extension/) client.
+This is intended for AWS Lambda functions and is a replacement for `createCloudWatchClient`.
+
+This client will only submit metrics as a [distribution](https://docs.datadoghq.com/metrics/distributions/) which enables globally accurate aggregations for percentiles (p50, p75, p90, etc).
+
+```typescript
+import { createLambdaExtensionClient } from 'seek-datadog-custom-metrics';
+
+// Expects `name` and `metrics` properties
+import config from '../config';
+
+// Returns a standard hot-shots StatsD instance
+const { metricsClient, withLambdaExtension } =
+  createLambdaExtensionClient(config);
+
+export const handler = withLambdaExtension((event, _ctx) => {
+  try {
+    logger.info('request');
+
+    await lambdaFunction(event);
+  } catch (err) {
+    logger.error({ err }, 'request');
+
+    metricsClient.increment('invocation_error');
+
+    throw new Error('invoke error');
+  }
+});
+```
+
 ### `createNoOpClient`
 
 `createNoOpClient` returns a no-op client.

diff --git a/package.json b/package.json
@@ -30,17 +30,24 @@
     "test:ci": "skuba test --coverage",
     "test:watch": "skuba test --watch"
   },
-  "dependencies": {},
+  "dependencies": {
+    "@types/aws-lambda": "^8.10.108"
+  },
   "devDependencies": {
     "@types/node": "16.18.12",
+    "datadog-lambda-js": "6.84.0",
     "dd-trace": "3.9.3",
     "hot-shots": "9.3.0",
     "skuba": "5.1.1"
   },
   "peerDependencies": {
+    "datadog-lambda-js": "6.x",
     "hot-shots": "6.x || 7.x || 8.x || 9.x"
   },
   "peerDependenciesMeta": {
+    "datadog-lambda-js": {
+      "optional": true
+    },
     "hot-shots": {
       "optional": true
     }

diff --git a/src/LambdaExtensionMetricsClient.ts b/src/LambdaExtensionMetricsClient.ts
@@ -0,0 +1,39 @@
+/**
+ * Abstract interface for recording metrics for AWS Lambda
+ */
+export interface LambdaExtensionMetricsClient {
+  /**
+   * Records a time in milliseconds
+   *
+   * @param name  - Name of the metric to record.
+   * @param value - Time in milliseconds. Fractional values are supported.
+   * @param tags  - Optional list of tags for the metric.
+   */
+  timing(name: string, value: number, tags?: string[]): void;
+
+  /**
+   * Measures the statistical distribution of a set of values
+   *
+   * @param name  - Name of the metric to record.
+   * @param value - Value to include in the statistical distribution.
+   * @param tags  - Optional list of tags for the metric.
+   */
+  distribution(name: string, value: number, tags?: string[]): void;
+
+  /**
+   * Increments a counter by one
+   *
+   * @param name  - Name of the metric to increment.
+   * @param tags  - Optional list of tags for the metric.
+   */
+  increment(name: string, tags?: string[]): void;
+
+  /**
+   * Increments a counter by the specified integer count
+   *
+   * @param name  - Name of the metric to increment.
+   * @param count - Number to increment the counter by.
+   * @param tags  - Optional list of tags for the metric.
+   */
+  increment(name: string, count: number, tags?: string[]): void;
+}
diff --git a/src/createCloudWatchClient.ts b/src/createCloudWatchClient.ts
@@ -18,8 +18,7 @@ const sanitiseTag = (tag: string): string => tag.replace(/\||@|,/g, '_');
  * Creates a new CloudWatch Datadog client configured for the given app
  *
  * @deprecated This depends on Datadog's deprecated CloudWatch log integration.
- * This has been superseded by the Datadog Lambda Extension which does not
- * support the `count` metric type required by our `MetricClient` interface.
+ * Consumers should migrate to the `createLambdaExtensionClient` function.
  *
  * @see {@link https://docs.datadoghq.com/serverless/libraries_integrations/extension/}
  * @see {@link https://docs.datadoghq.com/serverless/custom_metrics/#deprecated-cloudwatch-logs}

diff --git a/src/createLambdaExtensionClient.test.ts b/src/createLambdaExtensionClient.test.ts
@@ -0,0 +1,131 @@
+import * as datadogJS from 'datadog-lambda-js';
+
+import { createLambdaExtensionClient } from './createLambdaExtensionClient';
+
+const sendDistributionMetric = jest
+  .spyOn(datadogJS, 'sendDistributionMetric')
+  .mockReturnValue();
+
+const datadog = jest.spyOn(datadogJS, 'datadog').mockReturnValue(() => {});
+
+describe('createLambdaExtensionClient', () => {
+  const { metricsClient, withLambdaExtension } = createLambdaExtensionClient({
+    name: 'test',
+    metrics: true,
+  });
+
+  const tags = ['pipe|special|char', 'env:prod'];
+
+  afterEach(() => jest.resetAllMocks());
+
+  describe('withLambdaExtension', () => {
+    it('should call the `datadog` wrapper', () => {
+      withLambdaExtension(() => Promise.resolve({}));
+
+      expect(datadog).toHaveBeenCalledTimes(1);
+    });
+
+    it('should not call the `datadog` wrapper when metrics is turned off', () => {
+      const client = createLambdaExtensionClient({
+        name: 'test',
+        metrics: false,
+      });
+
+      client.withLambdaExtension(() => Promise.resolve({}));
+
+      expect(datadog).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('timing', () => {
+    it('should record integer timings without tags', () => {
+      metricsClient.timing('my_custom_metric', 100);
+
+      expect(sendDistributionMetric).toHaveBeenCalledTimes(1);
+      expect(sendDistributionMetric).toHaveBeenCalledWith(
+        'test.my_custom_metric',
+        100,
+      );
+    });
+
+    it('should record float timings with tags', () => {
+      metricsClient.timing('my_custom_metric', 1234.5, tags);
+
+      expect(sendDistributionMetric).toHaveBeenCalledTimes(1);
+      expect(sendDistributionMetric).toHaveBeenCalledWith(
+        'test.my_custom_metric',
+        1234.5,
+        'pipe_special_char',
+        'env:prod',
+      );
+    });
+  });
+
+  describe('distribution', () => {
+    it('should record integer values without tags', () => {
+      metricsClient.distribution('my_custom_metric', 100);
+
+      expect(sendDistributionMetric).toHaveBeenCalledTimes(1);
+      expect(sendDistributionMetric).toHaveBeenCalledWith(
+        'test.my_custom_metric',
+        100,
+      );
+    });
+
+    it('should record float values with tags', () => {
+      metricsClient.distribution('my_custom_metric', 1234.5, tags);
+
+      expect(sendDistributionMetric).toHaveBeenCalledTimes(1);
+      expect(sendDistributionMetric).toHaveBeenCalledWith(
+        'test.my_custom_metric',
+        1234.5,
+        'pipe_special_char',
+        'env:prod',
+      );
+    });
+  });
+
+  describe('increment', () => {
+    it('should increment with an implicit value and no tags', () => {
+      metricsClient.increment('my_custom_metric');
+
+      expect(sendDistributionMetric).toHaveBeenCalledTimes(1);
+      expect(sendDistributionMetric).toHaveBeenCalledWith(
+        'test.my_custom_metric',
+        1,
+      );
+    });
+
+    it('should increment with an implicit value and tags', () => {
+      metricsClient.increment('my_custom_metric', ['comma,special,char']);
+
+      expect(sendDistributionMetric).toHaveBeenCalledTimes(1);
+      expect(sendDistributionMetric).toHaveBeenCalledWith(
+        'test.my_custom_metric',
+        1,
+        'comma_special_char',
+      );
+    });
+
+    it('should increment with an explicit value and no tags', () => {
+      metricsClient.increment('my_custom_metric', 10);
+
+      expect(sendDistributionMetric).toHaveBeenCalledTimes(1);
+      expect(sendDistributionMetric).toHaveBeenCalledWith(
+        'test.my_custom_metric',
+        10,
+      );
+    });
+
+    it('should increment with an explicit value and tags', () => {
+      metricsClient.increment('my_custom_metric', 10, ['compound:tag']);
+
+      expect(sendDistributionMetric).toHaveBeenCalledTimes(1);
+      expect(sendDistributionMetric).toHaveBeenCalledWith(
+        'test.my_custom_metric',
+        10,
+        'compound:tag',
+      );
+    });
+  });
+});
diff --git a/src/createLambdaExtensionClient.ts b/src/createLambdaExtensionClient.ts
@@ -0,0 +1,100 @@
+import type { Context } from 'aws-lambda';
+import { datadog, sendDistributionMetric } from 'datadog-lambda-js';
+
+import { LambdaExtensionMetricsClient } from './LambdaExtensionMetricsClient';
+
+interface DatadogMetric {
+  name: string;
+  tags?: string[];
+  value: number;
+}
+
+interface DatadogConfig {
+  name: string;
+  metrics: boolean;
+}
+
+type Handler<Event, Output> = (
+  event: Event,
+  ctx: Readonly<Context>,
+) => Promise<Output>;
+
+interface LambdaExtensionClient {
+  metricsClient: LambdaExtensionMetricsClient;
+  /**
+   * Conditionally wraps your AWS lambda handler function based on the provided config.
+   *
+   * This is necessary for initialising metrics/tracing support.
+   */
+  // This also "fixes" its broken type definitions.
+  withLambdaExtension: <Event, Output = unknown>(
+    fn: Handler<Event, Output>,
+  ) => Handler<Event, Output>;
+}
+
+/**
+ * Replaces tag special characters with an underscore
+ */
+const sanitiseTag = (tag: string): string => tag.replace(/\||@|,/g, '_');
+
+/**
+ * Creates a new Datadog Lambda client configured for the given app.
+ *
+ * @see {@link https://docs.datadoghq.com/serverless/libraries_integrations/extension/}
+ *
+ * @param config - Application configuration
+ */
+export const createLambdaExtensionClient = (
+  config: DatadogConfig,
+): LambdaExtensionClient => {
+  const send = (metric: DatadogMetric) => {
+    const { value } = metric;
+
+    const sanitisedLambdaName = config.name.replace(new RegExp('-', 'g'), '_');
+    const name = `${sanitisedLambdaName}.${metric.name.toLowerCase()}`;
+
+    const tags = (metric.tags || []).map(sanitiseTag);
+
+    if (config.metrics) {
+      sendDistributionMetric(name, value, ...tags);
+    }
+  };
+
+  const sendCount = (
+    name: string,
+    countOrTags?: number | string[],
+    tagsIfCount?: string[],
+  ): void => {
+    let count: number;
+    let tags: string[] | undefined;
+
+    if (typeof countOrTags === 'number') {
+      count = countOrTags;
+      tags = tagsIfCount;
+    } else {
+      count = 1;
+      tags = countOrTags;
+    }
+
+    send({ name, tags, value: count });
+  };
+
+  return {
+    withLambdaExtension: <Event, Output = unknown>(
+      fn: Handler<Event, Output>,
+    ): Handler<Event, Output> =>
+      config.metrics ? (datadog(fn) as Handler<Event, Output>) : fn,
+
+    metricsClient: {
+      increment: sendCount,
+
+      distribution: (name: string, value: number, tags?: string[]): void => {
+        send({ name, tags, value });
+      },
+
+      timing: (name: string, value: number, tags?: string[]): void => {
+        send({ name, tags, value });
+      },
+    },
+  };
+};
diff --git a/src/createTimedSpan.ts b/src/createTimedSpan.ts
@@ -1,5 +1,7 @@
 import { MetricsClient } from './MetricsClient';
 
+type TimingMetricsClient = Pick<MetricsClient, 'increment' | 'timing'>;
+
 /**
  * Sends timing related metrics for an asynchronous operation
  *
@@ -12,7 +14,7 @@ import { MetricsClient } from './MetricsClient';
  * @param block         - Function returning the promise to time
  */
 export const createTimedSpan =
-  (metricsClient: MetricsClient) =>
+  (metricsClient: TimingMetricsClient) =>
   async <T>(name: string, block: () => PromiseLike<T>): Promise<T> => {
     const startTime = process.hrtime.bigint();
 

diff --git a/src/index.test.ts b/src/index.test.ts
@@ -1,5 +1,6 @@
 import {
   createCloudWatchClient,
+  createLambdaExtensionClient,
   createNoOpClient,
   createStatsDClient,
   createTimedSpan,
@@ -15,6 +16,10 @@ describe('index', () => {
     expect(createCloudWatchClient).toBeInstanceOf(Function);
   });
 
+  it('should export a createLambdaExtensionClient function', () => {
+    expect(createLambdaExtensionClient).toBeInstanceOf(Function);
+  });
+
   it('should export a createNoOpClient function', () => {
     expect(createNoOpClient).toBeInstanceOf(Function);
   });