swarm-stt: the speech-to-text core (hash cache, shared runs, ffmpeg and whisper-cli engine)

swarm-stt/src/service.rs

@@ -0,0 +1,242 @@
+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);
+    }
+}