diff --git a/recognization/code/live.py b/recognization/code/live.py index 385f379..3850448 100644 --- a/recognization/code/live.py +++ b/recognization/code/live.py @@ -13,6 +13,7 @@ lock so classification callbacks can land concurrently. from __future__ import annotations +import asyncio import collections import os import statistics @@ -169,6 +170,9 @@ class LiveSession: self._pending: Optional[Tuple[int, bytes]] = None self._frame_ready = threading.Event() + # One waiter, created on the caller's loop. None until the first async take. + self._waiter: Optional[asyncio.Event] = None + self._waiter_loop: Optional[asyncio.AbstractEventLoop] = None self._closed = False self._seq = 0 @@ -213,23 +217,77 @@ class LiveSession: self._frames_dropped += 1 self._pending = (int(frame_id), payload) self._frame_ready.set() + self._wake_waiter_locked() return True + def _pop_pending_locked(self) -> Optional[Tuple[int, bytes]]: + pending = self._pending + self._pending = None + self._frame_ready.clear() + return pending + + def _wake_waiter_locked(self) -> None: + """Set the async event from this thread. Caller holds self._lock.""" + waiter = self._waiter + loop = self._waiter_loop + if waiter is None or loop is None: + return + try: + running = asyncio.get_running_loop() + except RuntimeError: + running = None + if running is loop: + waiter.set() + return + try: + loop.call_soon_threadsafe(waiter.set) + except RuntimeError: + # Loop already closed. The sync event still wakes take_pending. + pass + def take_pending(self, timeout: float = 0.2) -> Optional[Tuple[int, bytes]]: """Blocking pop for the processing loop; None on timeout/close.""" if not self._frame_ready.wait(timeout): return None with self._lock: - pending = self._pending - self._pending = None - self._frame_ready.clear() - return pending + return self._pop_pending_locked() + + async def take_pending_async(self) -> Optional[Tuple[int, bytes]]: + """Awaitable pop. None means the session closed with nothing left. + + Bind the waiter to the running loop on first use. A frame that arrives + before that bind is still visible: the event starts set when one is queued. + """ + loop = asyncio.get_running_loop() + with self._lock: + if self._waiter is not None and self._waiter_loop is not loop: + raise RuntimeError("take_pending_async is bound to another event loop") + if self._waiter is None: + self._waiter = asyncio.Event() + self._waiter_loop = loop + if self._pending is not None or self._closed: + self._waiter.set() + waiter = self._waiter + + while True: + await waiter.wait() + with self._lock: + if self._pending is not None: + pending = self._pop_pending_locked() + if self._waiter is waiter and self._pending is None and not self._closed: + self._waiter.clear() + return pending + if self._closed: + return None + if self._waiter is waiter: + self._waiter.clear() def close(self) -> None: with self._lock: self._closed = True self._pending = None - self._frame_ready.set() + self._frame_ready.set() + self._wake_waiter_locked() @property def closed(self) -> bool: @@ -293,6 +351,8 @@ class LiveSession: self._latest_frame_id = None self._frame_origin = [0.0, 0.0] self._shift_map = None + # ponytail: do not touch _waiter here. asyncio.Event is loop-affine; + # take_pending_async clears a stale wake on its own loop. # ------------------------------------------------------------------ frame path def process_frame(self, frame_id: int, payload: bytes) -> Dict[str, Any]: diff --git a/recognization/code/test_live.py b/recognization/code/test_live.py index 1e9ed14..3fbdc7d 100644 --- a/recognization/code/test_live.py +++ b/recognization/code/test_live.py @@ -11,6 +11,7 @@ newest frame. from __future__ import annotations import os +import asyncio import sys import threading import time @@ -391,6 +392,39 @@ class ClassificationBudgetTests(unittest.TestCase): session.close() self.assertTrue(session.closed) self.assertIsNone(session.take_pending(timeout=0.05)) + + def test_async_take_keeps_latest_and_wakes_on_close(self): + """The server waiter must not miss a frame that beat registration, and close must wake it.""" + session, engine, _, _ = make_session() + + async def scenario(): + session.submit_frame(1, frame_bytes()) + session.submit_frame(2, frame_bytes()) + first = await asyncio.wait_for(session.take_pending_async(), timeout=1.0) + self.assertEqual(first[0], 2) + + waiting = asyncio.create_task(session.take_pending_async()) + await asyncio.sleep(0) + self.assertFalse(waiting.done(), "an empty queue must not complete the waiter") + session.reset() + session.submit_frame(3, frame_bytes()) + self.assertEqual((await asyncio.wait_for(waiting, timeout=1.0))[0], 3) + + cancelled = asyncio.create_task(session.take_pending_async()) + await asyncio.sleep(0) + cancelled.cancel() + with self.assertRaises(asyncio.CancelledError): + await cancelled + session.submit_frame(4, frame_bytes()) + self.assertEqual((await asyncio.wait_for(session.take_pending_async(), timeout=1.0))[0], 4) + + closing = asyncio.create_task(session.take_pending_async()) + await asyncio.sleep(0) + session.close() + self.assertIsNone(await asyncio.wait_for(closing, timeout=1.0)) + self.assertFalse(session.submit_frame(5, frame_bytes())) + + asyncio.run(scenario()) engine.shutdown()