From 7da9cea0f66a6af20936d0d47e4ca5d0b5c8f595 Mon Sep 17 00:00:00 2001
From: firestar5683 <168790843+firestar5683@users.noreply.github.com>
Date: Fri, 30 Oct 2026 20:41:23 -0500
Subject: [PATCH] Yield FLM evidence queues before fresh pose and trial
 delivery

---
 openpilot/starpilot/flm/fullrate_collector.py |  24 ++-
 openpilot/starpilot/flm/steerlearnerd.py      |  14 +-
 .../starpilot/flm/tests/_drain_fixture.py     |   2 +-
 .../tests/test_paired_drain_time_slices.py    | 139 ++++++++++++++++++
 prebuilt.json                                 |   2 +-
 5 files changed, 172 insertions(+), 9 deletions(-)
 create mode 100644 openpilot/starpilot/flm/tests/test_paired_drain_time_slices.py

diff --git a/openpilot/starpilot/flm/fullrate_collector.py b/openpilot/starpilot/flm/fullrate_collector.py
index 2f9849c60a..6a52260ed6 100644
--- a/openpilot/starpilot/flm/fullrate_collector.py
+++ b/openpilot/starpilot/flm/fullrate_collector.py
@@ -11,6 +11,7 @@ from collections import Counter, OrderedDict, deque
 from openpilot.starpilot.flm.torque_eligibility import asynchronous_hint, bounded_context_output, exact_unclipped_pair, consumed_sources_current
 from openpilot.starpilot.flm.native_refinement_wire import valid_inputs
 
+import time
 from dataclasses import dataclass, replace
 import math
 import re
@@ -923,9 +924,11 @@ class PairedDrainState:
   def __init__(self, control_socket, applied_socket):
     self.sockets = (control_socket, applied_socket)
     self.heads: tuple = (None, None)
+    self.counts = (0, 0)
+    self.exhausted = (False, False)
 
 
-def drain_pair(control_socket, control, applied_socket, applied, receive, collector, now_ns, state=None):
+def drain_pair(control_socket, control, applied_socket, applied, receive, collector, now_ns, state=None, *, budget_ns=None, clock=time.monotonic_ns):
   """Merge two bounded FIFO tapes by their original publisher headers.
 
   Only the two queue heads compete; never sort a producer's own tape. A held
@@ -937,17 +940,27 @@ def drain_pair(control_socket, control, applied_socket, applied, receive, collec
   A persistent state retains only the two already-read heads at an overrun.
   Their original wire and clocks are checked again on the next bounded drain.
   """
+  if budget_ns is not None and (state is None or type(budget_ns) is not int or budget_ns <= 0):
+    raise ValueError('time-sliced drain requires persistent state and a positive budget')
+  deadline = None if budget_ns is None else clock() + budget_ns
   sockets, callbacks = (control_socket, applied_socket), (control, applied)
   if state is None:
     state = PairedDrainState(*sockets)
   pending = list(state.heads)
   state.heads = (None, None)
-  counts = [0, 0]
+  counts = list(state.counts)
+  exhausted = list(state.exhausted)
+  state.counts = (0, 0)
+  state.exhausted = (False, False)
+  processed = 0
 
   def head(index):
+    if exhausted[index]:
+      return None
     event = pending[index][0] if pending[index] is not None else receive(sockets[index])
     pending[index] = None
     if event is None:
+      exhausted[index] = True
       return None
     _check_evidence_wire(event)
     stamp = event.logMonoTime
@@ -965,10 +978,17 @@ def drain_pair(control_socket, control, applied_socket, applied, receive, collec
     head(0)
     head(1)
     while pending[0] is not None or pending[1] is not None:
+      if processed and deadline is not None and clock() >= deadline:
+        # Yield scheduling only. Preserve original FIFO heads and the SAME
+        # per-stream overrun count; a slice cannot renew evidence or a batch.
+        state.heads, state.counts = tuple(pending), (counts[0], counts[1])
+        state.exhausted = (exhausted[0], exhausted[1])
+        return False
       index = 0 if pending[1] is None or (pending[0] is not None and pending[0][1] <= pending[1][1]) else 1
       event, _stamp = pending[index]
       callbacks[index](event, collector, now_ns)
       counts[index] += 1
+      processed += 1
       pending[index] = None
       head(index)
   except _DrainOverrun:
diff --git a/openpilot/starpilot/flm/steerlearnerd.py b/openpilot/starpilot/flm/steerlearnerd.py
index 6dcd479094..3d97d9657c 100644
--- a/openpilot/starpilot/flm/steerlearnerd.py
+++ b/openpilot/starpilot/flm/steerlearnerd.py
@@ -40,9 +40,11 @@ class _SourceDelivery:
     if sm.updated['deviceMotion']:
       self.pose_header = int(sm.logMonoTime['deviceMotion'])
 
-  def next(self, sm, *, pending_evidence=()):
+  def next(self, sm, *, pending_evidence=(), retained_evidence=False):
     if not self.pending:
-      if pending_evidence:
+      if retained_evidence:
+        sm.update(0)  # Process our retained FIFO heads before waiting again.
+      elif pending_evidence:
         # A received pose may still need a genuinely newer consumed state.
         # Poll existing tapes without consuming/reordering their FIFO heads.
         sockets = (sm.sock['deviceMotion'], *pending_evidence)
@@ -536,8 +538,9 @@ def run(
         pending_stamp = getattr(pending_poses[0], 'mono_ns', 0)
         pending_now = clock()
         pending_fresh = type(pending_stamp) is int and 0 < pending_stamp <= pending_now <= pending_stamp + MAX_GAP_NS
-      delivered_pose_header = (source_delivery.next(sm, pending_evidence=(control_socket, applied_socket))
-                               if pending_fresh else source_delivery.next(sm))
+      retained_evidence = any(head is not None for head in drain_state.heads)
+      delivered_pose_header = source_delivery.next(sm,
+        pending_evidence=(control_socket, applied_socket) if pending_fresh else (), retained_evidence=retained_evidence)
       source_rejected = False
       now = clock()
       # This is an ORIGINAL header from the real poll or an already consumed
@@ -697,7 +700,8 @@ def run(
               return
           _collect_applied_output(collector, applied, sm, CP, received_now)
 
-      drain_pair(control_socket, control, applied_socket, output, messaging.recv_one_or_none, sink, now, drain_state)
+      drain_pair(control_socket, control, applied_socket, output, messaging.recv_one_or_none, sink, now, drain_state,
+        budget_ns=5_000_000 if live_actor is not None else None)
       if collector is not None:
         # Decoding queued frames may outlast the sampled shared inputs. Consume
         # real pending deliveries before pose/native checks; keep their original
diff --git a/openpilot/starpilot/flm/tests/_drain_fixture.py b/openpilot/starpilot/flm/tests/_drain_fixture.py
index c584b70a7c..e2ca91d024 100644
--- a/openpilot/starpilot/flm/tests/_drain_fixture.py
+++ b/openpilot/starpilot/flm/tests/_drain_fixture.py
@@ -6,7 +6,7 @@ The actual FIFO merge and bounded wire receiver have their own transport tests.
 
 
 def separate_drain_fixture(single):
-  def paired(control_socket, control, applied_socket, applied, receive, collector, now_ns, _state=None):
+  def paired(control_socket, control, applied_socket, applied, receive, collector, now_ns, _state=None, *, budget_ns=None):
     single(control_socket, control, receive, collector, now_ns)
     single(applied_socket, applied, receive, collector, now_ns)
   return paired
diff --git a/openpilot/starpilot/flm/tests/test_paired_drain_time_slices.py b/openpilot/starpilot/flm/tests/test_paired_drain_time_slices.py
new file mode 100644
index 0000000000..301f7c7ece
--- /dev/null
+++ b/openpilot/starpilot/flm/tests/test_paired_drain_time_slices.py
@@ -0,0 +1,139 @@
+"""Time slices preserve FIFO, real receipt clocks and original overload limits."""
+from dataclasses import replace
+from collections import deque
+from types import SimpleNamespace
+import pytest
+from openpilot.starpilot.flm import fullrate_collector as module
+from openpilot.starpilot.flm.fullrate_collector import FullrateCollector, PairedDrainState
+from openpilot.starpilot.flm.learning_wire import read_applied, read_control
+from openpilot.starpilot.flm.tests.test_fullrate_pipeline import SyntheticAdapter, snapshots, serialized_control, serialized_output
+from openpilot.starpilot.flm.tests.test_paired_evidence_drain import BASE, small_event, transport
+
+
+class WorkClock:
+  def __init__(self):
+    self.ns = 0
+
+  def __call__(self):
+    return self.ns
+
+
+def test_yield_keeps_every_original_head_and_never_splits_the_maneuver():
+  rows = [[small_event(BASE + 2*i + side) for i in range(10)] for side in (0, 1)]
+  sockets, receive, _ = transport(*rows)
+  collector = FullrateCollector(SyntheticAdapter())
+  state = PairedDrainState(*sockets)
+  clock = WorkClock()
+  seen = []
+
+  def callback(event, current, received):
+    assert current is collector
+    seen.append((event, received))
+    clock.ns += 3_000_000
+
+  for tick in range(20):
+    before = len(seen)
+    done = module.drain_pair(sockets[0], callback, sockets[1], callback, receive, collector,
+      BASE+tick*10_000_000, state, budget_ns=5_000_000, clock=clock)
+    assert 0 < len(seen)-before <= 2
+    assert collector.dropped == 0
+    if done:
+      break
+    assert any(head is not None for head in state.heads)
+  assert done and len(seen) == 20
+  assert [event.logMonoTime for event, _ in seen] == list(range(BASE, BASE+20))
+  assert len({id(event)for event, _ in seen}) == 20
+  assert state.heads == (None, None) and state.counts == (0, 0)
+
+
+def test_a_stream_already_empty_in_this_batch_is_not_reopened_by_a_yield():
+  sockets = object(), object()
+  first, newer = small_event(BASE), small_event(BASE+10)
+  outputs = [small_event(BASE+20+i)for i in range(3)]
+  queues = {sockets[0]:deque([first]),sockets[1]:deque(outputs)}
+  collector = FullrateCollector(SyntheticAdapter())
+  state = PairedDrainState(*sockets)
+  clock = WorkClock()
+  seen = []
+
+  def callback(event, _current, _stamp):
+    seen.append(event)
+    clock.ns += 2_000_000
+
+  def consume():
+    return module.drain_pair(sockets[0],callback,sockets[1],callback,
+      lambda socket:queues[socket].popleft()if queues[socket]else None,
+      collector,BASE,state,budget_ns=1_000_000,clock=clock)
+
+  assert not consume() and seen == [first]
+  queues[sockets[0]].append(newer)
+  for _ in range(3):
+    done = consume()
+  assert done and seen == [first,*outputs]
+  assert consume() and seen == [first,*outputs,newer]
+
+
+def test_slicing_cannot_reset_the_original_per_stream_overrun_limit(monkeypatch):
+  monkeypatch.setattr(module, 'MAX_DRAIN', 4)
+  rows = [small_event(BASE+i)for i in range(7)]
+  sockets, receive, _ = transport(rows, [])
+  state = PairedDrainState(*sockets)
+  collector = FullrateCollector(SyntheticAdapter())
+  clock = WorkClock()
+  seen = []
+
+  def callback(event, _current, _received):
+    seen.append(event)
+    clock.ns += 2_000_000
+
+  for tick in range(10):
+    done = module.drain_pair(sockets[0], callback, sockets[1], callback, receive, collector,
+      BASE+tick, state, budget_ns=1_000_000, clock=clock)
+    if collector.dropped:
+      assert len(seen) == 4 and collector.reason == 'fullrate_overrun'
+      assert state.counts == (0, 0)
+      break
+  else:
+    raise AssertionError('A timed yield bypassed the original overrun guard')
+  for tick in range(10):
+    done = module.drain_pair(sockets[0], callback, sockets[1], callback, receive, collector,
+      BASE+100+tick, state, budget_ns=1_000_000, clock=clock)
+    if done:
+      break
+  assert done and seen == rows and collector.dropped == 1
+
+
+@pytest.mark.parametrize('delay,native,expected', [(10_000_000, True, 1), (200_000_000, True, 0), (10_000_000, False, 0)])
+def test_yielded_output_uses_new_receipt_time_and_unchanged_native_refusal(delay, native, expected):
+  control, output, _ = snapshots(BASE)
+  sockets, receive, _ = transport([serialized_control(control)], [serialized_output(replace(output,native_allowed=native))])
+  state = PairedDrainState(*sockets)
+  collector = FullrateCollector(SyntheticAdapter())
+  clock = WorkClock()
+  received = []
+
+  def consume_control(event, current, stamp):
+    received.append(stamp)
+    current.control(read_control(event), stamp)
+    clock.ns += 2_000_000
+
+  def consume_output(event, current, stamp):
+    received.append(stamp)
+    current.output(read_applied(event), stamp)
+    clock.ns += 2_000_000
+
+  assert not module.drain_pair(sockets[0], consume_control, sockets[1], consume_output,
+    receive, collector, BASE, state, budget_ns=1_000_000, clock=clock)
+  assert module.drain_pair(sockets[0], consume_control, sockets[1], consume_output,
+    receive, collector, BASE+delay, state, budget_ns=1_000_000, clock=clock)
+  assert received == [BASE, BASE+delay] and len(collector.ready) == expected
+
+
+def test_retained_head_wakes_source_receiver_without_waiting_for_another_pose():
+  from openpilot.starpilot.flm.steerlearnerd import _SourceDelivery
+  calls = []
+  sm = SimpleNamespace(update=lambda timeout:calls.append(timeout),
+    updated={'deviceMotion':True},logMonoTime={'deviceMotion':BASE})
+  receiver = _SourceDelivery()
+  assert receiver.next(sm,retained_evidence=True) == BASE
+  assert calls == [0] and not receiver.pending and receiver.pose_header == 0
diff --git a/prebuilt.json b/prebuilt.json
index ce34c55742..fbd91afd1d 100644
--- a/prebuilt.json
+++ b/prebuilt.json
@@ -1 +1 @@
-{"protected_inventory_sha256": "0cd90e4360c4ac51d61657f8ab32b534523b33df9481485c739f8c79fdb2c878", "schema_version": 1}
+{"protected_inventory_sha256": "2cf347236d514931aed950b4cb4821430b89d8fcda7e381449f69bbc604eff11", "schema_version": 1}