kcode_k1_audio_classification_driver/
lib.rs

1pub use kcode_k1_audio_fragment_runner::FragmentId;
2
3use kcode_k1_audio_fragment_runner as runner;
4use kcode_k1_audio_fragment_transactions::{self as transactions, FragmentStageV1};
5use kcode_k1_objects::K1Objects;
6use kcode_k1_peering::K1Peering;
7use kcode_speaker_v3_analysis::Analyzer;
8use std::collections::{HashMap, VecDeque};
9use std::future::Future;
10use std::pin::Pin;
11use std::sync::atomic::{AtomicBool, Ordering};
12use std::sync::{Arc, Mutex};
13use std::thread;
14
15pub const MAX_ACTIVE_ATTEMPTS: usize = 8;
16pub type EngineFuture<'a> = Pin<Box<dyn Future<Output = Result<(), String>> + 'a>>;
17
18pub trait FragmentEngine: Send + Sync + 'static {
19    fn run<'a>(
20        &'a self,
21        fragment_id: FragmentId,
22        ogg_bytes: &'a [u8],
23        is_active: &'a (dyn Fn() -> bool + Send + Sync),
24    ) -> EngineFuture<'a>;
25}
26
27#[derive(Clone, Copy, Debug, Eq, PartialEq)]
28pub enum StartOutcome {
29    Started,
30    Pending,
31    AlreadyActive,
32}
33
34pub struct AudioClassificationDriver {
35    inner: Arc<Inner>,
36}
37
38impl AudioClassificationDriver {
39    pub fn open(peering: Arc<K1Peering>, objects: Arc<K1Objects>, analyzer: Analyzer) -> Self {
40        let engine: Arc<dyn FragmentEngine> = Arc::new(RunnerEngine {
41            peering: peering.clone(),
42            analyzer: Arc::new(analyzer),
43        });
44        Self::with_engine(peering, objects, engine)
45    }
46
47    pub fn with_engine(
48        peering: Arc<K1Peering>,
49        objects: Arc<K1Objects>,
50        engine: Arc<dyn FragmentEngine>,
51    ) -> Self {
52        Self::with_resources(Arc::new(LiveResources { peering, objects }), engine)
53    }
54
55    pub fn start(&self, id: FragmentId) -> Result<StartOutcome, String> {
56        let (outcome, lane) = self.inner.reserve(id)?;
57        if let Some(lane) = lane {
58            self.inner.launch(lane)?;
59        }
60        Ok(outcome)
61    }
62
63    pub fn abort(&self, id: FragmentId) {
64        let mut state = lock(&self.inner.state);
65        if let Some(entry) = state.current.remove(&id) {
66            entry.active.store(false, Ordering::Release);
67        }
68        state.pending.retain(|(pending, _)| *pending != id);
69    }
70
71    pub fn ensure_healthy(&self) -> Result<(), String> {
72        match &lock(&self.inner.state).fault {
73            Some(error) => Err(format!("audio classification driver is unhealthy: {error}")),
74            None => Ok(()),
75        }
76    }
77
78    pub fn shutdown(&self) {
79        lock(&self.inner.state).stop(None);
80    }
81
82    fn with_resources(
83        resources: Arc<dyn FragmentResources>,
84        engine: Arc<dyn FragmentEngine>,
85    ) -> Self {
86        let state = State {
87            accepting: true,
88            fault: None,
89            next_generation: 0,
90            lanes: 0,
91            current: HashMap::new(),
92            pending: VecDeque::new(),
93        };
94        Self {
95            inner: Arc::new(Inner {
96                resources,
97                engine,
98                state: Mutex::new(state),
99            }),
100        }
101    }
102}
103
104impl Drop for AudioClassificationDriver {
105    fn drop(&mut self) {
106        self.shutdown();
107    }
108}
109
110struct RunnerEngine {
111    peering: Arc<K1Peering>,
112    analyzer: Arc<Analyzer>,
113}
114
115impl FragmentEngine for RunnerEngine {
116    fn run<'a>(
117        &'a self,
118        id: FragmentId,
119        bytes: &'a [u8],
120        active: &'a (dyn Fn() -> bool + Send + Sync),
121    ) -> EngineFuture<'a> {
122        Box::pin(runner::run_while_active(
123            &self.analyzer,
124            &self.peering,
125            id,
126            bytes,
127            active,
128        ))
129    }
130}
131
132trait FragmentResources: Send + Sync + 'static {
133    fn load(&self, id: FragmentId) -> Result<Option<(String, Vec<u8>)>, String>;
134    fn fail_queue(&self, id: FragmentId, error: String) -> Result<(), String>;
135}
136
137struct LiveResources {
138    peering: Arc<K1Peering>,
139    objects: Arc<K1Objects>,
140}
141
142impl FragmentResources for LiveResources {
143    fn load(&self, id: FragmentId) -> Result<Option<(String, Vec<u8>)>, String> {
144        self.objects
145            .load(id)
146            .map(|object| object.map(|object| (object.file_type, object.data)))
147    }
148
149    fn fail_queue(&self, id: FragmentId, error: String) -> Result<(), String> {
150        transactions::submit_failure(&self.peering, id, FragmentStageV1::Queue, None, error)
151            .map(|_| ())
152    }
153}
154
155struct State {
156    accepting: bool,
157    fault: Option<String>,
158    next_generation: u64,
159    lanes: usize,
160    current: HashMap<FragmentId, Entry>,
161    pending: VecDeque<(FragmentId, u64)>,
162}
163
164struct Entry {
165    generation: u64,
166    active: Arc<AtomicBool>,
167}
168
169struct Lane {
170    id: FragmentId,
171    generation: u64,
172    active: Arc<AtomicBool>,
173}
174
175impl State {
176    fn stop(&mut self, fault: Option<String>) {
177        self.accepting = false;
178        self.pending.clear();
179        for entry in self.current.values() {
180            entry.active.store(false, Ordering::Release);
181        }
182        self.current.clear();
183        if self.fault.is_none() {
184            self.fault = fault;
185        }
186    }
187
188    fn next_lane(&mut self) -> Option<Lane> {
189        if !self.accepting || self.lanes == MAX_ACTIVE_ATTEMPTS {
190            return None;
191        }
192        let (id, generation) = self.pending.pop_front()?;
193        let entry = self
194            .current
195            .get(&id)
196            .filter(|entry| entry.generation == generation)?;
197        self.lanes += 1;
198        Some(Lane {
199            id,
200            generation,
201            active: entry.active.clone(),
202        })
203    }
204}
205
206struct Inner {
207    resources: Arc<dyn FragmentResources>,
208    engine: Arc<dyn FragmentEngine>,
209    state: Mutex<State>,
210}
211
212impl Inner {
213    fn reserve(&self, id: FragmentId) -> Result<(StartOutcome, Option<Lane>), String> {
214        let mut state = lock(&self.state);
215        if let Some(error) = &state.fault {
216            return Err(format!("audio classification driver is unhealthy: {error}"));
217        }
218        if !state.accepting {
219            return Err("audio classification driver is shut down".into());
220        }
221        if state.current.contains_key(&id) {
222            return Ok((StartOutcome::AlreadyActive, None));
223        }
224        state.next_generation = state
225            .next_generation
226            .checked_add(1)
227            .ok_or_else(|| "audio classification generation overflow".to_owned())?;
228        let generation = state.next_generation;
229        let active = Arc::new(AtomicBool::new(true));
230        state.current.insert(
231            id,
232            Entry {
233                generation,
234                active: active.clone(),
235            },
236        );
237        if state.lanes == MAX_ACTIVE_ATTEMPTS {
238            state.pending.push_back((id, generation));
239            return Ok((StartOutcome::Pending, None));
240        }
241        state.lanes += 1;
242        Ok((
243            StartOutcome::Started,
244            Some(Lane {
245                id,
246                generation,
247                active,
248            }),
249        ))
250    }
251
252    fn launch(self: &Arc<Self>, lane: Lane) -> Result<(), String> {
253        let id = lane.id;
254        let generation = lane.generation;
255        let inner = self.clone();
256        match thread::Builder::new()
257            .name("k1-audio-fragment".to_owned())
258            .spawn(move || inner.run_lane(lane))
259        {
260            Ok(_) => Ok(()),
261            Err(error) => {
262                let error = format!("start audio classification lane: {error}");
263                self.finished(id, generation, Err(error.clone()));
264                Err(error)
265            }
266        }
267    }
268
269    fn run_lane(self: Arc<Self>, lane: Lane) {
270        let result = tokio::runtime::Builder::new_current_thread()
271            .build()
272            .map_err(|error| format!("create audio classification runtime: {error}"))
273            .and_then(|runtime| {
274                runtime.block_on(run_attempt(
275                    self.resources.as_ref(),
276                    self.engine.as_ref(),
277                    lane.id,
278                    lane.active.clone(),
279                ))
280            });
281        self.finished(lane.id, lane.generation, result);
282    }
283
284    fn finished(self: &Arc<Self>, id: FragmentId, generation: u64, result: Result<(), String>) {
285        let next = {
286            let mut state = lock(&self.state);
287            state.lanes = state.lanes.saturating_sub(1);
288            let active = state.current.get(&id).and_then(|entry| {
289                (entry.generation == generation).then(|| entry.active.load(Ordering::Acquire))
290            });
291            if active.is_some() {
292                state.current.remove(&id);
293            }
294            if let Some(error) = result.err().filter(|_| active == Some(true)) {
295                state.stop(Some(format!("runner persistence failure: {error}")));
296                None
297            } else {
298                state.next_lane()
299            }
300        };
301        if let Some(lane) = next {
302            let _ = self.launch(lane);
303        }
304    }
305}
306
307async fn run_attempt(
308    resources: &dyn FragmentResources,
309    engine: &dyn FragmentEngine,
310    id: FragmentId,
311    active: Arc<AtomicBool>,
312) -> Result<(), String> {
313    let object = match resources.load(id) {
314        Ok(Some(object)) if object.0 == "audio/ogg" => object,
315        Ok(Some(_)) => return fail_if_active(resources, id, &active, "wrong Object media type"),
316        Ok(None) => return fail_if_active(resources, id, &active, "audio Object is unavailable"),
317        Err(error) => {
318            return fail_if_active(
319                resources,
320                id,
321                &active,
322                &format!("load audio Object: {error}"),
323            );
324        }
325    };
326    let is_active = || active.load(Ordering::Acquire);
327    if is_active() {
328        engine.run(id, &object.1, &is_active).await?;
329    }
330    Ok(())
331}
332
333fn fail_if_active(
334    resources: &dyn FragmentResources,
335    id: FragmentId,
336    active: &AtomicBool,
337    error: &str,
338) -> Result<(), String> {
339    if active.load(Ordering::Acquire) {
340        resources.fail_queue(id, error.to_owned())?;
341    }
342    Ok(())
343}
344
345fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
346    mutex.lock().unwrap_or_else(|error| error.into_inner())
347}
348
349#[cfg(test)]
350mod tests {
351    use super::*;
352    use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
353    use std::time::{Duration, Instant};
354
355    type Action = (Option<Arc<AtomicBool>>, Result<(), String>);
356
357    #[derive(Clone)]
358    enum Load {
359        Missing,
360        Wrong,
361        Error,
362        Wait(Arc<AtomicBool>),
363    }
364
365    #[derive(Default)]
366    struct Harness {
367        actions: Mutex<HashMap<FragmentId, VecDeque<Action>>>,
368        loads: Mutex<HashMap<FragmentId, Load>>,
369        load_calls: AtomicUsize,
370        failures: AtomicUsize,
371        polls: AtomicUsize,
372    }
373
374    impl FragmentEngine for Harness {
375        fn run<'a>(
376            &'a self,
377            id: FragmentId,
378            _bytes: &'a [u8],
379            _active: &'a (dyn Fn() -> bool + Send + Sync),
380        ) -> EngineFuture<'a> {
381            let (gate, result) = lock(&self.actions)
382                .get_mut(&id)
383                .and_then(VecDeque::pop_front)
384                .unwrap_or((None, Ok(())));
385            Box::pin(async move {
386                self.polls.fetch_add(1, AtomicOrdering::SeqCst);
387                if let Some(gate) = gate {
388                    wait_gate(&gate);
389                }
390                result
391            })
392        }
393    }
394
395    impl FragmentResources for Harness {
396        fn load(&self, id: FragmentId) -> Result<Option<(String, Vec<u8>)>, String> {
397            self.load_calls.fetch_add(1, AtomicOrdering::SeqCst);
398            let load = lock(&self.loads).get(&id).cloned();
399            if let Some(Load::Wait(gate)) = &load {
400                wait_gate(gate);
401            }
402            match load {
403                None | Some(Load::Wait(_)) => Ok(Some(("audio/ogg".into(), vec![1]))),
404                Some(Load::Missing) => Ok(None),
405                Some(Load::Wrong) => Ok(Some(("text/plain".into(), vec![1]))),
406                Some(Load::Error) => Err("read failed".into()),
407            }
408        }
409
410        fn fail_queue(&self, _id: FragmentId, _error: String) -> Result<(), String> {
411            self.failures.fetch_add(1, AtomicOrdering::SeqCst);
412            Ok(())
413        }
414    }
415
416    fn wait_gate(gate: &AtomicBool) {
417        while !gate.load(Ordering::Acquire) {
418            thread::yield_now();
419        }
420    }
421
422    fn wait_for(condition: impl Fn() -> bool) {
423        let deadline = Instant::now() + Duration::from_secs(2);
424        while !condition() && Instant::now() < deadline {
425            thread::sleep(Duration::from_millis(2));
426        }
427        assert!(condition());
428    }
429
430    #[test]
431    fn isolation_admission_generation_failures_and_shutdown() {
432        let id = |value| FragmentId::from_bytes([value; 12]);
433        let gate = || Arc::new(AtomicBool::new(false));
434        let harness = Arc::new(Harness::default());
435        let driver = AudioClassificationDriver::with_resources(harness.clone(), harness.clone());
436        let (first, rest) = (gate(), gate());
437        lock(&harness.actions).insert(id(1), vec![(Some(first.clone()), Ok(()))].into());
438        for value in 2..=9 {
439            lock(&harness.actions).insert(id(value), vec![(Some(rest.clone()), Ok(()))].into());
440        }
441        assert!((1..=8).all(|value| driver.start(id(value)) == Ok(StartOutcome::Started)));
442        assert_eq!(driver.start(id(9)), Ok(StartOutcome::Pending));
443        wait_for(|| harness.polls.load(AtomicOrdering::SeqCst) == 8);
444        driver.abort(id(1));
445        thread::sleep(Duration::from_millis(10));
446        assert_eq!(harness.polls.load(AtomicOrdering::SeqCst), 8);
447        first.store(true, Ordering::Release);
448        wait_for(|| harness.polls.load(AtomicOrdering::SeqCst) == 9);
449        rest.store(true, Ordering::Release);
450
451        let harness = Arc::new(Harness::default());
452        let driver = AudioClassificationDriver::with_resources(harness.clone(), harness.clone());
453        let (old, new) = (gate(), gate());
454        lock(&harness.actions).insert(
455            id(1),
456            vec![
457                (Some(old.clone()), Err("stale".into())),
458                (Some(new.clone()), Ok(())),
459            ]
460            .into(),
461        );
462        assert_eq!(driver.start(id(1)), Ok(StartOutcome::Started));
463        wait_for(|| harness.polls.load(AtomicOrdering::SeqCst) == 1);
464        driver.abort(id(1));
465        assert_eq!(driver.start(id(1)), Ok(StartOutcome::Started));
466        wait_for(|| harness.polls.load(AtomicOrdering::SeqCst) == 2);
467        old.store(true, Ordering::Release);
468        wait_for(|| lock(&driver.inner.state).lanes == 1);
469        assert_eq!(driver.ensure_healthy(), Ok(()));
470        assert_eq!(driver.start(id(1)), Ok(StartOutcome::AlreadyActive));
471        new.store(true, Ordering::Release);
472        wait_for(|| lock(&driver.inner.state).lanes == 0);
473        lock(&harness.actions).insert(id(2), vec![(None, Err("current".into()))].into());
474        assert_eq!(driver.start(id(2)), Ok(StartOutcome::Started));
475        wait_for(|| driver.ensure_healthy().is_err());
476        assert!(driver.start(id(3)).is_err());
477
478        let harness = Arc::new(Harness::default());
479        let driver = AudioClassificationDriver::with_resources(harness.clone(), harness.clone());
480        for (value, load) in [(1, Load::Missing), (2, Load::Wrong), (3, Load::Error)] {
481            lock(&harness.loads).insert(id(value), load);
482            assert_eq!(driver.start(id(value)), Ok(StartOutcome::Started));
483        }
484        wait_for(|| harness.failures.load(AtomicOrdering::SeqCst) == 3);
485        assert_eq!(harness.polls.load(AtomicOrdering::SeqCst), 0);
486        let load_gate = gate();
487        lock(&harness.loads).insert(id(4), Load::Wait(load_gate.clone()));
488        assert_eq!(driver.start(id(4)), Ok(StartOutcome::Started));
489        wait_for(|| harness.load_calls.load(AtomicOrdering::SeqCst) == 4);
490        assert_eq!(driver.start(id(5)), Ok(StartOutcome::Started));
491        wait_for(|| harness.polls.load(AtomicOrdering::SeqCst) == 1);
492        let started = Instant::now();
493        driver.shutdown();
494        assert!(started.elapsed() < Duration::from_millis(100));
495        assert!(driver.start(id(6)).is_err());
496        load_gate.store(true, Ordering::Release);
497        wait_for(|| lock(&driver.inner.state).lanes == 0);
498    }
499}
kcode_k1_audio_classification_driver/lib.rs

kcode_k1_audio_classification_driver/
lib.rs
