koan_core/audio/

streaming.rs

//! `PartialFileSource` — a Read+Seek adapter over a file that is still downloading.
//!
//! The download thread writes the track to a `.part` file and publishes how far
//! it has got in an `AtomicU64`; the Symphonia decoder reads the same file
//! through a `PartialFileSource`, which blocks when the read position catches
//! up to the write head. Playback starts long before the transfer finishes and
//! seeking anywhere below the write head costs a `lseek`.
//!
//! Nothing is copied. An earlier design pumped the file into a shared `Vec<u8>`
//! so the decoder could read from memory, which cost as much RAM as the track
//! was long — half a gigabyte for a nine-hour recording, held for as long as it
//! played. The bytes are already on disk; the page cache is better at this.
//!
//! The open descriptor survives the download's final rename from `.part` to its
//! cache path, so a transfer landing mid-playback changes nothing for a reader.

use std::fs::File;
use std::io::{self, Read, Seek, SeekFrom};
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

/// Longest a read may block waiting for bytes before the transfer counts as
/// dead. A download that stops advancing must surface as an error, not park the
/// decode thread forever holding the ring buffer producer.
const STALL_LIMIT: Duration = Duration::from_secs(30);

/// How long to wait between checks for bytes that have not landed yet.
const POLL_INTERVAL: Duration = Duration::from_millis(10);

/// Where a download has got to, as the source needs to know it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamStatus {
    /// Bytes are still arriving.
    Downloading,
    /// Every byte landed. Reads past the end are a clean EOF.
    Complete,
    /// The transfer died before delivering everything. Reads past the written
    /// bytes fail rather than reporting EOF, which would silently truncate the
    /// track and look like a short file.
    Failed,
}

/// A `Read + Seek` view of a file that is still being written.
pub struct PartialFileSource {
    file: File,
    pos: u64,
    /// How many bytes the download has committed to disk so far.
    bytes_written: Arc<AtomicU64>,
    /// Total expected length, or 0 when the server sent no Content-Length.
    total: u64,
    status: Arc<dyn Fn() -> StreamStatus + Send + Sync>,
    stall_limit: Duration,
    /// Whether a read may wait at the write head for more of the download.
    ///
    /// Playback waits; probing does not. A container asked to describe itself
    /// reads whatever it needs to, and Ogg needs its last page — so a probe
    /// that waits waits for the whole transfer. Refusing instead turns that
    /// into an immediate answer of "not from what has arrived", which is
    /// something the caller can act on.
    wait_for_bytes: bool,
    /// Whether to state the advertised length.
    ///
    /// Saying nothing is what stops a container going looking for its tail:
    /// Ogg reads its final page only when the source claims both a length and
    /// seekability, and Symphonia scans for trailing metadata on the same
    /// terms — which no partial file can satisfy. So this is what every
    /// container mid-download ends up opened with, not only Ogg.
    ///
    /// It governs the end of the stream as well as the length, because they
    /// have to be the same end. See `Seek`.
    advertise_len: bool,
    /// Whether `SeekFrom::End` answers for the whole file or for what has
    /// arrived. Separate from `advertise_len` because Ogg needs the first
    /// without the second — see `ProbeMode`.
    whole_file_end: bool,
}

/// How much a probe may claim, and how far it may read.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProbeMode {
    /// State the length. A container that can describe itself from the bytes
    /// already downloaded does, and the result is fully seekable.
    Full,
    /// State no length, so a container that would go looking for its tail
    /// settles for what it can read from the front. Opens immediately, and
    /// stays seekable within what has arrived for anything that describes its
    /// frames from the front. What it gives up is whatever only the tail could
    /// give — for Ogg that is the duration, and with it seeking at all.
    ///
    /// The end is what has arrived. FLAC bisects between its first frame and
    /// the end it is given, so an end it cannot reach sends every probe into
    /// bytes that are not on disk.
    Lengthless,
    /// The same, except that the end stays the whole file's.
    ///
    /// Ogg takes the end it is handed as the end of the *stream*. Handed the
    /// write head it reports a track that is already over — nought
    /// milliseconds — and the decode thread reaches the end of it in a second
    /// and moves on to the next, over and over, so a large Opus file
    /// downloading never plays at all. It cannot seek mid-download under either
    /// answer; this is the difference between playing and not.
    LengthlessWholeEnd,
}

impl PartialFileSource {
    /// Open `path` for playback: reads wait at the write head for the download
    /// to catch up. `bytes_written` is the download's own counter and `total`
    /// its advertised length, 0 when it sent none.
    ///
    /// `mode` must be whatever the probe settled on. A container opened without
    /// a length to describe itself has to be decoded without one too — given a
    /// length it goes looking for its tail all over again, and this time on the
    /// decode thread, where the cost is silence rather than a busy player.
    pub fn open(
        path: &Path,
        bytes_written: Arc<AtomicU64>,
        total: u64,
        status: Arc<dyn Fn() -> StreamStatus + Send + Sync>,
        mode: ProbeMode,
    ) -> io::Result<Self> {
        let mut source = Self::with_stall_limit(path, bytes_written, total, status, STALL_LIMIT)?;
        source.advertise_len = mode == ProbeMode::Full;
        source.whole_file_end = mode != ProbeMode::Lengthless;
        Ok(source)
    }

    /// Open for a probe: never waits at the write head, so a container that
    /// cannot describe itself from what has arrived says so at once.
    pub fn open_for_probe(
        path: &Path,
        bytes_written: Arc<AtomicU64>,
        total: u64,
        status: Arc<dyn Fn() -> StreamStatus + Send + Sync>,
        mode: ProbeMode,
    ) -> io::Result<Self> {
        let mut source = Self::with_stall_limit(path, bytes_written, total, status, STALL_LIMIT)?;
        source.wait_for_bytes = false;
        source.advertise_len = mode == ProbeMode::Full;
        source.whole_file_end = mode != ProbeMode::Lengthless;
        Ok(source)
    }

    fn with_stall_limit(
        path: &Path,
        bytes_written: Arc<AtomicU64>,
        total: u64,
        status: Arc<dyn Fn() -> StreamStatus + Send + Sync>,
        stall_limit: Duration,
    ) -> io::Result<Self> {
        Ok(Self {
            file: File::open(path)?,
            pos: 0,
            bytes_written,
            total,
            status,
            stall_limit,
            wait_for_bytes: true,
            advertise_len: true,
            whole_file_end: true,
        })
    }

    /// Bytes known to be readable — what the download has written, or the whole
    /// file once it has landed.
    fn available(&self) -> u64 {
        let written = self.bytes_written.load(Ordering::Acquire);
        match (self.status)() {
            StreamStatus::Complete => self.file.metadata().map(|m| m.len()).unwrap_or(written),
            _ => written,
        }
    }

    /// Read straight from the file, tolerating a short read at the write head:
    /// `bytes_written` is published by the downloader as it goes and the data
    /// behind it can lag by a moment.
    fn read_available(&mut self, buf: &mut [u8], limit: u64) -> io::Result<usize> {
        let to_read = (limit as usize).min(buf.len());
        self.file.read(&mut buf[..to_read]).inspect(|n| {
            self.pos += *n as u64;
        })
    }
}

impl Read for PartialFileSource {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }

        let deadline = Instant::now() + self.stall_limit;
        loop {
            let available = self.available();
            if available > self.pos {
                let n = self.read_available(buf, available - self.pos)?;
                if n > 0 {
                    return Ok(n);
                }
                // The counter ran ahead of what is visible on disk. Fall
                // through and wait rather than reporting a false EOF.
            }

            match (self.status)() {
                StreamStatus::Failed => {
                    return Err(io::Error::new(
                        io::ErrorKind::BrokenPipe,
                        "stream download failed before delivering the whole track",
                    ));
                }
                // Everything landed and there is nothing past `pos`: real EOF.
                StreamStatus::Complete if available <= self.pos => return Ok(0),
                StreamStatus::Complete => {}
                StreamStatus::Downloading => {
                    // A server that sent a Content-Length has delivered it all.
                    if self.total > 0 && available >= self.total && self.pos >= self.total {
                        return Ok(0);
                    }
                }
            }

            if !self.wait_for_bytes {
                return Err(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    "past what the download has delivered",
                ));
            }
            if Instant::now() >= deadline {
                return Err(io::Error::new(
                    io::ErrorKind::TimedOut,
                    "stream download stalled",
                ));
            }
            std::thread::sleep(POLL_INTERVAL);
        }
    }
}

impl Seek for PartialFileSource {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        let target: i64 = match pos {
            SeekFrom::Start(n) => n as i64,
            SeekFrom::Current(n) => self.pos as i64 + n,
            // The end has to be the same end `byte_len` describes. A source
            // that states a length is asked about the whole file and answers
            // for it; one that states none has only what has arrived, and
            // answering with the advertised total there sends a reader
            // bisecting into bytes that are not on disk yet — which is the
            // whole file's worth of waiting for a FLAC seeked mid-download.
            SeekFrom::End(n) => {
                let len = if self.whole_file_end && self.total > 0 {
                    self.total
                } else {
                    self.available()
                };
                len as i64 + n
            }
        };

        if target < 0 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "seek before beginning of stream",
            ));
        }

        self.pos = self.file.seek(SeekFrom::Start(target as u64))?;
        Ok(self.pos)
    }
}

// Symphonia requires MediaSource: Read + Seek + Send + Any
impl symphonia::core::io::MediaSource for PartialFileSource {
    fn is_seekable(&self) -> bool {
        // Backward seeks and forward seeks below the write head are a `lseek`
        // on a file that is already there. A forward seek past it lands on a
        // read that blocks until the bytes arrive, which is the honest
        // behaviour — callers clamp to `seekable_ms` to avoid asking.
        //
        // Saying no would cost more than it saved: a reader told a stream is
        // unseekable does not stop seeking, it walks the whole file to the
        // target instead, and cannot go backwards at all.
        true
    }

    fn byte_len(&self) -> Option<u64> {
        (self.advertise_len && self.total > 0).then_some(self.total)
    }
}

#[cfg(test)]
mod tests {
    use std::io::Write;
    use std::sync::atomic::AtomicU8;

    use symphonia::core::io::MediaSource;

    use super::*;

    /// A file plus the counter and status a download would publish, so a test
    /// can advance either independently.
    struct Fixture {
        _dir: tempfile::TempDir,
        path: std::path::PathBuf,
        written: Arc<AtomicU64>,
        status: Arc<AtomicU8>,
    }

    const DOWNLOADING: u8 = 0;
    const COMPLETE: u8 = 1;
    const FAILED: u8 = 2;

    impl Fixture {
        fn new() -> Self {
            let dir = tempfile::tempdir().unwrap();
            let path = dir.path().join("track.opus.part");
            File::create(&path).unwrap();
            Self {
                _dir: dir,
                path,
                written: Arc::new(AtomicU64::new(0)),
                status: Arc::new(AtomicU8::new(DOWNLOADING)),
            }
        }

        /// Append bytes and publish them, as the downloader does per chunk.
        fn push(&self, chunk: &[u8]) {
            let mut f = std::fs::OpenOptions::new()
                .append(true)
                .open(&self.path)
                .unwrap();
            f.write_all(chunk).unwrap();
            f.flush().unwrap();
            self.written
                .fetch_add(chunk.len() as u64, Ordering::Release);
        }

        fn set(&self, status: u8) {
            self.status.store(status, Ordering::Release);
        }

        fn status_fn(&self) -> Arc<dyn Fn() -> StreamStatus + Send + Sync> {
            let status = self.status.clone();
            Arc::new(move || match status.load(Ordering::Acquire) {
                COMPLETE => StreamStatus::Complete,
                FAILED => StreamStatus::Failed,
                _ => StreamStatus::Downloading,
            })
        }

        fn source(&self, total: u64) -> PartialFileSource {
            self.source_with_stall(total, STALL_LIMIT)
        }

        fn source_with_stall(&self, total: u64, stall: Duration) -> PartialFileSource {
            let status = self.status.clone();
            PartialFileSource::with_stall_limit(
                &self.path,
                self.written.clone(),
                total,
                Arc::new(move || match status.load(Ordering::Acquire) {
                    COMPLETE => StreamStatus::Complete,
                    FAILED => StreamStatus::Failed,
                    _ => StreamStatus::Downloading,
                }),
                stall,
            )
            .unwrap()
        }
    }

    #[test]
    fn reads_what_has_landed() {
        let fx = Fixture::new();
        fx.push(b"hello streaming world");
        fx.set(COMPLETE);

        let mut out = Vec::new();
        fx.source(21).read_to_end(&mut out).unwrap();
        assert_eq!(out, b"hello streaming world");
    }

    #[test]
    fn read_stops_at_the_write_head_then_resumes() {
        let fx = Fixture::new();
        fx.push(b"abcd");
        let mut src = fx.source(10);

        let mut first = [0u8; 8];
        assert_eq!(src.read(&mut first).unwrap(), 4);
        assert_eq!(&first[..4], b"abcd");

        // The rest arrives while the reader is blocked on it.
        std::thread::spawn({
            let path = fx.path.clone();
            let written = fx.written.clone();
            move || {
                std::thread::sleep(Duration::from_millis(20));
                let mut f = std::fs::OpenOptions::new()
                    .append(true)
                    .open(&path)
                    .unwrap();
                f.write_all(b"efghij").unwrap();
                f.flush().unwrap();
                written.fetch_add(6, Ordering::Release);
            }
        });

        let mut rest = [0u8; 8];
        let n = src.read(&mut rest).unwrap();
        assert_eq!(&rest[..n], b"efghij");
    }

    #[test]
    fn seeks_freely_below_the_write_head() {
        let fx = Fixture::new();
        fx.push(b"0123456789");
        let mut src = fx.source(1_000_000);

        assert_eq!(src.seek(SeekFrom::Start(5)).unwrap(), 5);
        let mut out = [0u8; 3];
        src.read_exact(&mut out).unwrap();
        assert_eq!(&out, b"567");

        // Backwards, into bytes already read — no re-download, no buffer.
        assert_eq!(src.seek(SeekFrom::Start(1)).unwrap(), 1);
        src.read_exact(&mut out).unwrap();
        assert_eq!(&out, b"123");

        assert_eq!(src.seek(SeekFrom::Current(-2)).unwrap(), 2);
    }

    #[test]
    fn seek_from_end_uses_the_advertised_length() {
        let fx = Fixture::new();
        fx.push(b"0123456789");
        let mut src = fx.source(10);

        assert_eq!(src.seek(SeekFrom::End(0)).unwrap(), 10);
        assert_eq!(src.seek(SeekFrom::End(-3)).unwrap(), 7);

        let mut out = [0u8; 3];
        src.read_exact(&mut out).unwrap();
        assert_eq!(&out, b"789");
    }

    #[test]
    fn a_lengthless_source_ends_where_the_download_does() {
        // The end has to agree with `byte_len`. FLAC seeks by bisecting
        // between its first frame and `SeekFrom::End(0)`, so answering with the
        // advertised total aims the search at bytes that are not on disk yet:
        // every probe of the range waits at the write head, and a seek into a
        // half-downloaded track spends the stall limit before failing.
        let fx = Fixture::new();
        fx.push(b"0123456789");

        let mut src = PartialFileSource::open(
            &fx.path,
            fx.written.clone(),
            1_000,
            fx.status_fn(),
            ProbeMode::Lengthless,
        )
        .unwrap();
        assert_eq!(src.byte_len(), None);
        assert_eq!(src.seek(SeekFrom::End(0)).unwrap(), 10);

        // Stating a length means answering for the whole of it.
        let mut src = PartialFileSource::open(
            &fx.path,
            fx.written.clone(),
            1_000,
            fx.status_fn(),
            ProbeMode::Full,
        )
        .unwrap();
        assert_eq!(src.seek(SeekFrom::End(0)).unwrap(), 1_000);
    }

    /// Ogg's half of the same question, and the opposite answer.
    ///
    /// Handed the write head as the end, Ogg reports a stream that is already
    /// over: a large Opus file downloading opened, said nought milliseconds,
    /// and the decode thread finished it and moved on, over and over, so it
    /// never played at all.
    #[test]
    fn an_ogg_source_still_ends_at_the_whole_file() {
        let fx = Fixture::new();
        fx.push(b"0123456789");
        let mut src = PartialFileSource::open(
            &fx.path,
            fx.written.clone(),
            1_000,
            fx.status_fn(),
            ProbeMode::LengthlessWholeEnd,
        )
        .unwrap();
        // Still no length: the point of opening this way is that Ogg does not
        // go looking for a tail that has not arrived.
        assert_eq!(src.byte_len(), None);
        // But the end it is told about is the file's, not the download's.
        assert_eq!(src.seek(SeekFrom::End(0)).unwrap(), 1_000);
    }

    #[test]
    fn a_landed_download_ends_at_the_whole_file() {
        // The bound moves with the transfer, so it stops bounding anything
        // once every byte is there.
        let fx = Fixture::new();
        fx.push(b"0123456789");
        fx.set(COMPLETE);

        let mut src = PartialFileSource::open(
            &fx.path,
            fx.written.clone(),
            10,
            fx.status_fn(),
            ProbeMode::Lengthless,
        )
        .unwrap();
        assert_eq!(src.seek(SeekFrom::End(0)).unwrap(), 10);
    }

    #[test]
    fn seek_before_start_errors() {
        let fx = Fixture::new();
        fx.push(b"hello");
        assert!(fx.source(5).seek(SeekFrom::Current(-1)).is_err());
    }

    #[test]
    fn failed_download_errors_instead_of_reporting_eof() {
        let fx = Fixture::new();
        fx.push(b"partial");
        fx.set(FAILED);
        let mut src = fx.source(1000);

        let mut out = [0u8; 7];
        src.read_exact(&mut out).unwrap();
        assert_eq!(&out, b"partial");

        // Past the written bytes: an error, never a clean EOF — Ok(0) here
        // would end the track early and look like a short file.
        assert_eq!(
            src.read(&mut out).unwrap_err().kind(),
            io::ErrorKind::BrokenPipe
        );
    }

    #[test]
    fn failure_wakes_a_blocked_reader() {
        let fx = Fixture::new();
        let mut src = fx.source(1000);

        let status = fx.status.clone();
        std::thread::spawn(move || {
            std::thread::sleep(Duration::from_millis(20));
            status.store(FAILED, Ordering::Release);
        });

        let mut out = [0u8; 8];
        assert_eq!(
            src.read(&mut out).unwrap_err().kind(),
            io::ErrorKind::BrokenPipe
        );
    }

    #[test]
    fn a_probe_reads_only_what_has_arrived() {
        // The first attempt must fail at the write head rather than wait: a
        // container that would go looking for its tail has to be detected, not
        // waited for.
        let fx = Fixture::new();
        fx.push(b"0123456789");
        let mut src = PartialFileSource::open_for_probe(
            &fx.path,
            fx.written.clone(),
            1_000,
            fx.status_fn(),
            ProbeMode::Full,
        )
        .unwrap();

        let mut out = [0u8; 10];
        src.read_exact(&mut out).unwrap();
        assert_eq!(
            src.read(&mut out).unwrap_err().kind(),
            io::ErrorKind::UnexpectedEof,
            "past the write head is an answer, not a wait"
        );
    }

    #[test]
    fn playback_reads_wait_for_what_has_not_arrived() {
        // And the second attempt does wait, which is what lets a container
        // whose audio starts further in than the streaming threshold — a FLAC
        // with a large padding block, most of them — be opened at all.
        let fx = Fixture::new();
        fx.push(b"0123456789");
        let mut src = PartialFileSource::open(
            &fx.path,
            fx.written.clone(),
            1_000,
            fx.status_fn(),
            ProbeMode::Lengthless,
        )
        .unwrap();

        let mut out = [0u8; 10];
        src.read_exact(&mut out).unwrap();

        std::thread::spawn({
            let path = fx.path.clone();
            let written = fx.written.clone();
            move || {
                std::thread::sleep(Duration::from_millis(20));
                let mut f = std::fs::OpenOptions::new()
                    .append(true)
                    .open(&path)
                    .unwrap();
                f.write_all(b"abcde").unwrap();
                f.flush().unwrap();
                written.fetch_add(5, Ordering::Release);
            }
        });

        let n = src.read(&mut out).unwrap();
        assert_eq!(&out[..n], b"abcde", "it waited rather than giving up");
    }

    #[test]
    fn stalled_download_times_out() {
        // A download with a Content-Length that never arrives: the read must
        // give up rather than park the decode thread forever.
        let fx = Fixture::new();
        let mut src = fx.source_with_stall(1000, Duration::from_millis(20));
        let mut out = [0u8; 8];
        assert_eq!(
            src.read(&mut out).unwrap_err().kind(),
            io::ErrorKind::TimedOut
        );
    }

    #[test]
    fn completion_ends_the_read_at_the_true_length() {
        // A chunked transfer reports no total; completion is what says the file
        // is whole, and its length on disk is what is readable.
        let fx = Fixture::new();
        fx.push(b"chunked");
        fx.set(COMPLETE);

        let mut out = Vec::new();
        fx.source(0).read_to_end(&mut out).unwrap();
        assert_eq!(out, b"chunked");
    }

    #[test]
    fn survives_the_part_file_being_renamed() {
        // The download's final act is a rename. A reader that already has the
        // file open must not notice.
        let fx = Fixture::new();
        fx.push(b"0123456789");
        let mut src = fx.source(10);

        let mut out = [0u8; 4];
        src.read_exact(&mut out).unwrap();
        assert_eq!(&out, b"0123");

        std::fs::rename(&fx.path, fx.path.with_extension("")).unwrap();
        fx.set(COMPLETE);

        let mut rest = Vec::new();
        src.read_to_end(&mut rest).unwrap();
        assert_eq!(rest, b"456789");
    }

    #[test]
    fn byte_len_is_the_advertised_length_only() {
        let fx = Fixture::new();
        assert_eq!(fx.source(42).byte_len(), Some(42));
        // No Content-Length: the length is genuinely unknown, and claiming one
        // would have Symphonia compute a duration from it.
        assert_eq!(fx.source(0).byte_len(), None);
    }

    #[test]
    fn is_seekable_true() {
        let fx = Fixture::new();
        assert!(fx.source(0).is_seekable());
    }
}