[codex] test(telemetry): deduplicate fork-cloned heartbeats (#7304)

Co-authored-by: Munir Abdinur <munir.abdinur@datadoghq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 3 people

manifests/python.yml

@@ -2510,9 +2510,6 @@ manifest:
   tests/test_telemetry.py::Test_Telemetry: v1.16.0
   tests/test_telemetry.py::Test_Telemetry::test_api_still_v1: irrelevant
   tests/test_telemetry.py::Test_Telemetry::test_app_dependencies_loaded: irrelevant
-  tests/test_telemetry.py::Test_Telemetry::test_app_heartbeats_delays:
-    - weblog_declaration:
-        uwsgi-poc: flaky (APMAPI-2167)
   tests/test_telemetry.py::Test_Telemetry::test_app_product_change: missing_feature (Weblog GET/enable_product and app-product-change event is not implemented yet.)
   tests/test_telemetry.py::Test_Telemetry::test_app_started_is_first_message:
     - declaration: flaky (APMAPI-1858)

tests/test_telemetry.py

@@ -9,6 +9,7 @@
 
 from dateutil.parser import isoparse
 
+from tests.test_telemetry_heartbeat_utils import heartbeat_delays_by_runtime
 from utils import bug, context, features, interfaces, logger, rfc, scenarios, weblog
 from utils.interfaces._misc_validators import HeadersMatchValidator, HeadersPresenceValidator
 from utils.telemetry import get_lang_configs, load_telemetry_json
@@ -312,59 +313,25 @@ def save_data(data: dict, container: dict):
             raise Exception("The following telemetry messages were not forwarded by the agent")
 
     @staticmethod
-    def _get_heartbeat_delays_by_runtime() -> dict:
+    def _get_heartbeat_delays_by_runtime() -> dict[str, list[float]]:
         """Returns a dict where :
         The key is the runtime id
         The value is a list of delay observed on this runtime id
         """
 
-        telemetry_data = list(interfaces.library.get_telemetry_data())
-        heartbeats_by_runtime = defaultdict(list)
-
-        for data in telemetry_data:
-            if data["request"]["content"].get("request_type") == "app-heartbeat":
-                heartbeats_by_runtime[data["request"]["content"]["runtime_id"]].append(data)
-
-        delays_by_runtime = {}
-
-        # Measure only long-lived application runtimes, not the short-lived children the session-id
-        # tests spawn via /spawn_child (a forked child can emit under its parent's runtime_id before
-        # regenerating its own -- a sub-interval "duplicate" that flakes the check). Select by lifespan,
-        # not heartbeat count, so a slow-drifting worker with fewer heartbeats still gets measured.
-        def lifespan(heartbeats: list[Any]) -> float:
-            times = [isoparse(d["request"]["timestamp_start"]) for d in heartbeats]
-            return (max(times) - min(times)).total_seconds()
-
-        heartbeat_counts = {rid: len(hbs) for rid, hbs in heartbeats_by_runtime.items()}
-        lifespans = {rid: lifespan(hbs) for rid, hbs in heartbeats_by_runtime.items()}
-        longest = max(lifespans.values(), default=0.0)
-        measurable_runtimes = {
-            rid: hbs for rid, hbs in heartbeats_by_runtime.items() if len(hbs) > 2 and lifespans[rid] >= longest * 0.5
-        }
-        assert measurable_runtimes, (
+        delays_by_runtime, heartbeat_counts = heartbeat_delays_by_runtime(
+            interfaces.library.get_telemetry_data(flatten_message_batches=False),
+            # Require 2 intervals of lifespan so a short-lived process (e.g. a forked child
+            # from the session-id tests) can't skew the average with just a couple of
+            # samples. A normal long-lived runtime clears this easily.
+            min_lifespan=context.telemetry_heartbeat_interval * 2,
+        )
+        assert delays_by_runtime, (
             f"No runtime emitted enough heartbeats to check delays (runtimes seen: {heartbeat_counts})"
         )
 
-        for runtime_id, heartbeats in measurable_runtimes.items():
-            logger.debug(f"Heartbeats for runtime {runtime_id}:")
-
-            # In theory, it's sorted. Let be safe
-            heartbeats.sort(key=lambda data: isoparse(data["request"]["timestamp_start"]))
-
-            prev_message_time = None
-            delays: list[float] = []
-            for data in heartbeats:
-                curr_message_time = isoparse(data["request"]["timestamp_start"])
-                if prev_message_time is None:
-                    logger.debug(f"  * {data['log_filename']}: {curr_message_time}")
-                else:
-                    delay = (curr_message_time - prev_message_time).total_seconds()
-                    logger.debug(f"  * {data['log_filename']}: {curr_message_time} => {delay}s ellapsed")
-                    delays.append(delay)
-
-                prev_message_time = curr_message_time
-
-            delays_by_runtime[runtime_id] = delays
+        for runtime_id, delays in delays_by_runtime.items():
+            logger.debug(f"Heartbeat delays for runtime {runtime_id}: {delays}")
 
         return delays_by_runtime
 

tests/test_telemetry_heartbeat_utils.py

@@ -0,0 +1,71 @@
+import itertools
+from collections import defaultdict
+from collections.abc import Iterable
+from typing import Any
+
+from dateutil.parser import isoparse
+
+
+def heartbeat_delays_by_runtime(
+    telemetry_data: Iterable[dict[str, Any]],
+    *,
+    min_lifespan: float = 0.0,
+) -> tuple[dict[str, list[float]], dict[str, int]]:
+    """Return heartbeat delays and heartbeat counts, grouped by runtime ID.
+
+    Expects unflattened telemetry data (`get_telemetry_data(flatten_message_batches=False)`):
+    a flattened message-batch entry inherits the outer batch's seq_id, so heartbeats packed
+    in the same batch would otherwise collide on (runtime_id, seq_id) alone and get deduped
+    into one.
+
+    `min_lifespan` drops runtimes whose heartbeats span less time than that (e.g. a forked
+    child that exits after only 2-3 heartbeats). With so few samples, a real (non-duplicate)
+    heartbeat sent during shutdown can dominate the average delay -- dropping the runtime is
+    the only way to avoid measuring it. A long-lived runtime with the same anomaly is still
+    measured, and can still fail.
+    """
+    heartbeats_by_runtime: dict[str, dict[tuple[int, int], dict[str, Any]]] = defaultdict(dict)
+
+    for data in telemetry_data:
+        content: dict[str, Any] = data["request"]["content"]
+        runtime_id: str = content.get("runtime_id", "")
+        seq_id: int = content.get("seq_id", 0)
+
+        if content.get("request_type") == "message-batch":
+            entries = list(enumerate(content.get("payload", [])))
+        else:
+            entries = [(0, content)]
+
+        for batch_index, entry in entries:
+            if entry.get("request_type") != "app-heartbeat":
+                continue
+
+            # Dedup key (seq_id, batch_index) catches two cases: a real retry (identical
+            # resend), and a forked child's clone heartbeat. A child briefly shares its
+            # parent's runtime_id and copies its seq_id counter at fork time, so it can
+            # independently emit a heartbeat with the same seq_id. Not a guaranteed
+            # invariant, but the best signal we have to tell "fork clone" apart from
+            # "genuinely fast heartbeat" without tracking PIDs.
+            heartbeats_by_runtime[runtime_id].setdefault((seq_id, batch_index), data)
+
+    heartbeat_counts = {runtime_id: len(heartbeats) for runtime_id, heartbeats in heartbeats_by_runtime.items()}
+    delays_by_runtime: dict[str, list[float]] = {}
+
+    for runtime_id, heartbeats_by_key in heartbeats_by_runtime.items():
+        if len(heartbeats_by_key) <= 2:
+            continue
+
+        heartbeats = sorted(
+            heartbeats_by_key.values(),
+            key=lambda data: isoparse(data["request"]["timestamp_start"]),
+        )
+        times = [isoparse(data["request"]["timestamp_start"]) for data in heartbeats]
+        lifespan = (times[-1] - times[0]).total_seconds()
+        if lifespan < min_lifespan:
+            continue
+
+        delays_by_runtime[runtime_id] = [
+            (current - previous).total_seconds() for previous, current in itertools.pairwise(times)
+        ]
+
+    return delays_by_runtime, heartbeat_counts

tests/test_the_test/test_telemetry_heartbeat.py

@@ -0,0 +1,97 @@
+from datetime import datetime, timedelta, UTC
+from typing import Any
+
+import pytest
+
+from tests.test_telemetry_heartbeat_utils import heartbeat_delays_by_runtime
+
+
+pytestmark = pytest.mark.scenario("TEST_THE_TEST")
+
+BASE_TIME = datetime(2026, 1, 1, tzinfo=UTC)
+
+
+def _heartbeat(runtime_id: str, seq_id: int, offset: float) -> dict[str, Any]:
+    timestamp = BASE_TIME + timedelta(seconds=offset)
+    return {
+        "request": {
+            "timestamp_start": timestamp.isoformat(),
+            "content": {
+                "request_type": "app-heartbeat",
+                "runtime_id": runtime_id,
+                "seq_id": seq_id,
+            },
+        },
+        "log_filename": f"{runtime_id}-{seq_id}-{offset}.json",
+    }
+
+
+def test_fork_duplicate_does_not_affect_parent_cadence() -> None:
+    messages = [
+        _heartbeat("parent", 1, 0),
+        _heartbeat("parent", 2, 2),
+        _heartbeat("parent", 2, 2.05),
+        _heartbeat("parent", 3, 4),
+        _heartbeat("parent", 4, 6),
+        _heartbeat("child", 1, 2.1),
+        _heartbeat("child", 2, 4.1),
+    ]
+    delays, heartbeat_counts = heartbeat_delays_by_runtime(messages)
+
+    assert heartbeat_counts == {"parent": 4, "child": 2}
+    assert set(delays) == {"parent"}
+    assert delays["parent"] == pytest.approx([2.0, 2.0, 2.0])
+
+
+def test_distinct_fast_heartbeats_remain_measurable() -> None:
+    messages = [
+        _heartbeat("parent", 1, 0),
+        _heartbeat("parent", 2, 1),
+        _heartbeat("parent", 3, 2),
+        _heartbeat("parent", 4, 3),
+    ]
+    delays, heartbeat_counts = heartbeat_delays_by_runtime(messages)
+
+    assert heartbeat_counts == {"parent": 4}
+    assert delays["parent"] == pytest.approx([1.0, 1.0, 1.0])
+
+
+def test_short_lived_runtime_excluded_by_min_lifespan() -> None:
+    """A short-lived runtime with too few samples shouldn't be measured at all: a shutdown
+    heartbeat right before exit isn't a duplicate, so dedup can't drop it, and with only
+    2-3 samples it would skew the average.
+    """
+    messages = [
+        _heartbeat("parent", 1, 0),
+        _heartbeat("parent", 2, 2),
+        _heartbeat("parent", 3, 4),
+        _heartbeat("parent", 4, 6),
+        _heartbeat("parent", 5, 8),
+        # child lives ~2.05s: two normal-ish heartbeats, then an out-of-cadence one
+        # right before it exits
+        _heartbeat("child", 1, 0.1),
+        _heartbeat("child", 2, 2.05),
+        _heartbeat("child", 3, 2.12),
+    ]
+    delays, heartbeat_counts = heartbeat_delays_by_runtime(messages, min_lifespan=6.0)
+
+    assert heartbeat_counts == {"parent": 5, "child": 3}
+    assert set(delays) == {"parent"}
+
+
+def test_long_lived_runtime_anomaly_still_measured() -> None:
+    """min_lifespan only excludes runtimes too short-lived to measure. A long-lived runtime
+    with the same anomalous fast heartbeat is still measured, and can still fail.
+    """
+    messages = [
+        _heartbeat("parent", 1, 0),
+        _heartbeat("parent", 2, 2),
+        _heartbeat("parent", 3, 4),
+        _heartbeat("parent", 4, 6),
+        _heartbeat("parent", 5, 8),
+        _heartbeat("parent", 6, 8.05),  # anomalous fast heartbeat, but lifespan is 8.05s
+    ]
+    delays, heartbeat_counts = heartbeat_delays_by_runtime(messages, min_lifespan=6.0)
+
+    assert heartbeat_counts == {"parent": 6}
+    assert delays["parent"] == pytest.approx([2.0, 2.0, 2.0, 2.0, 0.05])

utils/scripts/libraries_and_scenarios_rules.yml

@@ -547,6 +547,12 @@ patterns:
         scenario_groups: null
         libraries: null
 
+    # This helper has no test items of its own, so it's invisible to the scenario map that
+    # drives dynamic selection below. List its consumers' scenarios explicitly.
+    tests/test_telemetry_heartbeat_utils.py:
+        scenario_groups: null
+        scenarios: [DEFAULT, TEST_THE_TEST]
+
     tests/fuzzer/*:
         scenario_groups: exotics
         libraries: null