truce_core/

audio_tap.rs

//! Audio tap: stream interleaved frames off the audio thread to a
//! background consumer.
//!
//! The audio thread pushes frames wait-free from `process`; a background
//! consumer drains them. This is the *stream + analyze* half of the
//! worker pattern (a spectrum analyzer, a loudness meter, an
//! oscilloscope), where [`crate::tasks`]'s discrete task pool is the
//! *construction offload* half.
//!
//! Two ways to drain:
//! - a `BackgroundTask` handler (in `truce_plugin`) on the shared pool,
//!   woken by a coalescing task each block - bounded threads, best when
//!   analysis is bursty or coalescable.
//! - a dedicated [`StreamWorker`] via [`AudioTap::spawn_worker`] - one
//!   owned thread that parks on the tap and drains sequentially. Consumer
//!   state lives thread-local (no lock), and it never stalls on unrelated
//!   pool work, at the cost of one thread per worker.
//!
//! The tap is a plain `#[skip]` field the plugin owns (usually
//! `Arc<AudioTap<S>>`), so both `process` (via `&params`) and the
//! handler (via `&params`) reach it - the same shared-`Arc` mechanism as
//! any audio -> worker channel. It composes with the pool rather than
//! introducing new threading:
//!
//! ```ignore
//! #[derive(Params)]
//! struct AnalyzerParams {
//!     #[skip]
//!     tap: Arc<AudioTap<f32>>,        // Default: a stereo tap at the default capacity
//!     // ... published spectrum atoms (also #[skip]) ...
//! }
//!
//! // process (audio thread):
//! params.tap.push_frames(&interleaved);
//! if let Some(t) = ctx.tasks::<Analyze>() { t.spawn_coalescing(Analyze); }
//!
//! // run (pool thread):
//! params.tap.drain_with(|chunk| { /* run the FFT, publish */ });
//! ```

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::thread::{self, JoinHandle, Thread};
use std::time::Duration;

use crossbeam_queue::ArrayQueue;

/// Safety-net wake interval for a dedicated [`StreamWorker`]. Normal
/// wakeups come from `push_frames` unparking the thread; this only bounds
/// how long a missed unpark or a pending shutdown can go unnoticed.
const WORKER_PARK_TIMEOUT: Duration = Duration::from_millis(100);

/// Frame capacity [`AudioTap::default`] builds: 32768 stereo frames, a
/// generous consumer-scheduling margin (~170 ms at 192 kHz, ~740 ms at
/// 44.1 kHz). Call [`AudioTap::new`] for a different size or channel count.
const DEFAULT_TAP_FRAMES: usize = 32 * 1024;
/// Channel count [`AudioTap::default`] builds. Stereo is the common case;
/// call [`AudioTap::new`] for mono or higher channel counts.
const DEFAULT_TAP_CHANNELS: usize = 2;

/// A lock-free audio tap. The audio thread is the sole producer
/// ([`Self::push_frames`], wait-free); a background consumer drains it
/// ([`Self::drain_with`]). Whole-frame drop-on-full keeps a drop from
/// desyncing channels.
pub struct AudioTap<S> {
    ring: ArrayQueue<S>,
    channels: usize,
    /// Reused pop buffer whose lock also serializes drains: with a shared
    /// pool a second worker could pick up a coalesced drain while the
    /// first is still running, and a per-sample ring must have one
    /// consumer at a time. `try_lock` makes the second worker bow out
    /// (the first is already draining the same data). Locked only on the
    /// pool thread, never on the audio thread.
    scratch: Mutex<Vec<S>>,
    /// The dedicated [`StreamWorker`] thread, if one drains this tap.
    /// `push_frames` unparks it (a lock-free atomic load + unpark) so the
    /// consumer wakes without polling. Unset when the tap is drained by
    /// the shared pool instead.
    waker: OnceLock<Thread>,
}

/// Builds a stereo tap at a default capacity (32768 frames, ~170 ms at
/// 192 kHz), so a plugin can hold an `Arc<AudioTap<S>>` as a `#[skip]`
/// field and still `#[derive(Default)]` its params. Use [`AudioTap::new`]
/// when the size or channel count needs to differ from the default.
impl<S: Copy + Send + 'static> Default for AudioTap<S> {
    fn default() -> Self {
        Self::new(DEFAULT_TAP_FRAMES, DEFAULT_TAP_CHANNELS)
    }
}

impl<S: Copy + Send + 'static> AudioTap<S> {
    /// Build a tap holding up to `frame_capacity` interleaved frames of
    /// `channels` each. Size `frame_capacity` for the worst realistic
    /// consumer-scheduling gap; drop-on-full is the safety net beyond it
    /// (e.g. the analyzer uses ~32k frames, ~170 ms at 192 kHz).
    ///
    /// # Panics
    ///
    /// Panics if `channels` is zero.
    #[must_use]
    pub fn new(frame_capacity: usize, channels: usize) -> Self {
        assert!(channels > 0, "AudioTap needs at least one channel");
        Self {
            ring: ArrayQueue::new(frame_capacity * channels),
            channels,
            scratch: Mutex::new(Vec::new()),
            waker: OnceLock::new(),
        }
    }

    /// Channel count each frame carries.
    #[must_use]
    pub fn channels(&self) -> usize {
        self.channels
    }

    /// Push interleaved frames from the audio thread. Wait-free. Drops
    /// whole frames (never a partial one, so a drop can't desync
    /// channels) once the ring is full.
    pub fn push_frames(&self, interleaved: &[S]) {
        for frame in interleaved.chunks_exact(self.channels) {
            // The producer is the sole pusher and the consumer only frees
            // space, so free room seen here is a lower bound - if a whole
            // frame fits now it still fits when we push it.
            if self.ring.capacity() - self.ring.len() < self.channels {
                break;
            }
            for &sample in frame {
                let _ = self.ring.push(sample);
            }
        }
        // Wake a dedicated worker, if one is attached. One unpark per
        // block; the park token means a push landing between the worker's
        // drain and its next park is never lost.
        if let Some(worker) = self.waker.get() {
            worker.unpark();
        }
    }

    /// Discard everything buffered. Called off the audio thread (e.g.
    /// from `reset` on a sample-rate change, so frames captured at the
    /// old rate aren't analyzed against the new one). Serialized with
    /// [`Self::drain_with`]; if a drain is mid-flight this is a no-op and
    /// the drain finishes the stale frames - `reset` should follow with
    /// the consumer's own state reset.
    pub fn clear(&self) {
        let Ok(_guard) = self.scratch.try_lock() else {
            return;
        };
        while self.ring.pop().is_some() {}
    }

    /// Drain everything buffered and hand it to `f` as one interleaved
    /// slice. Runs on the consumer (a pool worker); safe to call from a
    /// `BackgroundTask::run`. If another worker is already draining
    /// this tap it returns without double-draining - that worker sees the
    /// same data. `f` is not called when nothing is buffered.
    pub fn drain_with(&self, mut f: impl FnMut(&[S])) {
        let Ok(mut scratch) = self.scratch.try_lock() else {
            return;
        };
        scratch.clear();
        while let Some(sample) = self.ring.pop() {
            scratch.push(sample);
        }
        if !scratch.is_empty() {
            f(&scratch);
        }
    }

    /// Spawn a dedicated thread that drains this tap sequentially. The
    /// thread parks until [`Self::push_frames`] unparks it, then hands
    /// every buffered frame to `on_drain` as one interleaved slice, in
    /// order. Consumer state lives inside the closure - a single owner,
    /// so no lock. The returned [`StreamWorker`] joins the thread on drop.
    ///
    /// This is the dedicated-thread alternative to draining on the shared
    /// task pool: pick it when the consumer runs continuously and would
    /// otherwise contend with unrelated pool work. It spends one thread
    /// per worker, so unlike the pool it does not bound thread growth.
    ///
    /// Attach **at most one** worker per tap: the per-sample ring has a
    /// single consumer, and only one thread can be registered for the
    /// wake-on-push. Drain a tap with either a `StreamWorker` or the shared
    /// pool, not both.
    ///
    /// # Panics
    ///
    /// Panics if a worker is already attached to this tap, or if the OS
    /// refuses to spawn the worker thread.
    #[must_use]
    pub fn spawn_worker(
        self: Arc<Self>,
        name: &str,
        mut on_drain: impl FnMut(&[S]) + Send + 'static,
    ) -> StreamWorker {
        // Fail loud on a second attach rather than silently registering no
        // waker (which would fall back to the park timeout and race the
        // first worker on the drain lock). Checked before spawning so a
        // rejected call leaks no thread.
        assert!(
            self.waker.get().is_none(),
            "AudioTap already has a StreamWorker; attach at most one worker per tap",
        );
        let tap = Arc::clone(&self);
        let shutdown = Arc::new(AtomicBool::new(false));
        let shutdown_thread = Arc::clone(&shutdown);
        let handle = thread::Builder::new()
            .name(name.to_owned())
            .spawn(move || {
                while !shutdown_thread.load(Ordering::Acquire) {
                    tap.drain_with(&mut on_drain);
                    thread::park_timeout(WORKER_PARK_TIMEOUT);
                }
            })
            .expect("spawn truce stream worker");
        // Publish the thread so `push_frames` can unpark it. Set from the
        // spawner (not the worker itself) so it is in place before the
        // first push; a push before it lands just relies on the timeout.
        let _ = self.waker.set(handle.thread().clone());
        StreamWorker {
            shutdown,
            thread: handle.thread().clone(),
            handle: Some(handle),
        }
    }
}

/// Handle to a dedicated [`AudioTap`] consumer thread spawned by
/// [`AudioTap::spawn_worker`]. Dropping it asks the thread to stop and
/// joins it, so the worker's lifetime is tied to whatever owns the handle
/// (typically a `#[skip]` field alongside the tap).
pub struct StreamWorker {
    shutdown: Arc<AtomicBool>,
    thread: Thread,
    handle: Option<JoinHandle<()>>,
}

impl Drop for StreamWorker {
    fn drop(&mut self) {
        self.shutdown.store(true, Ordering::Release);
        // Wake the thread out of its park so it sees the flag now rather
        // than after the timeout.
        self.thread.unpark();
        if let Some(handle) = self.handle.take() {
            let _ = handle.join();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn round_trips_frames() {
        let tap = AudioTap::<f32>::new(16, 2);
        tap.push_frames(&[1.0, 2.0, 3.0, 4.0]); // two stereo frames
        let mut got = Vec::new();
        tap.drain_with(|chunk| got.extend_from_slice(chunk));
        assert_eq!(got, vec![1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn drop_on_full_stays_frame_aligned() {
        // Capacity 2 frames (4 samples). Push 4 frames; the last 2 drop.
        let tap = AudioTap::<i32>::new(2, 2);
        tap.push_frames(&[1, 1, 2, 2, 3, 3, 4, 4]);
        let mut got = Vec::new();
        tap.drain_with(|chunk| got.extend_from_slice(chunk));
        // Whatever survived is a whole number of frames (even length),
        // and never a split frame.
        assert_eq!(got.len() % 2, 0);
        assert_eq!(got, vec![1, 1, 2, 2]);
    }

    #[test]
    fn drain_of_empty_tap_does_not_call_f() {
        let tap = AudioTap::<f32>::new(4, 1);
        let mut called = false;
        tap.drain_with(|_| called = true);
        assert!(!called, "no callback when nothing is buffered");
    }

    #[test]
    fn stream_worker_drains_in_order() {
        use std::sync::mpsc;

        let tap = Arc::new(AudioTap::<i32>::new(64, 2));
        let (tx, rx) = mpsc::channel();
        let worker = tap
            .clone()
            .spawn_worker("test-stream-worker", move |chunk| {
                for &sample in chunk {
                    let _ = tx.send(sample);
                }
            });
        tap.push_frames(&[1, 2, 3, 4]);

        // The worker drains FIFO, so the samples arrive in push order.
        let mut got = Vec::new();
        for _ in 0..4 {
            got.push(
                rx.recv_timeout(Duration::from_secs(5))
                    .expect("worker drained the pushed frames"),
            );
        }
        assert_eq!(got, vec![1, 2, 3, 4]);
        drop(worker);
    }

    #[test]
    fn default_is_a_usable_stereo_tap() {
        let tap = AudioTap::<f32>::default();
        assert_eq!(tap.channels(), DEFAULT_TAP_CHANNELS);
        tap.push_frames(&[1.0, 2.0]);
        let mut got = Vec::new();
        tap.drain_with(|chunk| got.extend_from_slice(chunk));
        assert_eq!(got, vec![1.0, 2.0]);
    }

    #[test]
    #[should_panic(expected = "at most one worker per tap")]
    fn second_worker_panics() {
        let tap = Arc::new(AudioTap::<i32>::new(16, 2));
        let _first = tap.clone().spawn_worker("first", |_| {});
        // A second worker on the same tap is misuse and must fail loudly.
        let _second = tap.clone().spawn_worker("second", |_| {});
    }
}