hyperhive/swarm-stt/src/service.rs

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use tokio::sync::{OnceCell, Semaphore};

use crate::cache::{self, Cache};
use crate::engine::Engine;

/// What [`Service`] enforces around the engine.
#[derive(Debug, Clone)]
pub struct Config {
    /// The largest audio file taken, in bytes.
    pub max_bytes: usize,
}

/// Why a request got no transcript.
#[derive(Debug, thiserror::Error)]
pub enum Error {
    /// The audio is larger than [`Config::max_bytes`].
    #[error("the audio is {size} bytes, over the {max} byte limit")]
    TooLarge {
        /// The size received.
        size: usize,
        /// The limit.
        max: usize,
    },
    /// The audio has no bytes.
    #[error("the audio is empty")]
    Empty,
    /// The engine failed; the message says how.
    #[error("{0}")]
    Engine(String),
}

/// A transcript, and whether it came from the cache or from a run another
/// request had already started.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Reply {
    /// The text spoken in the audio (empty for silence).
    pub text: String,
    /// `true` when no engine run was started for this request.
    pub cached: bool,
}

type Flight = Arc<OnceCell<Result<String, String>>>;

/// Transcribes voice memos once each.
pub struct Service<E> {
    engine: E,
    cache: Cache,
    config: Config,
    flights: Mutex<HashMap<String, Flight>>,
    /// One engine run at a time: whisper uses every core it is given.
    engine_slot: Semaphore,
}

impl<E: Engine> Service<E> {
    /// A service running `engine`, caching in `cache`.
    #[must_use]
    pub fn new(engine: E, cache: Cache, config: Config) -> Self {
        Self {
            engine,
            cache,
            config,
            flights: Mutex::new(HashMap::new()),
            engine_slot: Semaphore::new(1),
        }
    }

    /// The transcript of `audio`: from the cache, from a run already in
    /// flight for the same bytes, or from a new engine run.
    pub async fn transcribe(&self, audio: &[u8]) -> Result<Reply, Error> {
        if audio.is_empty() {
            return Err(Error::Empty);
        }
        if audio.len() > self.config.max_bytes {
            return Err(Error::TooLarge {
                size: audio.len(),
                max: self.config.max_bytes,
            });
        }
        let key = cache::key(audio);
        if let Some(text) = self.cache.get(&key).await {
            return Ok(Reply { text, cached: true });
        }
        let flight = Arc::clone(
            self.flights
                .lock()
                .expect("the flights lock is never held across a panic")
                .entry(key.clone())
                .or_default(),
        );
        let mut ran = false;
        let outcome = flight
            .get_or_init(|| async {
                ran = true;
                self.run(&key, audio).await
            })
            .await
            .clone();
        // The first caller to finish clears the entry; a later request finds
        // the cache or starts a new run.
        self.flights
            .lock()
            .expect("the flights lock is never held across a panic")
            .remove(&key);
        outcome
            .map(|text| Reply { text, cached: !ran })
            .map_err(Error::Engine)
    }

    async fn run(&self, key: &str, audio: &[u8]) -> Result<String, String> {
        let _slot = self
            .engine_slot
            .acquire()
            .await
            .map_err(|_| "the engine is shut down".to_owned())?;
        let text = self.engine.transcribe(audio).await?;
        if let Err(e) = self.cache.put(key, &text).await {
            tracing::warn!(error = %e, "caching a transcript failed");
        }
        Ok(text)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::time::Duration;

    use async_trait::async_trait;

    use super::{Config, Error, Reply, Service};
    use crate::cache::Cache;
    use crate::engine::Engine;

    /// Answers `heard: <bytes>` after `delay`, counting its runs.
    struct Fake {
        runs: Arc<AtomicUsize>,
        delay: Duration,
        fail: bool,
    }

    #[async_trait]
    impl Engine for Fake {
        async fn transcribe(&self, audio: &[u8]) -> Result<String, String> {
            self.runs.fetch_add(1, Ordering::SeqCst);
            tokio::time::sleep(self.delay).await;
            if self.fail {
                Err("engine broke".to_owned())
            } else {
                Ok(format!("heard: {}", String::from_utf8_lossy(audio)))
            }
        }
    }

    fn service(
        dir: &std::path::Path,
        delay: Duration,
        fail: bool,
    ) -> (Service<Fake>, Arc<AtomicUsize>) {
        let runs = Arc::new(AtomicUsize::new(0));
        let fake = Fake {
            runs: Arc::clone(&runs),
            delay,
            fail,
        };
        let cache = Cache::new(dir, Duration::from_hours(24));
        (Service::new(fake, cache, Config { max_bytes: 16 }), runs)
    }

    #[tokio::test]
    async fn the_same_memo_is_transcribed_once_and_then_served_from_the_cache() {
        let dir = tempfile::tempdir().unwrap();
        let (service, runs) = service(dir.path(), Duration::ZERO, false);
        let first = service.transcribe(b"memo").await.unwrap();
        let second = service.transcribe(b"memo").await.unwrap();
        assert_eq!(
            first,
            Reply {
                text: "heard: memo".into(),
                cached: false
            }
        );
        assert_eq!(
            second,
            Reply {
                text: "heard: memo".into(),
                cached: true
            }
        );
        assert_eq!(runs.load(Ordering::SeqCst), 1);
    }

    #[tokio::test(start_paused = true)]
    async fn requests_arriving_together_for_one_memo_share_one_run() {
        let dir = tempfile::tempdir().unwrap();
        let (service, runs) = service(dir.path(), Duration::from_secs(5), false);
        let (a, b, c) = tokio::join!(
            service.transcribe(b"memo"),
            service.transcribe(b"memo"),
            service.transcribe(b"memo")
        );
        let replies = [a.unwrap(), b.unwrap(), c.unwrap()];
        assert!(replies.iter().all(|r| r.text == "heard: memo"));
        assert_eq!(replies.iter().filter(|r| !r.cached).count(), 1);
        assert_eq!(runs.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn different_memos_get_their_own_runs() {
        let dir = tempfile::tempdir().unwrap();
        let (service, runs) = service(dir.path(), Duration::ZERO, false);
        service.transcribe(b"one").await.unwrap();
        service.transcribe(b"two").await.unwrap();
        assert_eq!(runs.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn empty_and_oversized_audio_never_reach_the_engine() {
        let dir = tempfile::tempdir().unwrap();
        let (service, runs) = service(dir.path(), Duration::ZERO, false);
        assert!(matches!(service.transcribe(b"").await, Err(Error::Empty)));
        let big = [0u8; 17];
        assert!(matches!(
            service.transcribe(&big).await,
            Err(Error::TooLarge { size: 17, max: 16 })
        ));
        assert_eq!(runs.load(Ordering::SeqCst), 0);
    }

    #[tokio::test]
    async fn a_failed_run_is_not_cached_and_the_next_request_runs_again() {
        let dir = tempfile::tempdir().unwrap();
        let (service, runs) = service(dir.path(), Duration::ZERO, true);
        let err = service.transcribe(b"memo").await.unwrap_err();
        assert!(matches!(err, Error::Engine(ref m) if m == "engine broke"));
        service.transcribe(b"memo").await.unwrap_err();
        assert_eq!(runs.load(Ordering::SeqCst), 2);
    }
}