#![forbid(unsafe_code)]
use std::{
collections::BTreeMap,
fmt,
future::Future,
path::Path,
pin::Pin,
str::FromStr,
sync::{Arc, Mutex},
task::{Context, Poll},
time::Duration,
};
use chrono::{DateTime, Utc};
use kcode_kennedy_session_kweb_contracts::DecodedKwebTool;
use kcode_kweb_db::{NodeId, Provenance, TransactionId};
use kcode_kweb_manager::{
ErrorKind as ManagerErrorKind, KmapOperation, KmapOwnerSelection, KmapWrite, KwebManager,
};
use rusqlite::{Connection, OptionalExtension, params};
use serde::{Deserialize, Serialize};
use tokio::sync::Notify;
#[derive(Debug)]
pub enum Error {
InvalidInput(String),
Conflict(String),
Internal(String),
}
impl fmt::Display for Error {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InvalidInput(value) | Self::Conflict(value) | Self::Internal(value) => {
formatter.write_str(value)
}
}
}
}
impl std::error::Error for Error {}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct Submission {
pub idempotency_id: String,
pub operation: DecodedKwebTool,
pub expected_revisions: BTreeMap<String, String>,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum State {
Queued,
Applying,
Committed,
Failed,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct Record {
pub submission: Submission,
pub state: State,
pub submitted_at: DateTime<Utc>,
pub attempt_count: u32,
pub next_attempt_at: Option<DateTime<Utc>>,
pub transaction_id: Option<String>,
pub created_node_id: Option<String>,
pub error: Option<String>,
}
#[derive(Clone)]
pub struct CommandLane {
inner: Arc<Inner>,
}
pub struct CommandRuntime(Pin<Box<dyn Future<Output = Result<(), Error>> + Send>>);
impl Future for CommandRuntime {
type Output = Result<(), Error>;
fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
self.0.as_mut().poll(context)
}
}
struct Inner {
connection: Mutex<Connection>,
manager: KwebManager,
notify: Notify,
}
pub fn open(
path: impl AsRef<Path>,
manager: KwebManager,
) -> Result<(CommandLane, CommandRuntime), Error> {
let connection = Connection::open(path).map_err(internal)?;
connection.execute_batch(
"PRAGMA busy_timeout=15000; PRAGMA journal_mode=WAL; PRAGMA synchronous=FULL;
CREATE TABLE IF NOT EXISTS kmap_admin_commands(
id TEXT PRIMARY KEY CHECK(length(id)=32), operation_json TEXT NOT NULL,
expected_json TEXT NOT NULL, state TEXT NOT NULL, submitted_ms INTEGER NOT NULL,
attempt_count INTEGER NOT NULL DEFAULT 0, next_attempt_ms INTEGER,
transaction_id TEXT, created_node_id TEXT, error TEXT
);
UPDATE kmap_admin_commands SET state='queued',next_attempt_ms=NULL WHERE state='applying';"
).map_err(internal)?;
let inner = Arc::new(Inner {
connection: Mutex::new(connection),
manager,
notify: Notify::new(),
});
let lane = CommandLane {
inner: inner.clone(),
};
Ok((lane, CommandRuntime(Box::pin(run(inner)))))
}
impl CommandLane {
pub fn submit(&self, submission: Submission) -> Result<Record, Error> {
validate_id(&submission.idempotency_id)?;
typed_request(&submission, Utc::now())?;
let operation_json = serde_json::to_string(&submission.operation).map_err(internal)?;
let expected_json =
serde_json::to_string(&submission.expected_revisions).map_err(internal)?;
let connection = self
.inner
.connection
.lock()
.map_err(|_| internal("Kmap command mutex is poisoned"))?;
if let Some(existing) = load(&connection, &submission.idempotency_id)? {
if existing.submission == submission {
return Ok(existing);
}
return Err(Error::Conflict(
"idempotency_id belongs to another Kmap command".into(),
));
}
let submitted_ms = Utc::now().timestamp_millis();
connection.execute(
"INSERT INTO kmap_admin_commands(id,operation_json,expected_json,state,submitted_ms)
VALUES(?1,?2,?3,'queued',?4)",
params![submission.idempotency_id, operation_json, expected_json, submitted_ms],
).map_err(internal)?;
let record = load(&connection, &submission.idempotency_id)?
.ok_or_else(|| internal("durable Kmap command disappeared"))?;
drop(connection);
self.inner.notify.notify_one();
Ok(record)
}
pub fn snapshot(&self) -> Result<Vec<Record>, Error> {
let connection = self
.inner
.connection
.lock()
.map_err(|_| internal("Kmap command mutex is poisoned"))?;
let mut statement = connection
.prepare(
"SELECT id,operation_json,expected_json,state,submitted_ms,attempt_count,
next_attempt_ms,transaction_id,created_node_id,error
FROM kmap_admin_commands ORDER BY submitted_ms,id",
)
.map_err(internal)?;
let rows = statement.query_map([], row).map_err(internal)?;
rows.collect::<Result<Vec<_>, _>>().map_err(internal)
}
}
async fn run(inner: Arc<Inner>) -> Result<(), Error> {
loop {
let Some(record) = claim(&inner)? else {
tokio::select! {
_ = inner.notify.notified() => {},
_ = tokio::time::sleep(Duration::from_secs(1)) => {},
}
continue;
};
let manager = inner.manager.clone();
let submission = record.submission.clone();
let request = typed_request(&submission, record.submitted_at)?;
let result = tokio::task::spawn_blocking(move || manager.apply_kmap_operation(request))
.await
.map_err(|error| internal(format!("Kmap command task failed: {error}")))?;
match result {
Ok(outcome) => finish(
&inner,
&submission.idempotency_id,
State::Committed,
Some(outcome.transaction_id.to_string()),
outcome.created_node_id.map(|id| id.to_string()),
None,
None,
)?,
Err(error) if error.kind() == ManagerErrorKind::Unavailable => {
let seconds = 2_u64.saturating_pow(record.attempt_count.min(6));
let next = Utc::now() + chrono::Duration::seconds(seconds.min(60) as i64);
finish(
&inner,
&submission.idempotency_id,
State::Queued,
None,
None,
Some(error.to_string()),
Some(next),
)?;
}
Err(error) => finish(
&inner,
&submission.idempotency_id,
State::Failed,
None,
None,
Some(error.to_string()),
None,
)?,
}
}
}
fn claim(inner: &Inner) -> Result<Option<Record>, Error> {
let connection = inner
.connection
.lock()
.map_err(|_| internal("Kmap command mutex is poisoned"))?;
let command = connection
.query_row(
"SELECT id,next_attempt_ms FROM kmap_admin_commands WHERE state='queued'
ORDER BY submitted_ms,id LIMIT 1",
[],
|row| Ok((row.get::<_, String>(0)?, row.get::<_, Option<i64>>(1)?)),
)
.optional()
.map_err(internal)?;
let Some((id, next_attempt_ms)) = command else {
return Ok(None);
};
if next_attempt_ms.is_some_and(|next| next > Utc::now().timestamp_millis()) {
return Ok(None);
}
connection
.execute(
"UPDATE kmap_admin_commands SET state='applying',attempt_count=attempt_count+1,
next_attempt_ms=NULL,error=NULL WHERE id=?1 AND state='queued'",
[&id],
)
.map_err(internal)?;
load(&connection, &id)
}
fn finish(
inner: &Inner,
id: &str,
state: State,
transaction: Option<String>,
created: Option<String>,
error: Option<String>,
next: Option<DateTime<Utc>>,
) -> Result<(), Error> {
let connection = inner
.connection
.lock()
.map_err(|_| internal("Kmap command mutex is poisoned"))?;
connection
.execute(
"UPDATE kmap_admin_commands SET state=?2,transaction_id=?3,created_node_id=?4,
error=?5,next_attempt_ms=?6 WHERE id=?1",
params![
id,
state_text(state),
transaction,
created,
error,
next.map(|value| value.timestamp_millis())
],
)
.map_err(internal)?;
Ok(())
}
fn load(connection: &Connection, id: &str) -> Result<Option<Record>, Error> {
connection.query_row(
"SELECT id,operation_json,expected_json,state,submitted_ms,attempt_count,next_attempt_ms,
transaction_id,created_node_id,error FROM kmap_admin_commands WHERE id=?1",
[id], row,
).optional().map_err(internal)
}
fn row(row: &rusqlite::Row<'_>) -> rusqlite::Result<Record> {
let operation_json: String = row.get(1)?;
let expected_json: String = row.get(2)?;
let state: String = row.get(3)?;
let submitted_ms: i64 = row.get(4)?;
let next_ms: Option<i64> = row.get(6)?;
let parse_error = |error| {
rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(error))
};
Ok(Record {
submission: Submission {
idempotency_id: row.get(0)?,
operation: serde_json::from_str(&operation_json).map_err(parse_error)?,
expected_revisions: serde_json::from_str(&expected_json).map_err(parse_error)?,
},
state: parse_state(&state)?,
submitted_at: time(submitted_ms)?,
attempt_count: row.get::<_, i64>(5)? as u32,
next_attempt_at: next_ms.map(time).transpose()?,
transaction_id: row.get(7)?,
created_node_id: row.get(8)?,
error: row.get(9)?,
})
}
fn typed_request(submission: &Submission, submitted_at: DateTime<Utc>) -> Result<KmapWrite, Error> {
let node = |value: &str| {
NodeId::from_str(value).map_err(|error| Error::InvalidInput(error.to_string()))
};
let owner = |value: &str| match value {
"self" => Ok(KmapOwnerSelection::SelfNode),
"unowned" => Ok(KmapOwnerSelection::Unowned),
other => node(other).map(KmapOwnerSelection::Node),
};
let nodes = |values: &[String]| {
values
.iter()
.map(|value| node(value))
.collect::<Result<Vec<_>, _>>()
};
let operation = match &submission.operation {
DecodedKwebTool::ConnectNodes(ids) => KmapOperation::ConnectNodes(nodes(ids)?),
DecodedKwebTool::ConsolidateFanout {
parent,
fanout,
aggregator,
} => KmapOperation::ConsolidateFanout {
parent: node(parent)?,
fanout: nodes(fanout)?,
aggregator: node(aggregator)?,
},
DecodedKwebTool::SetFixedConnection {
parent,
child,
slot,
} => KmapOperation::SetFixedConnection {
parent: node(parent)?,
child: child.as_deref().map(node).transpose()?,
slot: *slot,
},
DecodedKwebTool::CreateNode {
parents,
owner: value,
short_name,
short_description,
long_description,
} => KmapOperation::CreateNode {
parents: nodes(parents)?,
owner: owner(value)?,
short_name: short_name.clone(),
short_description: short_description.clone(),
long_description: long_description.clone(),
},
DecodedKwebTool::UpdateNode {
id,
owner: value,
short_name,
short_description,
long_description,
} => KmapOperation::UpdateNode {
id: node(id)?,
owner: owner(value)?,
short_name: short_name.clone(),
short_description: short_description.clone(),
long_description: long_description.clone(),
},
};
let expected_revisions = submission
.expected_revisions
.iter()
.map(|(id, revision)| {
Ok((
node(id)?,
TransactionId::from_str(revision)
.map_err(|error| Error::InvalidInput(error.to_string()))?,
))
})
.collect::<Result<BTreeMap<_, _>, Error>>()?;
Ok(KmapWrite {
operation,
expected_revisions,
provenance: Provenance {
author: "kennedy-administrator".into(),
source: "kennedy-memory-explorer".into(),
source_created_at: submitted_at,
data: format!(
"Kennedy administrator command {}",
submission.idempotency_id
),
},
})
}
fn state_text(state: State) -> &'static str {
match state {
State::Queued => "queued",
State::Applying => "applying",
State::Committed => "committed",
State::Failed => "failed",
}
}
fn parse_state(value: &str) -> rusqlite::Result<State> {
match value {
"queued" => Ok(State::Queued),
"applying" => Ok(State::Applying),
"committed" => Ok(State::Committed),
"failed" => Ok(State::Failed),
_ => Err(rusqlite::Error::InvalidColumnType(
3,
"state".into(),
rusqlite::types::Type::Text,
)),
}
}
fn time(value: i64) -> rusqlite::Result<DateTime<Utc>> {
DateTime::from_timestamp_millis(value)
.ok_or_else(|| rusqlite::Error::IntegralValueOutOfRange(4, value))
}
fn validate_id(value: &str) -> Result<(), Error> {
if value.len() == 32
&& value
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
{
Ok(())
} else {
Err(Error::InvalidInput(
"idempotency_id must be 32 lowercase hexadecimal characters".into(),
))
}
}
fn internal(error: impl fmt::Display) -> Error {
Error::Internal(error.to_string())
}