Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 Comment from Claude working with Aaron Silverman:

[P2] Move this section above "Opening Issues", add the variable to the configuration table, and drop the release-status paragraph. The README ships to PyPI.

  • The section is after ## License and the link definitions.
  • The variable is missing from the "Configuration" table.
  • pyproject.toml sets readme = "README.md", so the "pending coordinated release" paragraph goes to PyPI and becomes stale.
  • With the current pins, Python 3.8 and 3.9 do not get this behavior: pyproject.toml pins them to ddtrace>=3.19.1,<4, and the tracer PR targets ddtrace 4.x. The README does not say so.
Suggested README — assumes the rename and backoff comments are taken
--- a/README.md
+++ b/README.md
@@ -31,8 +31,17 @@ Besides the environment variables supported by dd-trace-py, the datadog-lambda-p
 | DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH | Determines the level of detail captured from AWS Lambda payloads, which are then assigned as tags for the `aws.lambda` span. It specifies the nesting depth of the JSON payload structure to process. Once the specified maximum depth is reached, the tag's value is set to the stringified value of any nested elements beyond this level.  <br> For example, given the input payload: <pre>{<br>  "lv1" : {<br>    "lv2": {<br>      "lv3": "val"<br>    }<br>  }<br>}</pre> If the depth is set to `2`, the resulting tag's key is set to `function.request.lv1.lv2` and the value is `{\"lv3\": \"val\"}`. <br> If the depth is set to `0`, the resulting tag's key is set to `function.request` and value is `{\"lv1\":{\"lv2\":{\"lv3\": \"val\"}}}` | `10` |
 | DD_EXCEPTION_REPLAY_ENABLED | When set to `true`, the Lambda will run with Error Tracking Exception Replay enabled, capturing local variables. | `false` |
 | DD_DSM_EXCHANGE_NAME | If you are using Datadog Data Streams Monitoring and propagating context through Amazon Event Bridge, this variable specifies the name of the EventBridge EventBus. The name of the bus is *not* passed through in the event payload and therefore cannot be automatically inferred at consume time. | `false` |
+| DD_FEATURE_FLAGS_FLUSH_TIMEOUT_MS | Maximum time, in milliseconds, that the wrapper waits at the end of an invocation to send pending feature flag evaluation and exposure events. `0` disables the wait. See [Feature flag telemetry at invocation end](#feature-flag-telemetry-at-invocation-end). | `500` |
 
 
+## Feature flag telemetry at invocation end
+
+The wrapper tries to send pending feature flag evaluation and exposure events before it returns a response or rethrows a handler exception. No handler code or provider shutdown is required, and providers stay available for warm invocations.
+
+`DD_FEATURE_FLAGS_FLUSH_TIMEOUT_MS` limits the added wait. The wait is also limited to the remaining invocation time minus 100 ms. After three timeouts in a row that use the full limit, the wrapper skips the wait for 60 seconds.
+
+Delivery is best effort. A timeout stops the wait but does not cancel a request in flight, which can complete after a warm thaw. Abrupt termination can still lose buffered events. This requires ddtrace `<version>` or later, so it is available on Python 3.10 and later.
+
 ## Opening Issues
 
 If you encounter a bug with this package, we want to hear about it. Before opening a new issue, search the existing issues to avoid duplicates.
@@ -76,27 +85,3 @@ 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.

If you don't take those, keep the current name and drop the backoff sentence. Replace <version> with the first ddtrace release that contains DataDog/dd-trace-py#21135.


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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we can remove this from the README. We don't have any other content like this for other products. We usually just expect people to use our public docs site instead.

3 changes: 3 additions & 0 deletions datadog_lambda/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.0, float

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 Comment from Claude working with Aaron Silverman:

[P2] Rename the variable to DD_FEATURE_FLAGS_FLUSH_TIMEOUT_MS, to match the tracer's plural feature flag prefix.

The tracer's DD_FEATURE_FLAGS_ENABLED, DD_FEATURE_FLAGS_CONFIGURATION_SOURCE and DD_FEATURE_FLAGS_CONFIGURATION_SOURCE_AGENTLESS_* settings use the plural prefix (ddtrace/internal/settings/openfeature.py). Other settings there use DD_FFE_, DD_FLAGGING_ or DD_EXPERIMENTAL_FLAGGING_, so a singular DD_FEATURE_FLAG_ would add another spelling. We want one prefix for feature flag variables across repositories. A name is hard to change after it ships, because a rename usually needs an alias and a deprecation period.

Suggested change — config, README, runbook and tests
--- a/README.md
+++ b/README.md
@@ -85,7 +85,7 @@ 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
+`DD_FEATURE_FLAGS_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
--- a/datadog_lambda/config.py
+++ b/datadog_lambda/config.py
@@ -86,7 +86,7 @@ class Config:
 
     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.0, float
+        "DD_FEATURE_FLAGS_FLUSH_TIMEOUT_MS", 500.0, float
     )
     flush_to_log = _get_env("DD_FLUSH_TO_LOG", "false", as_bool)
     logs_injection = _get_env("DD_LOGS_INJECTION", "true", as_bool)
--- a/tests/integration/FEATURE_FLAG_FLUSH.md
+++ b/tests/integration/FEATURE_FLAG_FLUSH.md
@@ -30,7 +30,7 @@ Verify that the selected profile can access the approved account before deployin
    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`.
+   service, `DD_ENV` to the test environment and `DD_FEATURE_FLAGS_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.
--- a/tests/test_feature_flags.py
+++ b/tests/test_feature_flags.py
@@ -72,9 +72,9 @@ class FeatureFlagFlushTest(unittest.TestCase):
                 feature_flags, "config", Config()
             ), patch.object(feature_flags, "_telemetry") as telemetry:
                 if value is None:
-                    os.environ.pop("DD_FEATURE_FLAG_FLUSH_TIMEOUT_MS", None)
+                    os.environ.pop("DD_FEATURE_FLAGS_FLUSH_TIMEOUT_MS", None)
                 else:
-                    os.environ["DD_FEATURE_FLAG_FLUSH_TIMEOUT_MS"] = value
+                    os.environ["DD_FEATURE_FLAGS_FLUSH_TIMEOUT_MS"] = value
                 telemetry.FLUSH_API_VERSION = 1
                 telemetry.flush.return_value = "completed"
                 feature_flags.flush_feature_flag_telemetry(

The full suite passes with the rename.

)
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)
Expand Down
38 changes: 38 additions & 0 deletions datadog_lambda/feature_flags.py
Original file line number Diff line number Diff line change
@@ -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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Once we get this PR more ready to merge, I'd like to run a before/after test on the cold start duration. We've found in the past that adding imports, when the code ultimately isn't being used here, can add significant amounts of time to the beginning of the first lambda invocation. From what I can tell, I don't think that will be the case here, but it will be good to know for certain. I'll run that testing once we get to a good point.



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.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 Comment from Claude working with Aaron Silverman:

[P3] Return early when context has no get_remaining_time_in_millis, and log the unsupported-tracer skip.

This comment says: "No remaining-time information means no safe invocation budget." The next line calls context.get_remaining_time_in_millis() without a check. A missing method raises AttributeError, and the outer except logs it at debug level as "Feature flag telemetry flush failed", with a traceback. The skip is correct, but the log reads like a failure.

Also, line 20 returns without a log when FLUSH_API_VERSION is not 1. If the tracer moves to contract version 2, the flush stops with no log.

Suggested change and tests
--- a/datadog_lambda/feature_flags.py
+++ b/datadog_lambda/feature_flags.py
@@ -17,14 +17,19 @@ SAFETY_MARGIN_MS = 100.0
 
 def flush_feature_flag_telemetry(context):
     try:
-        if _telemetry is None or getattr(_telemetry, "FLUSH_API_VERSION", None) != 1:
+        if _telemetry is None:
+            return
+        if getattr(_telemetry, "FLUSH_API_VERSION", None) != 1:
+            logger.debug("Skipping feature flag telemetry flush: unsupported tracer")
             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)
+        get_remaining_ms = getattr(context, "get_remaining_time_in_millis", None)
+        if get_remaining_ms is None:
+            return
+        budget_ms = min(cap_ms, get_remaining_ms() - SAFETY_MARGIN_MS)
         if budget_ms <= 0:
             logger.debug(
                 "Skipping feature flag telemetry flush: no invocation time remains"
--- a/tests/test_feature_flags.py
+++ b/tests/test_feature_flags.py
@@ -92,6 +92,26 @@ class FeatureFlagFlushTest(unittest.TestCase):
             )
             telemetry.flush.assert_called_once()
 
+    def test_unsupported_tracer_is_logged(self):
+        with patch.object(
+            feature_flags, "_telemetry", Mock(FLUSH_API_VERSION=2)
+        ), patch.object(feature_flags, "logger") as log:
+            feature_flags.flush_feature_flag_telemetry(
+                Mock(get_remaining_time_in_millis=lambda: 1000)
+            )
+            log.debug.assert_called_once_with(
+                "Skipping feature flag telemetry flush: unsupported tracer"
+            )
+
+    def test_missing_context_skips_without_failure_log(self):
+        with patch.object(feature_flags, "_telemetry") as telemetry, patch.object(
+            feature_flags, "logger"
+        ) as log:
+            telemetry.FLUSH_API_VERSION = 1
+            feature_flags.flush_feature_flag_telemetry(None)
+            telemetry.flush.assert_not_called()
+            log.debug.assert_not_called()
+
     def test_incomplete_diagnostic(self):
         with patch.object(feature_flags, "_telemetry") as telemetry, patch.object(
             feature_flags, "logger"

The full suite passes. Each new test fails without its change.

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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 Comment from Claude working with Aaron Silverman:

[P1] Back off after repeated timeouts that use the full cap. Without a backoff, a slow or unreachable intake adds up to the whole cap, 500 ms by default, to every invocation that has unsent flag events.

  • For a positive budget, flush() in fix(openfeature): add bounded invocation-end EVP flush dd-trace-py#21135 returns timed_out only after its deadline passes, so each timeout costs the full budget.
  • The tracer's direct-intake route has no cooldown. Only agentless-local routes have one (FeatureFlagEVPRouteSelector.send() in ddtrace/internal/openfeature/_evp_transport.py).
  • Before this change, the wrapper did not wait for feature flag intake at invocation end. During an intake incident, functions with unsent flag events can pay up to the cap on each invocation, and Lambda bills that time.

Datadog's serverless libraries already back off in this situation. datadogpy, which this wrapper uses to send metrics without the extension, stops sending for 300 s after 3 timeouts. datadog-lambda-go opens a circuit breaker for 60 s after more than 4 failures in 30 s. The suggested fix takes the threshold from datadogpy and the window from datadog-lambda-go: after 3 full-cap timeouts with no completed or failed flush between them, it skips the wait for 60 s. A timeout from a short remaining-time budget neither counts nor resets the streak. During the window, invocations skip the invocation-end flush, and the background writers keep sending on their own schedule. After the window, the next invocation flushes again.

Keep the 500 ms default and the 100 ms margin. The industry standard for serverless SDKs is to flush before the invocation returns, with caps of 1 s or longer, so 500 ms is already conservative. A reserve of about 100 ms of remaining time matches published guidance.

Precedent — Datadog libraries that already back off
  • datadogpy: _max_timeouts = 3 and _backoff_period = 300 in datadog/api/__init__.py. api_client.py logs "Max number of datadog timeouts exceeded, backing off for %s seconds" when the timeout counter reaches the maximum.
  • datadog-lambda-go: defaultCircuitBreakerInterval = 30s, defaultCircuitBreakerTimeout = 60s and defaultCircuitBreakerTotalFailures = 4 in internal/metrics/constants.go. The breaker trips when TotalFailures > 4.

A timeout caused by a short remaining-time budget does not count, because it says nothing about the intake. Any outcome other than timed_out resets the count.

Suggested fix and tests — each test fails when its branch is removed
--- a/datadog_lambda/feature_flags.py
+++ b/datadog_lambda/feature_flags.py
@@ -2,6 +2,7 @@
 
 import logging
 import math
+import time
 
 from datadog_lambda.config import config
 
@@ -13,9 +14,15 @@ except ImportError:
 
 logger = logging.getLogger(__name__)
 SAFETY_MARGIN_MS = 100.0
+# After repeated full-cap timeouts, stop adding the wait to every invocation for a while.
+MAX_CONSECUTIVE_TIMEOUTS = 3
+TIMEOUT_BACKOFF_S = 60.0
+_consecutive_timeouts = 0
+_skip_until = 0.0
 
 
 def flush_feature_flag_telemetry(context):
+    global _consecutive_timeouts, _skip_until
     try:
         if _telemetry is None or getattr(_telemetry, "FLUSH_API_VERSION", None) != 1:
             return
@@ -30,7 +37,17 @@ def flush_feature_flag_telemetry(context):
                 "Skipping feature flag telemetry flush: no invocation time remains"
             )
             return
+        if time.time() < _skip_until:
+            logger.debug("Skipping feature flag telemetry flush: backing off")
+            return
         outcome = _telemetry.flush(budget_ms / 1000.0)
+        if outcome != "timed_out":
+            _consecutive_timeouts = 0
+        elif budget_ms >= cap_ms:
+            _consecutive_timeouts += 1
+            if _consecutive_timeouts >= MAX_CONSECUTIVE_TIMEOUTS:
+                _consecutive_timeouts = 0
+                _skip_until = time.time() + TIMEOUT_BACKOFF_S
         if outcome != "completed":
             logger.debug("Feature flag telemetry flush incomplete: %s", outcome)
     except Exception:
--- a/tests/test_feature_flags.py
+++ b/tests/test_feature_flags.py
@@ -9,6 +9,37 @@ from datadog_lambda.config import Config
 
 
 class FeatureFlagFlushTest(unittest.TestCase):
+    def setUp(self):
+        for name, value in [("_consecutive_timeouts", 0), ("_skip_until", 0.0)]:
+            patcher = patch.object(feature_flags, name, value)
+            patcher.start()
+            self.addCleanup(patcher.stop)
+
+    def _flush_with_outcomes(self, outcomes, remaining_ms=30000):
+        context = Mock(get_remaining_time_in_millis=lambda: remaining_ms)
+        with patch.object(feature_flags, "_telemetry") as telemetry:
+            telemetry.FLUSH_API_VERSION = 1
+            telemetry.flush.side_effect = outcomes
+            for _ in outcomes:
+                feature_flags.flush_feature_flag_telemetry(context)
+            return telemetry.flush.call_count
+
+    def test_repeated_full_budget_timeouts_back_off(self):
+        with patch.object(feature_flags.time, "time", return_value=1000.0) as now:
+            self.assertEqual(self._flush_with_outcomes(["timed_out"] * 4), 3)
+            now.return_value = 1000.0 + feature_flags.TIMEOUT_BACKOFF_S
+            self.assertEqual(self._flush_with_outcomes(["completed"]), 1)
+
+    def test_success_resets_timeout_count(self):
+        outcomes = ["timed_out", "timed_out", "completed", "timed_out", "timed_out"]
+        self.assertEqual(self._flush_with_outcomes(outcomes), 5)
+        self.assertEqual(feature_flags._skip_until, 0.0)
+
+    def test_short_budget_timeouts_do_not_back_off(self):
+        self.assertEqual(
+            self._flush_with_outcomes(["timed_out"] * 4, remaining_ms=200), 4
+        )
+
     def test_budget(self):
         for remaining, cap, expected in [
             (5000, 500, 0.5),

The change adds module state, so FeatureFlagFlushTest needs the setUp shown, which resets it between tests. The full suite passes (1000 tests). Each new test fails when I remove its branch: the backoff, the full-cap check, or the reset on success.

if outcome != "completed":
logger.debug("Feature flag telemetry flush incomplete: %s", outcome)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I checked your other PR, do we even need to return something from the _telemetry.flush? If all we're doing is logging, then why not just log from dd-trace-py?

except Exception:
# Never replace the response or the original handler exception.
logger.debug("Feature flag telemetry flush failed", exc_info=True)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be very honest, this method just seems really over engineered as it is right now. Why the check of FLUSH_API_VERSION? And if we route the flushed telemetry data through the lambda extension (instead of agentless), the outbound flush http request will be very fast and wouldn't need all this extra overhead.

Instead, this PR could simply be these two additions to datadog_lambda/wrapper.py and the config addition:

# added to datadog_lambda/config.py
feature_flags_enabled = _get_env("DD_FEATURE_FLAGS_ENABLED", "true", as_bool)
# at top of file
_feature_flags = None
if config.feature_flags_enabled:
    try:
        from ddtrace.internal.openfeature import _telemetry as _feature_flags
    except ImportError:
        pass
# at end of self._after
if _feature_flags is not None:
    _feature_flags.flush()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, if we just call periodic on the writer directly, this PR could become just 5d3d3e3 and you wouldn't even need the ddtrace PR.

24 changes: 17 additions & 7 deletions datadog_lambda/wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from importlib import import_module
from time import time_ns

from datadog_lambda.feature_flags import flush_feature_flag_telemetry

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should put this import alongside where we do all the other feature imports from ddtrace (ex llmobs below).

You'll want to put the import behind a config check, so that way we're not doing the import unless the user has configured to enable the feature.

Then, where you call the actual flush_feature_flag_telemetry, you'll do another check of the config value.

Then, add DD_FEATURE_FLAGS_ENABLED to the config, alongside the DD_FEATURE_FLAG_FLUSH_TIMEOUT_MS you already added. This will allow users to turn this off and potentially avoid the import.

from datadog_lambda.extension import should_use_extension, flush_extension
from datadog_lambda.cold_start import (
set_cold_start,
Expand Down Expand Up @@ -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:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 Comment from Claude working with Aaron Silverman:

[P2] Add a test for this branch, and mention the change in the PR description. A handler's SystemExit or KeyboardInterrupt now propagates.

Since #835 (released in v8.126.0), this branch swallowed SystemExit, KeyboardInterrupt and any other BaseException that is not an Exception and whose class name does not contain BlockingException. The wrapper then returned self.response, which _before resets to None. Unless AppSec replaced the response, sys.exit(1) in a handler produced a successful null response. This else: raise restores the behaviour from before #835. That is the right behaviour, but handlers that call sys.exit() will now fail the invocation, and no test in tests/ raises SystemExit or KeyboardInterrupt.

Suggested test — fails with main's behaviour
--- a/tests/test_feature_flags.py
+++ b/tests/test_feature_flags.py
@@ -105,6 +105,25 @@ class FeatureFlagFlushTest(unittest.TestCase):
                 "Feature flag telemetry flush incomplete: %s", "timed_out"
             )
 
+    def test_non_blocking_base_exception_propagates_and_flushes(self):
+        for raised in [SystemExit(3), KeyboardInterrupt()]:
+            decorated = object.__new__(wrapper._LambdaDecorator)
+            decorated.func = Mock(side_effect=raised)
+            decorated.response = None
+            decorated.blocking_response = None
+            decorated.span = None
+            decorated._before = Mock()
+            decorated._after = Mock()
+            context = Mock(get_remaining_time_in_millis=lambda: 30000)
+            with self.subTest(raised=type(raised).__name__), patch.object(
+                wrapper, "flush_feature_flag_telemetry"
+            ) as flush:
+                with self.assertRaises(type(raised)) as caught:
+                    decorated({}, context)
+                self.assertIs(caught.exception, raised)
+                decorated._after.assert_called_once()
+                flush.assert_called_once_with(context)
+
     def test_wrapper_preserves_result_and_exception_when_cleanup_fails(self):
         for raises in [False, True]:
             original = ValueError("handler error")

With else: raise removed, which is main's behaviour, both subtests fail with SystemExit not raised and KeyboardInterrupt not raised. The full suite passes with the test added.

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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_after already catches and logs everything, so the outer try/except is a no-op. Can we call flush_feature_flag_telemetry(context) inside _after alongside LLMObs.flush(), before flush_stats/flush_extension? That matches the other flushes. If an early _after failure skipping flushes is a concern, it applies to all of them and belongs in a separate change. The BaseException re-raise and the submit_errors_metric hardening also look unrelated to this PR.

The other changes here are also unnecessary and can be reverted.

if self.blocking_response:
return self.blocking_response
return self.response
Expand Down
69 changes: 69 additions & 0 deletions tests/integration/FEATURE_FLAG_FLUSH.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
# 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. Registration waits for provider initialization
before the first invocation, bounded by the provider's default 10-second timeout.

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==0.10.0` (which provides
`set_provider_and_wait`). 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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems more like a PR description than something we'd want to commit to the repo. I think we should remove it.

49 changes: 49 additions & 0 deletions tests/integration/feature_flags.py
Original file line number Diff line number Diff line change
@@ -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_and_wait(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
Loading
Loading