hyperhive/swarm-stt/src/main.rs

//! The swarm's speech-to-text service.
//!
//! It connects to the queue as its own OIDC client, hears every agent's
//! request subject (`swarm_wire::stt`), and answers each audio file with its
//! transcript. [`swarm_stt::Service`] does the work: the cache, the shared
//! 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, Mutex};

use std::time::{Duration, Instant};

use anyhow::{Context as _, Result, bail};
use futures_util::StreamExt as _;
use metrics::{Metrics, Outcome};
use swarm_queue_client::{ClientSecret, QueueConfig};
use swarm_secret_client::stored_secret::StoredSecret;
use swarm_stt::engine::Whisper;
use swarm_stt::{Config, Error, Service, cache::Cache};
use swarm_wire::stt::{Response, subject_wildcard};
use tokio::sync::Semaphore;

mod metrics;

/// The cert-auth role this service logs in to the store as. `swarm-bao.nix`
/// writes it, matching the leaf's subject.
const CERT_ROLE: &str = "swarm-stt";

/// The environment prefix `swarm-stt.nix` sets.
const PREFIX: &str = "SWARM_STT";

/// The queue group every instance joins, so a request reaches one of them.
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();
    swarm_secret_client::install_crypto_provider();
    let var = |name: &str| {
        let full = format!("{PREFIX}_{name}");
        std::env::var(&full).with_context(|| format!("{full} is unset"))
    };
    let number = |name: &str, default: u64| -> Result<u64> {
        let full = format!("{PREFIX}_{name}");
        match std::env::var(&full) {
            Ok(v) => v.parse().with_context(|| format!("{full} is not a number")),
            Err(_) => Ok(default),
        }
    };

    let Settings {
        service,
        cache,
        waiting,
    } = settings(&var, &number)?;
    let per_agent: Arc<Mutex<HashSet<String>>> = Arc::default();
    let queue = connect_queue(&var).await?;
    let metrics = Arc::new(Metrics::from_env());

    tokio::spawn(async move {
        loop {
            let removed = cache.purge().await;
            if removed > 0 {
                tracing::info!(removed, "removed expired transcripts");
            }
            tokio::time::sleep(PURGE_EVERY).await;
        }
    });

    let mut requests = queue
        .queue_subscribe(subject_wildcard(), QUEUE_GROUP.to_owned())
        .await
        .context("subscribing to the request subject")?;
    tracing::info!(
        subject = subject_wildcard(),
        "swarm-stt: answering requests"
    );
    while let Some(request) = requests.next().await {
        let (queue, service) = (queue.clone(), Arc::clone(&service));
        let metrics = Arc::clone(&metrics);
        // 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 = match turned_away {
                Ok(held) => {
                    let count = metrics.hold();
                    let started = Instant::now();
                    let (response, outcome) = answer(&service, &request.payload).await;
                    metrics.request(outcome, Some(started.elapsed()));
                    drop((held, count));
                    response
                }
                Err(message) => {
                    metrics.request(Outcome::Busy, None);
                    Response::Failed {
                        message: message.to_owned(),
                    }
                }
            };
            match serde_json::to_vec(&response) {
                Ok(body) => {
                    if let Err(e) = queue.publish(reply, body.into()).await {
                        tracing::warn!(error = %e, "publishing a reply failed");
                    }
                }
                Err(e) => tracing::error!(error = %e, "encoding a reply failed"),
            }
        });
    }
    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, and what came of it.
async fn answer(service: &Service<Whisper>, audio: &[u8]) -> (Response, Outcome) {
    match service.transcribe(audio).await {
        Ok(reply) => {
            let outcome = if reply.cached {
                Outcome::Cached
            } else {
                Outcome::Transcribed
            };
            (
                Response::Ok {
                    text: reply.text,
                    cached: reply.cached,
                },
                outcome,
            )
        }
        Err(Error::Engine(message)) => {
            tracing::warn!(%message, "a transcription failed");
            (Response::Failed { message }, Outcome::Failed)
        }
        Err(e) => (
            Response::Failed {
                message: e.to_string(),
            },
            Outcome::Refused,
        ),
    }
}

async fn client_secret(client_id: &str) -> Result<String> {
    use swarm_secret_client::builder::{ClientSecret, client_secret_path};

    let path = client_secret_path(client_id)?;
    let store = swarm_secret_client::SecretStore::from_env(CERT_ROLE)
        .await
        .context("logging in to the secret store")?;
    let stored: Option<ClientSecret> = store
        .read_optional(&path)
        .await
        .with_context(|| format!("reading {path}"))?;
    match stored {
        Some(s) if !s.value.trim().is_empty() => Ok(s.value.trim().to_owned()),
        _ => bail!(
            "nothing at {path}: authelia's mint, in its container, has not \
             stored this service's client secret"
        ),
    }
}

#[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());
    }
}