kcode-k1-audio-classification 0.5.2

pub use kcode_k1_audio_classification_projection::{
    ExecutedAnalysis, FragmentId, FragmentStageV1, FragmentStatus, LlmJobState, LlmJobStatus,
    OverallState, SpeakerLabelV1, StageState, StageStatus,
};
pub use kcode_k1_audio_fragment_transactions::PersonId;

use kcode_k1_audio_classification_coordinator::AudioClassificationCoordinator;
use kcode_k1_audio_classification_format::{AudioClassificationEventV3, QueueV2, encode_event};
use kcode_k1_audio_fragment_submit as fragment_submit;
use kcode_k1_audio_fragment_transactions as transactions;
use kcode_k1_objects::K1Objects;
use kcode_k1_peering::K1Peering;
use kcode_k1_txn_ordering::{K1TxnOrdering, SubsystemId};
use kcode_speaker_v3_analysis::Analyzer;
use std::collections::HashSet;
use std::path::Path;
use std::sync::{Arc, Mutex};

const SUBSYSTEM_NAME: &str = "audio-classification";

pub struct AudioClassification {
    coordinator: AudioClassificationCoordinator,
    peering: Arc<K1Peering>,
    objects: Arc<K1Objects>,
    reservations: Reservations,
}

impl AudioClassification {
    pub fn open(
        root: &Path,
        ordering: Arc<K1TxnOrdering>,
        peering: Arc<K1Peering>,
        objects: Arc<K1Objects>,
        analyzer: Analyzer,
    ) -> Result<Self, String> {
        let coordinator = AudioClassificationCoordinator::open(
            root,
            ordering,
            peering.clone(),
            objects.clone(),
            analyzer,
        )?;
        Ok(Self {
            coordinator,
            peering,
            objects,
            reservations: Reservations::default(),
        })
    }

    pub fn submit(&self, ogg_bytes: &[u8]) -> Result<FragmentId, String> {
        self.coordinator.ensure_healthy()?;
        match fragment_submit::submit(&self.objects, &self.peering, ogg_bytes) {
            Ok(id) => Ok(id),
            Err(error) => {
                self.fault_if_committed(&error);
                Err(error)
            }
        }
    }

    pub fn status(&self, fragment_id: FragmentId) -> Result<Option<FragmentStatus>, String> {
        self.coordinator.status(fragment_id)
    }

    pub fn retry(&self, fragment_id: FragmentId) -> Result<(), String> {
        self.coordinator.ensure_healthy()?;
        self.require_state(
            fragment_id,
            OverallState::Failed,
            "retry requires Failed state",
        )?;
        self.submit_reserved(fragment_id, Operation::Retry, || {
            queue_existing(&self.peering, fragment_id)
        })
    }

    pub fn discard(&self, fragment_id: FragmentId) -> Result<(), String> {
        self.coordinator.ensure_healthy()?;
        if self.known_state(fragment_id)? == OverallState::Discarded {
            return Ok(());
        }
        self.submit_reserved(fragment_id, Operation::Discard, || {
            transactions::submit_discard(&self.peering, fragment_id).map(|_| ())
        })
    }

    pub fn submit_labels(
        &self,
        fragment_id: FragmentId,
        labels: Vec<SpeakerLabelV1>,
    ) -> Result<(), String> {
        self.coordinator.ensure_healthy()?;
        let interim = self.coordinator.validate_labels(fragment_id, &labels)?;
        self.submit_reserved(fragment_id, Operation::Labels, || {
            transactions::submit_label_confirmation(&self.peering, fragment_id, interim, labels)
                .map(|_| ())
        })
    }

    fn known_state(&self, id: FragmentId) -> Result<OverallState, String> {
        self.coordinator
            .status(id)?
            .map(|status| status.state)
            .ok_or_else(|| "unknown audio fragment".to_owned())
    }

    fn require_state(
        &self,
        id: FragmentId,
        state: OverallState,
        error: &str,
    ) -> Result<(), String> {
        if self.known_state(id)? == state {
            Ok(())
        } else {
            Err(error.to_owned())
        }
    }

    fn submit_reserved(
        &self,
        id: FragmentId,
        operation: Operation,
        submit: impl FnOnce() -> Result<(), String>,
    ) -> Result<(), String> {
        if !self.reservations.reserve(id, operation) {
            return Err(operation.active_error().to_owned());
        }
        let result = submit();
        if result.is_ok()
            || result
                .as_ref()
                .is_err_and(|error| !is_committed_error(error))
        {
            self.reservations.release(id, operation);
        }
        if let Err(error) = &result {
            self.fault_if_committed(error);
        }
        result
    }

    fn fault_if_committed(&self, error: &str) {
        if is_committed_error(error) {
            self.coordinator.fault_ambiguous_commitment();
        }
    }
}

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

#[derive(Clone, Copy, Eq, Hash, PartialEq)]
enum Operation {
    Retry,
    Discard,
    Labels,
}

impl Operation {
    fn active_error(self) -> &'static str {
        match self {
            Self::Retry => "retry is already active",
            Self::Discard => "discard is already active",
            Self::Labels => "label submission is already active",
        }
    }
}

#[derive(Default)]
struct Reservations(Mutex<HashSet<(FragmentId, Operation)>>);

impl Reservations {
    fn reserve(&self, id: FragmentId, operation: Operation) -> bool {
        lock(&self.0).insert((id, operation))
    }

    fn release(&self, id: FragmentId, operation: Operation) {
        lock(&self.0).remove(&(id, operation));
    }
}

fn queue_existing(peering: &K1Peering, fragment_id: FragmentId) -> Result<(), String> {
    let payload = encode_event(&AudioClassificationEventV3::Queue(QueueV2 {
        audio_object_id: fragment_id,
    }))
    .map_err(|error| format!("encode retry queue: {error}"))?;
    peering.submit_txn(subsystem_id()?, &payload).map(|_| ())
}

fn subsystem_id() -> Result<SubsystemId, String> {
    SubsystemId::from_str(SUBSYSTEM_NAME)
}

fn is_committed_error(error: &str) -> bool {
    error.to_ascii_lowercase().contains("committed")
}

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::*;

    fn fragment_id(byte: u8) -> FragmentId {
        FragmentId::from_bytes([byte; 12])
    }

    #[test]
    fn retry_queue_uses_event_v5() {
        let bytes = encode_event(&AudioClassificationEventV3::Queue(QueueV2 {
            audio_object_id: fragment_id(1),
        }))
        .expect("encode queue");
        assert_eq!(bytes.first(), Some(&5));
    }

    #[test]
    fn reservations_are_scoped_by_fragment_and_operation() {
        let reservations = Reservations::default();
        let first = fragment_id(1);
        let second = fragment_id(2);
        assert!(reservations.reserve(first, Operation::Retry));
        assert!(!reservations.reserve(first, Operation::Retry));
        assert!(reservations.reserve(first, Operation::Discard));
        assert!(reservations.reserve(second, Operation::Retry));
        reservations.release(first, Operation::Retry);
        assert!(reservations.reserve(first, Operation::Retry));
    }

    #[test]
    fn only_committed_errors_cross_the_ambiguous_boundary() {
        assert!(is_committed_error("transaction COMMITTED as abc"));
        assert!(!is_committed_error("submission unavailable"));
    }

    #[test]
    fn subsystem_identity_is_valid() {
        subsystem_id().expect("audio-classification subsystem ID");
    }
}