Skip to main content

kcode_k1_audio_classification/
lib.rs

1pub use kcode_k1_audio_classification_projection::{
2    ExecutedAnalysis, FragmentId, FragmentStageV1, FragmentStatus, LlmJobState, LlmJobStatus,
3    OverallState, SpeakerLabelV1, StageState, StageStatus,
4};
5
6use kcode_k1_audio_classification_coordinator::AudioClassificationCoordinator;
7use kcode_k1_audio_classification_format::{AudioClassificationEventV3, QueueV2, encode_event};
8use kcode_k1_audio_fragment_submit as fragment_submit;
9use kcode_k1_audio_fragment_transactions as transactions;
10use kcode_k1_objects::K1Objects;
11use kcode_k1_peering::K1Peering;
12use kcode_k1_txn_ordering::{K1TxnOrdering, SubsystemId};
13use kcode_speaker_v3_analysis::Analyzer;
14use std::collections::HashSet;
15use std::path::Path;
16use std::sync::{Arc, Mutex};
17
18const SUBSYSTEM_NAME: &str = "audio-classification";
19
20pub struct AudioClassification {
21    coordinator: AudioClassificationCoordinator,
22    peering: Arc<K1Peering>,
23    objects: Arc<K1Objects>,
24    reservations: Reservations,
25}
26
27impl AudioClassification {
28    pub fn open(
29        root: &Path,
30        ordering: Arc<K1TxnOrdering>,
31        peering: Arc<K1Peering>,
32        objects: Arc<K1Objects>,
33        analyzer: Analyzer,
34    ) -> Result<Self, String> {
35        let coordinator = AudioClassificationCoordinator::open(
36            root,
37            ordering,
38            peering.clone(),
39            objects.clone(),
40            analyzer,
41        )?;
42        Ok(Self::from_coordinator(coordinator, peering, objects))
43    }
44
45    pub fn submit(&self, ogg_bytes: &[u8]) -> Result<FragmentId, String> {
46        self.coordinator.ensure_healthy()?;
47        match fragment_submit::submit(&self.objects, &self.peering, ogg_bytes) {
48            Ok(id) => Ok(id),
49            Err(error) => {
50                self.fault_if_committed(&error);
51                Err(error)
52            }
53        }
54    }
55
56    pub fn status(&self, fragment_id: FragmentId) -> Result<Option<FragmentStatus>, String> {
57        self.coordinator.status(fragment_id)
58    }
59
60    pub fn retry(&self, fragment_id: FragmentId) -> Result<(), String> {
61        self.coordinator.ensure_healthy()?;
62        self.require_state(
63            fragment_id,
64            OverallState::Failed,
65            "retry requires Failed state",
66        )?;
67        self.submit_reserved(fragment_id, Operation::Retry, || {
68            queue_existing(&self.peering, fragment_id)
69        })
70    }
71
72    pub fn discard(&self, fragment_id: FragmentId) -> Result<(), String> {
73        self.coordinator.ensure_healthy()?;
74        if self.known_state(fragment_id)? == OverallState::Discarded {
75            return Ok(());
76        }
77        self.submit_reserved(fragment_id, Operation::Discard, || {
78            transactions::submit_discard(&self.peering, fragment_id).map(|_| ())
79        })
80    }
81
82    pub fn submit_labels(
83        &self,
84        fragment_id: FragmentId,
85        labels: Vec<SpeakerLabelV1>,
86    ) -> Result<(), String> {
87        self.coordinator.ensure_healthy()?;
88        let interim = self.coordinator.validate_labels(fragment_id, &labels)?;
89        self.submit_reserved(fragment_id, Operation::Labels, || {
90            transactions::submit_label_confirmation(&self.peering, fragment_id, interim, labels)
91                .map(|_| ())
92        })
93    }
94
95    fn from_coordinator(
96        coordinator: AudioClassificationCoordinator,
97        peering: Arc<K1Peering>,
98        objects: Arc<K1Objects>,
99    ) -> Self {
100        Self {
101            coordinator,
102            peering,
103            objects,
104            reservations: Reservations::default(),
105        }
106    }
107
108    fn known_state(&self, id: FragmentId) -> Result<OverallState, String> {
109        self.coordinator
110            .status(id)?
111            .map(|status| status.state)
112            .ok_or_else(|| "unknown audio fragment".to_owned())
113    }
114
115    fn require_state(
116        &self,
117        id: FragmentId,
118        state: OverallState,
119        error: &str,
120    ) -> Result<(), String> {
121        if self.known_state(id)? == state {
122            Ok(())
123        } else {
124            Err(error.to_owned())
125        }
126    }
127
128    fn submit_reserved(
129        &self,
130        id: FragmentId,
131        operation: Operation,
132        submit: impl FnOnce() -> Result<(), String>,
133    ) -> Result<(), String> {
134        if !self.reservations.reserve(id, operation) {
135            return Err(operation.active_error().to_owned());
136        }
137        let result = submit();
138        if result.is_ok()
139            || result
140                .as_ref()
141                .is_err_and(|error| !is_committed_error(error))
142        {
143            self.reservations.release(id, operation);
144        }
145        if let Err(error) = &result {
146            self.fault_if_committed(error);
147        }
148        result
149    }
150
151    fn fault_if_committed(&self, error: &str) {
152        if is_committed_error(error) {
153            self.coordinator.fault_ambiguous_commitment();
154        }
155    }
156
157    #[cfg(test)]
158    fn inject_error_burst(&self, id: FragmentId, errors: Vec<String>) -> Result<(), String> {
159        self.coordinator.inject_errors(id, errors)
160    }
161}
162
163impl Drop for AudioClassification {
164    fn drop(&mut self) {
165        self.coordinator.shutdown();
166    }
167}
168
169#[derive(Clone, Copy, Eq, Hash, PartialEq)]
170enum Operation {
171    Retry,
172    Discard,
173    Labels,
174}
175
176impl Operation {
177    fn active_error(self) -> &'static str {
178        match self {
179            Self::Retry => "retry is already active",
180            Self::Discard => "discard is already active",
181            Self::Labels => "label submission is already active",
182        }
183    }
184}
185
186#[derive(Default)]
187struct Reservations(Mutex<HashSet<(FragmentId, Operation)>>);
188
189impl Reservations {
190    fn reserve(&self, id: FragmentId, operation: Operation) -> bool {
191        lock(&self.0).insert((id, operation))
192    }
193
194    fn release(&self, id: FragmentId, operation: Operation) {
195        lock(&self.0).remove(&(id, operation));
196    }
197}
198
199fn queue_existing(peering: &K1Peering, fragment_id: FragmentId) -> Result<(), String> {
200    let payload = encode_event(&AudioClassificationEventV3::Queue(QueueV2 {
201        audio_object_id: fragment_id,
202    }))
203    .map_err(|error| format!("encode retry queue: {error}"))?;
204    peering.submit_txn(subsystem_id()?, &payload).map(|_| ())
205}
206
207fn subsystem_id() -> Result<SubsystemId, String> {
208    SubsystemId::from_str(SUBSYSTEM_NAME)
209}
210
211fn is_committed_error(error: &str) -> bool {
212    error.to_ascii_lowercase().contains("committed")
213}
214
215fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
216    mutex.lock().unwrap_or_else(|error| error.into_inner())
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222    use kcode_k1_audio_classification_test_adapter::{Facade, FacadeBuilder, run_all};
223
224    struct TestFacade(AudioClassification);
225    struct Builder;
226
227    impl Facade for TestFacade {
228        fn submit(&self, bytes: &[u8]) -> Result<FragmentId, String> {
229            AudioClassification::submit(&self.0, bytes)
230        }
231
232        fn status(&self, fragment_id: FragmentId) -> Result<Option<FragmentStatus>, String> {
233            AudioClassification::status(&self.0, fragment_id)
234        }
235
236        fn retry(&self, fragment_id: FragmentId) -> Result<(), String> {
237            AudioClassification::retry(&self.0, fragment_id)
238        }
239
240        fn discard(&self, fragment_id: FragmentId) -> Result<(), String> {
241            AudioClassification::discard(&self.0, fragment_id)
242        }
243
244        fn submit_labels(
245            &self,
246            fragment_id: FragmentId,
247            labels: Vec<SpeakerLabelV1>,
248        ) -> Result<(), String> {
249            AudioClassification::submit_labels(&self.0, fragment_id, labels)
250        }
251
252        fn inject_error_burst(
253            &self,
254            fragment_id: FragmentId,
255            errors: Vec<String>,
256        ) -> Result<(), String> {
257            AudioClassification::inject_error_burst(&self.0, fragment_id, errors)
258        }
259    }
260
261    impl FacadeBuilder for Builder {
262        type Facade = TestFacade;
263
264        fn build(
265            &self,
266            coordinator: AudioClassificationCoordinator,
267            peering: Arc<K1Peering>,
268            objects: Arc<K1Objects>,
269        ) -> Self::Facade {
270            TestFacade(AudioClassification::from_coordinator(
271                coordinator,
272                peering,
273                objects,
274            ))
275        }
276    }
277
278    #[test]
279    fn published_conformance() {
280        run_all(&Builder).expect("audio classification conformance");
281    }
282}