fix(traces): retry a refused batch within a timed flush, and start the budget once the lock is held

A flush with budget left stopped at the first retriable failure, so a
30 s shutdown flush made one attempt and then discarded the backlog. A
caller-driven flush now waits out the backoff and retries while budget
remains, with one last attempt at the deadline; timer flushes and
untimed flushes are unchanged. The budget starts once no other flush is
in flight, and a flush that never gets the lock says so at debug. A
size-1 413 restores the batch size it halved from, and the ramp doubles
instead of adding one. The resource is encoded once per exporter.

Author: turnipdabeets

posthog/test/tracing/test_export.py

@@ -1,4 +1,5 @@
 import threading
+import time
 from types import SimpleNamespace
 from unittest import mock
 
@@ -152,8 +153,9 @@ def test_a_spent_budget_still_ships_one_batch(self):
         assert queued(pipeline) == []
 
     def test_returns_without_draining_when_another_flush_holds_the_lock_past_the_deadline(
-        self,
+        self, caplog
     ):
+        caplog.set_level("DEBUG", logger="posthog")
         pipeline, sender, _ = make_traces()
         pipeline.start_span("a").end()
         timer = FakeTimer.instances[-1]
@@ -165,6 +167,31 @@ def test_returns_without_draining_when_another_flush_holds_the_lock_past_the_dea
         assert sender.payloads == []
         assert pipeline._exporter._flush_timer is timer
         assert not timer.cancelled
+        assert "another flush was still in flight" in caplog.text
+
+    def test_the_budget_starts_once_the_lock_is_held(self, clock):
+        pipeline, sender, _ = make_traces(max_export_batch_size=1)
+        exporter = pipeline._exporter
+
+        def send_slowly(client, payload):
+            sender.payloads.append(payload)
+            clock["now"] += 0.1
+            return SendOutcome("ok")
+
+        exporter._send = send_slowly
+        for _ in range(2):
+            pipeline.start_span("a").end()
+        exporter._flush_lock.acquire()
+
+        def release_after_a_while():
+            time.sleep(0.05)
+            clock["now"] += 1.0
+            exporter._flush_lock.release()
+
+        threading.Thread(target=release_after_a_while).start()
+        pipeline.flush(timeout=0.5)
+        assert len(sender.payloads) == 2
+        assert queued(pipeline) == []
 
     def test_re_arms_after_a_timer_that_failed_to_start(self):
         class FlakyTimer(FakeTimer):
@@ -263,9 +290,19 @@ def test_ramps_the_batch_size_back_up_after_a_413_shrink(self):
         for i in range(8):
             pipeline.start_span(str(i)).end()
         pipeline.flush()
-        assert pipeline._exporter._max_export_batch_size == 6
+        assert pipeline._exporter._max_export_batch_size == 8
         assert [len(b) for b in sender.batches()] == [8, 4, 4]
 
+    def test_restores_the_batch_size_once_a_413_isolates_the_oversized_span(self):
+        sender = FakeSender(*([SendOutcome("too-large")] * 4), SendOutcome("ok"))
+        pipeline, _, _ = make_traces(sender=sender, max_export_batch_size=8)
+        for i in range(8):
+            pipeline.start_span(str(i)).end()
+        pipeline.flush()
+        assert [len(b) for b in sender.batches()] == [8, 4, 2, 1, 7]
+        assert pipeline._exporter._max_export_batch_size == 8
+        assert queued(pipeline) == []
+
     def test_a_batch_measured_too_large_locally_splits_only_that_drain(self):
         sender = FakeSender(TOO_LARGE_LOCALLY, SendOutcome("ok"))
         pipeline, _, _ = make_traces(sender=sender, max_export_batch_size=8)
@@ -911,3 +948,88 @@ def test_a_forked_child_drops_the_inherited_queue_and_timer(self):
         )
         pipeline.start_span("child-span").end()
         assert [r.name for r in queued(pipeline)] == ["child-span"]
+
+
+def waits_advance(clock, pipeline):
+    """Make the exporter's backoff wait move the fake clock instead of sleeping."""
+    waited = []
+
+    def wait(seconds):
+        waited.append(seconds)
+        clock["now"] += seconds
+        return False
+
+    pipeline._exporter._wait_for_retry = wait
+    return waited
+
+
+class TestRetryWithinBudget:
+    def test_retries_a_retriable_failure_after_its_backoff(self, clock):
+        sender = FakeSender(SendOutcome("retry-later"), SendOutcome("ok"))
+        pipeline, _, _ = make_traces(sender=sender)
+        waited = waits_advance(clock, pipeline)
+        pipeline.start_span("a").end()
+        pipeline.flush(timeout=30)
+        assert len(sender.payloads) == 2
+        assert waited == [5]
+        assert queued(pipeline) == []
+
+    def test_makes_a_last_attempt_at_the_deadline(self, clock):
+        sender = FakeSender(SendOutcome("retry-later"))
+        pipeline, _, _ = make_traces(sender=sender)
+        waited = waits_advance(clock, pipeline)
+        pipeline.start_span("a").end()
+        pipeline.flush(timeout=12)
+        # Attempts at 0, 5 and 12: the second backoff of 10 is cut to the budget.
+        assert len(sender.payloads) == 3
+        assert waited == [5, 7]
+        assert len(queued(pipeline)) == 1
+
+    def test_honours_a_retry_after_within_the_budget(self, clock):
+        sender = FakeSender(SendOutcome("retry-later", 60), SendOutcome("ok"))
+        pipeline, _, _ = make_traces(sender=sender)
+        waited = waits_advance(clock, pipeline)
+        pipeline.start_span("a").end()
+        pipeline.flush(timeout=40)
+        assert waited == [MAX_RETRY_AFTER_SECONDS]
+        assert queued(pipeline) == []
+
+    def test_a_timer_flush_does_not_retry(self, clock):
+        sender = FakeSender(SendOutcome("retry-later"))
+        pipeline, _, _ = make_traces(sender=sender)
+        waited = waits_advance(clock, pipeline)
+        pipeline.start_span("a").end()
+        FakeTimer.instances[-1].fire()
+        assert len(sender.payloads) == 1
+        assert waited == []
+
+    def test_a_flush_without_a_timeout_does_not_retry(self, clock):
+        sender = FakeSender(SendOutcome("retry-later"))
+        pipeline, _, _ = make_traces(sender=sender)
+        waited = waits_advance(clock, pipeline)
+        pipeline.start_span("a").end()
+        pipeline.flush()
+        assert len(sender.payloads) == 1
+        assert waited == []
+
+    def test_close_cuts_the_wait_short(self, clock):
+        sender = FakeSender(SendOutcome("retry-later"))
+        pipeline, _, _ = make_traces(sender=sender)
+        exporter = pipeline._exporter
+        pipeline.start_span("a").end()
+
+        def close_during_wait(seconds):
+            exporter.close()
+            return True
+
+        exporter._wait_for_retry = close_during_wait
+        pipeline.flush(timeout=30)
+        assert len(sender.payloads) == 1
+
+    def test_the_real_wait_returns_when_close_is_called(self):
+        pipeline, _, _ = make_traces()
+        exporter = pipeline._exporter
+        threading.Thread(target=lambda: (time.sleep(0.05), exporter.close())).start()
+        started = time.monotonic()
+        assert exporter._wait_for_retry(5) is True
+        assert time.monotonic() - started < 2

posthog/tracing/_export.py

@@ -20,6 +20,7 @@
     build_otlp_span,
     build_resource_attributes,
     build_traces_payload,
+    to_resource_key_value_list,
 )
 from ._transport import SendOutcome, send_traces_batch
 
@@ -101,20 +102,28 @@ def __init__(
         self._config = config
         self._drops = drops
         self._send = send
-        self._resource_attributes = build_resource_attributes(
-            config.service_name,
-            config.service_version,
-            config.environment,
-            config.resource_attributes,
+        # Encoded once: the resource is the same for every batch.
+        self._resource = to_resource_key_value_list(
+            build_resource_attributes(
+                config.service_name,
+                config.service_version,
+                config.environment,
+                config.resource_attributes,
+            )
         )
 
         self._lock = threading.Lock()
         self._flush_lock = threading.Lock()
+        # Set by close(), so a flush waiting out a backoff returns at once.
+        self._closing = threading.Event()
         self._closed = False
         self._queue: List[SpanRecord] = []
         self._flush_timer: Optional[threading.Timer] = None
         self._flush_timer_fires_at = 0.0
         self._max_export_batch_size = config.max_export_batch_size
+        # The size before a server 413 halved it, restored once the oversized
+        # span is isolated and dropped.
+        self._batch_size_before_halving: Optional[int] = None
         self._consecutive_failures = 0
         # Drawn once per failure, so the timer and the retry-budget charge see
         # the same delay.
@@ -151,15 +160,24 @@ def flush(
     ) -> None:
         """Drain the queue: one pass over what was queued, then one follow-up pass.
 
-        With a ``timeout``, no request starts once it is spent, except the
-        first. A request already in flight is bounded by the client's timeout.
+        With a ``timeout``, the budget starts once no other flush is in flight;
+        one still in flight after that long a wait is left to finish and
+        nothing is sent here. No request starts once the budget is spent,
+        except the first, and a retriable failure is retried after its backoff
+        while budget remains, once more at the deadline. Without a timeout a
+        retriable failure is left to the timer. A request already in flight is
+        bounded by the client's timeout.
         """
-        deadline = None if timeout is None else time.monotonic() + timeout
         # -1 is Lock.acquire's unbounded form.
         if not self._flush_lock.acquire(
             timeout=-1 if timeout is None else max(0.0, timeout)
         ):
+            log.debug(
+                "Skipping a span flush: another flush was still in flight after %ss",
+                timeout,
+            )
             return
+        deadline = None if timeout is None else time.monotonic() + timeout
         try:
             with self._lock:
                 # A timer that waited behind another flush may have been
@@ -168,25 +186,50 @@ def flush(
                     return
                 self._clear_timer_locked()
             try:
-                removed, stop = self._drain(deadline)
-                if (
-                    removed
-                    and not stop
-                    and self._queue
-                    and (deadline is None or time.monotonic() < deadline)
-                ):
-                    self._drain(deadline)
+                self._drain_within_budget(deadline, retry=_timer is None)
             finally:
                 with self._lock:
                     self._rearm_after_pass_locked()
         finally:
             self._flush_lock.release()
             self._drops.warn_if_due(force=True)
 
+    def _drain_within_budget(self, deadline: Optional[float], retry: bool) -> None:
+        while True:
+            removed, stop = self._drain(deadline)
+            if not stop:
+                if (
+                    removed
+                    and self._queue
+                    and (deadline is None or time.monotonic() < deadline)
+                ):
+                    self._drain(deadline)
+                return
+            if deadline is None or not retry:
+                return
+            wait = self._retry_wait_locked()
+            remaining = deadline - time.monotonic()
+            if wait is None or remaining <= 0:
+                return
+            if self._wait_for_retry(min(wait, remaining)):
+                return
+
+    def _retry_wait_locked(self) -> Optional[float]:
+        """The backoff to wait out before retrying, or ``None`` when there is nothing to retry."""
+        with self._lock:
+            if self._closed or not self._queue or not self._consecutive_failures:
+                return None
+            return self._next_flush_delay_locked()
+
+    def _wait_for_retry(self, seconds: float) -> bool:
+        """Wait out a backoff; ``True`` when close() cut the wait short."""
+        return self._closing.wait(seconds)
+
     def close(self) -> None:
         """Stop exporting and discard what is still queued. Called at shutdown."""
         with self._lock:
             self._closed = True
+            self._closing.set()
             self._clear_timer_locked()
             discarded = len(self._queue)
             self._queue = []
@@ -212,10 +255,12 @@ def reinit_after_fork(self) -> None:
         # in the child, so they are replaced rather than acquired.
         self._lock = threading.Lock()
         self._flush_lock = threading.Lock()
+        self._closing = threading.Event()
         self._flush_timer = None
         self._flush_timer_fires_at = 0.0
         self._queue = []
         self._max_export_batch_size = self._config.max_export_batch_size
+        self._batch_size_before_halving = None
         self._retry_after.reset()
         self._end_failure_sequence_locked()
 
@@ -281,7 +326,7 @@ def _drain(self, deadline: Optional[float]) -> Tuple[int, bool]:
             sent_any = True
             try:
                 outcome = self._send(
-                    self._client, build_traces_payload(spans, self._resource_attributes)
+                    self._client, build_traces_payload(spans, self._resource)
                 )
             except Exception:
                 log.debug("Span batch send failed", exc_info=True)
@@ -312,20 +357,29 @@ def _apply_outcome_locked(
         if outcome.kind == "ok":
             del self._queue[:size]
             self._end_failure_sequence_locked()
+            self._batch_size_before_halving = None
             if self._max_export_batch_size < self._config.max_export_batch_size:
-                self._max_export_batch_size += 1
+                self._max_export_batch_size = min(
+                    self._config.max_export_batch_size, self._max_export_batch_size * 2
+                )
             return size, False, None
 
         if outcome.kind == "too-large":
             if size == 1:
                 del self._queue[:1]
                 self._end_failure_sequence_locked()
                 self._drops.record(1, "it is too large for the ingestion endpoint")
+                # The oversized span is gone; the batches after it are not suspect.
+                if self._batch_size_before_halving is not None:
+                    self._max_export_batch_size = self._batch_size_before_halving
+                    self._batch_size_before_halving = None
                 return 1, False, None
             # Halve the refused batch, not the configured size: a shallow queue
             # would otherwise resend the same body.
             halved = max(1, size // 2)
             if not outcome.measured_locally:
+                if self._batch_size_before_halving is None:
+                    self._batch_size_before_halving = self._max_export_batch_size
                 self._max_export_batch_size = halved
             self._reset_head_batch_budget_locked()
             return (