use std::{
collections::HashSet,
sync::{
Arc, Mutex,
mpsc::{self, SyncSender},
},
thread,
};
use anyhow::Context;
use chrono::Utc;
use kcode_audio_speaker_review::{
ChunkConfirmation, ConfirmationState, CorrectionPacket, ObservationKey, apply_confirmation,
restore_training, validate_confirmation_coverage,
};
use kcode_speaker_system::SpeechClassifier;
use rusqlite::{Connection, OptionalExtension, params};
use crate::{Error, StoredCorrectionPacket};
#[derive(Clone)]
pub(crate) struct Lane {
sender: SyncSender<Command>,
}
struct Command {
confirmation: ChunkConfirmation,
reply: SyncSender<Result<CorrectionPacket, Error>>,
}
impl Lane {
pub(crate) fn start(
db: Arc<Mutex<Connection>>,
classifier: Arc<SpeechClassifier>,
) -> std::io::Result<Self> {
let (sender, receiver) = mpsc::sync_channel::<Command>(1);
thread::Builder::new()
.name("audio-speaker-confirmation".into())
.spawn(move || {
while let Ok(command) = receiver.recv() {
let result = process(
&db,
&command.confirmation,
|packet, confirmation, legacy_keys| {
apply_confirmation(&classifier, packet, confirmation, legacy_keys)
},
|packet, legacy_keys| restore_training(&classifier, packet, legacy_keys),
);
let _ = command.reply.send(result);
}
})?;
Ok(Self { sender })
}
pub(crate) fn confirm(
&self,
confirmation: ChunkConfirmation,
) -> Result<CorrectionPacket, Error> {
let (reply, result) = mpsc::sync_channel(1);
self.sender
.send(Command {
confirmation,
reply,
})
.map_err(|_| Error::internal("speaker-confirmation worker stopped"))?;
result
.recv()
.map_err(|_| Error::internal("speaker-confirmation worker stopped"))?
}
}
fn process<Apply, Restore>(
db: &Arc<Mutex<Connection>>,
confirmation: &ChunkConfirmation,
apply: Apply,
restore: Restore,
) -> Result<CorrectionPacket, Error>
where
Apply: FnOnce(
&mut CorrectionPacket,
&ChunkConfirmation,
&HashSet<ObservationKey>,
) -> anyhow::Result<()>,
Restore: FnOnce(&CorrectionPacket, &HashSet<ObservationKey>) -> Vec<String>,
{
let recording_id = confirmation.recording_id;
let packet_json = read_review_packet(db, recording_id)?;
let mut stored = StoredCorrectionPacket::decode(&packet_json).map_err(Error::internal)?;
validate_confirmation_coverage(&stored.packet, confirmation).map_err(Error::invalid)?;
let legacy_keys = stored.legacy_observation_keys().map_err(Error::internal)?;
let previous = stored.packet.clone();
apply(&mut stored.packet, confirmation, &legacy_keys).map_err(Error::internal)?;
let updated_json = stored.encode().map_err(Error::internal)?;
let all_signed = stored.packet.confirmation_state == ConfirmationState::Confirmed;
let persistence = persist_packet(db, recording_id, &packet_json, &updated_json, all_signed);
match persistence {
Ok(1) => Ok(stored.packet),
Ok(_) => {
report_restore_errors(recording_id, restore(&previous, &legacy_keys));
Err(Error::conflict(
"Speaker review changed while confirmation was being applied.",
))
}
Err(error) => {
report_restore_errors(recording_id, restore(&previous, &legacy_keys));
Err(Error::internal(error))
}
}
}
fn read_review_packet(
db: &Arc<Mutex<Connection>>,
recording_id: uuid::Uuid,
) -> Result<String, Error> {
let stored = {
let db = db.lock().map_err(Error::internal)?;
db.query_row(
"SELECT status,correction_packet_json,final_transcript
FROM audio_recordings WHERE id=?1",
[recording_id.to_string()],
|row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, Option<String>>(1)?,
row.get::<_, Option<String>>(2)?,
))
},
)
.optional()
.map_err(Error::internal)?
};
let Some((status, packet_json, transcript)) = stored else {
return Err(Error::not_found());
};
if status != "ready_for_ingress" || transcript.is_some() {
return Err(Error::conflict(
"Speaker confirmation requires a review-ready recording.",
));
}
packet_json.ok_or_else(|| Error::conflict("The completed recording has no correction packet."))
}
fn persist_packet(
db: &Arc<Mutex<Connection>>,
recording_id: uuid::Uuid,
previous_json: &str,
updated_json: &str,
all_signed: bool,
) -> anyhow::Result<usize> {
let db = db
.lock()
.map_err(|_| anyhow::anyhow!("AudioIngress database lock was poisoned"))?;
db.execute(
"UPDATE audio_recordings
SET correction_packet_json=?1,
status=CASE WHEN ?2 THEN 'reconciling' ELSE status END,
attempt_count=CASE WHEN ?2 THEN 0 ELSE attempt_count END,
next_attempt_at=NULL,last_error=NULL,updated_at=?3
WHERE id=?4 AND status='ready_for_ingress' AND final_transcript IS NULL
AND correction_packet_json=?5",
params![
updated_json,
all_signed,
Utc::now().to_rfc3339(),
recording_id.to_string(),
previous_json,
],
)
.context("persisting speaker confirmation")
}
fn report_restore_errors(recording_id: uuid::Uuid, errors: Vec<String>) {
if !errors.is_empty() {
tracing::error!(
%recording_id,
?errors,
"Could not fully restore classifier state after packet persistence failed"
);
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::Utc;
use kcode_audio_speaker_review::{CorrectionChunk, ParsedChunk};
use uuid::Uuid;
fn packet(id: Uuid) -> CorrectionPacket {
CorrectionPacket {
recording_id: id,
user_id: "user".into(),
sha256: "a".repeat(64),
original_filename: "voice.wav".into(),
size_bytes: 1,
recorded_at: Utc::now(),
chunk_count: 2,
chunks: (0..2)
.map(|index| CorrectionChunk {
chunk_index: index,
chunk_count: 2,
audio_start_ms: index as u64 * 1_000,
audio_end_ms: (index as u64 + 1) * 1_000,
raw_gemini_response: "raw".into(),
parsed: ParsedChunk {
clip_valid: false,
clip_validity_reason: Some("no profile".into()),
speakers: vec![],
},
observations: vec![],
signed_off: false,
})
.collect(),
confirmation_state: ConfirmationState::Unconfirmed,
}
}
fn database(packet: &CorrectionPacket) -> Arc<Mutex<Connection>> {
let db = Connection::open_in_memory().unwrap();
db.execute_batch(
"CREATE TABLE audio_recordings(
id TEXT PRIMARY KEY,status TEXT,correction_packet_json TEXT,final_transcript TEXT,
attempt_count INTEGER,next_attempt_at TEXT,last_error TEXT,updated_at TEXT
);",
)
.unwrap();
db.execute(
"INSERT INTO audio_recordings VALUES(?1,'ready_for_ingress',?2,NULL,0,NULL,NULL,'now')",
params![
packet.recording_id.to_string(),
serde_json::to_string(packet).unwrap()
],
)
.unwrap();
Arc::new(Mutex::new(db))
}
fn confirmation(id: Uuid, chunk_index: usize) -> ChunkConfirmation {
ChunkConfirmation {
recording_id: id,
chunk_index,
observations: vec![],
}
}
#[test]
fn capability_and_compensation_callbacks_run_without_the_audio_lock() {
let id = Uuid::new_v4();
let packet = packet(id);
let db = database(&packet);
let apply_db = Arc::clone(&db);
let restore_db = Arc::clone(&db);
let error = process(
&db,
&confirmation(id, 0),
move |packet, _, _| {
let connection = apply_db
.try_lock()
.expect("audio DB released for capability");
connection
.execute(
"UPDATE audio_recordings SET correction_packet_json='changed' WHERE id=?1",
[id.to_string()],
)
.unwrap();
packet.chunks[0].signed_off = true;
Ok(())
},
move |_, _| {
let _guard = restore_db
.try_lock()
.expect("audio DB released for compensation");
vec![]
},
)
.unwrap_err();
assert_eq!(error.kind(), crate::ErrorKind::Conflict);
}
#[test]
fn lane_serializes_competing_chunk_confirmations_without_lost_updates() {
let id = Uuid::new_v4();
let db = database(&packet(id));
let classifier_path =
std::env::temp_dir().join(format!("audio-confirmation-lane-{id}.sqlite3"));
let classifier = Arc::new(SpeechClassifier::open(&classifier_path).unwrap());
let lane = Lane::start(Arc::clone(&db), classifier).unwrap();
let first = lane.clone();
let second = lane.clone();
let a = thread::spawn(move || first.confirm(confirmation(id, 0)).unwrap());
let b = thread::spawn(move || second.confirm(confirmation(id, 1)).unwrap());
a.join().unwrap();
b.join().unwrap();
let (status, json): (String, String) = db
.lock()
.unwrap()
.query_row(
"SELECT status,correction_packet_json FROM audio_recordings WHERE id=?1",
[id.to_string()],
|row| Ok((row.get(0)?, row.get(1)?)),
)
.unwrap();
let stored = StoredCorrectionPacket::decode(&json).unwrap();
assert!(stored.packet.chunks.iter().all(|chunk| chunk.signed_off));
assert_eq!(
stored.packet.confirmation_state,
ConfirmationState::Confirmed
);
assert_eq!(status, "reconciling");
drop(lane);
let _ = std::fs::remove_file(classifier_path);
}
}