kcode_k1_audio_classification_driver/

lib.rs

pub use kcode_k1_audio_fragment_runner::FragmentId;

use kcode_k1_audio_fragment_runner as runner;
use kcode_k1_audio_fragment_transactions::{self as transactions, FragmentStageV1};
use kcode_k1_objects::K1Objects;
use kcode_k1_peering::K1Peering;
use kcode_speaker_v3_analysis::Analyzer;
use std::collections::{HashMap, VecDeque};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;

pub const MAX_ACTIVE_ATTEMPTS: usize = 8;
pub type EngineFuture<'a> = Pin<Box<dyn Future<Output = Result<(), String>> + 'a>>;

pub trait FragmentEngine: Send + Sync + 'static {
    fn run<'a>(
        &'a self,
        fragment_id: FragmentId,
        ogg_bytes: &'a [u8],
        is_active: &'a (dyn Fn() -> bool + Send + Sync),
    ) -> EngineFuture<'a>;
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StartOutcome {
    Started,
    Pending,
    AlreadyActive,
}

pub struct AudioClassificationDriver {
    inner: Arc<Inner>,
}

impl AudioClassificationDriver {
    pub fn open(peering: Arc<K1Peering>, objects: Arc<K1Objects>, analyzer: Analyzer) -> Self {
        let engine: Arc<dyn FragmentEngine> = Arc::new(RunnerEngine {
            peering: peering.clone(),
            analyzer: Arc::new(analyzer),
        });
        Self::with_engine(peering, objects, engine)
    }

    pub fn with_engine(
        peering: Arc<K1Peering>,
        objects: Arc<K1Objects>,
        engine: Arc<dyn FragmentEngine>,
    ) -> Self {
        Self::with_resources(Arc::new(LiveResources { peering, objects }), engine)
    }

    pub fn start(&self, id: FragmentId) -> Result<StartOutcome, String> {
        let (outcome, lane) = self.inner.reserve(id)?;
        lane.map_or(Ok(()), |lane| self.inner.launch(lane))?;
        Ok(outcome)
    }

    pub fn abort(&self, id: FragmentId) {
        let mut state = lock(&self.inner.state);
        if let Some(entry) = state.current.remove(&id) {
            entry.active.store(false, Ordering::Release);
        }
        state.pending.retain(|(pending, _)| *pending != id);
    }

    pub fn ensure_healthy(&self) -> Result<(), String> {
        match &lock(&self.inner.state).fault {
            Some(error) => Err(format!("audio classification driver is unhealthy: {error}")),
            None => Ok(()),
        }
    }

    pub fn shutdown(&self) {
        lock(&self.inner.state).stop(None);
    }

    fn with_resources(
        resources: Arc<dyn FragmentResources>,
        engine: Arc<dyn FragmentEngine>,
    ) -> Self {
        Self {
            inner: Arc::new(Inner {
                resources,
                engine,
                state: Mutex::new(State {
                    accepting: true,
                    ..State::default()
                }),
            }),
        }
    }
}

impl Drop for AudioClassificationDriver {
    fn drop(&mut self) {
        self.shutdown();
    }
}

struct RunnerEngine {
    peering: Arc<K1Peering>,
    analyzer: Arc<Analyzer>,
}

impl FragmentEngine for RunnerEngine {
    fn run<'a>(
        &'a self,
        id: FragmentId,
        bytes: &'a [u8],
        active: &'a (dyn Fn() -> bool + Send + Sync),
    ) -> EngineFuture<'a> {
        Box::pin(runner::run_while_active(
            &self.analyzer,
            &self.peering,
            id,
            bytes,
            active,
        ))
    }
}

trait FragmentResources: Send + Sync + 'static {
    fn load(&self, id: FragmentId) -> Result<Option<(String, Vec<u8>)>, String>;
    fn fail_queue(&self, id: FragmentId, error: String) -> Result<(), String>;
}

struct LiveResources {
    peering: Arc<K1Peering>,
    objects: Arc<K1Objects>,
}

impl FragmentResources for LiveResources {
    fn load(&self, id: FragmentId) -> Result<Option<(String, Vec<u8>)>, String> {
        self.objects
            .load(id)
            .map(|object| object.map(|object| (object.file_type, object.data)))
    }

    fn fail_queue(&self, id: FragmentId, error: String) -> Result<(), String> {
        transactions::submit_failure(&self.peering, id, FragmentStageV1::Queue, None, error)
            .map(|_| ())
    }
}

#[derive(Default)]
struct State {
    accepting: bool,
    fault: Option<String>,
    next_generation: u64,
    lanes: usize,
    current: HashMap<FragmentId, Entry>,
    pending: VecDeque<(FragmentId, u64)>,
}

struct Entry {
    generation: u64,
    active: Arc<AtomicBool>,
}

struct Lane {
    id: FragmentId,
    generation: u64,
    active: Arc<AtomicBool>,
}

impl State {
    fn stop(&mut self, fault: Option<String>) {
        self.accepting = false;
        self.pending.clear();
        for entry in self.current.drain().map(|(_, entry)| entry) {
            entry.active.store(false, Ordering::Release);
        }
        self.fault = self.fault.take().or(fault);
    }

    fn next_lane(&mut self) -> Option<Lane> {
        if !self.accepting || self.lanes == MAX_ACTIVE_ATTEMPTS {
            return None;
        }
        let (id, generation) = self.pending.pop_front()?;
        let entry = self
            .current
            .get(&id)
            .filter(|entry| entry.generation == generation)?;
        self.lanes += 1;
        Some(Lane {
            id,
            generation,
            active: entry.active.clone(),
        })
    }
}

struct Inner {
    resources: Arc<dyn FragmentResources>,
    engine: Arc<dyn FragmentEngine>,
    state: Mutex<State>,
}

impl Inner {
    fn reserve(&self, id: FragmentId) -> Result<(StartOutcome, Option<Lane>), String> {
        let mut state = lock(&self.state);
        if let Some(error) = &state.fault {
            return Err(format!("audio classification driver is unhealthy: {error}"));
        }
        if !state.accepting {
            return Err("audio classification driver is shut down".into());
        }
        if state.current.contains_key(&id) {
            return Ok((StartOutcome::AlreadyActive, None));
        }
        state.next_generation = state
            .next_generation
            .checked_add(1)
            .ok_or_else(|| "audio classification generation overflow".to_owned())?;
        let generation = state.next_generation;
        let active = Arc::new(AtomicBool::new(true));
        state.current.insert(
            id,
            Entry {
                generation,
                active: active.clone(),
            },
        );
        if state.lanes == MAX_ACTIVE_ATTEMPTS {
            state.pending.push_back((id, generation));
            return Ok((StartOutcome::Pending, None));
        }
        state.lanes += 1;
        Ok((
            StartOutcome::Started,
            Some(Lane {
                id,
                generation,
                active,
            }),
        ))
    }

    fn launch(self: &Arc<Self>, lane: Lane) -> Result<(), String> {
        let id = lane.id;
        let generation = lane.generation;
        let inner = self.clone();
        if let Err(error) = thread::Builder::new()
            .name("k1-audio-fragment".to_owned())
            .spawn(move || inner.run_lane(lane))
        {
            let error = format!("start audio classification lane: {error}");
            self.finished(id, generation, Err(error.clone()));
            return Err(error);
        }
        Ok(())
    }

    fn run_lane(self: Arc<Self>, lane: Lane) {
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .map_err(|error| format!("create audio classification runtime: {error}"))
                .and_then(|runtime| {
                    runtime.block_on(run_attempt(
                        self.resources.as_ref(),
                        self.engine.as_ref(),
                        lane.id,
                        lane.active.clone(),
                    ))
                })
        }))
        .unwrap_or_else(|_| Err("audio classification lane panicked".into()));
        self.finished(lane.id, lane.generation, result);
    }

    fn finished(self: &Arc<Self>, id: FragmentId, generation: u64, result: Result<(), String>) {
        let next = {
            let mut state = lock(&self.state);
            state.lanes = state.lanes.saturating_sub(1);
            let active = state.current.get(&id).and_then(|entry| {
                (entry.generation == generation).then(|| entry.active.load(Ordering::Acquire))
            });
            if active.is_some() {
                state.current.remove(&id);
            }
            if let Some(error) = result.err().filter(|_| active == Some(true)) {
                state.stop(Some(format!("runner persistence failure: {error}")));
                None
            } else {
                state.next_lane()
            }
        };
        if let Some(lane) = next {
            let _ = self.launch(lane);
        }
    }
}

async fn run_attempt(
    resources: &dyn FragmentResources,
    engine: &dyn FragmentEngine,
    id: FragmentId,
    active: Arc<AtomicBool>,
) -> Result<(), String> {
    let object = match resources.load(id) {
        Ok(Some(object)) if object.0 == "audio/ogg" => object,
        Ok(Some(_)) => return fail_if_active(resources, id, &active, "wrong Object media type"),
        Ok(None) => return fail_if_active(resources, id, &active, "audio Object is unavailable"),
        Err(error) => {
            let error = format!("load audio Object: {error}");
            return fail_if_active(resources, id, &active, &error);
        }
    };
    let is_active = || active.load(Ordering::Acquire);
    if is_active() {
        engine.run(id, &object.1, &is_active).await?;
    }
    Ok(())
}

fn fail_if_active(
    resources: &dyn FragmentResources,
    id: FragmentId,
    active: &AtomicBool,
    error: &str,
) -> Result<(), String> {
    if active.load(Ordering::Acquire) {
        resources.fail_queue(id, error.to_owned())?;
    }
    Ok(())
}

fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
    mutex.lock().unwrap_or_else(|error| error.into_inner())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
    use std::time::{Duration, Instant};

    type Action = (Option<Arc<AtomicBool>>, Result<(), String>);

    #[derive(Clone)]
    enum Load {
        Missing,
        Wrong,
        Error,
        Wait(Arc<AtomicBool>),
    }

    #[derive(Default)]
    struct Harness {
        actions: Mutex<HashMap<FragmentId, VecDeque<Action>>>,
        loads: Mutex<HashMap<FragmentId, Load>>,
        load_calls: AtomicUsize,
        failures: AtomicUsize,
        polls: AtomicUsize,
    }

    impl FragmentEngine for Harness {
        fn run<'a>(
            &'a self,
            id: FragmentId,
            _bytes: &'a [u8],
            _active: &'a (dyn Fn() -> bool + Send + Sync),
        ) -> EngineFuture<'a> {
            if id == FragmentId::from_bytes([0; 12]) {
                return Box::pin(async {
                    let mut command =
                        tokio::process::Command::new(std::env::current_exe().unwrap());
                    command.arg("--list").stdout(std::process::Stdio::null());
                    tokio::time::timeout(Duration::from_secs(1), command.status())
                        .await
                        .unwrap()
                        .map(|_| ())
                        .map_err(|error| error.to_string())
                });
            }
            let (gate, result) = lock(&self.actions)
                .get_mut(&id)
                .and_then(VecDeque::pop_front)
                .unwrap_or((None, Ok(())));
            Box::pin(async move {
                self.polls.fetch_add(1, AtomicOrdering::SeqCst);
                if let Some(gate) = gate {
                    wait_gate(&gate);
                }
                result.inspect_err(|error| assert_ne!(error, "panic"))
            })
        }
    }

    impl FragmentResources for Harness {
        fn load(&self, id: FragmentId) -> Result<Option<(String, Vec<u8>)>, String> {
            self.load_calls.fetch_add(1, AtomicOrdering::SeqCst);
            let load = lock(&self.loads).get(&id).cloned();
            if let Some(Load::Wait(gate)) = &load {
                wait_gate(gate);
            }
            match load {
                None | Some(Load::Wait(_)) => Ok(Some(("audio/ogg".into(), vec![1]))),
                Some(Load::Missing) => Ok(None),
                Some(Load::Wrong) => Ok(Some(("text/plain".into(), vec![1]))),
                Some(Load::Error) => Err("read failed".into()),
            }
        }

        fn fail_queue(&self, _id: FragmentId, _error: String) -> Result<(), String> {
            self.failures.fetch_add(1, AtomicOrdering::SeqCst);
            Ok(())
        }
    }

    fn wait_gate(gate: &AtomicBool) {
        while !gate.load(Ordering::Acquire) {
            thread::yield_now();
        }
    }

    fn wait_for(condition: impl Fn() -> bool) {
        let deadline = Instant::now() + Duration::from_secs(2);
        while !condition() && Instant::now() < deadline {
            thread::sleep(Duration::from_millis(2));
        }
        assert!(condition());
    }

    #[test]
    fn runtime_isolation_admission_generation_failures_and_shutdown() {
        let id = |value| FragmentId::from_bytes([value; 12]);
        let gate = || Arc::new(AtomicBool::new(false));
        let harness = Arc::new(Harness::default());
        let driver = AudioClassificationDriver::with_resources(harness.clone(), harness.clone());
        assert_eq!(driver.start(id(0)), Ok(StartOutcome::Started));
        wait_for(|| lock(&driver.inner.state).lanes == 0);
        let (first, rest) = (gate(), gate());
        lock(&harness.actions).insert(id(1), vec![(Some(first.clone()), Ok(()))].into());
        for value in 2..=9 {
            lock(&harness.actions).insert(id(value), vec![(Some(rest.clone()), Ok(()))].into());
        }
        assert!((1..=8).all(|value| driver.start(id(value)) == Ok(StartOutcome::Started)));
        assert_eq!(driver.start(id(9)), Ok(StartOutcome::Pending));
        wait_for(|| harness.polls.load(AtomicOrdering::SeqCst) == 8);
        driver.abort(id(1));
        thread::sleep(Duration::from_millis(10));
        assert_eq!(harness.polls.load(AtomicOrdering::SeqCst), 8);
        first.store(true, Ordering::Release);
        wait_for(|| harness.polls.load(AtomicOrdering::SeqCst) == 9);
        rest.store(true, Ordering::Release);

        let harness = Arc::new(Harness::default());
        let driver = AudioClassificationDriver::with_resources(harness.clone(), harness.clone());
        let (old, new) = (gate(), gate());
        lock(&harness.actions).insert(
            id(1),
            vec![
                (Some(old.clone()), Err("panic".into())),
                (Some(new.clone()), Ok(())),
            ]
            .into(),
        );
        assert_eq!(driver.start(id(1)), Ok(StartOutcome::Started));
        wait_for(|| harness.polls.load(AtomicOrdering::SeqCst) == 1);
        driver.abort(id(1));
        assert_eq!(driver.start(id(1)), Ok(StartOutcome::Started));
        wait_for(|| harness.polls.load(AtomicOrdering::SeqCst) == 2);
        old.store(true, Ordering::Release);
        wait_for(|| lock(&driver.inner.state).lanes == 1);
        assert_eq!(driver.ensure_healthy(), Ok(()));
        assert_eq!(driver.start(id(1)), Ok(StartOutcome::AlreadyActive));
        new.store(true, Ordering::Release);
        wait_for(|| lock(&driver.inner.state).lanes == 0);
        lock(&harness.actions).insert(id(2), vec![(None, Err("panic".into()))].into());
        assert_eq!(driver.start(id(2)), Ok(StartOutcome::Started));
        wait_for(|| driver.ensure_healthy().is_err());
        assert_eq!(lock(&driver.inner.state).lanes, 0);
        assert!(driver.start(id(3)).is_err());

        let harness = Arc::new(Harness::default());
        let driver = AudioClassificationDriver::with_resources(harness.clone(), harness.clone());
        for (value, load) in [(1, Load::Missing), (2, Load::Wrong), (3, Load::Error)] {
            lock(&harness.loads).insert(id(value), load);
            assert_eq!(driver.start(id(value)), Ok(StartOutcome::Started));
        }
        wait_for(|| harness.failures.load(AtomicOrdering::SeqCst) == 3);
        assert_eq!(harness.polls.load(AtomicOrdering::SeqCst), 0);
        let load_gate = gate();
        lock(&harness.loads).insert(id(4), Load::Wait(load_gate.clone()));
        assert_eq!(driver.start(id(4)), Ok(StartOutcome::Started));
        wait_for(|| harness.load_calls.load(AtomicOrdering::SeqCst) == 4);
        assert_eq!(driver.start(id(5)), Ok(StartOutcome::Started));
        wait_for(|| harness.polls.load(AtomicOrdering::SeqCst) == 1);
        let started = Instant::now();
        driver.shutdown();
        assert!(started.elapsed() < Duration::from_millis(100));
        assert!(driver.start(id(6)).is_err());
        load_gate.store(true, Ordering::Release);
        wait_for(|| lock(&driver.inner.state).lanes == 0);
    }
}