fix(traces): keep the scheduled flush when a replacement timer fails to start, and count an unencodable span once

The replacement timer now starts before the old one is cancelled, so a
thread that cannot be created leaves the earlier flush in place. Records
that fail to encode leave the queue before the send is settled, so a
failed send no longer counts them a second time.

Author: turnipdabeets

posthog/test/tracing/test_export.py

@@ -180,10 +180,12 @@ def start(self):
             pipeline, sender, _ = make_traces(max_export_batch_size=3)
             for _ in range(3):
                 pipeline.start_span("a").end()
-            assert pipeline._exporter._flush_timer is None
+            interval_timer = pipeline._exporter._flush_timer
+            assert interval_timer.delay == 5 and not interval_timer.cancelled
             pipeline.start_span("b").end()
             timer = pipeline._exporter._flush_timer
-            assert timer is not None and timer.started and timer.delay == 0
+            assert timer is not interval_timer and timer.started and timer.delay == 0
+            assert interval_timer.cancelled
             timer.fire()
         assert queued(pipeline) == []
         assert [len(b) for b in sender.batches()] == [3, 1]
@@ -358,6 +360,26 @@ def test_never_returns_a_span_it_failed_to_encode(self):
         assert sender.payloads == []
         assert queued(pipeline) == []
 
+    def test_counts_a_span_it_failed_to_encode_once(self, caplog):
+        caplog.set_level("WARNING", logger="posthog")
+        sender = FakeSender(SendOutcome("fatal"))
+        pipeline, _, _ = make_traces(sender=sender, max_export_batch_size=2)
+        pipeline.start_span("bad").end()
+        pipeline.start_span("fine").end()
+        encode = export_module.build_otlp_span
+
+        def encode_unless_bad(record):
+            if record.name == "bad":
+                raise RuntimeError("bad")
+            return encode(record)
+
+        with mock.patch.object(export_module, "build_otlp_span", encode_unless_bad):
+            pipeline.flush()
+        assert [[s["name"] for s in b] for b in sender.batches()] == [["fine"]]
+        assert queued(pipeline) == []
+        (record,) = [r for r in caplog.records if "Dropping" in r.getMessage()]
+        assert "Dropping 2 span(s)" in record.getMessage()
+
 
 class TestDroppedBatchesEndTheFailureSequence:
     @pytest.mark.parametrize(

posthog/tracing/_export.py

@@ -255,16 +255,22 @@ def _drain(self, deadline: Optional[float]) -> Tuple[int, bool]:
                     and not self._retry_after.is_open()
                 )
 
-            spans = self._encode(batch)
-            if not spans:
+            spans, failed = self._encode(batch)
+            if failed:
+                # Out of the queue before the send is settled, so a failed
+                # send does not count them a second time.
                 with self._lock:
                     if self._is_closed():
                         return removed, True
-                    del self._queue[:size]
-                    self._reset_head_batch_budget_locked()
-                remaining -= size
-                removed += size
-                continue
+                    for index in reversed(failed):
+                        del self._queue[index]
+                    if not spans:
+                        self._reset_head_batch_budget_locked()
+                size -= len(failed)
+                remaining -= len(failed)
+                removed += len(failed)
+                if not spans:
+                    continue
 
             # The first request is exempt, so a flush called with no budget left
             # (a serverless handler, say) still ships a batch.
@@ -361,15 +367,18 @@ def _is_closed(self) -> bool:
         # A method, so the check after each send is not narrowed away.
         return self._closed
 
-    def _encode(self, batch: List[SpanRecord]) -> List[dict]:
+    def _encode(self, batch: List[SpanRecord]) -> Tuple[List[dict], List[int]]:
+        """The batch as OTLP spans, and the indexes of records that could not be encoded."""
         encoded: List[dict] = []
-        for record in batch:
+        failed: List[int] = []
+        for index, record in enumerate(batch):
             try:
                 encoded.append(build_otlp_span(record))
             except Exception:
                 log.debug("Failed to encode a span; dropping it", exc_info=True)
                 self._drops.record(1, "its attributes could not be encoded")
-        return encoded
+                failed.append(index)
+        return encoded, failed
 
     def _discard_if_disabled_locked(self) -> bool:
         if not getattr(self._client, "disabled", False):
@@ -431,16 +440,15 @@ def _rearm_after_pass_locked(self) -> None:
         self._arm_timer_locked(delay, replace=True)
 
     def _arm_timer_locked(self, delay: float, replace: bool = False) -> None:
-        if self._flush_timer is not None:
-            if not replace:
-                return
-            self._flush_timer.cancel()
-            self._flush_timer = None
+        if self._flush_timer is not None and not replace:
+            return
         timer = threading.Timer(delay, lambda: self._timer_flush(timer))
         timer.daemon = True
-        # Registered only once started, so a failed start leaves nothing
-        # behind and the next span end arms again.
+        # Started before the old timer is cancelled, so a failed start keeps
+        # whatever flush was already scheduled.
         timer.start()
+        if self._flush_timer is not None:
+            self._flush_timer.cancel()
         self._flush_timer = timer
         self._flush_timer_fires_at = time.monotonic() + delay