fix: stop reading private aws-cdk-lib Function.environment field (#621)

* [SVLS-9359] stop reading private aws-cdk-lib Function.environment field

* [SVLS-9359] add setEnvironment method to seed construct-tracked env vars

Author: ava-silver

README.md

@@ -280,6 +280,60 @@ _Note_: The descriptions use the npm package parameters, but they also apply to
 | `llmObsMlApp`                | `llm_obs_ml_app`                | The name of your LLM application, service, or project, under which all traces and spans are grouped. This helps distinguish between different applications or experiments. See [Application naming guidelines](https://docs.datadoghq.com/llm_observability/sdk/?tab=nodejs#application-naming-guidelines) for allowed characters and other constraints. To override this value for a given root span, see [Tracing multiple applications](https://docs.datadoghq.com/llm_observability/sdk/?tab=nodejs#tracing-multiple-applications).  Required if `llmObsEnabled` is `true` |
 | `llmObsAgentlessEnabled`     | `llm_obs_agentless_enabled`     | Only required if you are not using the Datadog Lambda Extension, in which case this should be set to `true`.  Defaults to `false`.                                                                                                                                                                                                                                                                                                                                                                                                                                             |
 
+#### Setting `DD_*` environment variables
+
+To configure Datadog variables for every instrumented function, set the matching field on `DatadogLambdaProps` (for example, `enableDatadogTracing`, `logLevel`, `env`, or `tags`).
+
+To override a value on a single function, use one of:
+
+- `datadogLambda.setEnvironment(func, key, value)` before `datadogLambda.addLambdaFunctions()`, to override a construct default while letting the construct finish instrumenting the function.
+- `func.addEnvironment(key, value)` after `datadogLambda.addLambdaFunctions()`, to override the value set during instrumentation.
+
+When more than one source sets the same key, the following order applies (highest precedence first):
+
+1. `func.addEnvironment(key, value)` called after `datadogLambda.addLambdaFunctions()`.
+2. `DatadogLambdaProps` fields dedicated to that key. These fields overwrite a value for the same key set through `datadogLambda.setEnvironment()`:
+   - Unified service tagging: `env`, `service`, `version`
+   - Cold-start tracing: `enableColdStartTracing`, `minColdStartTraceDuration`, `coldStartTraceSkipLibs`
+   - Other tracer settings: `enableProfiling`, `encodeAuthorizerContext`, `decodeAuthorizerContext`, `apmFlushDeadline`
+   - LLM Observability: `llmObsEnabled`, `llmObsMlApp`, `llmObsAgentlessEnabled`
+   - Transport: `site`, `apiKey`, `apiKeySecretArn`, `apiKeySsmArn`, `apiKmsKey`, `flushMetricsToLogs`
+3. `datadogLambda.setEnvironment(func, key, value)` called before `datadogLambda.addLambdaFunctions()`.
+4. Construct defaults, which apply only when nothing else set the key: `enableDatadogTracing`, `datadogAppSecMode`, `enableMergeXrayTraces`, `injectLogContext`, `enableDatadogLogs`, `captureLambdaPayload`, `captureCloudServicePayload`, `logLevel`
+
+`datadogLambda.addLambdaFunctions` merges the `DD_TAGS` environment variable from three sources, in order:
+
+1. `DatadogLambdaProps.tags`.
+2. Per-function tags from `datadogLambda.setEnvironment(func, 'DD_TAGS', ...)`.
+3. `git.commit.sha` and `git.repository_url` from source code integration.
+
+If the same tag key appears in more than one source, the later source wins.
+
+The following example shows these rules in practice:
+
+```typescript
+const myFunction = new lambda.Function(this, 'MyFunction', {
+  // ...
+});
+
+const datadogLambda = new DatadogLambda(this, 'DatadogLambda', {
+  // ...
+  tags: 'env:prod,team:platform',
+});
+
+datadogLambda.setEnvironment(myFunction, 'DD_TRACE_ENABLED', 'false');
+datadogLambda.setEnvironment(myFunction, 'DD_TAGS', 'service:worker,team:payments');
+datadogLambda.addLambdaFunctions([myFunction]);
+
+myFunction.addEnvironment('DD_LOG_LEVEL', 'debug');
+```
+
+Final values on `myFunction`:
+
+- `DD_TRACE_ENABLED=false`, from `datadogLambda.setEnvironment`, overriding the default.
+- `DD_LOG_LEVEL=debug`, from `myFunction.addEnvironment`, overriding the construct.
+- `DD_TAGS=env:prod,service:worker,team:payments,git.commit.sha:...,git.repository_url:...`. `team:payments` replaces `team:platform`, and source code integration appends the git tags.
+
 #### Default layer versions
 
 When you don't pass a `*LayerVersion` or `*LayerArn`, the construct uses a default layer version bundled with the package. These defaults track the latest released Datadog Lambda layers at the time the construct version was published, and are exposed via the `DatadogDefaultLayerVersions` class so you can reference them directly in any language:

src/datadog-lambda.ts

@@ -13,6 +13,7 @@ import * as logs from "aws-cdk-lib/aws-logs";
 import { ISecret, Secret } from "aws-cdk-lib/aws-secretsmanager";
 import { Construct } from "constructs";
 import log from "loglevel";
+import { hasTrackedEnv, mergeTrackedGitTags, setTrackedEnv } from "./env-tracker";
 import {
   applyLayers,
   redirectHandlers,
@@ -181,6 +182,25 @@ export class DatadogLambda extends Construct {
     }
   }
 
+  /**
+   * Pre-set a Datadog environment variable on `lambdaFunction`. Call before
+   * `addLambdaFunctions([lambdaFunction])`.
+   *
+   * Precedence, highest first:
+   *   1. `func.addEnvironment()` called after `addLambdaFunctions()`.
+   *   2. `DatadogLambdaProps` fields dedicated to `key` (for example, `env` for `DD_ENV`).
+   *   3. This method.
+   *   4. Construct defaults (for example, `enableDatadogTracing` for `DD_TRACE_ENABLED`).
+   *
+   * `addLambdaFunctions` merges `DD_TAGS` from `DatadogLambdaProps.tags`, per-function
+   * tags from this method, and git tags from source code integration, in that order.
+   * 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);
+  }
+
   public overrideGitMetadata(gitCommitSha: string, gitRepoUrl?: string): void {
     if (gitCommitSha) {
       this.gitCommitShaOverride = gitCommitSha;
@@ -191,23 +211,18 @@ export class DatadogLambda extends Construct {
 
     // If any lambdas have already been added, override the commit sha and url
     if (this.lambdas) {
-      this.lambdas.forEach((lambdaFunction: any) => {
-        const existingTags = lambdaFunction.environment.map.get(DD_TAGS);
-        if (existingTags === undefined) {
+      this.lambdas.forEach((lambdaFunction: LambdaFunction) => {
+        if (!hasTrackedEnv(lambdaFunction, DD_TAGS)) {
           return;
         }
-        const tags = existingTags.value.split(",");
-        if (gitCommitSha) {
-          const index = tags.findIndex((val: string) => val.split(":")[0] === "git.commit.sha");
-          tags[index] = `git.commit.sha:${gitCommitSha}`;
-        }
+        const gitTags = [
+          gitCommitSha ? `git.commit.sha:${gitCommitSha}` : undefined,
+          gitRepoUrl ? `git.repository_url:${gitRepoUrl}` : undefined,
+        ].filter((tag): tag is string => tag !== undefined);
 
-        if (gitRepoUrl) {
-          const index = tags.findIndex((val: string) => val.split(":")[0] === "git.repository_url");
-          tags[index] = `git.repository_url:${gitRepoUrl}`;
+        if (gitTags.length > 0) {
+          mergeTrackedGitTags(lambdaFunction, gitTags.join(","));
         }
-
-        lambdaFunction.addEnvironment(DD_TAGS, tags.join(","));
       });
     }
   }

src/env-tracker.ts

@@ -0,0 +1,112 @@
+/*
+ * Unless explicitly stated otherwise all files in this repository are licensed
+ * under the Apache License Version 2.0.
+ *
+ * This product includes software developed at Datadog (https://www.datadoghq.com/).
+ * Copyright 2021 Datadog, Inc.
+ */
+
+import { LambdaFunction } from "./interfaces";
+
+const DD_TAGS = "DD_TAGS";
+
+type Tags = Map<string, string>;
+
+interface TrackedEnvironment {
+  readonly values: Map<string, string>;
+  readonly propTags: Tags;
+  readonly functionTags: Tags;
+  readonly gitTags: Tags;
+  tagsSet: boolean;
+}
+
+// Bookkeeping for env vars this library writes to Lambda functions. aws-cdk-lib does not
+// expose Function.environment publicly, so we mirror our own writes here and read from
+// this map instead of the private field.
+//
+// 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).
+//
+// 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();
+
+export function setTrackedEnv(lam: LambdaFunction, key: string, value: string): void {
+  if (key === DD_TAGS) {
+    const tracked = getOrCreateTrackedEnvironment(lam);
+    replaceTags(tracked.functionTags, value);
+    writeTags(lam, tracked);
+    return;
+  }
+
+  getOrCreateTrackedEnvironment(lam).values.set(key, value);
+  lam.addEnvironment(key, value);
+}
+
+export function setTrackedPropTags(lam: LambdaFunction, value: string): void {
+  const tracked = getOrCreateTrackedEnvironment(lam);
+  replaceTags(tracked.propTags, value);
+  writeTags(lam, tracked);
+}
+
+export function mergeTrackedGitTags(lam: LambdaFunction, value: string): void {
+  const tracked = getOrCreateTrackedEnvironment(lam);
+  mergeTags(tracked.gitTags, value);
+  writeTags(lam, tracked);
+}
+
+export function hasTrackedEnv(lam: LambdaFunction, key: string): boolean {
+  const tracked = ddEnvTracker.get(lam);
+  return key === DD_TAGS ? (tracked?.tagsSet ?? false) : (tracked?.values.has(key) ?? false);
+}
+
+function getOrCreateTrackedEnvironment(lam: LambdaFunction): TrackedEnvironment {
+  let tracked = ddEnvTracker.get(lam);
+  if (!tracked) {
+    tracked = {
+      values: new Map(),
+      propTags: new Map(),
+      functionTags: new Map(),
+      gitTags: new Map(),
+      tagsSet: false,
+    };
+    ddEnvTracker.set(lam, tracked);
+  }
+  return tracked;
+}
+
+function replaceTags(tags: Tags, value: string): void {
+  tags.clear();
+  mergeTags(tags, value);
+}
+
+function mergeTags(tags: Tags, value: string): void {
+  for (const tag of value.split(",")) {
+    // Split only the key because tag values can contain colons.
+    const separator = tag.indexOf(":");
+    const key = separator > 0 ? tag.slice(0, separator) : tag;
+    tags.delete(key);
+    tags.set(key, tag);
+  }
+}
+
+function writeTags(lam: LambdaFunction, tracked: TrackedEnvironment): void {
+  tracked.tagsSet = true;
+  lam.addEnvironment(DD_TAGS, serializeTags(tracked));
+}
+
+function serializeTags(tracked: TrackedEnvironment): string {
+  const tags: Tags = new Map();
+  for (const source of [tracked.propTags, tracked.functionTags, tracked.gitTags]) {
+    for (const [key, tag] of source) {
+      // Move replaced tags to the position of the higher-precedence source.
+      tags.delete(key);
+      tags.set(key, tag);
+    }
+  }
+  return [...tags.values()].join(",");
+}