preserve singleton behavior

Author: ava-silver

src/env-tracker.ts

@@ -6,6 +6,7 @@
  * Copyright 2021 Datadog, Inc.
  */
 
+import { Node } from "constructs";
 import { LambdaFunction } from "./interfaces";
 
 const DD_TAGS = "DD_TAGS";
@@ -26,14 +27,14 @@ interface TrackedEnvironment {
 //
 // Not exported from index.ts -- internal to the package.
 //
-// WeakMap so functions can be garbage-collected when their stack goes out of scope (for
-// example, between test cases).
+// WeakMap so construct nodes can be garbage-collected when their stack goes out of scope
+// (for example, between test cases).
 //
 // Env vars set via func.addEnvironment() outside this library are invisible here and will
 // be overwritten if the library writes the same key. Configure DD_* vars via
 // DatadogLambdaProps or datadogLambda.setEnvironment(), or call func.addEnvironment()
 // after datadogLambda.addLambdaFunctions().
-const ddEnvTracker: WeakMap<LambdaFunction, TrackedEnvironment> = new WeakMap();
+const ddEnvTracker: WeakMap<Node, TrackedEnvironment> = new WeakMap();
 
 export function setTrackedEnv(lam: LambdaFunction, key: string, value: string): void {
   if (key === DD_TAGS) {
@@ -60,12 +61,12 @@ export function mergeTrackedGitTags(lam: LambdaFunction, value: string): void {
 }
 
 export function hasTrackedEnv(lam: LambdaFunction, key: string): boolean {
-  const tracked = ddEnvTracker.get(lam);
+  const tracked = ddEnvTracker.get(lam.permissionsNode);
   return key === DD_TAGS ? (tracked?.tagsSet ?? false) : (tracked?.values.has(key) ?? false);
 }
 
 function getOrCreateTrackedEnvironment(lam: LambdaFunction): TrackedEnvironment {
-  let tracked = ddEnvTracker.get(lam);
+  let tracked = ddEnvTracker.get(lam.permissionsNode);
   if (!tracked) {
     tracked = {
       values: new Map(),
@@ -74,7 +75,7 @@ function getOrCreateTrackedEnvironment(lam: LambdaFunction): TrackedEnvironment
       gitTags: new Map(),
       tagsSet: false,
     };
-    ddEnvTracker.set(lam, tracked);
+    ddEnvTracker.set(lam.permissionsNode, tracked);
   }
   return tracked;
 }

src/forwarder.ts

@@ -34,7 +34,7 @@ export function addForwarder(
 ): void {
   const forwarder = getForwarder(scope, forwarderArn);
   const forwarderDestination = new LambdaDestination(forwarder, { addPermissions: createForwarderPermissions });
-  const subscriptionFilterName = generateSubscriptionFilterName(Names.uniqueId(lam), forwarderArn);
+  const subscriptionFilterName = generateSubscriptionFilterName(Names.nodeUniqueId(lam.permissionsNode), forwarderArn);
   log.debug(`Adding log subscription ${subscriptionFilterName} for ${lam.functionName}`);
   lam.logGroup.addSubscriptionFilter(subscriptionFilterName, {
     destination: forwarderDestination,

test/datadog-lambda.spec.ts

@@ -1194,6 +1194,38 @@ describe("setEnvironment", () => {
     });
   });
 
+  it("shares tracked values between singleton declarations with the same UUID", () => {
+    const app = new App();
+    const stack = new Stack(app, "stack");
+    const first = new lambda.SingletonFunction(stack, "First", {
+      runtime: lambda.Runtime.NODEJS_18_X,
+      code: lambda.Code.fromInline("test"),
+      handler: "hello.handler",
+      uuid: "f21a2d33-8ef9-4e38-8af6-9dc39dd63f12",
+    });
+    const second = new lambda.SingletonFunction(stack, "Second", {
+      runtime: lambda.Runtime.NODEJS_18_X,
+      code: lambda.Code.fromInline("test"),
+      handler: "hello.handler",
+      uuid: "f21a2d33-8ef9-4e38-8af6-9dc39dd63f12",
+    });
+    const datadogLambda = new DatadogLambda(stack, "Datadog", {
+      nodeLayerVersion: NODE_LAYER_VERSION,
+      enableDatadogTracing: true,
+      sourceCodeIntegration: false,
+    });
+    datadogLambda.setEnvironment(first, "DD_TRACE_ENABLED", "false");
+    datadogLambda.addLambdaFunctions([second], stack);
+
+    Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
+      Environment: {
+        Variables: {
+          DD_TRACE_ENABLED: "false",
+        },
+      },
+    });
+  });
+
   it("lets an unconditional construct setting override a tracked value", () => {
     const app = new App();
     const stack = new Stack(app, "stack");

test/index.spec.ts

@@ -1,7 +1,8 @@
-import { App, Stack, NestedStack } from "aws-cdk-lib";
+import { App, Names, Stack, NestedStack } from "aws-cdk-lib";
 import { Match, Template } from "aws-cdk-lib/assertions";
 import * as lambda from "aws-cdk-lib/aws-lambda";
 import { findDatadogSubscriptionFilters } from "./test-utils";
+import { generateSubscriptionFilterName } from "../src/forwarder";
 import {
   DatadogLambda,
   DD_ACCOUNT_ID,
@@ -61,6 +62,12 @@ describe("addLambdaFunctions", () => {
     expect(pythonLambdaSubscriptionFilters).toHaveLength(1);
     expect(singletonLambdaSubscriptionFilters).toHaveLength(1);
     expect(nodeLambdaSubscriptionFilters[0].destinationArn).toEqual(pythonLambdaSubscriptionFilters[0].destinationArn);
+    expect(singletonLambdaSubscriptionFilters[0].id).toEqual(
+      generateSubscriptionFilterName(
+        Names.nodeUniqueId(singletonLambda.permissionsNode),
+        "arn:test:forwarder:sa-east-1:12345678:1",
+      ),
+    );
 
     const singletonResource = singletonLambda.permissionsNode.defaultChild as lambda.CfnFunction;
     const singletonLogicalId = stack.getLogicalId(singletonResource);