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

kcode_k1_audio_classification/
lib.rs
