use kcode_k1_audio_classification_format::decode_event;
pub use kcode_k1_audio_classification_projection_reducer::ProjectionEffect;
use kcode_k1_audio_classification_projection_reducer::{event_fragment_id, reduce};
#[cfg(feature = "testkit")]
use kcode_k1_audio_classification_projection_state::append_error;
pub use kcode_k1_audio_classification_projection_state::{
ExecutedAnalysis, FragmentId, FragmentStageV1, FragmentStatus, LlmJobState, LlmJobStatus,
OverallState, SpeakerLabelV1, StageState, StageStatus,
};
use kcode_k1_audio_classification_projection_state::{
actionable_state, interrupted_stage, validate_status,
};
use kcode_k1_transaction::Transaction;
use kcode_k1_txn_ordering::{K1TxnOrdering, TxId};
use rusqlite::{Connection, Error as SqlError, ErrorCode, OptionalExtension, params};
use std::ffi::OsString;
use std::fmt::Display;
use std::fs;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::{Mutex, MutexGuard};
const DATABASE_NAME: &str = "audio-classification.sqlite3";
const SUBSYSTEM: &str = "audio-classification";
const SCHEMA_VERSION: i64 = 1;
const METADATA_SQL: &str = "CREATE TABLE metadata(singleton INTEGER PRIMARY KEY CHECK(singleton = 1), schema_version INTEGER NOT NULL, last_applied_txid BLOB)";
const FRAGMENTS_SQL: &str = "CREATE TABLE fragments(fragment_id BLOB PRIMARY KEY, actionable_state INTEGER NOT NULL, encoded_status BLOB NOT NULL)";
const INDEX_SQL: &str = "CREATE INDEX fragments_actionable_state ON fragments(actionable_state)";
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AppliedEvent {
pub fragment_id: FragmentId,
pub effect: ProjectionEffect,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct InterruptedFragment {
pub fragment_id: FragmentId,
pub stage: FragmentStageV1,
}
pub struct Projection {
connection: Mutex<Connection>,
}
impl Projection {
pub fn open(root: &Path, ordering: &K1TxnOrdering) -> Result<(Self, Option<TxId>), String> {
fs::create_dir_all(root).map_err(text)?;
let path = root.join(DATABASE_NAME);
let exists = path.try_exists().map_err(text)?;
let (connection, cursor) = if exists {
match open_existing(&path) {
Ok(value) => value,
Err(OpenIssue::Recoverable) => recreate(&path)?,
Err(OpenIssue::Fatal(error)) => return Err(error),
}
} else {
remove_sidecars(&path)?;
create_database(&path)?
};
if let Some(cursor) = cursor {
let bytes = ordering.get_txn(cursor).map_err(text)?;
let valid = bytes.as_deref().is_some_and(|bytes| {
Transaction::parse(bytes).is_ok_and(|value| value.subsystem().as_str() == SUBSYSTEM)
});
if !valid {
drop(connection);
let (connection, cursor) = recreate(&path)?;
return Ok((
Self {
connection: Mutex::new(connection),
},
cursor,
));
}
}
Ok((
Self {
connection: Mutex::new(connection),
},
cursor,
))
}
pub fn apply(&self, callback_txid: TxId, payload: &[u8]) -> Result<AppliedEvent, String> {
let event = decode_event(payload).map_err(text)?;
let fragment_id = event_fragment_id(&event);
let mut connection = self.lock()?;
let transaction = connection.transaction().map_err(text)?;
let reduction = reduce(
load_status(&transaction, fragment_id)?,
&event,
callback_txid,
)?;
write_status(&transaction, fragment_id, &reduction.status)?;
let updated = transaction
.execute(
"UPDATE metadata SET last_applied_txid = ?1 WHERE singleton = 1",
params![&callback_txid.as_bytes()[..]],
)
.map_err(text)?;
if updated != 1 {
return Err("projection metadata row is missing".to_string());
}
transaction.commit().map_err(text)?;
Ok(AppliedEvent {
fragment_id,
effect: reduction.effect,
})
}
pub fn status(&self, fragment_id: FragmentId) -> Result<Option<FragmentStatus>, String> {
load_status(&self.lock()?, fragment_id)
}
pub fn queued(&self) -> Result<Vec<FragmentId>, String> {
Ok(actionable_statuses(&self.lock()?, 1)?
.into_iter()
.map(|value| value.0)
.collect())
}
pub fn running(&self) -> Result<Vec<InterruptedFragment>, String> {
Ok(actionable_statuses(&self.lock()?, 2)?
.into_iter()
.map(|(fragment_id, status)| InterruptedFragment {
fragment_id,
stage: interrupted_stage(&status),
})
.collect())
}
pub fn validate_labels(
&self,
fragment_id: FragmentId,
labels: &[SpeakerLabelV1],
) -> Result<TxId, String> {
let status = self.status(fragment_id)?.ok_or("unknown fragment")?;
kcode_k1_audio_classification_projection_state::validate_labels(&status, labels)
}
pub fn clear(&self) -> Result<(), String> {
let mut connection = self.lock()?;
let transaction = connection.transaction().map_err(text)?;
transaction
.execute("DELETE FROM fragments", [])
.map_err(text)?;
transaction
.execute(
"UPDATE metadata SET last_applied_txid = NULL WHERE singleton = 1",
[],
)
.map_err(text)?;
transaction.commit().map_err(text)
}
#[cfg(feature = "testkit")]
pub fn inject_errors(
&self,
fragment_id: FragmentId,
errors: Vec<String>,
) -> Result<(), String> {
let mut connection = self.lock()?;
let transaction = connection.transaction().map_err(text)?;
let mut status = load_status(&transaction, fragment_id)?.ok_or("unknown fragment")?;
for error in errors {
append_error(&mut status, error);
}
write_status(&transaction, fragment_id, &status)?;
transaction.commit().map_err(text)
}
fn lock(&self) -> Result<MutexGuard<'_, Connection>, String> {
self.connection.lock().map_err(text)
}
}
fn load_status<C: Deref<Target = Connection>>(
connection: &C,
fragment_id: FragmentId,
) -> Result<Option<FragmentStatus>, String> {
let stored = connection
.query_row(
"SELECT actionable_state, encoded_status FROM fragments WHERE fragment_id = ?1",
params![&fragment_id.as_bytes()[..]],
|row| Ok((row.get::<_, i64>(0)?, row.get::<_, Vec<u8>>(1)?)),
)
.optional()
.map_err(text)?;
stored
.map(|(actionable, bytes)| decode_status(&bytes, actionable))
.transpose()
}
fn write_status(
connection: &Connection,
fragment_id: FragmentId,
status: &FragmentStatus,
) -> Result<(), String> {
let encoded = postcard::to_allocvec(status).map_err(text)?;
connection.execute(
"INSERT INTO fragments(fragment_id, actionable_state, encoded_status) VALUES(?1, ?2, ?3) ON CONFLICT(fragment_id) DO UPDATE SET actionable_state = excluded.actionable_state, encoded_status = excluded.encoded_status",
params![&fragment_id.as_bytes()[..], actionable_state(status.state), encoded],
).map_err(text)?;
Ok(())
}
fn actionable_statuses<C: Deref<Target = Connection>>(
connection: &C,
actionable: i64,
) -> Result<Vec<(FragmentId, FragmentStatus)>, String> {
let mut statement = connection
.prepare("SELECT fragment_id, encoded_status FROM fragments WHERE actionable_state = ?1")
.map_err(text)?;
let rows = statement
.query_map(params![actionable], |row| {
Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, Vec<u8>>(1)?))
})
.map_err(text)?;
rows.map(|row| {
let (id, encoded) = row.map_err(text)?;
Ok((fragment_id(&id)?, decode_status(&encoded, actionable)?))
})
.collect()
}
fn fragment_id(bytes: &[u8]) -> Result<FragmentId, String> {
let bytes: [u8; 12] = bytes.try_into().map_err(text)?;
Ok(FragmentId::from_bytes(bytes))
}
fn decode_status(bytes: &[u8], actionable: i64) -> Result<FragmentStatus, String> {
let status: FragmentStatus = postcard::from_bytes(bytes).map_err(text)?;
if postcard::to_allocvec(&status).map_err(text)? != bytes {
return Err("stored fragment status is noncanonical".to_string());
}
validate_status(&status, actionable)?;
Ok(status)
}
fn create_database(path: &Path) -> Result<(Connection, Option<TxId>), String> {
let connection = Connection::open(path).map_err(text)?;
configure(&connection).map_err(text)?;
connection.execute_batch(&format!(
"{METADATA_SQL};{FRAGMENTS_SQL};{INDEX_SQL};INSERT INTO metadata(singleton, schema_version, last_applied_txid) VALUES(1, {SCHEMA_VERSION}, NULL);"
)).map_err(text)?;
Ok((connection, None))
}
fn open_existing(path: &Path) -> Result<(Connection, Option<TxId>), OpenIssue> {
let connection = Connection::open(path).map_err(classify)?;
configure(&connection).map_err(classify)?;
let check: String = connection
.query_row("PRAGMA quick_check", [], |row| row.get(0))
.map_err(classify)?;
if check != "ok" {
return Err(OpenIssue::Recoverable);
}
validate_schema(&connection)?;
let cursor = validate_rows(&connection)?;
Ok((connection, cursor))
}
fn configure(connection: &Connection) -> Result<(), SqlError> {
connection.pragma_update(None, "journal_mode", "WAL")?;
connection.pragma_update(None, "synchronous", "FULL")
}
fn validate_schema(connection: &Connection) -> Result<(), OpenIssue> {
let mut statement = connection.prepare(
"SELECT type, name, sql FROM sqlite_schema WHERE name NOT LIKE 'sqlite_%' ORDER BY type, name",
).map_err(classify)?;
let rows = statement
.query_map([], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)))
.map_err(classify)?;
let schema = rows
.collect::<Result<Vec<(String, String, String)>, _>>()
.map_err(classify)?;
let expected = vec![
(
"index".into(),
"fragments_actionable_state".into(),
INDEX_SQL.into(),
),
("table".into(), "fragments".into(), FRAGMENTS_SQL.into()),
("table".into(), "metadata".into(), METADATA_SQL.into()),
];
if schema != expected {
return Err(OpenIssue::Recoverable);
}
Ok(())
}
fn validate_rows(connection: &Connection) -> Result<Option<TxId>, OpenIssue> {
let mut metadata = connection
.prepare("SELECT singleton, schema_version, last_applied_txid FROM metadata")
.map_err(classify)?;
let rows = metadata
.query_map([], |row| {
Ok((
row.get::<_, i64>(0)?,
row.get::<_, i64>(1)?,
row.get::<_, Option<Vec<u8>>>(2)?,
))
})
.map_err(classify)?;
let rows = rows.collect::<Result<Vec<_>, _>>().map_err(classify)?;
if rows.len() != 1 || rows[0].0 != 1 || rows[0].1 != SCHEMA_VERSION {
return Err(OpenIssue::Recoverable);
}
let mut fragments = connection
.prepare("SELECT fragment_id, actionable_state, encoded_status FROM fragments")
.map_err(classify)?;
let values = fragments
.query_map([], |row| {
Ok((
row.get::<_, Vec<u8>>(0)?,
row.get::<_, i64>(1)?,
row.get::<_, Vec<u8>>(2)?,
))
})
.map_err(classify)?;
for value in values {
let (id, actionable, encoded) = value.map_err(classify)?;
fragment_id(&id).map_err(|_| OpenIssue::Recoverable)?;
decode_status(&encoded, actionable).map_err(|_| OpenIssue::Recoverable)?;
}
rows[0]
.2
.as_deref()
.map(|bytes| {
let bytes: [u8; 12] = bytes.try_into().map_err(|_| OpenIssue::Recoverable)?;
Ok(TxId::from_bytes(bytes))
})
.transpose()
}
fn classify(error: SqlError) -> OpenIssue {
if matches!(
error.sqlite_error_code(),
Some(ErrorCode::DatabaseCorrupt | ErrorCode::NotADatabase)
) {
OpenIssue::Recoverable
} else {
OpenIssue::Fatal(error.to_string())
}
}
fn recreate(path: &Path) -> Result<(Connection, Option<TxId>), String> {
remove_sidecars(path)?;
remove_if_present(path)?;
create_database(path)
}
fn remove_sidecars(path: &Path) -> Result<(), String> {
for suffix in ["-wal", "-shm"] {
remove_if_present(&path_with_suffix(path, suffix))?;
}
Ok(())
}
fn remove_if_present(path: &Path) -> Result<(), String> {
match fs::remove_file(path) {
Ok(()) => Ok(()),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(error) => Err(error.to_string()),
}
}
fn path_with_suffix(path: &Path, suffix: &str) -> PathBuf {
let mut value = OsString::from(path.as_os_str());
value.push(suffix);
PathBuf::from(value)
}
fn text(error: impl Display) -> String {
error.to_string()
}
enum OpenIssue {
Recoverable,
Fatal(String),
}
#[cfg(test)]
mod tests {
use super::*;
use kcode_k1_audio_classification_format::{
AudioClassificationEventV3, DiscardedV2, FailedV2, ProgressUpdateV1, ProgressV1, QueueV2,
encode_event,
};
use kcode_k1_transaction::{GENESIS_PARENT, SubsystemId, build_signed_transaction};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
static NEXT: AtomicU64 = AtomicU64::new(0);
const LEGACY: &[u8] = &[
0, 0, 2, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0,
];
struct Root(PathBuf);
impl Root {
fn new() -> Self {
let id = NEXT.fetch_add(1, Ordering::Relaxed);
let path =
std::env::temp_dir().join(format!("audio-projection-{}-{id}", std::process::id()));
fs::create_dir_all(&path).unwrap();
Self(path)
}
fn path(&self, name: &str) -> PathBuf {
self.0.join(name)
}
}
impl Drop for Root {
fn drop(&mut self) {
let _ = fs::remove_dir_all(&self.0);
}
}
fn id(value: u8) -> TxId {
TxId::from_bytes([value; 12])
}
fn queue(fragment_id: FragmentId) -> AudioClassificationEventV3 {
AudioClassificationEventV3::Queue(QueueV2 {
audio_object_id: fragment_id,
})
}
fn apply(
projection: &Projection,
callback: TxId,
event: AudioClassificationEventV3,
) -> AppliedEvent {
projection
.apply(callback, &encode_event(&event).unwrap())
.unwrap()
}
fn canonical(ordering: &K1TxnOrdering, subsystem: &str) -> TxId {
let bytes = build_signed_transaction(
ordering.tip().unwrap_or(GENESIS_PARENT),
1,
[1; 32],
SubsystemId::from_str(subsystem).unwrap(),
b"payload",
|_| Ok([2; 64]),
)
.unwrap();
let id = TxId::for_transaction(&bytes);
ordering.submit_txn(&bytes).unwrap();
id
}
#[test]
fn legacy_bytes_open_without_rebuild() {
let root = Root::new();
let ordering = K1TxnOrdering::open(&root.path("ordering")).unwrap();
let directory = root.path("projection");
fs::create_dir_all(&directory).unwrap();
let connection = Connection::open(directory.join(DATABASE_NAME)).unwrap();
connection.execute_batch(&format!(
"{METADATA_SQL};{FRAGMENTS_SQL};{INDEX_SQL};INSERT INTO metadata VALUES(1, 1, NULL);"
)).unwrap();
connection
.execute(
"INSERT INTO fragments VALUES(?1, 1, ?2)",
params![&id(1).as_bytes()[..], LEGACY],
)
.unwrap();
drop(connection);
let (projection, cursor) = Projection::open(&directory, &ordering).unwrap();
let status: FragmentStatus = projection.status(id(1)).unwrap().unwrap();
assert_eq!(
(cursor, status.state, status.queue.state),
(None, OverallState::Queued, StageState::Succeeded)
);
assert_eq!(postcard::to_allocvec(&status).unwrap(), LEGACY);
}
#[test]
fn callback_is_atomic_and_startup_recovers_cursors() {
let root = Root::new();
let ordering = K1TxnOrdering::open(&root.path("ordering")).unwrap();
let directory = root.path("projection");
let (projection, cursor) = Projection::open(&directory, &ordering).unwrap();
assert_eq!(cursor, None);
let tail = canonical(&ordering, SUBSYSTEM);
apply(&projection, tail, queue(id(2)));
drop(projection);
assert_eq!(
Projection::open(&directory, &ordering).unwrap().1,
Some(tail)
);
for bytes in [vec![1], id(99).as_bytes().to_vec()] {
let (projection, _) = Projection::open(&directory, &ordering).unwrap();
projection
.connection
.lock()
.unwrap()
.execute("UPDATE metadata SET last_applied_txid = ?1", params![bytes])
.unwrap();
drop(projection);
assert_eq!(Projection::open(&directory, &ordering).unwrap().1, None);
}
let wrong = canonical(&ordering, "wrong-subsystem");
let (projection, _) = Projection::open(&directory, &ordering).unwrap();
projection
.connection
.lock()
.unwrap()
.execute(
"UPDATE metadata SET last_applied_txid = ?1",
params![&wrong.as_bytes()[..]],
)
.unwrap();
drop(projection);
assert_eq!(Projection::open(&directory, &ordering).unwrap().1, None);
let atomic = root.path("atomic");
let (projection, _) = Projection::open(&atomic, &ordering).unwrap();
Connection::open(atomic.join(DATABASE_NAME)).unwrap().execute_batch(
"CREATE TRIGGER block_cursor BEFORE UPDATE ON metadata BEGIN SELECT RAISE(ABORT, 'blocked'); END;",
).unwrap();
assert!(
projection
.apply(id(8), &encode_event(&queue(id(8))).unwrap())
.is_err()
);
assert_eq!(projection.status(id(8)).unwrap(), None);
let cursor: Option<Vec<u8>> = projection
.connection
.lock()
.unwrap()
.query_row("SELECT last_applied_txid FROM metadata", [], |row| {
row.get(0)
})
.unwrap();
assert_eq!(cursor, None);
}
#[test]
fn retry_discard_and_independent_roots() {
let root = Root::new();
let ordering = K1TxnOrdering::open(&root.path("ordering")).unwrap();
let (projection, _) = Projection::open(&root.path("projection"), &ordering).unwrap();
let fragment = id(3);
apply(&projection, id(4), queue(fragment));
apply(
&projection,
id(5),
AudioClassificationEventV3::Failed(FailedV2 {
fragment_id: fragment,
stage: FragmentStageV1::Queue,
llm_job_sequence: None,
error: "failed".into(),
}),
);
assert_eq!(
apply(&projection, id(6), queue(fragment)).effect,
ProjectionEffect::Start
);
assert_eq!(
projection.status(fragment).unwrap().unwrap().attempt_count,
0
);
apply(
&projection,
id(7),
AudioClassificationEventV3::Progress(ProgressV1 {
fragment_id: fragment,
update: ProgressUpdateV1::LlmJobStarted {
sequence: 1,
stage: FragmentStageV1::Transcript,
name: "retry".into(),
},
}),
);
assert_eq!(
projection.status(fragment).unwrap().unwrap().attempt_count,
1
);
assert_eq!(
apply(
&projection,
id(8),
AudioClassificationEventV3::Discarded(DiscardedV2 {
fragment_id: fragment,
})
)
.effect,
ProjectionEffect::Abort
);
let discarded = projection.status(fragment).unwrap().unwrap();
assert_eq!(
apply(&projection, id(9), queue(fragment)).effect,
ProjectionEffect::None
);
assert_eq!(projection.status(fragment).unwrap().unwrap(), discarded);
let cursor: Vec<u8> = projection
.connection
.lock()
.unwrap()
.query_row("SELECT last_applied_txid FROM metadata", [], |row| {
row.get(0)
})
.unwrap();
assert_eq!(cursor, id(9).as_bytes());
let first = Arc::new(projection);
let (second, _) = Projection::open(&root.path("second"), &ordering).unwrap();
let held = first.connection.lock().unwrap();
let blocked = Arc::clone(&first);
let (started_tx, started_rx) = std::sync::mpsc::channel();
let (done_tx, done_rx) = std::sync::mpsc::channel();
let worker = std::thread::spawn(move || {
started_tx.send(()).unwrap();
blocked.status(id(1)).unwrap();
done_tx.send(()).unwrap();
});
started_rx.recv().unwrap();
assert!(second.queued().unwrap().is_empty());
assert!(done_rx.try_recv().is_err());
drop(held);
worker.join().unwrap();
}
}