LogPusher: flush() blocks until every prior entry is processed

Adds a per-entry monotonic seq to the queue. `flush(timeout=None)` snapshots
_pushed_seq and waits on the condvar until the drain has advanced
_processed_seq past it, where "processed" means either sent or
overflow-dropped. Dropped entries advance the marker so a flush doesn't
wait forever on entries that will never reach the server.

Makes task_attempt._cleanup's `self._log_pusher.flush()` actually wait for
final logs to ship, restoring the pre-rewrite contract. close() stays
best-effort per directive.

Also replaces the _wait_for polling in test_failures_always_deliver_via_retry
with deterministic flush(timeout=...) synchronization — removes the xdist-
load flake that showed up as assert 2==1 / 6==1 on the seed/seed-then-push
race.

Co-authored-by: Russell Power <rjpower@users.noreply.github.com>

Author: github-actions[bot]

lib/iris/src/iris/log_server/client.py

@@ -55,11 +55,16 @@
 
 
 class LogPusher:
-    """Non-blocking buffered client for pushing log entries to a remote LogService.
+    """Buffered client for pushing log entries to a remote LogService.
 
-    ``push`` appends to an in-memory queue; a background thread drains it
-    in per-key batches. Send failures re-buffer and back off exponentially
-    — only the ``MAX_LOG_BUFFER_SIZE`` overflow path drops entries.
+    ``push`` is non-blocking: it appends to an in-memory queue and returns.
+    A background thread drains the queue in per-key batches. Send failures
+    re-buffer and back off exponentially — only the ``MAX_LOG_BUFFER_SIZE``
+    overflow path drops entries.
+
+    ``flush`` blocks until every entry enqueued before the call has been
+    processed (sent or overflow-dropped). Use the ``timeout`` argument to
+    bound the wait — by default ``flush`` waits indefinitely.
 
     ``server_url`` is passed to ``resolver`` (default: identity) to obtain
     the actual http address. Retryable failures invalidate the cached RPC
@@ -87,12 +92,20 @@ def __init__(
 
         # All shared state is guarded by _cond. The drain thread is the
         # only owner of _client, so no separate client lock. ``_queue`` is
-        # a single FIFO of (key, entry); the drain thread groups by key
-        # just before sending. Trimming on overflow is one popleft.
+        # a single FIFO of (seq, key, entry); the drain thread groups by
+        # key just before sending. Trimming on overflow is one popleft.
+        # ``seq`` is a monotonic per-entry counter used by blocking flush.
         self._cond = threading.Condition()
-        self._queue: deque[tuple[str, logging_pb2.LogEntry]] = deque()
+        self._queue: deque[tuple[int, str, logging_pb2.LogEntry]] = deque()
         self._closed = False
 
+        # Monotonic counters for blocking flush(). ``_pushed_seq`` advances
+        # on every entry enqueued. ``_processed_seq`` advances when the
+        # drain thread acks an entry as either successfully sent or
+        # overflow-dropped — both terminal states from flush's POV.
+        self._pushed_seq = 0
+        self._processed_seq = 0
+
         # Built lazily by the drain thread on first send; invalidated on
         # any failure so the next attempt re-resolves.
         self._client: LogServiceClientSync | None = None
@@ -115,27 +128,53 @@ def push(self, key: str, entries: list[logging_pb2.LogEntry]) -> None:
             if self._closed:
                 return
             for e in entries:
-                self._queue.append((key, e))
+                self._pushed_seq += 1
+                self._queue.append((self._pushed_seq, key, e))
             self._trim_oldest_locked()
             if len(self._queue) >= self._batch_size:
-                self._cond.notify()
+                self._cond.notify_all()
 
-    def flush(self) -> None:
-        """Poke the drain thread to send whatever is buffered now.
+    def flush(self, timeout: float | None = None) -> bool:
+        """Block until every entry enqueued before this call has been processed.
 
-        Non-blocking. For draining on shutdown, use ``close``.
+        "Processed" means either successfully sent or overflow-dropped —
+        both terminal states. Returns ``True`` if the drain caught up,
+        ``False`` on timeout. ``timeout=None`` waits indefinitely.
+
+        For shutdown drain, prefer ``close`` (best-effort, won't block on
+        a stuck server).
         """
         with self._cond:
-            if self._queue:
-                self._cond.notify()
+            target = self._pushed_seq
+            if target == 0 or self._processed_seq >= target:
+                return True
+            self._cond.notify_all()
+            deadline = (time.monotonic() + timeout) if timeout is not None else None
+            while self._processed_seq < target:
+                if self._closed:
+                    return self._processed_seq >= target
+                if deadline is None:
+                    # Re-check periodically so a wedged drain still surfaces.
+                    self._cond.wait(timeout=1.0)
+                else:
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        return False
+                    self._cond.wait(timeout=remaining)
+            return True
 
     def close(self) -> None:
-        """Stop the drain thread after one best-effort drain, close the RPC client."""
+        """Stop the drain thread after one best-effort drain, close the RPC client.
+
+        Best-effort: if a send is in flight when ``close()`` returns the
+        join timeout, we still close the cached client. Use ``flush()``
+        first if you need to guarantee final delivery.
+        """
         with self._cond:
             if self._closed:
                 return
             self._closed = True
-            self._cond.notify()
+            self._cond.notify_all()
         # Join the drain thread; it will send what it can and exit.
         self._thread.join(timeout=max(self._flush_interval * 2, 10.0))
         if self._client is not None:
@@ -150,28 +189,38 @@ def close(self) -> None:
     # ------------------------------------------------------------------
 
     def _trim_oldest_locked(self) -> None:
-        """Drop oldest entries until under ``_max_buffer_size``."""
+        """Drop oldest entries until under ``_max_buffer_size``.
+
+        Dropped entries advance ``_processed_seq`` so blocking ``flush``
+        doesn't wait forever on entries that will never reach the server.
+        """
         dropped = 0
+        max_dropped_seq = 0
         while len(self._queue) > self._max_buffer_size:
-            self._queue.popleft()
+            seq, _key, _entry = self._queue.popleft()
+            if seq > max_dropped_seq:
+                max_dropped_seq = seq
             dropped += 1
         if dropped:
             logger.warning(
                 "LogPusher buffer overflow: dropped %d oldest entries (cap=%d)",
                 dropped,
                 self._max_buffer_size,
             )
+            if max_dropped_seq > self._processed_seq:
+                self._processed_seq = max_dropped_seq
+                self._cond.notify_all()
 
-    def _take_queue_locked(self) -> list[tuple[str, logging_pb2.LogEntry]]:
+    def _take_queue_locked(self) -> list[tuple[int, str, logging_pb2.LogEntry]]:
         """Drain the entire queue, preserving arrival order."""
         items = list(self._queue)
         self._queue.clear()
         return items
 
-    def _rebuffer_at_head_locked(self, items: list[tuple[str, logging_pb2.LogEntry]]) -> None:
+    def _rebuffer_at_head_locked(self, items: list[tuple[int, str, logging_pb2.LogEntry]]) -> None:
         """Put unsent items back at the head of the queue (original order)."""
-        for pair in reversed(items):
-            self._queue.appendleft(pair)
+        for triple in reversed(items):
+            self._queue.appendleft(triple)
         self._trim_oldest_locked()
 
     # ------------------------------------------------------------------
@@ -194,7 +243,11 @@ def _run(self) -> None:
                     return
                 items = self._take_queue_locked()
 
-            unsent = self._send_items(items)
+            sent_max_seq, unsent = self._send_items(items)
+            with self._cond:
+                if sent_max_seq > self._processed_seq:
+                    self._processed_seq = sent_max_seq
+                    self._cond.notify_all()
             if not unsent:
                 self._backoff.reset()
                 continue
@@ -216,44 +269,50 @@ def _run(self) -> None:
 
     def _send_items(
         self,
-        items: list[tuple[str, logging_pb2.LogEntry]],
-    ) -> list[tuple[str, logging_pb2.LogEntry]]:
+        items: list[tuple[int, str, logging_pb2.LogEntry]],
+    ) -> tuple[int, list[tuple[int, str, logging_pb2.LogEntry]]]:
         """Group ``items`` by key (stable on first occurrence) and push one
-        RPC per key. On any failure, return every item from that key onward
-        so the caller can re-buffer it at the head of the queue.
+        RPC per key. Returns ``(max_sent_seq, unsent_items)``.
 
+        On any failure, every item from that key onward is returned as
+        unsent so the caller can re-buffer it at the head of the queue.
         Every failure mode — resolver error, retryable RPC error, or
         non-retryable RPC error — re-buffers so no log entries are silently
         dropped. Retryable errors additionally invalidate the cached client
         so the next attempt re-resolves the endpoint.
         """
-        groups: dict[str, list[logging_pb2.LogEntry]] = {}
-        for key, entry in items:
-            groups.setdefault(key, []).append(entry)
+        groups: dict[str, list[tuple[int, logging_pb2.LogEntry]]] = {}
+        for seq, key, entry in items:
+            groups.setdefault(key, []).append((seq, entry))
 
         sent_keys: set[str] = set()
-        for key, entries in groups.items():
+        max_sent_seq = 0
+        for key, seq_entries in groups.items():
             try:
                 client = self._get_client()
             except Exception as exc:
                 logger.warning("LogPusher: endpoint resolution failed: %s", exc)
-                return [p for p in items if p[0] not in sent_keys]
+                return max_sent_seq, [p for p in items if p[1] not in sent_keys]
             try:
+                entries = [e for _s, e in seq_entries]
                 client.push_logs(logging_pb2.PushLogsRequest(key=key, entries=entries))
                 sent_keys.add(key)
+                for seq, _e in seq_entries:
+                    if seq > max_sent_seq:
+                        max_sent_seq = seq
             except Exception as exc:
                 retryable = is_retryable_error(exc)
                 logger.warning(
                     "LogPusher: send failure for key=%s (%d entries, retryable=%s): %s",
                     key,
-                    len(entries),
+                    len(seq_entries),
                     retryable,
                     exc,
                 )
                 if retryable:
                     self._invalidate(str(exc))
-                return [p for p in items if p[0] not in sent_keys]
-        return []
+                return max_sent_seq, [p for p in items if p[1] not in sent_keys]
+        return max_sent_seq, []
 
     def _build_client(self, address: str) -> LogServiceClientSync:
         return LogServiceClientSync(

lib/iris/tests/test_remote_log_handler.py

@@ -117,7 +117,11 @@ def _wait_for(predicate, timeout: float = 5.0) -> None:
 
 
 def test_log_pusher_buffers_and_flushes_on_demand(tracked_log_service_client):
-    """Entries buffered below batch_size are drained on flush()."""
+    """Entries buffered below batch_size are drained on flush().
+
+    flush() blocks until every entry enqueued before the call has shipped,
+    so the assertions can run immediately without polling.
+    """
     pusher = LogPusher(
         "http://h:1",
         batch_size=1000,
@@ -127,15 +131,62 @@ def test_log_pusher_buffers_and_flushes_on_demand(tracked_log_service_client):
         entry = logging_pb2.LogEntry(source="test", data="line1")
         pusher.push("key-a", [entry, entry, entry])
         pusher.push("key-b", [entry])
-        pusher.flush()
-        _wait_for(lambda: len(tracked_log_service_client) == 1 and len(tracked_log_service_client[0].pushes) >= 2)
+        assert pusher.flush(timeout=5.0)
 
         totals = {p.key: len(p.entries) for p in tracked_log_service_client[0].pushes}
         assert totals == {"key-a": 3, "key-b": 1}
     finally:
         pusher.close()
 
 
+def test_log_pusher_flush_is_blocking(tracked_log_service_client):
+    """flush() returns only after every previously-pushed entry has been sent."""
+    pusher = LogPusher(
+        "http://h:1",
+        batch_size=1000,
+        flush_interval=999.0,
+    )
+    try:
+        entry = logging_pb2.LogEntry(source="test", data="line")
+        pusher.push("k", [entry, entry])
+        # No polling — flush must block until shipped.
+        assert pusher.flush(timeout=5.0) is True
+        assert len(tracked_log_service_client[0].pushes) == 1
+        assert len(tracked_log_service_client[0].pushes[0].entries) == 2
+    finally:
+        pusher.close()
+
+
+def test_log_pusher_flush_timeout_returns_false(monkeypatch):
+    """flush(timeout=...) returns False when the drain can't catch up in time.
+
+    Seeds a non-retryable error so the drain rebuffers and enters the
+    backoff window; flush is given less time than the backoff interval.
+    """
+    created: list[_FakeLogServiceClient] = []
+
+    def factory(address, timeout_ms=10_000, interceptors=()):
+        c = _FakeLogServiceClient(address, timeout_ms=timeout_ms, interceptors=interceptors)
+        created.append(c)
+        return c
+
+    monkeypatch.setattr(client_mod, "LogServiceClientSync", factory)
+
+    pusher = LogPusher("http://h:1", batch_size=1, flush_interval=999.0)
+    try:
+        entry = logging_pb2.LogEntry(source="test", data="primer")
+        pusher.push("k", [entry])
+        # Wait for the cached client to exist, then seed a non-retryable
+        # error so the next send rebuffers and the drain enters backoff.
+        assert pusher.flush(timeout=5.0) is True
+        created[0].errors.append(ConnectError(Code.NOT_FOUND, "missing"))
+        pusher.push("k", [logging_pb2.LogEntry(source="test", data="stuck")])
+        # Backoff is 0.5s; a 0.05s flush cannot catch up.
+        assert pusher.flush(timeout=0.05) is False
+    finally:
+        pusher.close()
+
+
 def test_log_pusher_flushes_at_batch_size(tracked_log_service_client):
     """Reaching batch_size wakes the drain thread without waiting for a timer."""
     pusher = LogPusher(
@@ -305,8 +356,9 @@ def resolver(_url):
         # for resolver_raises, which has no client to seed).
         pusher.push("k", [_entry("a")])
         if scenario != "resolver_raises":
-            _wait_for(lambda: created and created[0].pushes)
-            # Seed the cached client with the scenario-appropriate error.
+            # Block until "a" has shipped, so seeding the next error is
+            # race-free with the drain thread's next iteration.
+            assert pusher.flush(timeout=5.0)
             err = (
                 ConnectError(Code.NOT_FOUND, "missing")
                 if scenario == "non_retryable"
@@ -316,15 +368,19 @@ def resolver(_url):
 
         pusher.push("k", [_entry("b")])
 
-        # "b" must eventually land somewhere.
+        # Wait deterministically for "b" to be processed (sent or dropped).
+        assert pusher.flush(timeout=10.0)
+
+        # "b" must have landed somewhere — the buffer-overflow path is not
+        # exercised here, so processed implies delivered.
         def delivered():
             return any(any(e.data == "b" for p in c.pushes for e in p.entries) for c in created)
 
-        _wait_for(delivered, timeout=10.0)
+        assert delivered(), "entry 'b' was never delivered to any client"
 
         if scenario.startswith("retryable"):
             # Retryable RPC failure invalidated the first client; second built.
-            _wait_for(lambda: len(created) >= 2)
+            assert len(created) >= 2
             assert created[0].closed is True
         elif scenario == "resolver_raises":
             # Resolver raised on first call → no client yet. Second call