diff --git a/datadog_lambda/config.py b/datadog_lambda/config.py index 6487c9e7..dbee6c2b 100644 --- a/datadog_lambda/config.py +++ b/datadog_lambda/config.py @@ -85,6 +85,10 @@ 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_flags_enabled = _get_env("DD_FEATURE_FLAGS_ENABLED", "true", as_bool) + feature_flag_flush_timeout_ms = _get_env( + "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) merge_xray_traces = _get_env("DD_MERGE_XRAY_TRACES", "false", as_bool) diff --git a/datadog_lambda/feature_flags.py b/datadog_lambda/feature_flags.py new file mode 100644 index 00000000..de3d3bc1 --- /dev/null +++ b/datadog_lambda/feature_flags.py @@ -0,0 +1,61 @@ +"""Compatibility boundary for invocation-end feature flag telemetry.""" + +import logging +import math +from time import monotonic + +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 +MAX_CONSECUTIVE_TIMEOUTS = 3 +TIMEOUT_BACKOFF_S = 60.0 +# Warm invocations share this state; a new process starts with no cooldown. +_consecutive_timeouts = 0 +_skip_until = 0.0 + + +def flush_feature_flag_telemetry(context): + global _consecutive_timeouts, _skip_until + try: + 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. + 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" + ) + return + if monotonic() < _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: + # A short invocation budget says nothing about intake health. + _consecutive_timeouts += 1 + if _consecutive_timeouts >= MAX_CONSECUTIVE_TIMEOUTS: + _consecutive_timeouts = 0 + _skip_until = monotonic() + TIMEOUT_BACKOFF_S + 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) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 73057983..acaff5ce 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -77,6 +77,10 @@ if config.llmobs_enabled: from ddtrace.llmobs import LLMObs +flush_feature_flag_telemetry = None +if config.feature_flags_enabled: + from datadog_lambda.feature_flags import flush_feature_flag_telemetry + if config.exception_replay_enabled: from ddtrace.debugging._exception.replay import SpanExceptionHandler @@ -400,6 +404,12 @@ def _after(self, event, context): except Exception as e: logger.debug("Failed to create cold start spans. %s", e) + if ( + config.feature_flags_enabled + and flush_feature_flag_telemetry is not None + ): + flush_feature_flag_telemetry(context) + if not config.flush_to_log or should_use_extension: from datadog_lambda.metric import flush_stats diff --git a/tests/integration/feature_flags.py b/tests/integration/feature_flags.py new file mode 100644 index 00000000..837e0bc9 --- /dev/null +++ b/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_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 diff --git a/tests/test_config.py b/tests/test_config.py index 9c5da63a..bf24b1f7 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -101,6 +101,7 @@ def _test_as_list(env_key, conf_key, default): *_test_as_bool("DD_MERGE_XRAY_TRACES", "merge_xray_traces", default=False), *_test_as_bool("DD_PROFILING_ENABLED", "profiling_enabled", default=False), *_test_as_bool("DD_LLMOBS_ENABLED", "llmobs_enabled", default=False), + *_test_as_bool("DD_FEATURE_FLAGS_ENABLED", "feature_flags_enabled", default=True), *_test_as_bool( "DD_EXCEPTION_REPLAY_ENABLED", "exception_replay_enabled", default=False ), diff --git a/tests/test_feature_flags.py b/tests/test_feature_flags.py new file mode 100644 index 00000000..d55479b4 --- /dev/null +++ b/tests/test_feature_flags.py @@ -0,0 +1,236 @@ +import os +import subprocess +import sys +import textwrap +import unittest +from unittest.mock import Mock, patch + +from datadog_lambda import feature_flags +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 test_full_budget_timeouts_back_off_and_retry_after_expiry(self): + with patch.object( + feature_flags, "monotonic", return_value=1000.0 + ) as now, patch.object(feature_flags, "_telemetry") as telemetry, patch.object( + Config, "feature_flag_flush_timeout_ms", 250 + ): + telemetry.FLUSH_API_VERSION = 1 + telemetry.flush.return_value = "timed_out" + context = Mock(get_remaining_time_in_millis=lambda: 30000) + for _ in range(3): + feature_flags.flush_feature_flag_telemetry(context) + self.assertEqual(telemetry.flush.call_count, 3) + telemetry.flush.assert_called_with(0.25) + for timestamp in (1000.0, 1030.0, 1059.999): + now.return_value = timestamp + feature_flags.flush_feature_flag_telemetry(context) + self.assertEqual(telemetry.flush.call_count, 3) + now.return_value = 1060.0 + telemetry.flush.return_value = "completed" + feature_flags.flush_feature_flag_telemetry(context) + self.assertEqual(telemetry.flush.call_count, 4) + # Recovery starts a new streak, not an immediate second cooldown. + telemetry.flush.return_value = "timed_out" + for _ in range(2): + feature_flags.flush_feature_flag_telemetry(context) + telemetry.flush.return_value = "completed" + feature_flags.flush_feature_flag_telemetry(context) + self.assertEqual(telemetry.flush.call_count, 7) + + def test_completed_and_failed_flushes_reset_timeout_streak(self): + for outcome in ("completed", "failed"): + with self.subTest(outcome=outcome), patch.object( + feature_flags, "_telemetry" + ) as telemetry, patch.object( + feature_flags, "monotonic", return_value=1000.0 + ): + telemetry.FLUSH_API_VERSION = 1 + outcomes = [ + "timed_out", + "timed_out", + outcome, + "timed_out", + "timed_out", + "completed", + ] + telemetry.flush.side_effect = outcomes + for _ in outcomes: + feature_flags.flush_feature_flag_telemetry( + Mock(get_remaining_time_in_millis=lambda: 30000) + ) + self.assertEqual(telemetry.flush.call_count, len(outcomes)) + + def test_short_budget_timeouts_neither_count_nor_reset_streak(self): + with patch.object(feature_flags, "_telemetry") as telemetry, patch.object( + feature_flags, "monotonic", return_value=1000.0 + ): + telemetry.FLUSH_API_VERSION = 1 + telemetry.flush.return_value = "timed_out" + for remaining in (30000, 200, 200, 200, 30000, 30000): + feature_flags.flush_feature_flag_telemetry( + Mock(get_remaining_time_in_millis=lambda: remaining) + ) + self.assertEqual(telemetry.flush.call_count, 6) + feature_flags.flush_feature_flag_telemetry( + Mock(get_remaining_time_in_millis=lambda: 30000) + ) + self.assertEqual(telemetry.flush.call_count, 6) + + 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 self.subTest(telemetry=telemetry), patch.object( + feature_flags, "_telemetry", telemetry + ): + context = Mock() + context.get_remaining_time_in_millis.return_value = 1000 + feature_flags.flush_feature_flag_telemetry(context) + context.get_remaining_time_in_millis.assert_not_called() + if telemetry is not None: + telemetry.flush.assert_not_called() + + def test_timeout_environment_parsing_preserves_flush(self): + for value, expected in [ + (None, 0.5), + ("250.5", 0.2505), + ("", 0.5), + ("garbage", 0.5), + ("500ms", 0.5), + ]: + with self.subTest(value=value), patch.dict(os.environ), patch.object( + feature_flags, "config", Config() + ), patch.object(feature_flags, "_telemetry") as telemetry: + if value is None: + os.environ.pop("DD_FEATURE_FLAGS_FLUSH_TIMEOUT_MS", None) + else: + 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( + Mock(get_remaining_time_in_millis=lambda: 1000) + ) + telemetry.flush.assert_called_once_with(expected) + + 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_missing_remaining_time_skips_without_failure_log(self): + for context in (None, object()): + with self.subTest(context=context), 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(context) + telemetry.flush.assert_not_called() + log.debug.assert_not_called() + + def test_unsupported_tracer_skip_is_logged(self): + with patch.object( + feature_flags, "_telemetry", Mock(FLUSH_API_VERSION=2) + ) as telemetry, patch.object(feature_flags, "logger") as log: + feature_flags.flush_feature_flag_telemetry( + Mock(get_remaining_time_in_millis=lambda: 1000) + ) + telemetry.flush.assert_not_called() + log.debug.assert_called_once_with( + "Skipping feature flag telemetry flush: unsupported tracer" + ) + + def test_disabled_feature_flags_do_not_import_flush_modules(self): + code = textwrap.dedent( + """ + import sys + + class BlockFlushImports: + def find_spec(self, fullname, *args): + if fullname in ( + "datadog_lambda.feature_flags", + "ddtrace.internal.openfeature._telemetry", + ): + raise AssertionError("unexpected import: " + fullname) + + sys.meta_path.insert(0, BlockFlushImports()) + from datadog_lambda import wrapper + assert not wrapper.config.feature_flags_enabled + assert "datadog_lambda.feature_flags" not in sys.modules + assert "ddtrace.internal.openfeature._telemetry" not in sys.modules + """ + ) + result = subprocess.run( + [sys.executable, "-c", code], + env={ + **os.environ, + "DD_FEATURE_FLAGS_ENABLED": "false", + "DD_INSTRUMENTATION_TELEMETRY_ENABLED": "false", + }, + capture_output=True, + text=True, + timeout=30, + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + + 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" + ) diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 26a53b0c..5521e0f6 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -4,6 +4,8 @@ import unittest import importlib import sys +import types +from contextlib import ExitStack import pytest from unittest.mock import MagicMock, patch, call, ANY @@ -13,6 +15,7 @@ import datadog_lambda.xray as xray from datadog_lambda.config import config +from datadog_lambda import feature_flags from datadog_lambda.metric import lambda_metric from datadog_lambda.thread_stats_writer import ThreadStatsWriter from ddtrace.trace import Span, tracer @@ -77,6 +80,225 @@ def setUp(self): self.mock_dd_lambda_layer_tag = patcher.start() self.addCleanup(patcher.stop) + @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 = MagicMock(configuration_source="agentless") + selector.select.return_value = route + selector.send.side_effect = lambda route, send: send(route) + connection = MagicMock() + 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 = wrapper.datadog_lambda_wrapper(handler) + context = get_mock_context() + context.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) + + def test_real_decorator_calls_feature_flag_flush_contract(self): + import ddtrace.internal.openfeature as openfeature_pkg + from datadog_lambda import feature_flags + + calls = [] + fake = types.ModuleType("ddtrace.internal.openfeature._telemetry") + fake.FLUSH_API_VERSION = 1 + fake.flush = lambda timeout: calls.append(timeout) or "completed" + self.addCleanup(importlib.reload, feature_flags) + with patch.dict(sys.modules, {fake.__name__: fake}), patch.object( + openfeature_pkg, "_telemetry", fake, create=True + ), patch.object(type(config), "feature_flags_enabled", True): + importlib.reload(feature_flags) + context = get_mock_context() + context.get_remaining_time_in_millis = lambda: 30000 + + handler = wrapper.datadog_lambda_wrapper( + lambda event, context: {"ok": True} + ) + self.assertEqual(handler({}, context), {"ok": True}) + + error = ValueError("handler error") + + def raising(event, context): + raise error + + with self.assertRaises(ValueError) as caught: + wrapper.datadog_lambda_wrapper(raising)({}, context) + self.assertIs(caught.exception, error) + self.assertEqual(calls, [0.5, 0.5]) + + def test_feature_flag_outcomes_allow_later_flushes(self): + for outcome in ("completed", "failed", "timed_out", "exception", "cooldown"): + for raises in (False, True): + with self.subTest(outcome=outcome, raises=raises), ExitStack() as stack: + for name, value in ( + ("feature_flags_enabled", True), + ("trace_enabled", False), + ("cold_start_tracing", False), + ("llmobs_enabled", True), + ("exception_replay_enabled", True), + ("local_test", True), + ): + stack.enter_context(patch.object(type(config), name, value)) + stack.enter_context( + patch.object(wrapper, "should_use_extension", True) + ) + stack.enter_context( + patch.object(wrapper, "SpanExceptionHandler", create=True) + ) + stack.enter_context( + patch.object(feature_flags, "_consecutive_timeouts", 0) + ) + stack.enter_context( + patch.object( + feature_flags, + "_skip_until", + 1060.0 if outcome == "cooldown" else 0.0, + ) + ) + stack.enter_context( + patch.object(feature_flags, "monotonic", return_value=1000.0) + ) + telemetry = stack.enter_context( + patch.object(feature_flags, "_telemetry") + ) + telemetry.FLUSH_API_VERSION = 1 + events = [] + + def flush_flags(timeout): + events.append("flags") + if outcome == "exception": + raise RuntimeError("telemetry unavailable") + return outcome + + telemetry.flush.side_effect = flush_flags + stack.enter_context( + patch( + "datadog_lambda.metric.flush_stats", + side_effect=lambda context: events.append("metrics"), + ) + ) + stack.enter_context( + patch.object( + wrapper, + "flush_extension", + side_effect=lambda: events.append("extension"), + ) + ) + llmobs = stack.enter_context( + patch.object(wrapper, "LLMObs", create=True) + ) + llmobs.flush.side_effect = lambda: events.append("llmobs") + replay = stack.enter_context( + patch.object(wrapper, "SignalUploader", create=True) + ) + replay._instance.periodic.side_effect = lambda: events.append( + "replay" + ) + error = ValueError("handler error") + + def handler(event, context): + if raises: + raise error + return {"ok": True} + + decorated = wrapper.datadog_lambda_wrapper(handler) + context = get_mock_context() + context.get_remaining_time_in_millis = lambda: 30000 + if raises: + with self.assertRaises(ValueError) as caught: + decorated({}, context) + self.assertIs(caught.exception, error) + else: + self.assertEqual(decorated({}, context), {"ok": True}) + expected = [] if outcome == "cooldown" else ["flags"] + self.assertEqual( + events, expected + ["metrics", "extension", "llmobs", "replay"] + ) + + def test_early_cleanup_failure_skips_feature_flag_flush(self): + with patch.object(type(config), "feature_flags_enabled", True), patch.object( + wrapper, "flush_feature_flag_telemetry" + ) as flush, patch( + "datadog_lambda.metric.submit_batch_item_failures_metric", + side_effect=RuntimeError("earlier cleanup failure"), + ): + for raises in (False, True): + with self.subTest(raises=raises): + error = ValueError("handler error") + + def handler(event, context): + if raises: + raise error + return {"ok": True} + + decorated = wrapper.datadog_lambda_wrapper(handler) + if raises: + with self.assertRaises(ValueError) as caught: + decorated({}, get_mock_context()) + self.assertIs(caught.exception, error) + else: + self.assertEqual( + decorated({}, get_mock_context()), {"ok": True} + ) + flush.assert_not_called() + + def test_disabled_feature_flags_skip_flush(self): + with patch.object(type(config), "feature_flags_enabled", False), patch.object( + wrapper, "flush_feature_flag_telemetry" + ) as flush: + context = get_mock_context() + context.get_remaining_time_in_millis = lambda: 30000 + handler = wrapper.datadog_lambda_wrapper( + lambda event, context: {"ok": True} + ) + self.assertEqual(handler({}, context), {"ok": True}) + + error = ValueError("handler error") + + def raising(event, context): + raise error + + with self.assertRaises(ValueError) as caught: + wrapper.datadog_lambda_wrapper(raising)({}, context) + self.assertIs(caught.exception, error) + flush.assert_not_called() + @patch("datadog_lambda.config.Config.trace_enabled", False) def test_datadog_lambda_wrapper(self): @wrapper.datadog_lambda_wrapper