Read ahead so a stalled source cannot speed up playback

services/media.py

@@ -20,12 +20,15 @@
 
 import os
 import sys
+import queue
+import threading
 import subprocess
 import tempfile
 import asyncio
 import logging
 from typing import Optional
 
+import discord
 import yt_dlp
 
 log = logging.getLogger("loopify.media")
@@ -340,6 +343,103 @@ def spawn_stream(track: dict) -> AudioStream:
     return AudioStream.launch(cmd)
 
 
+# Discord consumes 20 ms of 48 kHz stereo 16-bit PCM per frame.
+FRAME_SIZE = 3840
+FRAME_SECONDS = 0.02
+# How much audio to keep ready. 5 s is ~960 KB — nothing against 15 GB of RAM,
+# and long enough to cover the gaps yt-dlp leaves when YouTube throttles a
+# download after its initial burst.
+READ_AHEAD_SECONDS = 5.0
+
+
+class BufferedAudioSource(discord.AudioSource):
+    """
+    Keeps a few seconds of audio ready so a stalled source never delays a frame.
+
+    ``discord.py``'s player paces itself against a wall clock::
+
+        next_time = self._start + DELAY * self.loops
+        delay = max(0, DELAY + (next_time - time.perf_counter()))
+        time.sleep(delay)
+
+    If ``source.read()`` blocks — FFmpeg waiting on bytes from yt-dlp — the
+    player falls behind that schedule. ``delay`` then clamps to zero and it
+    sends frames as fast as it can until it catches up, which is audible as the
+    track briefly speeding up before settling. Reading ahead in a background
+    thread means a stall shorter than the buffer never reaches the player.
+
+    The queue is **bounded**: an unbounded one in front of a fast source would
+    pull a whole track into RAM and undo the backpressure that the yt-dlp pipe
+    exists to provide.
+    """
+
+    def __init__(self, source, *, seconds: float = READ_AHEAD_SECONDS) -> None:
+        self._source = source
+        self.capacity_frames = max(1, int(seconds / FRAME_SECONDS))
+        self._queue: "queue.Queue" = queue.Queue(maxsize=self.capacity_frames)
+        self._stop = threading.Event()
+        self._thread = threading.Thread(
+            target=self._fill, name="loopify-readahead", daemon=True,
+        )
+        self._thread.start()
+
+    @property
+    def buffered_frames(self) -> int:
+        return self._queue.qsize()
+
+    def is_opus(self) -> bool:
+        return self._source.is_opus()
+
+    def _fill(self) -> None:
+        """Pull from the wrapped source until it ends, fails, or we stop."""
+        try:
+            while not self._stop.is_set():
+                data = self._source.read()
+                if not data:
+                    break
+                # Block when full — that is the backpressure. Time out so the
+                # thread still notices _stop while the consumer is gone.
+                while not self._stop.is_set():
+                    try:
+                        self._queue.put(data, timeout=0.1)
+                        break
+                    except queue.Full:
+                        continue
+        except Exception as e:
+            log.warning("Read-ahead stopped: %s", e)
+        finally:
+            # A sentinel unblocks a consumer waiting on an exhausted source.
+            try:
+                self._queue.put_nowait(b"")
+            except queue.Full:
+                pass
+
+    def read(self) -> bytes:
+        if self._stop.is_set():
+            return b""
+        try:
+            return self._queue.get(timeout=1.0)
+        except queue.Empty:
+            # The producer is wedged; ending is better than stalling playback.
+            return b""
+
+    def cleanup(self) -> None:
+        """Stop the reader thread and tear down the wrapped source."""
+        if self._stop.is_set():
+            return
+        self._stop.set()
+        # Drain so a producer blocked on a full queue can see _stop and exit.
+        try:
+            while True:
+                self._queue.get_nowait()
+        except queue.Empty:
+            pass
+        try:
+            self._source.cleanup()
+        except Exception:
+            pass
+
+
 def make_pipe_source(stdin, *, volume: float = 0.5, ffmpeg_filter: str = "",
                      seek_seconds: float = 0.0):
     """
@@ -351,10 +451,15 @@ def make_pipe_source(stdin, *, volume: float = 0.5, ffmpeg_filter: str = "",
     read-and-discard rather than a real seek, but it costs almost nothing
     because yt-dlp delivers at network speed rather than in realtime.
     """
-    import discord
     options = f"-vn -af {ffmpeg_filter}" if ffmpeg_filter else "-vn"
     before = f"-ss {seek_seconds:.3f}" if seek_seconds > 0 else None
     source = discord.FFmpegPCMAudio(
         stdin, pipe=True, before_options=before, options=options,
     )
-    return discord.PCMVolumeTransformer(source, volume=volume)
+    # Order matters. The read-ahead goes around FFmpeg, which is what stalls,
+    # and the volume transformer stays outermost so `MusicPlayer.set_volume`
+    # still finds a PCMVolumeTransformer on `voice_client.source`. Volume is a
+    # cheap multiply, so applying it on the consumer side costs nothing.
+    return discord.PCMVolumeTransformer(
+        BufferedAudioSource(source), volume=volume,
+    )

tests/test_buffered_source.py

@@ -0,0 +1,254 @@
+"""
+``services.media.BufferedAudioSource`` - read-ahead in front of FFmpeg.
+
+discord.py's player paces itself against a wall clock and, if a read blocks,
+catches up by sending frames with no delay - audible as the track briefly
+speeding up. These tests drive sources that stall on purpose and assert the
+buffer absorbs it.
+"""
+
+import threading
+import time
+
+import discord
+import pytest
+
+from services.media import BufferedAudioSource, FRAME_SIZE
+
+
+class FakeSource(discord.AudioSource):
+    """A source that yields ``frames`` frames, optionally stalling partway."""
+
+    def __init__(self, frames: int, stall_at: int = -1, stall_for: float = 0.0):
+        self.frames = frames
+        self.stall_at = stall_at
+        self.stall_for = stall_for
+        self.served = 0
+        self.cleaned = False
+
+    def read(self) -> bytes:
+        if self.served == self.stall_at:
+            time.sleep(self.stall_for)
+        if self.served >= self.frames:
+            return b""
+        self.served += 1
+        return bytes([self.served % 256]) * FRAME_SIZE
+
+    def cleanup(self) -> None:
+        self.cleaned = True
+
+
+def drain(source) -> list:
+    frames = []
+    while True:
+        chunk = source.read()
+        if not chunk:
+            return frames
+        frames.append(chunk)
+
+
+# -- the bug this exists to prevent ------------------------------------
+
+def test_a_stalling_source_still_delivers_frames_on_time():
+    """
+    The player must not have to wait on a stalled source. With enough
+    read-ahead, a one-second gap is invisible to the caller.
+    """
+    inner = FakeSource(frames=200, stall_at=5, stall_for=1.0)
+    source = BufferedAudioSource(inner, seconds=4.0)
+    try:
+        # Give the background thread a moment to fill past the stall point.
+        deadline = time.monotonic() + 5
+        while source.buffered_frames < 50 and time.monotonic() < deadline:
+            time.sleep(0.01)
+
+        started = time.monotonic()
+        for _ in range(40):
+            assert source.read(), "the buffer ran dry during a stall"
+        elapsed = time.monotonic() - started
+
+        assert elapsed < 0.2, f"reads blocked for {elapsed:.2f}s; buffer did not absorb the stall"
+    finally:
+        source.cleanup()
+
+
+def test_every_frame_arrives_in_order_and_none_are_dropped():
+    inner = FakeSource(frames=120)
+    source = BufferedAudioSource(inner, seconds=1.0)
+    try:
+        frames = drain(source)
+    finally:
+        source.cleanup()
+
+    assert len(frames) == 120, "frames were dropped"
+    assert frames[0][0] == 1
+    assert frames[-1][0] == 120 % 256
+    assert all(len(f) == FRAME_SIZE for f in frames)
+
+
+def test_the_end_of_the_source_ends_playback():
+    inner = FakeSource(frames=3)
+    source = BufferedAudioSource(inner, seconds=1.0)
+    try:
+        assert len(drain(source)) == 3
+        assert source.read() == b"", "must keep reporting end, not block"
+    finally:
+        source.cleanup()
+
+
+def test_an_empty_source_ends_immediately():
+    source = BufferedAudioSource(FakeSource(frames=0), seconds=1.0)
+    try:
+        assert source.read() == b""
+    finally:
+        source.cleanup()
+
+
+# -- bounded memory ----------------------------------------------------
+
+def test_the_buffer_is_bounded():
+    """
+    An unbounded queue in front of a fast source would pull an entire track
+    into RAM, undoing the backpressure the yt-dlp pipe exists to provide.
+    """
+    inner = FakeSource(frames=100_000)
+    source = BufferedAudioSource(inner, seconds=2.0)
+    try:
+        time.sleep(0.5)                     # let it fill
+        cap = source.capacity_frames
+        assert cap == pytest.approx(100, abs=1), "2s at 20ms/frame is ~100 frames"
+        assert source.buffered_frames <= cap
+        assert inner.served <= cap + 2, "the producer ran past the buffer cap"
+    finally:
+        source.cleanup()
+
+
+def test_seconds_translates_to_frames():
+    source = BufferedAudioSource(FakeSource(frames=1), seconds=5.0)
+    try:
+        assert source.capacity_frames == 250      # 5s / 20ms
+    finally:
+        source.cleanup()
+
+
+# -- teardown ----------------------------------------------------------
+
+def test_cleanup_stops_the_thread_and_the_inner_source():
+    inner = FakeSource(frames=100_000)
+    source = BufferedAudioSource(inner, seconds=1.0)
+    time.sleep(0.2)
+    source.cleanup()
+
+    deadline = time.monotonic() + 3
+    while source._thread.is_alive() and time.monotonic() < deadline:
+        time.sleep(0.01)
+
+    assert not source._thread.is_alive(), "the reader thread outlived cleanup()"
+    assert inner.cleaned, "the wrapped source was not cleaned up"
+
+
+def test_cleanup_is_idempotent():
+    source = BufferedAudioSource(FakeSource(frames=10), seconds=1.0)
+    source.cleanup()
+    source.cleanup()
+    source.cleanup()
+
+
+def test_reading_after_cleanup_reports_end_rather_than_blocking():
+    source = BufferedAudioSource(FakeSource(frames=100_000), seconds=1.0)
+    source.cleanup()
+    started = time.monotonic()
+    assert source.read() == b""
+    assert time.monotonic() - started < 1.0, "read() blocked after cleanup"
+
+
+def test_no_reader_threads_are_left_behind():
+    before = threading.active_count()
+    for _ in range(5):
+        s = BufferedAudioSource(FakeSource(frames=1000), seconds=1.0)
+        time.sleep(0.05)
+        s.cleanup()
+
+    deadline = time.monotonic() + 5
+    while threading.active_count() > before and time.monotonic() < deadline:
+        time.sleep(0.05)
+    assert threading.active_count() <= before
+
+
+# -- a source that raises ----------------------------------------------
+
+def test_a_source_that_raises_ends_playback_instead_of_hanging():
+    class Exploding(discord.AudioSource):
+        def read(self):
+            raise OSError("pipe died")
+
+        def cleanup(self):
+            pass
+
+    source = BufferedAudioSource(Exploding(), seconds=1.0)
+    try:
+        started = time.monotonic()
+        assert source.read() == b""
+        assert time.monotonic() - started < 2.0
+    finally:
+        source.cleanup()
+
+
+def test_it_reports_opus_the_same_as_the_wrapped_source():
+    """discord.py asks this to decide whether to encode."""
+    source = BufferedAudioSource(FakeSource(frames=1), seconds=1.0)
+    try:
+        assert source.is_opus() is False
+    finally:
+        source.cleanup()
+
+
+# -- how make_pipe_source composes the chain ---------------------------
+
+def test_the_volume_transformer_stays_outermost(tmp_path):
+    """
+    MusicPlayer.set_volume does `isinstance(vc.source, PCMVolumeTransformer)`.
+    Wrapping the buffer around the outside would silently break volume control.
+    """
+    from services import media
+
+    path = tmp_path / "silence.raw"
+    path.write_bytes(b"\x00" * 4096)
+    with open(path, "rb") as handle:
+        source = media.make_pipe_source(handle, volume=0.7)
+        try:
+            assert isinstance(source, discord.PCMVolumeTransformer)
+            assert source.volume == pytest.approx(0.7)
+            assert isinstance(source.original, BufferedAudioSource)
+        finally:
+            source.cleanup()
+
+
+def test_cleanup_propagates_through_the_whole_chain(tmp_path):
+    from services import media
+
+    path = tmp_path / "silence.raw"
+    path.write_bytes(b"\x00" * 4096)
+    with open(path, "rb") as handle:
+        source = media.make_pipe_source(handle)
+        buffered = source.original
+        source.cleanup()
+
+    deadline = time.monotonic() + 3
+    while buffered._thread.is_alive() and time.monotonic() < deadline:
+        time.sleep(0.01)
+    assert not buffered._thread.is_alive(), "cleanup did not reach the read-ahead thread"
+
+
+def test_volume_can_still_be_changed_mid_playback(tmp_path):
+    from services import media
+
+    path = tmp_path / "silence.raw"
+    path.write_bytes(b"\x00" * 4096)
+    with open(path, "rb") as handle:
+        source = media.make_pipe_source(handle, volume=0.5)
+        try:
+            source.volume = 0.9
+            assert source.volume == pytest.approx(0.9)
+        finally:
+            source.cleanup()