Merge pull request #63 from vishnuvisnu/support-custom-timestamp

Added logic to support custom timestamp

Author: DarcyRaynerDD

src/index.spec.ts

@@ -2,7 +2,13 @@ import http from "http";
 import nock from "nock";
 
 import { Context, Handler } from "aws-lambda";
-import { datadog, getTraceHeaders, sendDistributionMetric, TraceHeaders } from "./index";
+import {
+  datadog,
+  getTraceHeaders,
+  sendDistributionMetric,
+  TraceHeaders,
+  sendDistributionMetricWithDate,
+} from "./index";
 import { incrementErrorsMetric, incrementInvocationsMetric } from "./metrics/enhanced-metrics";
 import { MetricsListener } from "./metrics/listener";
 import { LogLevel, setLogLevel } from "./utils";
@@ -163,6 +169,28 @@ describe("datadog", () => {
     expect(nock.isDone()).toBeTruthy();
   });
 
+  it("reads site keys from the environment using custom timestamp", async () => {
+    const site = "datadoghq.com";
+    const siteEnvVar = "DD_SITE";
+    const apiKey = "12345";
+    process.env[siteEnvVar] = site;
+
+    nock("https://api.datadoghq.com")
+      .post(`/api/v1/distribution_points?api_key=${apiKey}`, (request: any) => request.series[0].metric === "my-dist")
+      .reply(200, {});
+
+    const wrapped = datadog(
+      async () => {
+        sendDistributionMetricWithDate("my-dist", 100, new Date(), "first-tag", "second-tag");
+        return "";
+      },
+      { apiKey, forceWrap: true },
+    );
+    await wrapped({}, {} as any, () => {});
+
+    expect(nock.isDone()).toBeTruthy();
+  });
+
   it("makes the current trace headers available", async () => {
     let traceHeaders: Partial<TraceHeaders> = {};
     const event = {

src/index.ts

@@ -117,6 +117,23 @@ export function datadog<TEvent, TResult>(
   return wrappedFunc;
 }
 
+/**
+ * Sends a Distribution metric asynchronously to the Datadog API.
+ * @param name The name of the metric to send.
+ * @param value The value of the metric
+ * @param metricTime The timesamp associated with this metric data point.
+ * @param tags The tags associated with the metric. Should be of the format "tag:value".
+ */
+export function sendDistributionMetricWithDate(name: string, value: number, metricTime: Date, ...tags: string[]) {
+  tags = [...tags, getRuntimeTag()];
+
+  if (currentMetricsListener !== undefined) {
+    currentMetricsListener.sendDistributionMetricWithDate(name, value, metricTime, ...tags);
+  } else {
+    logError("handler not initialized");
+  }
+}
+
 /**
  * Sends a Distribution metric asynchronously to the Datadog API.
  * @param name The name of the metric to send.

src/metrics/enhanced-metrics.ts

@@ -65,7 +65,7 @@ export function getEnhancedMetricTags(context: Context): string[] {
  */
 function incrementEnhancedMetric(metricName: string, context: Context) {
   // Always write enhanced metrics to standard out
-  writeMetricToStdout(`${ENHANCED_LAMBDA_METRICS_NAMESPACE}.${metricName}`, 1, getEnhancedMetricTags(context));
+  writeMetricToStdout(`${ENHANCED_LAMBDA_METRICS_NAMESPACE}.${metricName}`, 1, new Date(), getEnhancedMetricTags(context));
 }
 
 export function incrementInvocationsMetric(context: Context): void {

src/metrics/listener.spec.ts

@@ -100,4 +100,25 @@ describe("MetricsListener", () => {
 
     expect(spy).toHaveBeenCalledWith(`{"e":1487076708,"m":"my-metric","t":["tag:a","tag:b"],"v":10}\n`);
   });
+
+  it("logs metrics when logForwarding is enabled with custom timestamp", 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.sendDistributionMetricWithDate("my-metric", 10, new Date(1584983836 * 1000), "tag:a", "tag:b");
+    await listener.onCompleteInvocation();
+
+    expect(spy).toHaveBeenCalledWith(`{"e":1584983836,"m":"my-metric","t":["tag:a","tag:b"],"v":10}\n`);
+  });
 });

src/metrics/listener.ts

@@ -83,12 +83,12 @@ export class MetricsListener {
     this.currentProcessor = undefined;
   }
 
-  public sendDistributionMetric(name: string, value: number, ...tags: string[]) {
+  public sendDistributionMetricWithDate(name: string, value: number, metricTime: Date, ...tags: string[]) {
     if (this.config.logForwarding) {
-      writeMetricToStdout(name, value, tags);
+      writeMetricToStdout(name, value, metricTime, tags);
       return;
     }
-    const dist = new Distribution(name, [{ timestamp: new Date(), value }], ...tags);
+    const dist = new Distribution(name, [{ timestamp: metricTime, value }], ...tags);
 
     if (this.currentProcessor !== undefined) {
       this.currentProcessor.then((processor) => {
@@ -99,6 +99,10 @@ export class MetricsListener {
     }
   }
 
+  public sendDistributionMetric(name: string, value: number, ...tags: string[]) {
+    this.sendDistributionMetricWithDate(name, value, new Date(Date.now()), ...tags);
+  }
+
   private async createProcessor(config: MetricsConfig, apiKey: Promise<string>) {
     const key = await apiKey;
     const url = `https://api.${config.siteURL}`;

src/metrics/metric-log.spec.ts

@@ -1,14 +1,13 @@
 import { buildMetricLog } from "./metric-log";
 
 describe("buildMetricLog", () => {
-  jest.spyOn(Date, "now").mockImplementation(() => 1487076708123);
   it("handles empty tag list", () => {
-    expect(buildMetricLog("my.test.metric", 1337, [])).toStrictEqual(
+    expect(buildMetricLog("my.test.metric", 1337, new Date(1487076708123), [])).toStrictEqual(
       '{"e":1487076708.123,"m":"my.test.metric","t":[],"v":1337}\n',
     );
   });
   it("writes timestamp in Unix seconds", () => {
-    expect(buildMetricLog("my.test.metric", 1337, ["region:us", "account:dev", "team:serverless"])).toStrictEqual(
+    expect(buildMetricLog("my.test.metric", 1337, new Date(1487076708123), ["region:us", "account:dev", "team:serverless"])).toStrictEqual(
       '{"e":1487076708.123,"m":"my.test.metric","t":["region:us","account:dev","team:serverless"],"v":1337}\n',
     );
   });

src/metrics/metric-log.ts

@@ -1,8 +1,8 @@
 // Builds the string representation of the metric that will be written to logs
-export function buildMetricLog(name: string, value: number, tags: string[]) {
+export function buildMetricLog(name: string, value: number, metricTime: Date, tags: string[]) {
   return `${JSON.stringify({
     // Date.now() returns Unix time in milliseconds, we convert to seconds for DD API submission
-    e: Date.now() / 1000,
+    e: metricTime.getTime() / 1000,
     m: name,
     t: tags,
     v: value,
@@ -15,8 +15,8 @@ export function buildMetricLog(name: string, value: number, tags: string[]) {
  * @param value Metric datapoint's value
  * @param tags Tags to apply to the metric
  */
-export function writeMetricToStdout(name: string, value: number, tags: string[]) {
+export function writeMetricToStdout(name: string, value: number, metricTime: Date, tags: string[]) {
   // We use process.stdout.write, because console.log will prepend metadata to the start
   // of the log that log forwarder doesn't know how to read.
-  process.stdout.write(buildMetricLog(name, value, tags));
+  process.stdout.write(buildMetricLog(name, value, metricTime, tags));
 }