SVLS-9158: Add tag exclusion for customer metrics enrichment

bottlecap/src/bin/bottlecap/main.rs

@@ -1219,8 +1219,25 @@ async fn start_dogstatsd(
 ) {
     // Start aggregator service and handle
     let start_time = Instant::now();
+    let enrichment_tags = if config.custom_metrics_exclude_tags.is_empty() {
+        tags_provider.get_tags_string()
+    } else {
+        debug!(
+            "Excluding tags from custom metrics: {:?}",
+            config.custom_metrics_exclude_tags
+        );
+        tags_provider
+            .get_tags_vec()
+            .into_iter()
+            .filter(|tag| {
+                let key = tag.split(':').next().unwrap_or("");
+                !config.custom_metrics_exclude_tags.iter().any(|e| e == key)
+            })
+            .collect::<Vec<_>>()
+            .join(",")
+    };
     let (aggregator_service, aggregator_handle) = MetricsAggregatorService::new(
-        SortedTags::parse(&tags_provider.get_tags_string()).unwrap_or(EMPTY_TAGS),
+        SortedTags::parse(&enrichment_tags).unwrap_or(EMPTY_TAGS),
         CONTEXTS,
     )
     .expect("can't create metrics service");

bottlecap/src/config/env.rs

@@ -283,6 +283,14 @@ pub struct EnvConfig {
     #[serde(deserialize_with = "deserialize_optional_string")]
     pub statsd_metric_namespace: Option<String>,
 
+    /// @env `DD_LAMBDA_CUSTOMER_METRICS_EXCLUDE_TAGS`
+    ///
+    /// Comma-separated list of tag keys to exclude from custom `DogStatsD` metrics
+    /// enrichment. Use this to drop auto-injected tags (e.g. `function_arn,region`)
+    /// from custom metrics to reduce billing.
+    #[serde(deserialize_with = "deserialize_array_from_comma_separated_string")]
+    pub lambda_customer_metrics_exclude_tags: Vec<String>,
+
     /// @env `DD_DOGSTATSD_SO_RCVBUF`
     /// Size of the receive buffer for `DogStatsD` UDP packets, in bytes (`SO_RCVBUF`).
     /// Increase to reduce packet loss under high-throughput metric bursts.
@@ -585,6 +593,13 @@ fn merge_config(config: &mut Config, env_config: &EnvConfig) {
         config.statsd_metric_namespace = parse_metric_namespace(namespace);
     }
 
+    merge_vec!(
+        config,
+        custom_metrics_exclude_tags,
+        env_config,
+        lambda_customer_metrics_exclude_tags
+    );
+
     // DogStatsD
     merge_option!(config, env_config, dogstatsd_so_rcvbuf);
     merge_option!(config, env_config, dogstatsd_buffer_size);
@@ -869,6 +884,11 @@ mod tests {
             );
             jail.set_env("DD_OTLP_CONFIG_LOGS_ENABLED", "true");
 
+            jail.set_env(
+                "DD_LAMBDA_CUSTOMER_METRICS_EXCLUDE_TAGS",
+                "function_arn,region",
+            );
+
             // DogStatsD
             jail.set_env("DD_DOGSTATSD_SO_RCVBUF", "1048576");
             jail.set_env("DD_DOGSTATSD_BUFFER_SIZE", "65507");
@@ -1029,6 +1049,7 @@ mod tests {
                 otlp_config_traces_probabilistic_sampler_sampling_percentage: Some(50),
                 otlp_config_logs_enabled: true,
                 statsd_metric_namespace: None,
+                custom_metrics_exclude_tags: vec!["function_arn".to_string(), "region".to_string()],
                 dogstatsd_so_rcvbuf: Some(1_048_576),
                 dogstatsd_buffer_size: Some(65507),
                 dogstatsd_queue_size: Some(2048),
@@ -1242,6 +1263,49 @@ mod tests {
         });
     }
 
+    #[test]
+    fn test_custom_metrics_exclude_tags_from_env() {
+        figment::Jail::expect_with(|jail| {
+            jail.clear_env();
+            jail.set_env(
+                "DD_LAMBDA_CUSTOMER_METRICS_EXCLUDE_TAGS",
+                "function_arn,region,account_id",
+            );
+
+            let mut config = Config::default();
+            let env_config_source = EnvConfigSource;
+            env_config_source
+                .load(&mut config)
+                .expect("Failed to load config");
+
+            assert_eq!(
+                config.custom_metrics_exclude_tags,
+                vec![
+                    "function_arn".to_string(),
+                    "region".to_string(),
+                    "account_id".to_string()
+                ]
+            );
+            Ok(())
+        });
+    }
+
+    #[test]
+    fn test_custom_metrics_exclude_tags_defaults_to_empty() {
+        figment::Jail::expect_with(|jail| {
+            jail.clear_env();
+
+            let mut config = Config::default();
+            let env_config_source = EnvConfigSource;
+            env_config_source
+                .load(&mut config)
+                .expect("Failed to load config");
+
+            assert!(config.custom_metrics_exclude_tags.is_empty());
+            Ok(())
+        });
+    }
+
     #[test]
     fn test_dogstatsd_config_defaults_to_none() {
         figment::Jail::expect_with(|jail| {

bottlecap/src/config/mod.rs

@@ -304,6 +304,7 @@ pub struct Config {
     // Metrics
     pub metrics_config_compression_level: i32,
     pub statsd_metric_namespace: Option<String>,
+    pub custom_metrics_exclude_tags: Vec<String>,
     /// Size of the receive buffer for `DogStatsD` UDP packets, in bytes (`SO_RCVBUF`).
     /// Increase to reduce packet loss under high-throughput metric bursts.
     /// If None, uses the OS default.
@@ -436,6 +437,8 @@ impl Default for Config {
             metrics_config_compression_level: 3,
             statsd_metric_namespace: None,
 
+            custom_metrics_exclude_tags: vec![],
+
             // DogStatsD
             // Defaults to None, which uses the OS default.
             dogstatsd_so_rcvbuf: None,

bottlecap/src/config/yaml.rs

@@ -1033,6 +1033,7 @@ api_security_sample_delay: 60 # Seconds
                 apm_filter_tags_regex_require: None,
                 apm_filter_tags_regex_reject: None,
                 statsd_metric_namespace: None,
+                custom_metrics_exclude_tags: vec![],
                 dogstatsd_so_rcvbuf: Some(1_048_576),
                 dogstatsd_buffer_size: Some(65507),
                 dogstatsd_queue_size: Some(2048),

integration-tests/bin/app.ts

@@ -6,6 +6,7 @@ import {Otlp} from '../lib/stacks/otlp';
 import {Snapstart} from '../lib/stacks/snapstart';
 import {LambdaManagedInstancesStack} from '../lib/stacks/lmi';
 import {AuthStack} from '../lib/stacks/auth';
+import {CustomMetrics} from '../lib/stacks/custom-metrics';
 import {AuthRoleStack} from '../lib/auth-role';
 import {ACCOUNT, getIdentifier, REGION} from '../config';
 import {CapacityProviderStack} from "../lib/capacity-provider";
@@ -40,6 +41,9 @@ const stacks = [
     new AuthStack(app, `integ-${identifier}-auth`, {
         env,
     }),
+    new CustomMetrics(app, `integ-${identifier}-custom-metrics`, {
+        env,
+    }),
 ]
 
 // Tag all stacks so we can easily clean them up

integration-tests/lambda/custom-metrics-node/index.js

@@ -0,0 +1,28 @@
+const dgram = require("dgram");
+
+exports.handler = async (event, context) => {
+  console.log("Sending custom DogStatsD metric");
+
+  await new Promise((resolve, reject) => {
+    const client = dgram.createSocket("udp4");
+    const message = Buffer.from("custom.exclude_tags_test:1|c");
+    client.send(message, 8125, "127.0.0.1", (err) => {
+      client.close();
+      if (err) reject(err);
+      else resolve();
+    });
+  });
+
+  // Give the extension a moment to receive the UDP packet
+  await new Promise((resolve) => setTimeout(resolve, 100));
+
+  console.log("Custom metric sent");
+
+  return {
+    statusCode: 200,
+    body: JSON.stringify({
+      message: "Success",
+      requestId: context.awsRequestId,
+    }),
+  };
+};

integration-tests/lib/stacks/custom-metrics.ts

@@ -0,0 +1,69 @@
+import * as cdk from "aws-cdk-lib";
+import * as lambda from "aws-cdk-lib/aws-lambda";
+import { Construct } from "constructs";
+import {
+  createLogGroup,
+  defaultDatadogEnvVariables,
+  defaultDatadogSecretPolicy,
+  getExtensionLayer,
+  getDefaultNodeLayer,
+  defaultNodeRuntime,
+} from "../util";
+
+export class CustomMetrics extends cdk.Stack {
+  constructor(scope: Construct, id: string, props: cdk.StackProps) {
+    super(scope, id, props);
+
+    const extensionLayer = getExtensionLayer(this);
+    const nodeLayer = getDefaultNodeLayer(this);
+
+    // Lambda function WITHOUT tag exclusion — all enrichment tags present
+    const unfilteredFunctionName = `${id}-unfiltered-lambda`;
+    const unfilteredFunction = new lambda.Function(
+      this,
+      unfilteredFunctionName,
+      {
+        runtime: defaultNodeRuntime,
+        architecture: lambda.Architecture.ARM_64,
+        handler: "/opt/nodejs/node_modules/datadog-lambda-js/handler.handler",
+        code: lambda.Code.fromAsset("./lambda/custom-metrics-node"),
+        functionName: unfilteredFunctionName,
+        timeout: cdk.Duration.seconds(30),
+        memorySize: 256,
+        environment: {
+          ...defaultDatadogEnvVariables,
+          DD_SERVICE: unfilteredFunctionName,
+          DD_TRACE_ENABLED: "true",
+          DD_LAMBDA_HANDLER: "index.handler",
+        },
+        logGroup: createLogGroup(this, unfilteredFunctionName),
+      },
+    );
+    unfilteredFunction.addToRolePolicy(defaultDatadogSecretPolicy);
+    unfilteredFunction.addLayers(extensionLayer);
+    unfilteredFunction.addLayers(nodeLayer);
+
+    // Lambda function WITH tag exclusion — function_arn and region excluded
+    const filteredFunctionName = `${id}-filtered-lambda`;
+    const filteredFunction = new lambda.Function(this, filteredFunctionName, {
+      runtime: defaultNodeRuntime,
+      architecture: lambda.Architecture.ARM_64,
+      handler: "/opt/nodejs/node_modules/datadog-lambda-js/handler.handler",
+      code: lambda.Code.fromAsset("./lambda/custom-metrics-node"),
+      functionName: filteredFunctionName,
+      timeout: cdk.Duration.seconds(30),
+      memorySize: 256,
+      environment: {
+        ...defaultDatadogEnvVariables,
+        DD_SERVICE: filteredFunctionName,
+        DD_TRACE_ENABLED: "true",
+        DD_LAMBDA_HANDLER: "index.handler",
+        DD_LAMBDA_CUSTOMER_METRICS_EXCLUDE_TAGS: "function_arn,region",
+      },
+      logGroup: createLogGroup(this, filteredFunctionName),
+    });
+    filteredFunction.addToRolePolicy(defaultDatadogSecretPolicy);
+    filteredFunction.addLayers(extensionLayer);
+    filteredFunction.addLayers(nodeLayer);
+  }
+}

integration-tests/tests/custom-metrics.test.ts

@@ -0,0 +1,96 @@
+import { hasMetricWithTag } from "./utils/datadog";
+import { forceColdStart, invokeLambda } from "./utils/lambda";
+import { getIdentifier, DEFAULT_DATADOG_INDEXING_WAIT_MS } from "../config";
+
+const identifier = getIdentifier();
+const stackName = `integ-${identifier}-custom-metrics`;
+
+const CUSTOM_METRIC_NAME = "custom.exclude_tags_test";
+const EXCLUDED_TAGS = ["function_arn", "region"];
+const KEPT_TAGS = ["functionname", "account_id"];
+
+describe("Customer Metrics Exclude Tags Integration Tests", () => {
+  let invocationStartTime: number;
+  let metricsEndTime: number;
+
+  const unfilteredFunctionName = `${stackName}-unfiltered-lambda`;
+  const filteredFunctionName = `${stackName}-filtered-lambda`;
+
+  beforeAll(async () => {
+    const functionNames = [unfilteredFunctionName, filteredFunctionName];
+
+    await Promise.all(functionNames.map((fn) => forceColdStart(fn)));
+
+    // Back up the query window by 60s so the metric bucket (which Datadog
+    // aligns to the rollup interval boundary, often before the invocation)
+    // falls inside the range we pass to /api/v1/query.
+    invocationStartTime = Date.now() - 60_000;
+
+    await Promise.all(functionNames.map((fn) => invokeLambda(fn)));
+
+    await new Promise((resolve) =>
+      setTimeout(resolve, DEFAULT_DATADOG_INDEXING_WAIT_MS),
+    );
+
+    metricsEndTime = Date.now();
+
+    console.log("Lambdas invoked and indexing wait complete");
+  }, 900000);
+
+  describe("unfiltered function (no DD_LAMBDA_CUSTOMER_METRICS_EXCLUDE_TAGS)", () => {
+    it.each(EXCLUDED_TAGS)(
+      "should have %s tag on custom metric",
+      async (tag) => {
+        const hasTag = await hasMetricWithTag(
+          CUSTOM_METRIC_NAME,
+          unfilteredFunctionName,
+          `${tag}:*`,
+          invocationStartTime,
+          metricsEndTime,
+        );
+        expect(hasTag).toBe(true);
+      },
+    );
+
+    it.each(KEPT_TAGS)("should have %s tag on custom metric", async (tag) => {
+      const hasTag = await hasMetricWithTag(
+        CUSTOM_METRIC_NAME,
+        unfilteredFunctionName,
+        `${tag}:*`,
+        invocationStartTime,
+        metricsEndTime,
+      );
+      expect(hasTag).toBe(true);
+    });
+  });
+
+  describe("filtered function (DD_LAMBDA_CUSTOMER_METRICS_EXCLUDE_TAGS=function_arn,region)", () => {
+    it.each(EXCLUDED_TAGS)(
+      "should NOT have %s tag on custom metric",
+      async (tag) => {
+        const hasTag = await hasMetricWithTag(
+          CUSTOM_METRIC_NAME,
+          filteredFunctionName,
+          `${tag}:*`,
+          invocationStartTime,
+          metricsEndTime,
+        );
+        expect(hasTag).toBe(false);
+      },
+    );
+
+    it.each(KEPT_TAGS)(
+      "should still have %s tag on custom metric",
+      async (tag) => {
+        const hasTag = await hasMetricWithTag(
+          CUSTOM_METRIC_NAME,
+          filteredFunctionName,
+          `${tag}:*`,
+          invocationStartTime,
+          metricsEndTime,
+        );
+        expect(hasTag).toBe(true);
+      },
+    );
+  });
+});

integration-tests/tests/utils/datadog.ts

@@ -320,3 +320,34 @@ async function getMetrics(
     value: p[1],
   }));
 }
+
+/**
+ * Query a custom metric filtered by function name and an additional tag filter.
+ * Use tagFilter like "function_arn:*" to check if a tag exists on the metric.
+ * Returns true if the query returns data points, false otherwise.
+ */
+export async function hasMetricWithTag(
+  metricName: string,
+  functionName: string,
+  tagFilter: string,
+  fromTime: number,
+  toTime: number,
+): Promise<boolean> {
+  const baseFunctionName = getServiceName(functionName).toLowerCase();
+  const query = `avg:${metricName}{functionname:${baseFunctionName},${tagFilter}}`;
+
+  console.log(`Querying metric with tag filter: ${query}`);
+
+  const response = await datadogClient.get('/api/v1/query', {
+    params: {
+      query,
+      from: Math.floor(fromTime / 1000),
+      to: Math.floor(toTime / 1000),
+    },
+  });
+
+  const series = response.data.series || [];
+  const hasData = series.some((s: any) => Array.isArray(s.pointlist) && s.pointlist.length > 0);
+  console.log(`Tag filter query returned ${series.length} series, hasData=${hasData}`);
+  return hasData;
+}