Fix a 5 second delay for the first listener on an idle serve worker

sendspin/serve/worker.py

@@ -9,9 +9,11 @@
 import asyncio
 import logging
 import multiprocessing as mp
+from collections import deque
 from multiprocessing.sharedctypes import Synchronized
 from contextlib import suppress
 
+from aiosendspin.clock import RawMonotonicClock
 from aiosendspin.noise import Identity
 from aiosendspin.server import (
     ClientAddedEvent,
@@ -59,6 +61,8 @@ def __init__(
         self._server: SendspinPlayerServer | None = None
         self._active_group: SendspinGroup | None = None
         self._stream: PushStream | None = None
+        self._clock = RawMonotonicClock()
+        self._recent_chunks: deque[AudioChunk] = deque()
 
     async def run(self) -> None:
         """Main worker loop: start server, consume audio queue."""
@@ -69,6 +73,7 @@ async def run(self) -> None:
             self._status_queue.put(WorkerListening(worker_id=self.worker_id, port=self.port))
 
             # Consume audio chunks from coordinator
+            fed_stream: PushStream | None = None
             while True:
                 msg = await loop.run_in_executor(None, self._audio_queue.get)
 
@@ -77,16 +82,15 @@ async def run(self) -> None:
                     break
 
                 if isinstance(msg, AudioChunk):
+                    self._keep_recent_chunk(msg)
                     stream = self._get_stream()
                     if stream is None:
                         continue
-                    fmt = AudioFormat(
-                        sample_rate=msg.sample_rate,
-                        bit_depth=msg.bit_depth,
-                        channels=msg.channels,
-                    )
                     try:
-                        stream.prepare_audio(msg.pcm_bytes, fmt)
+                        if stream is not fed_stream:
+                            fed_stream = stream
+                            self._prepare_recent_audio(stream)
+                        stream.prepare_audio(msg.pcm_bytes, _chunk_format(msg))
                         await stream.commit_audio(play_start_us=msg.play_start_us)
                     except StreamStoppedError:
                         # All clients disconnected — stream was stopped.
@@ -158,6 +162,28 @@ def _report_client_count(self) -> None:
         count = len(self._server.connected_clients)
         self._status_queue.put(WorkerClientCount(worker_id=self.worker_id, count=count))
 
+    def _keep_recent_chunk(self, chunk: AudioChunk) -> None:
+        """Keep the chunks that haven't finished playing, for a stream that starts later."""
+        now_us = self._clock.now_us()
+        while self._recent_chunks and _chunk_end_us(self._recent_chunks[0]) <= now_us:
+            self._recent_chunks.popleft()
+        # Restarting the audio source begins a new timeline that the kept chunks don't lead into.
+        if self._recent_chunks and not (
+            0 <= chunk.play_start_us - _chunk_end_us(self._recent_chunks[-1]) <= 1
+        ):
+            self._recent_chunks.clear()
+        self._recent_chunks.append(chunk)
+
+    def _prepare_recent_audio(self, stream: PushStream) -> None:
+        """Queue the kept chunks before the newest one, so a new stream starts near real time."""
+        *history, _ = self._recent_chunks
+        for chunk in history:
+            stream.prepare_historical_audio(
+                chunk.pcm_bytes,
+                _chunk_format(chunk),
+                start_time_us=history[0].play_start_us,
+            )
+
     def _get_stream(self) -> PushStream | None:
         """Get the active PushStream, or None if no clients are connected."""
         return self._stream
@@ -172,6 +198,19 @@ async def _shutdown_server(self) -> None:
                 await self._server.close()
 
 
+def _chunk_format(chunk: AudioChunk) -> AudioFormat:
+    return AudioFormat(
+        sample_rate=chunk.sample_rate,
+        bit_depth=chunk.bit_depth,
+        channels=chunk.channels,
+    )
+
+
+def _chunk_end_us(chunk: AudioChunk) -> int:
+    samples = len(chunk.pcm_bytes) // ((chunk.bit_depth // 8) * chunk.channels)
+    return chunk.play_start_us + samples * 1_000_000 // chunk.sample_rate
+
+
 def worker_main(
     worker_id: int,
     port: int,

tests/serve/test_worker.py

@@ -3,9 +3,10 @@
 from __future__ import annotations
 
 import multiprocessing as mp
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
 
 import pytest
+from aiosendspin.clock import ManualClock
 
 from sendspin.serve.ipc import AudioChunk, Shutdown, WorkerListening
 from sendspin.serve.worker import ServeWorker
@@ -113,6 +114,91 @@ async def test_worker_processes_audio_chunk(
     mock_stream.commit_audio.assert_called_once_with(play_start_us=1_000_000)
 
 
+def _chunk(play_start_us: int) -> AudioChunk:
+    # 4800 frames of 48 kHz stereo s16 is 100 ms.
+    return AudioChunk(
+        pcm_bytes=b"\x00" * 4800 * 4,
+        sample_rate=48000,
+        bit_depth=16,
+        channels=2,
+        play_start_us=play_start_us,
+    )
+
+
+async def _run_until_stream_appears(
+    worker: ServeWorker, audio_queue: mp.Queue, chunks: list[AudioChunk]
+) -> MagicMock:
+    """Feed chunks to a worker whose stream exists only from the last chunk on."""
+    for chunk in chunks:
+        audio_queue.put(chunk)
+    audio_queue.put(Shutdown())
+
+    mock_stream = MagicMock()
+    mock_stream.commit_audio = AsyncMock()
+    streams = [None] * (len(chunks) - 1) + [mock_stream]
+
+    with (
+        patch.object(worker, "_start_server", new_callable=AsyncMock),
+        patch.object(worker, "_get_stream", side_effect=streams),
+    ):
+        await worker.run()
+    return mock_stream
+
+
+@pytest.mark.asyncio
+async def test_worker_replays_recent_chunks_to_new_stream(
+    audio_queue: mp.Queue,
+    status_queue: mp.Queue,
+    total_listeners: mp.Value,
+) -> None:
+    """A first listener gets the unplayed chunks received before its stream."""
+    worker = ServeWorker(
+        worker_id=0,
+        port=8928,
+        audio_queue=audio_queue,
+        status_queue=status_queue,
+        total_listeners=total_listeners,
+    )
+    start_us = 1_000_000
+    worker._clock = ManualClock(start_us)
+    # The first chunk has finished playing.
+    chunks = [_chunk(start_us + i * 100_000) for i in range(-1, 3)]
+
+    stream = await _run_until_stream_appears(worker, audio_queue, chunks)
+
+    assert stream.mock_calls == [
+        call.prepare_historical_audio(chunks[1].pcm_bytes, ANY, start_time_us=start_us),
+        call.prepare_historical_audio(chunks[2].pcm_bytes, ANY, start_time_us=start_us),
+        call.prepare_audio(chunks[3].pcm_bytes, ANY),
+        call.commit_audio(play_start_us=chunks[3].play_start_us),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_worker_skips_chunks_from_previous_timeline(
+    audio_queue: mp.Queue,
+    status_queue: mp.Queue,
+    total_listeners: mp.Value,
+) -> None:
+    """Chunks the new timeline doesn't continue from are not replayed behind it."""
+    worker = ServeWorker(
+        worker_id=0,
+        port=8928,
+        audio_queue=audio_queue,
+        status_queue=status_queue,
+        total_listeners=total_listeners,
+    )
+    start_us = 1_000_000
+    worker._clock = ManualClock(start_us)
+    chunks = [_chunk(start_us + 3_000_000), _chunk(start_us), _chunk(start_us + 100_000)]
+
+    stream = await _run_until_stream_appears(worker, audio_queue, chunks)
+
+    stream.prepare_historical_audio.assert_called_once_with(
+        chunks[1].pcm_bytes, ANY, start_time_us=start_us
+    )
+
+
 def test_worker_clears_group_when_last_client_disconnects(
     audio_queue: mp.Queue,
     status_queue: mp.Queue,