fix: flush feature flag telemetry at invocation end

README.md

@@ -76,3 +76,27 @@ Unless explicitly stated otherwise all files in this repository are licensed und
 This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2019 Datadog, Inc.
 
 [1]: https://www.datadoghq.com/blog/troubleshoot-lambda-function-request-response-payloads/
+
+## Feature flag telemetry at invocation end
+
+With a tracer and Lambda layer that support automatic feature flag flushing,
+the wrapper attempts to send pending evaluation and exposure events before
+returning a response or rethrowing a handler exception. No additional handler
+code or provider shutdown is required; providers and polling remain available
+for warm invocations.
+
+`DD_FEATURE_FLAG_FLUSH_TIMEOUT_MS` sets the maximum added wait in milliseconds
+(default: `500`; `0` disables waiting). The actual budget is the smaller of that
+cap and the invocation's remaining time minus a 100 ms safety margin. The wrapper
+skips waiting when that budget is exhausted or unavailable. Incomplete flushes
+are reported in debug logs without changing the handler's result.
+
+Delivery is best effort. A timeout stops waiting but does not cancel or retry
+in-flight requests; they may complete after a warm thaw. Abrupt termination can
+still lose buffered events. Endpoint acceptance does not guarantee immediate
+visibility in the UI. Older tracers safely skip the integration.
+
+Release status: this capability requires the tracer's internal flush contract
+version 1. The first supported tracer/library/layer combination is pending
+coordinated release and deployed validation; an older published layer does not
+provide this behavior merely by setting the timeout variable.

datadog_lambda/config.py

@@ -85,6 +85,9 @@ def _resolve_env(self, key, default=None, cast=None, depends_on_tracing=False):
     enhanced_metrics_enabled = _get_env("DD_ENHANCED_METRICS", "true", as_bool)
 
     flush_in_thread = _get_env("DD_FLUSH_IN_THREAD", "false", as_bool)
+    feature_flag_flush_timeout_ms = _get_env(
+        "DD_FEATURE_FLAG_FLUSH_TIMEOUT_MS", "500", float
+    )
     flush_to_log = _get_env("DD_FLUSH_TO_LOG", "false", as_bool)
     logs_injection = _get_env("DD_LOGS_INJECTION", "true", as_bool)
     merge_xray_traces = _get_env("DD_MERGE_XRAY_TRACES", "false", as_bool)

datadog_lambda/feature_flags.py

@@ -0,0 +1,38 @@
+"""Compatibility boundary for invocation-end feature flag telemetry."""
+
+import logging
+import math
+
+from datadog_lambda.config import config
+
+try:
+    from ddtrace.internal.openfeature import _telemetry
+except ImportError:
+    _telemetry = None
+
+
+logger = logging.getLogger(__name__)
+SAFETY_MARGIN_MS = 100.0
+
+
+def flush_feature_flag_telemetry(context):
+    try:
+        if _telemetry is None or getattr(_telemetry, "FLUSH_API_VERSION", None) != 1:
+            return
+        cap_ms = config.feature_flag_flush_timeout_ms
+        if not math.isfinite(cap_ms) or cap_ms <= 0:
+            return
+        # No remaining-time information means no safe invocation budget.
+        remaining_ms = context.get_remaining_time_in_millis()
+        budget_ms = min(cap_ms, remaining_ms - SAFETY_MARGIN_MS)
+        if budget_ms <= 0:
+            logger.debug(
+                "Skipping feature flag telemetry flush: no invocation time remains"
+            )
+            return
+        outcome = _telemetry.flush(budget_ms / 1000.0)
+        if outcome != "completed":
+            logger.debug("Feature flag telemetry flush incomplete: %s", outcome)
+    except Exception:
+        # Never replace the response or the original handler exception.
+        logger.debug("Feature flag telemetry flush failed", exc_info=True)

datadog_lambda/wrapper.py

@@ -9,6 +9,7 @@
 from importlib import import_module
 from time import time_ns
 
+from datadog_lambda.feature_flags import flush_feature_flag_telemetry
 from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import (
     set_cold_start,
@@ -191,18 +192,27 @@ def __call__(self, event, context, **kwargs):
                 return self.blocking_response
             self.response = self.func(event, context, **kwargs)
         except Exception:
-            from datadog_lambda.metric import submit_errors_metric
-
-            submit_errors_metric(context)
-
-            if self.span:
-                self.span.set_traceback()
+            try:
+                from datadog_lambda.metric import submit_errors_metric
+
+                submit_errors_metric(context)
+                if self.span:
+                    self.span.set_traceback()
+            except Exception as e:
+                logger.error(format_err_with_traceback(e))
             raise
         except BaseException as e:
             if "BlockingException" in type(e).__name__:
                 self.blocking_response = get_asm_blocked_response(self.event_source)
+            else:
+                raise
         finally:
-            self._after(event, context)
+            try:
+                self._after(event, context)
+            except Exception as e:
+                logger.error(format_err_with_traceback(e))
+            finally:
+                flush_feature_flag_telemetry(context)
         if self.blocking_response:
             return self.blocking_response
         return self.response

tests/integration/FEATURE_FLAG_FLUSH.md

@@ -0,0 +1,67 @@
+# Feature flag invocation-end smoke test
+
+Status: local wrapper/tracer integration is covered in `tests/test_feature_flags.py`.
+The deployed ingestion test has not been run. Do not claim a supported published
+layer until the release record below is populated.
+
+## Existing examples and access
+
+The smoke handler initializes its provider and client once at module scope and
+reuses them across invocations.
+
+This repository's `scripts/run_integration_tests.sh` documents
+`sso-serverless-sandbox-account-admin`. `tests/integration/serverless.yml` uses
+`eu-west-1`, a sandbox execution role, and a dedicated deployment bucket. Its
+IAM restrictions require names beginning with `integration-tests`. Use that
+approved sandbox or obtain an explicitly approved replacement account/role.
+Verify that the selected profile can access the approved account before deploying.
+
+## Assemble and deploy
+
+1. Build a Python 3.11 layer containing both candidate changes: the tracer's
+   internal flush contract version 1 and this Lambda library. The ordinary layer
+   build installs the declared tracer dependency, so verify the zip actually
+   contains the candidate tracer instead of a published version without the fix.
+2. Package `feature_flags.py` and `openfeature-sdk`. Configure its handler as
+   `feature_flags.handle`. It uses the ordinary Datadog wrapper and has no flush
+   or shutdown calls. Reuse the approved sandbox's IAM and deployment settings
+   from `serverless.yml`, with a unique `integration-tests-*` function name.
+3. Configure `FFL_SMOKE_FLAG_KEY` with an existing string test flag that produces
+   exposure events for the test subjects. Set `DD_SERVICE` to a unique smoke-test
+   service, `DD_ENV` to the test environment and `DD_FEATURE_FLAG_FLUSH_TIMEOUT_MS=500`.
+   Keep evaluation telemetry and exposure intake enabled.
+4. For direct delivery, use `DD_FEATURE_FLAGS_CONFIGURATION_SOURCE=agentless`,
+   the test organization's `DD_API_KEY`/`DD_SITE`, and no local Agent endpoint.
+   Supply the API key through the sandbox's approved secret mechanism.
+5. Validate Agent delivery separately with a reachable Agent that advertises the
+   relevant EVP routes. The Lambda extension does not acquire new intake routes
+   from this change; do not treat extension presence as proof of Agent EVP support.
+
+## Validate
+
+Invoke the same function three times with distinct synthetic subjects, for
+example `{"subject":"ffl3412-<run>-1"}`, then a different subject with
+`"raise":true`, then a third success. Distinct subjects avoid intentional exposure
+deduplication. Confirm the second invocation returns the intentional handler
+error, the other two succeed, and logs show the same `sandbox_id` with increasing
+invocation counts. A changed sandbox ID means warm reuse was not demonstrated.
+
+Check both evaluation and exposure ingestion for every subject at the configured
+Datadog site, including the raising invocation. Capture the invocation request
+IDs and ingestion evidence. `evaluation_ms` should be below the background flush
+intervals; measure invocation duration separately because cleanup adds wait.
+Endpoint acceptance can precede UI visibility. No further invocation or deliberate
+sleep in the handler should be needed to trigger delivery.
+
+Repeat with a slow/unavailable endpoint and a short remaining-time budget; verify
+the original response/error is retained and the extra caller wait stays within
+the configured cap or remaining time minus 100 ms. Inspect debug diagnostics for
+incomplete flushes. Remove only the smoke-test stack/resources created by the run.
+
+## Release record
+
+Record the first supported **published** tracer version, Lambda library version,
+and layer ARN/version together, plus runtime, architecture, configuration source,
+service/environment, request IDs, and ingestion evidence. All are pending deployed
+validation and release coordination. Release the tracer first, then the wrapper
+and layer that contain it. Older tracer versions safely skip the capability.

tests/integration/feature_flags.py

@@ -0,0 +1,49 @@
+"""Short-invocation smoke handler; use an exposure-enabled string test flag."""
+
+import json
+import os
+import time
+import uuid
+
+from datadog_lambda.wrapper import datadog_lambda_wrapper
+from ddtrace.openfeature import DataDogProvider
+from openfeature import api
+from openfeature.evaluation_context import EvaluationContext
+
+provider = DataDogProvider()
+api.set_provider(provider)
+client = api.get_client("lambda-flush-smoke")
+sandbox_id = uuid.uuid4().hex
+invocations = 0
+
+
+@datadog_lambda_wrapper
+def handle(event, context):
+    global invocations
+    invocations += 1
+    started = time.monotonic()
+    subject = event["subject"]
+    details = client.get_string_details(
+        os.environ["FFL_SMOKE_FLAG_KEY"],
+        "smoke-default",
+        EvaluationContext(targeting_key=subject),
+    )
+    result = {
+        "request_id": context.aws_request_id,
+        "sandbox_id": sandbox_id,
+        "invocation": invocations,
+        "subject": subject,
+        "value": details.value,
+        "variant": details.variant,
+        "error_code": str(details.error_code) if details.error_code else None,
+        "evaluation_ms": (time.monotonic() - started) * 1000,
+    }
+    # Only synthetic test subjects belong in this smoke-test log.
+    print(json.dumps(result))
+    if details.error_code:
+        raise RuntimeError(
+            "Smoke flag did not resolve; verify test flag and configuration"
+        )
+    if event.get("raise"):
+        raise RuntimeError("intentional smoke-test handler failure")
+    return result

tests/test_feature_flags.py

@@ -0,0 +1,160 @@
+import importlib.util
+import unittest
+from unittest.mock import Mock, patch
+
+from datadog_lambda import feature_flags
+from datadog_lambda import wrapper
+
+
+class FeatureFlagFlushTest(unittest.TestCase):
+    def test_budget(self):
+        for remaining, cap, expected in [
+            (5000, 500, 0.5),
+            (200, 500, 0.1),
+            (101, 500, 0.001),
+        ]:
+            with self.subTest(remaining=remaining), patch.object(
+                feature_flags, "_telemetry"
+            ) as telemetry:
+                telemetry.FLUSH_API_VERSION = 1
+                telemetry.flush.return_value = "completed"
+                with patch.object(
+                    type(feature_flags.config), "feature_flag_flush_timeout_ms", cap
+                ):
+                    feature_flags.flush_feature_flag_telemetry(
+                        Mock(get_remaining_time_in_millis=lambda: remaining)
+                    )
+                telemetry.flush.assert_called_once_with(expected)
+
+    def test_no_budget(self):
+        for remaining, cap in [
+            (100, 500),
+            (0, 500),
+            (1000, 0),
+            (1000, -1),
+            (1000, float("nan")),
+        ]:
+            with self.subTest(remaining=remaining, cap=cap), patch.object(
+                feature_flags, "_telemetry"
+            ) as telemetry:
+                telemetry.FLUSH_API_VERSION = 1
+                with patch.object(
+                    type(feature_flags.config), "feature_flag_flush_timeout_ms", cap
+                ):
+                    feature_flags.flush_feature_flag_telemetry(
+                        Mock(get_remaining_time_in_millis=lambda: remaining)
+                    )
+                telemetry.flush.assert_not_called()
+
+    def test_old_tracer(self):
+        for telemetry in [None, Mock(FLUSH_API_VERSION=2)]:
+            with patch.object(feature_flags, "_telemetry", telemetry):
+                feature_flags.flush_feature_flag_telemetry(None)
+
+    def test_missing_context_and_telemetry_errors_are_safe(self):
+        with patch.object(feature_flags, "_telemetry") as telemetry:
+            telemetry.FLUSH_API_VERSION = 1
+            telemetry.flush.side_effect = RuntimeError("unavailable")
+            feature_flags.flush_feature_flag_telemetry(None)
+            feature_flags.flush_feature_flag_telemetry(
+                Mock(get_remaining_time_in_millis=lambda: 1000)
+            )
+            telemetry.flush.assert_called_once()
+
+    def test_incomplete_diagnostic(self):
+        with patch.object(feature_flags, "_telemetry") as telemetry, patch.object(
+            feature_flags, "logger"
+        ) as log:
+            telemetry.FLUSH_API_VERSION = 1
+            telemetry.flush.return_value = "timed_out"
+            feature_flags.flush_feature_flag_telemetry(
+                Mock(get_remaining_time_in_millis=lambda: 1000)
+            )
+            log.debug.assert_called_with(
+                "Feature flag telemetry flush incomplete: %s", "timed_out"
+            )
+
+    def test_wrapper_preserves_result_and_exception_when_cleanup_fails(self):
+        for raises in [False, True]:
+            original = ValueError("handler error")
+            handler = (
+                Mock(side_effect=original)
+                if raises
+                else Mock(return_value={"ok": True})
+            )
+            # Isolate invocation control flow from tracing/metrics initialization.
+            decorated = object.__new__(wrapper._LambdaDecorator)
+            decorated.func = handler
+            decorated.blocking_response = None
+            decorated.span = None
+            decorated._before = Mock()
+            decorated._after = Mock(side_effect=RuntimeError("cleanup error"))
+            context = Mock(get_remaining_time_in_millis=lambda: 30000)
+            with patch.object(wrapper, "flush_feature_flag_telemetry") as flush, patch(
+                "datadog_lambda.metric.submit_errors_metric",
+                side_effect=RuntimeError("metrics error"),
+            ):
+                if raises:
+                    with self.assertRaises(ValueError) as caught:
+                        decorated({}, context)
+                    self.assertIs(caught.exception, original)
+                else:
+                    self.assertEqual(decorated({}, context), {"ok": True})
+                flush.assert_called_once_with(context)
+
+
+class FeatureFlagTracerIntegrationTest(unittest.TestCase):
+    @unittest.skipUnless(
+        feature_flags._telemetry is not None
+        and importlib.util.find_spec("openfeature") is not None,
+        "requires tracer flush contract v1 and the optional OpenFeature SDK",
+    )
+    def test_short_warm_invocations_flush_both_writers(self):
+        from ddtrace.internal.openfeature._flagevaluation_writer import (
+            FlagEvaluationWriter,
+            _EvalEvent,
+        )
+        from ddtrace.internal.openfeature._evp_transport import EVPRoute
+        from ddtrace.internal.openfeature.writer import ExposureWriter
+
+        route = EVPRoute("http://agent:8126", "/evp_proxy/v2", {})
+        selector = Mock(configuration_source="agentless")
+        selector.select.return_value = route
+        selector.send.side_effect = lambda route, send: send(route)
+        connection = Mock()
+        connection.getresponse.return_value.status = 202
+        evaluations = FlagEvaluationWriter(
+            route_selector=selector, connection_factory=lambda *a, **k: connection
+        )
+        exposures = ExposureWriter(
+            enabled=True,
+            route_selector=selector,
+            connection_factory=lambda *a, **k: connection,
+        )
+        original = ValueError("handler failure")
+
+        def handler(event, context):
+            evaluations.enqueue(
+                _EvalEvent("flag", "on", "allocation", "subject", {}, False, "", 1)
+            )
+            exposures.enqueue({"timestamp": 1})
+            if event.get("raise"):
+                raise original
+            return {"ok": True}
+
+        decorated = object.__new__(wrapper._LambdaDecorator)
+        decorated.func = handler
+        decorated.blocking_response = None
+        decorated.span = None
+        decorated._before = Mock()
+        decorated._after = Mock(side_effect=RuntimeError("unrelated cleanup"))
+        context = Mock(get_remaining_time_in_millis=lambda: 30000)
+        with patch("datadog_lambda.metric.submit_errors_metric"):
+            self.assertEqual(decorated({}, context), {"ok": True})
+            self.assertEqual(connection.request.call_count, 2)
+            with self.assertRaises(ValueError) as caught:
+                decorated({"raise": True}, context)
+            self.assertIs(caught.exception, original)
+            self.assertEqual(connection.request.call_count, 4)
+            self.assertEqual(decorated({}, context), {"ok": True})
+            self.assertEqual(connection.request.call_count, 6)