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);
}
}