swarm-stt: reply deadline agrees with the limits, one request per agent, short errors to the sender

swarm-stt/src/main.rs

@@ -6,8 +6,9 @@
 //! runs and the engine. This file only reads the settings and carries
 //! requests between the queue and the service.
 
+use std::collections::HashSet;
 use std::path::PathBuf;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 use std::time::Duration;
 
@@ -33,6 +34,82 @@ const QUEUE_GROUP: &str = "swarm-stt";
 /// How often expired transcripts are removed.
 const PURGE_EVERY: Duration = Duration::from_hours(1);
 
+type Var<'a> = &'a dyn Fn(&str) -> Result<String>;
+type Number<'a> = &'a dyn Fn(&str, u64) -> Result<u64>;
+
+/// The service and the limits around it, from the environment.
+struct Settings {
+    service: Arc<Service<Whisper>>,
+    cache: Cache,
+    waiting: Arc<Semaphore>,
+}
+
+fn settings(var: Var<'_>, number: Number<'_>) -> Result<Settings> {
+    let state = PathBuf::from(var("STATE_DIR")?);
+    let cache = Cache::new(
+        state.join("cache"),
+        Duration::from_secs(number("CACHE_TTL_SECS", 24 * 60 * 60)?),
+    );
+    let engine_timeout_secs = number("RUN_TIMEOUT_SECS", 240)?;
+    let engine = Whisper {
+        ffmpeg: PathBuf::from(var("FFMPEG")?),
+        whisper: PathBuf::from(var("WHISPER_CLI")?),
+        model: PathBuf::from(var("MODEL")?),
+        scratch: state.join("scratch"),
+        max_secs: number("MAX_AUDIO_SECS", 300)?,
+        timeout: Duration::from_secs(engine_timeout_secs),
+    };
+    let max_bytes = usize::try_from(number("MAX_BYTES", 6 * 1024 * 1024)?)
+        .context("SWARM_STT_MAX_BYTES does not fit in memory")?;
+    let service = Arc::new(Service::new(engine, cache.clone(), Config { max_bytes }));
+    // Requests held at once, running or waiting for the engine's one slot:
+    // each holds its audio in memory, so past this the service says it is busy.
+    let max_waiting = number("MAX_WAITING", 3)?;
+    let worst_wait = Duration::from_secs(max_waiting.saturating_mul(engine_timeout_secs));
+    if worst_wait >= swarm_wire::stt::REPLY_DEADLINE {
+        bail!(
+            "SWARM_STT_MAX_WAITING ({max_waiting}) times SWARM_STT_RUN_TIMEOUT_SECS \
+             ({engine_timeout_secs}) must stay under the {}s reply deadline, or a queued \
+             request is answered after its sender has given up",
+            swarm_wire::stt::REPLY_DEADLINE.as_secs()
+        );
+    }
+    let waiting = Arc::new(Semaphore::new(
+        usize::try_from(max_waiting).context("SWARM_STT_MAX_WAITING is too large")?,
+    ));
+    Ok(Settings {
+        service,
+        cache,
+        waiting,
+    })
+}
+
+/// The queue connection, as the service's own OIDC client.
+async fn connect_queue(var: Var<'_>) -> Result<async_nats::Client> {
+    let client_id = var("OIDC_CLIENT_ID")?;
+    let secret =
+        swarm_secret_client::read_with_startup_retry("reading the queue client secret", || {
+            client_secret(&client_id)
+        })
+        .await?;
+    let reread_id = client_id.clone();
+    let secret = StoredSecret::new(client_id.clone(), secret, move || {
+        let client_id = reread_id.clone();
+        async move { client_secret(&client_id).await.map_err(Into::into) }
+    });
+    swarm_queue_client::connect(QueueConfig {
+        url: var("NATS_URL")?,
+        token_endpoint: var("OIDC_TOKEN_ENDPOINT")?,
+        client_id,
+        client_secret: ClientSecret::Stored(secret),
+        ca_file: std::env::var(format!("{PREFIX}_OIDC_CA_FILE"))
+            .ok()
+            .map(PathBuf::from),
+    })
+    .await
+    .map_err(Into::into)
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     hive_log::init();
@@ -49,49 +126,13 @@ async fn main() -> Result<()> {
         }
     };
 
-    let state = PathBuf::from(var("STATE_DIR")?);
-    let cache = Cache::new(
-        state.join("cache"),
-        Duration::from_secs(number("CACHE_TTL_SECS", 24 * 60 * 60)?),
-    );
-    let engine = Whisper {
-        ffmpeg: PathBuf::from(var("FFMPEG")?),
-        whisper: PathBuf::from(var("WHISPER_CLI")?),
-        model: PathBuf::from(var("MODEL")?),
-        scratch: state.join("scratch"),
-        max_secs: number("MAX_AUDIO_SECS", 300)?,
-        timeout: Duration::from_secs(number("RUN_TIMEOUT_SECS", 300)?),
-    };
-    let max_bytes = usize::try_from(number("MAX_BYTES", 6 * 1024 * 1024)?)
-        .context("SWARM_STT_MAX_BYTES does not fit in memory")?;
-    let service = Arc::new(Service::new(engine, cache.clone(), Config { max_bytes }));
-    // Requests held at once, running or waiting for the engine's one slot:
-    // each holds its audio in memory, so past this the service says it is busy.
-    let waiting = Arc::new(Semaphore::new(
-        usize::try_from(number("MAX_WAITING", 8)?).context("SWARM_STT_MAX_WAITING is too large")?,
-    ));
-
-    let client_id = var("OIDC_CLIENT_ID")?;
-    let secret =
-        swarm_secret_client::read_with_startup_retry("reading the queue client secret", || {
-            client_secret(&client_id)
-        })
-        .await?;
-    let reread_id = client_id.clone();
-    let secret = StoredSecret::new(client_id.clone(), secret, move || {
-        let client_id = reread_id.clone();
-        async move { client_secret(&client_id).await.map_err(Into::into) }
-    });
-    let queue = swarm_queue_client::connect(QueueConfig {
-        url: var("NATS_URL")?,
-        token_endpoint: var("OIDC_TOKEN_ENDPOINT")?,
-        client_id,
-        client_secret: ClientSecret::Stored(secret),
-        ca_file: std::env::var(format!("{PREFIX}_OIDC_CA_FILE"))
-            .ok()
-            .map(PathBuf::from),
-    })
-    .await?;
+    let Settings {
+        service,
+        cache,
+        waiting,
+    } = settings(&var, &number)?;
+    let per_agent: Arc<Mutex<HashSet<String>>> = Arc::default();
+    let queue = connect_queue(&var).await?;
 
     tokio::spawn(async move {
         loop {
@@ -113,20 +154,38 @@ async fn main() -> Result<()> {
     );
     while let Some(request) = requests.next().await {
         let (queue, service) = (queue.clone(), Arc::clone(&service));
-        let permit = Arc::clone(&waiting).try_acquire_owned().ok();
+        // Held at once overall, and one per asking agent, so one agent
+        // cannot fill every slot. The agent is the subject's last token.
+        let agent = request
+            .subject
+            .as_str()
+            .rsplit('.')
+            .next()
+            .unwrap_or_default()
+            .to_owned();
+        let turned_away = match (
+            Arc::clone(&waiting).try_acquire_owned(),
+            take_slot(&per_agent, &agent),
+        ) {
+            (Ok(permit), Some(slot)) => Ok((permit, slot)),
+            (Ok(_), None) => Err("you already have a transcription running; wait for it"),
+            (Err(_), _) => Err("the transcription service is busy, try again shortly"),
+        };
         tokio::spawn(async move {
             let Some(reply) = request.reply.clone() else {
                 tracing::warn!(subject = %request.subject, "a request with no reply subject");
                 return;
             };
-            let response = if permit.is_some() {
-                answer(&service, &request.payload).await
-            } else {
-                Response::Failed {
-                    message: "the transcription service is busy, try again shortly".to_owned(),
+            let response = match turned_away {
+                Ok(held) => {
+                    let response = answer(&service, &request.payload).await;
+                    drop(held);
+                    response
                 }
+                Err(message) => Response::Failed {
+                    message: message.to_owned(),
+                },
             };
-            drop(permit);
             match serde_json::to_vec(&response) {
                 Ok(body) => {
                     if let Err(e) = queue.publish(reply, body.into()).await {
@@ -140,6 +199,33 @@ async fn main() -> Result<()> {
     bail!("the request subscription ended")
 }
 
+/// One agent's place among the requests in flight; dropping it frees the place.
+struct Slot {
+    agents: Arc<Mutex<HashSet<String>>>,
+    agent: String,
+}
+
+impl Drop for Slot {
+    fn drop(&mut self) {
+        self.agents
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .remove(&self.agent);
+    }
+}
+
+/// `agent`'s place, or `None` when it already has a request in flight.
+fn take_slot(agents: &Arc<Mutex<HashSet<String>>>, agent: &str) -> Option<Slot> {
+    let fresh = agents
+        .lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner)
+        .insert(agent.to_owned());
+    fresh.then(|| Slot {
+        agents: Arc::clone(agents),
+        agent: agent.to_owned(),
+    })
+}
+
 /// The wire answer to one request's audio.
 async fn answer(service: &Service<Whisper>, audio: &[u8]) -> Response {
     match service.transcribe(audio).await {
@@ -176,3 +262,27 @@ async fn client_secret(client_id: &str) -> Result<String> {
         ),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashSet;
+    use std::sync::{Arc, Mutex};
+
+    use super::take_slot;
+
+    #[test]
+    fn an_agent_holds_one_slot_at_a_time_and_dropping_it_frees_it() {
+        let agents: Arc<Mutex<HashSet<String>>> = Arc::default();
+        let first = take_slot(&agents, "iris").expect("a free agent gets a slot");
+        assert!(
+            take_slot(&agents, "iris").is_none(),
+            "a second request is turned away"
+        );
+        assert!(
+            take_slot(&agents, "argus").is_some(),
+            "another agent is not"
+        );
+        drop(first);
+        assert!(take_slot(&agents, "iris").is_some());
+    }
+}