chore: remove singleton private access

src/datadog-lambda.ts

@@ -81,16 +81,14 @@ export class DatadogLambda extends Construct {
     // defined in DefaultDatadogProps (if not set by user)
     const baseProps: DatadogLambdaStrictProps = handleSettingPropDefaults(this.props);
 
-    const extractedLambdaFunctions = extractSingletonFunctions(lambdaFunctions);
-
-    if (extractedLambdaFunctions.length === 0) {
+    if (lambdaFunctions.length === 0) {
       return;
     }
 
-    const region = extractedLambdaFunctions[0].env.region;
+    const region = lambdaFunctions[0].env.region;
     log.debug(`Using region: ${region}`);
 
-    for (const lambdaFunction of extractedLambdaFunctions) {
+    for (const lambdaFunction of lambdaFunctions) {
       if (this.props.apiKeySecret !== undefined) {
         grantReadLambda(this.props.apiKeySecret, lambdaFunction);
       } else if (
@@ -197,8 +195,7 @@ export class DatadogLambda extends Construct {
    * On duplicate tag keys, the later source wins.
    */
   public setEnvironment(lambdaFunction: LambdaFunction, key: string, value: string): void {
-    const [extractedLambdaFunction] = extractSingletonFunctions([lambdaFunction]);
-    setTrackedEnv(extractedLambdaFunction, key, value);
+    setTrackedEnv(lambdaFunction, key, value);
   }
 
   public overrideGitMetadata(gitCommitSha: string, gitRepoUrl?: string): void {
@@ -237,8 +234,7 @@ export class DatadogLambda extends Construct {
     // @ts-ignore
     gitRepoUrl?: string,
   ): void {
-    const extractedLambdaFunctions = extractSingletonFunctions(lambdaFunctions);
-    setGitEnvironmentVariables(extractedLambdaFunctions, this.gitCommitShaOverride, this.gitRepoUrlOverride);
+    setGitEnvironmentVariables(lambdaFunctions, this.gitCommitShaOverride, this.gitRepoUrlOverride);
   }
 
   public addForwarderToNonLambdaLogGroups(logGroups: logs.ILogGroup[]) {
@@ -255,30 +251,32 @@ export class DatadogLambda extends Construct {
   }
 }
 
-export function addCdkConstructVersionTag(lambdaFunction: lambda.Function): void {
+export function addCdkConstructVersionTag(lambdaFunction: LambdaFunction): void {
   log.debug(`Adding CDK Construct version tag: ${versionJson.version}`);
-  Tags.of(lambdaFunction).add(TagKeys.CDK, `v${versionJson.version}`, {
+  Tags.of(
+    lambdaFunction instanceof lambda.SingletonFunction ? lambdaFunction.permissionsNode.defaultChild! : lambdaFunction,
+  ).add(TagKeys.CDK, `v${versionJson.version}`, {
     includeResourceTypes: ["AWS::Lambda::Function"],
   });
 }
 
-function setTagsForFunction(lambdaFunction: lambda.Function, props: DatadogLambdaProps): void {
+function setTagsForFunction(lambdaFunction: LambdaFunction, props: DatadogLambdaProps): void {
   if (props.forwarderArn) {
     setTags(lambdaFunction, props);
   }
 }
 
-function grantReadLambda(secret: ISecret, lambdaFunction: lambda.Function): void {
+function grantReadLambda(secret: ISecret, lambdaFunction: LambdaFunction): void {
   secret.grantRead(lambdaFunction);
   secret.encryptionKey?.grantDecrypt(lambdaFunction);
 }
 
-function grantReadLambdaFromSecretArn(construct: Construct, arn: string, lambdaFunction: lambda.Function): void {
+function grantReadLambdaFromSecretArn(construct: Construct, arn: string, lambdaFunction: LambdaFunction): void {
   const secret = Secret.fromSecretPartialArn(construct, "DatadogApiKeySecret", arn);
   secret.grantRead(lambdaFunction);
 }
 
-function grantReadLambdaFromSsmParameterArn(parameterArn: string, lambdaFunction: lambda.Function): void {
+function grantReadLambdaFromSsmParameterArn(parameterArn: string, lambdaFunction: LambdaFunction): void {
   // Grant IAM permissions to support both String and SecureString SSM parameters
   // For SecureString parameters, the Datadog Extension will decrypt at runtime using KMS
   const stack = Stack.of(lambdaFunction);
@@ -305,21 +303,6 @@ function grantReadLambdaFromSsmParameterArn(parameterArn: string, lambdaFunction
   );
 }
 
-function extractSingletonFunctions(lambdaFunctions: LambdaFunction[]): lambda.Function[] {
-  // extract lambdaFunction property from Singleton Function
-  // using bracket notation here since lambdaFunction is a private property
-  const extractedLambdaFunctions: lambda.Function[] = lambdaFunctions.map((fn) => {
-    // eslint-disable-next-line dot-notation
-    return isSingletonFunction(fn) ? fn["lambdaFunction"] : fn;
-  });
-
-  return extractedLambdaFunctions;
-}
-
-function isSingletonFunction(fn: LambdaFunction): fn is lambda.SingletonFunction {
-  return fn.hasOwnProperty("lambdaFunction");
-}
-
 export function validateProps(props: DatadogLambdaProps, apiKeyArnOverride = false): void {
   log.debug("Validating props...");
 

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

@@ -15,6 +15,7 @@ import * as sfn from "aws-cdk-lib/aws-stepfunctions";
 import { Construct } from "constructs";
 import log from "loglevel";
 import { SUBSCRIPTION_FILTER_PREFIX } from "./index";
+import { LambdaFunction } from "./interfaces";
 
 function getForwarder(scope: Construct, forwarderArn: string) {
   const forwarderConstructId = generateForwarderConstructId(forwarderArn);
@@ -27,13 +28,13 @@ function getForwarder(scope: Construct, forwarderArn: string) {
 
 export function addForwarder(
   scope: Construct,
-  lam: lambda.Function,
+  lam: LambdaFunction,
   forwarderArn: string,
   createForwarderPermissions: boolean,
 ): 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,

src/interfaces.ts

@@ -115,10 +115,6 @@ export interface Runtime {
   readonly name: string;
 }
 
-export interface Node {
-  readonly defaultChild: any;
-}
-
 export type LambdaFunction = lambda.Function | lambda.SingletonFunction;
 
 export interface DatadogStepFunctionsProps {

src/layer.ts

@@ -21,6 +21,7 @@ import {
   LAYER_PREFIX,
   EXTENSION_LAYER_PREFIX,
 } from "./index";
+import { LambdaFunction } from "./interfaces";
 import { DatadogDefaultLayerVersions } from "./layer-versions";
 
 const layers: Map<string, lambda.ILayerVersion> = new Map();
@@ -38,7 +39,7 @@ const DEFAULT_LAYER_VERSION_BY_RUNTIME: { [paramRuntime: string]: number } = {
 export function applyLayers(
   scope: Construct,
   region: string,
-  lam: lambda.Function,
+  lam: LambdaFunction,
   pythonLayerVersion?: number,
   pythonLayerArn?: string,
   nodeLayerVersion?: number,
@@ -176,7 +177,7 @@ export function applyLayers(
 export function applyExtensionLayer(
   scope: Construct,
   region: string,
-  lam: lambda.Function,
+  lam: LambdaFunction,
   extensionLayerVersion?: number,
   extensionLayerArn?: string,
   useLayersFromAccount?: string,
@@ -268,7 +269,7 @@ function addLayer(
   layerArn: string,
   isExtensionLayer: boolean,
   scope: Construct,
-  lam: lambda.Function,
+  lam: LambdaFunction,
   runtime: string,
 ): void {
   const layerId = generateLayerId(isExtensionLayer, lam.functionArn, runtime);

src/redirect.ts

@@ -18,6 +18,7 @@ import {
   JS_HANDLER,
   PYTHON_HANDLER,
 } from "./constants";
+import { LambdaFunction } from "./interfaces";
 
 /**
  * To avoid modifying code in the user's lambda handler, redirect the handler to a Datadog
@@ -27,7 +28,7 @@ import {
  *
  * Unchanged aside from parameter type
  */
-export function redirectHandlers(lam: lambda.Function, addLayers: boolean, useExtension: boolean): void {
+export function redirectHandlers(lam: LambdaFunction, addLayers: boolean, useExtension: boolean): void {
   log.debug(`Wrapping Lambda function handlers with Datadog handler...`);
 
   const runtime: string = lam.runtime.name;
@@ -40,7 +41,8 @@ export function redirectHandlers(lam: lambda.Function, addLayers: boolean, useEx
     return;
   }
 
-  const cfnFuntion = lam.node.defaultChild as lambda.CfnFunction;
+  const cfnFuntion = (lam instanceof lambda.SingletonFunction ? lam.permissionsNode : lam.node)
+    .defaultChild as lambda.CfnFunction;
   if (cfnFuntion === undefined) {
     log.debug("Unable to get Lambda Function handler");
     return;

src/tag.ts

@@ -11,29 +11,31 @@ import * as lambda from "aws-cdk-lib/aws-lambda";
 import * as sfn from "aws-cdk-lib/aws-stepfunctions";
 import log from "loglevel";
 import { TagKeys, DatadogLambdaProps, DatadogStepFunctionsProps } from "./index";
+import { LambdaFunction } from "./interfaces";
 
 const versionJson = require("../version.json");
 
 export function setTags(
-  resource: lambda.Function | sfn.StateMachine,
+  resource: LambdaFunction | sfn.StateMachine,
   props: DatadogLambdaProps | DatadogStepFunctionsProps,
 ): void {
   log.debug(`Adding datadog tags`);
+  const taggable = resource instanceof lambda.SingletonFunction ? resource.permissionsNode.defaultChild! : resource;
   if (props.env) {
-    Tags.of(resource).add(TagKeys.ENV, props.env);
+    Tags.of(taggable).add(TagKeys.ENV, props.env);
   }
   if (props.service) {
-    Tags.of(resource).add(TagKeys.SERVICE, props.service);
+    Tags.of(taggable).add(TagKeys.SERVICE, props.service);
   }
   if (props.version) {
-    Tags.of(resource).add(TagKeys.VERSION, props.version);
+    Tags.of(taggable).add(TagKeys.VERSION, props.version);
   }
   if (props.tags) {
     const tagsArray = props.tags.split(",");
     tagsArray.forEach((tag: string) => {
       const [key, value] = tag.split(":");
       if (key && value) {
-        Tags.of(resource).add(key, value);
+        Tags.of(taggable).add(key, value);
       }
     });
   }

src/transport.ts

@@ -6,9 +6,9 @@
  * Copyright 2020-2026 Datadog, Inc.
  */
 
-import * as lambda from "aws-cdk-lib/aws-lambda";
 import log from "loglevel";
 import { runtimeLookup, RuntimeType } from "./index";
+import { LambdaFunction } from "./interfaces";
 
 export const API_KEY_ENV_VAR = "DD_API_KEY";
 export const API_KEY_SECRET_ARN_ENV_VAR = "DD_API_KEY_SECRET_ARN";
@@ -75,7 +75,7 @@ export class Transport {
     this.apiKmsKey = apiKmsKey;
   }
 
-  applyEnvVars(lam: lambda.Function) {
+  applyEnvVars(lam: LambdaFunction) {
     log.debug(`Setting Datadog transport environment variables...`);
     lam.addEnvironment(FLUSH_METRICS_TO_LOGS_ENV_VAR, this.flushMetricsToLogs.toString());
     if (this.site !== undefined && this.flushMetricsToLogs === false) {

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,20 @@ 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);
+    const singletonProperties =
+      Template.fromStack(stack).findResources("AWS::Lambda::Function")[singletonLogicalId].Properties;
+    expect(singletonProperties.Tags).toEqual(
+      expect.arrayContaining([expect.objectContaining({ Key: "dd_cdk_construct" })]),
+    );
   });
 
   it("Throws an error when a customer redundantly calls the addLambdaFunctions function on the same lambda function(s) and forwarder", () => {

test/test-utils.ts

@@ -4,13 +4,9 @@ import { Construct } from "constructs";
 import { SUBSCRIPTION_FILTER_PREFIX } from "../src/index";
 
 export const findDatadogSubscriptionFilters = (baseConstruct: Construct) => {
-  // extract lambdaFunction property from Singleton Function
-  // using bracket notation here since lambdaFunction is a private property
-  const baseConstructModified: Construct = isSingletonFunction(baseConstruct)
-    ? baseConstruct["lambdaFunction"] // eslint-disable-line dot-notation
-    : baseConstruct;
+  const node = baseConstruct instanceof lambda.SingletonFunction ? baseConstruct.permissionsNode : baseConstruct.node;
 
-  return baseConstructModified.node
+  return node
     .findAll()
     .filter((construct) => construct.node.id.startsWith(SUBSCRIPTION_FILTER_PREFIX))
     .map((construct) => {
@@ -27,7 +23,3 @@ export const findDatadogSubscriptionFilters = (baseConstruct: Construct) => {
     })
     .reduce((acc, subscriptionFilters) => acc.concat(subscriptionFilters), []);
 };
-
-function isSingletonFunction(fn: Construct): fn is lambda.SingletonFunction {
-  return fn.hasOwnProperty("lambdaFunction");
-}