use rusqlite::Connection;
use tokio::sync::{mpsc, oneshot};
use khive_storage::error::StorageError;
use crate::error::SqliteError;
use crate::pool::ConnectionPool;
type WriteOp<R> = Box<dyn FnOnce(&Connection) -> Result<R, StorageError> + Send>;
pub struct WriteRequest<R: Send + 'static> {
op: WriteOp<R>,
reply: oneshot::Sender<Result<R, StorageError>>,
top_level: bool,
}
mod sealed {
pub trait Sealed {}
}
impl<R: Send + 'static> sealed::Sealed for WriteRequest<R> {}
pub trait AnyWriteRequest: sealed::Sealed + Send {
fn execute_and_reply(self: Box<Self>, conn: &Connection);
fn execute_and_reply_top_level(self: Box<Self>, conn: &Connection);
fn reply_error(self: Box<Self>, err: StorageError);
fn is_top_level(&self) -> bool;
}
impl<R: Send + 'static> AnyWriteRequest for WriteRequest<R> {
fn execute_and_reply(self: Box<Self>, conn: &Connection) {
let outcome = (self.op)(conn);
let final_result = match outcome {
Ok(value) => match conn.execute_batch("COMMIT") {
Ok(()) => Ok(value),
Err(e) => {
let _ = conn.execute_batch("ROLLBACK");
Err(StorageError::Pool {
operation: "writer_task_commit".into(),
message: e.to_string(),
})
}
},
Err(e) => {
let _ = conn.execute_batch("ROLLBACK");
Err(e)
}
};
let _ = self.reply.send(final_result);
}
fn execute_and_reply_top_level(self: Box<Self>, conn: &Connection) {
let outcome = (self.op)(conn);
let _ = self.reply.send(outcome);
}
fn reply_error(self: Box<Self>, err: StorageError) {
let _ = self.reply.send(Err(err));
}
fn is_top_level(&self) -> bool {
self.top_level
}
}
#[derive(Clone, Debug)]
pub struct WriterTaskHandle {
tx: mpsc::Sender<Box<dyn AnyWriteRequest + Send>>,
}
impl WriterTaskHandle {
async fn enqueue<R, F>(
&self,
op: F,
) -> Result<oneshot::Receiver<Result<R, StorageError>>, StorageError>
where
R: Send + 'static,
F: FnOnce(&Connection) -> Result<R, StorageError> + Send + 'static,
{
self.enqueue_inner(op, false).await
}
async fn enqueue_inner<R, F>(
&self,
op: F,
top_level: bool,
) -> Result<oneshot::Receiver<Result<R, StorageError>>, StorageError>
where
R: Send + 'static,
F: FnOnce(&Connection) -> Result<R, StorageError> + Send + 'static,
{
let (reply_tx, reply_rx) = oneshot::channel();
let request = WriteRequest {
op: Box::new(op),
reply: reply_tx,
top_level,
};
self.tx
.send(Box::new(request))
.await
.map_err(|_| StorageError::Internal("writer task channel closed".to_string()))?;
Ok(reply_rx)
}
pub async fn send<R, F>(&self, op: F) -> Result<R, StorageError>
where
R: Send + 'static,
F: FnOnce(&Connection) -> Result<R, StorageError> + Send + 'static,
{
let reply_rx = self.enqueue(op).await?;
reply_rx.await.map_err(|_| {
StorageError::Internal("writer task dropped before replying".to_string())
})?
}
pub async fn send_with_timeout<R, F>(
&self,
op: F,
timeout: std::time::Duration,
) -> Result<R, StorageError>
where
R: Send + 'static,
F: FnOnce(&Connection) -> Result<R, StorageError> + Send + 'static,
{
let reply_rx = match tokio::time::timeout(timeout, self.enqueue(op)).await {
Ok(Ok(reply_rx)) => reply_rx,
Ok(Err(e)) => return Err(e),
Err(_elapsed) => {
return Err(StorageError::WriteQueueFull {
timeout_ms: timeout.as_millis() as u64,
})
}
};
reply_rx.await.map_err(|_| {
StorageError::Internal("writer task dropped before replying".to_string())
})?
}
pub async fn send_top_level<R, F>(&self, op: F) -> Result<R, StorageError>
where
R: Send + 'static,
F: FnOnce(&Connection) -> Result<R, StorageError> + Send + 'static,
{
let reply_rx = self.enqueue_inner(op, true).await?;
reply_rx.await.map_err(|_| {
StorageError::Internal("writer task dropped before replying".to_string())
})?
}
pub fn queue_depth(&self) -> usize {
self.tx.max_capacity() - self.tx.capacity()
}
pub fn capacity(&self) -> usize {
self.tx.max_capacity()
}
}
pub fn spawn(pool: &ConnectionPool, capacity: usize) -> Result<WriterTaskHandle, SqliteError> {
let conn = pool.open_standalone_writer()?;
let (tx, rx) = mpsc::channel(capacity.max(1));
tokio::spawn(run_writer_task(conn, rx));
Ok(WriterTaskHandle { tx })
}
async fn run_writer_task(
mut conn: Connection,
mut rx: mpsc::Receiver<Box<dyn AnyWriteRequest + Send>>,
) {
while let Some(request) = rx.recv().await {
let outcome = tokio::task::spawn_blocking(move || {
if request.is_top_level() {
request.execute_and_reply_top_level(&conn);
return conn;
}
let _tx_handle =
khive_storage::tx_registry::register(Some("writer_task_tx".to_string()));
match conn.execute_batch("BEGIN IMMEDIATE") {
Ok(()) => request.execute_and_reply(&conn),
Err(e) => {
tracing::warn!(
error = %e,
"writer task: BEGIN IMMEDIATE failed; replying an \
error without running the request's operation"
);
request.reply_error(StorageError::Pool {
operation: "writer_task_begin".into(),
message: e.to_string(),
});
}
}
conn
})
.await;
match outcome {
Ok(returned_conn) => conn = returned_conn,
Err(join_err) => {
tracing::error!(
error = %join_err,
"writer task blocking closure panicked; writer task is exiting"
);
return;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::pool::PoolConfig;
use serial_test::serial;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
fn file_pool(path: &std::path::Path) -> ConnectionPool {
let cfg = PoolConfig {
path: Some(path.to_path_buf()),
..PoolConfig::default()
};
ConnectionPool::new(cfg).expect("pool open")
}
#[tokio::test]
#[serial(tx_registry)]
async fn begin_immediate_failure_replies_error_without_running_op() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("writer_task_begin_failure.db");
let cfg = PoolConfig {
path: Some(path.clone()),
busy_timeout: Duration::from_millis(150),
..PoolConfig::default()
};
let pool = ConnectionPool::new(cfg).unwrap();
{
let writer = pool.try_writer().unwrap();
writer
.conn()
.execute_batch("CREATE TABLE t (id INTEGER PRIMARY KEY, v TEXT)")
.unwrap();
}
let handle = spawn(&pool, 8).expect("writer task should spawn on a file-backed pool");
let lock_holder = pool.try_writer().unwrap();
lock_holder.conn().execute_batch("BEGIN IMMEDIATE").unwrap();
let op_ran = Arc::new(AtomicBool::new(false));
let op_ran_clone = Arc::clone(&op_ran);
let result = handle
.send(move |conn| {
op_ran_clone.store(true, Ordering::SeqCst);
conn.execute("INSERT INTO t (id, v) VALUES (99, 'should-not-land')", [])
.map_err(|e| StorageError::Pool {
operation: "test_insert".into(),
message: e.to_string(),
})
})
.await;
assert!(
matches!(
&result,
Err(StorageError::Pool { operation, .. }) if operation == "writer_task_begin"
),
"expected a writer_task_begin Pool error on BEGIN IMMEDIATE \
failure, got {result:?}"
);
assert!(
!op_ran.load(Ordering::SeqCst),
"the request's operation closure must never run when BEGIN \
IMMEDIATE fails — running it would land a partial write in \
autocommit mode for a request the caller is told failed"
);
lock_holder.conn().execute_batch("ROLLBACK").unwrap();
drop(lock_holder);
let reader = pool.reader().expect("reader");
let count: i64 = reader
.conn()
.query_row("SELECT COUNT(*) FROM t WHERE id = 99", [], |row| row.get(0))
.unwrap();
assert_eq!(
count, 0,
"no row must have landed from the request whose BEGIN IMMEDIATE failed"
);
}
#[tokio::test]
#[serial(tx_registry)]
async fn writer_task_executes_op_and_commits() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("writer_task_commit.db");
let pool = file_pool(&path);
{
let writer = pool.try_writer().unwrap();
writer
.conn()
.execute_batch("CREATE TABLE t (id INTEGER PRIMARY KEY, v TEXT)")
.unwrap();
}
let handle = spawn(&pool, 8).expect("writer task should spawn on a file-backed pool");
let affected = handle
.send(|conn| {
conn.execute("INSERT INTO t (id, v) VALUES (1, 'hello')", [])
.map_err(|e| StorageError::Pool {
operation: "test_insert".into(),
message: e.to_string(),
})
})
.await
.expect("op should succeed");
assert_eq!(affected, 1);
let reader = pool.reader().expect("reader");
let v: String = reader
.conn()
.query_row("SELECT v FROM t WHERE id = 1", [], |row| row.get(0))
.expect("row must be committed and visible to a reader");
assert_eq!(v, "hello");
}
#[test]
fn spawn_fails_on_in_memory_pool() {
let cfg = PoolConfig {
path: None,
..PoolConfig::default()
};
let pool = ConnectionPool::new(cfg).unwrap();
let result = spawn(&pool, 8);
assert!(
result.is_err(),
"in-memory pools must reject spawn, not panic"
);
}
#[tokio::test]
async fn full_channel_applies_backpressure_not_immediate_error() {
let (tx, _rx) = mpsc::channel::<Box<dyn AnyWriteRequest + Send>>(1);
let handle = WriterTaskHandle { tx };
let first = tokio::spawn({
let handle = handle.clone();
async move {
let _ = handle.send(|_conn| Ok::<(), StorageError>(())).await;
}
});
tokio::time::sleep(Duration::from_millis(20)).await;
let second = tokio::time::timeout(
Duration::from_millis(100),
handle.send(|_conn| Ok::<(), StorageError>(())),
)
.await;
assert!(
second.is_err(),
"a full channel must apply backpressure (send suspends) rather \
than erroring immediately — no try_send escape hatch per ADR-067"
);
first.abort();
}
#[tokio::test]
async fn send_with_timeout_maps_full_channel_to_write_queue_full() {
let (tx, _rx) = mpsc::channel::<Box<dyn AnyWriteRequest + Send>>(1);
let handle = WriterTaskHandle { tx };
let first = tokio::spawn({
let handle = handle.clone();
async move {
let _ = handle.send(|_conn| Ok::<(), StorageError>(())).await;
}
});
tokio::time::sleep(Duration::from_millis(20)).await;
let result = handle
.send_with_timeout(
|_conn| Ok::<(), StorageError>(()),
Duration::from_millis(50),
)
.await;
match result {
Err(StorageError::WriteQueueFull { timeout_ms }) => assert_eq!(timeout_ms, 50),
other => panic!("expected WriteQueueFull, got {other:?}"),
}
first.abort();
}
#[tokio::test]
#[serial(tx_registry)]
async fn send_with_timeout_returns_op_result_when_op_outlives_the_timeout() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("writer_task_slow_op.db");
let pool = file_pool(&path);
{
let writer = pool.try_writer().unwrap();
writer
.conn()
.execute_batch("CREATE TABLE t (id INTEGER PRIMARY KEY, v TEXT)")
.unwrap();
}
let handle = spawn(&pool, 8).expect("writer task should spawn on a file-backed pool");
let result = handle
.send_with_timeout(
|conn| {
std::thread::sleep(Duration::from_millis(150));
conn.execute("INSERT INTO t (id, v) VALUES (1, 'slow')", [])
.map_err(|e| StorageError::Pool {
operation: "test_insert".into(),
message: e.to_string(),
})
},
Duration::from_millis(20),
)
.await;
let affected = result.expect(
"an accepted request must return its real result even when the \
op takes longer than the enqueue timeout, not WriteQueueFull",
);
assert_eq!(affected, 1);
let reader = pool.reader().expect("reader");
let v: String = reader
.conn()
.query_row("SELECT v FROM t WHERE id = 1", [], |row| row.get(0))
.expect("the slow op's write must have committed");
assert_eq!(v, "slow");
}
#[tokio::test]
async fn dropped_receiver_maps_send_to_internal_error() {
let (tx, rx) = mpsc::channel::<Box<dyn AnyWriteRequest + Send>>(4);
drop(rx);
let handle = WriterTaskHandle { tx };
let result = handle.send(|_conn| Ok::<(), StorageError>(())).await;
match result {
Err(StorageError::Internal(_)) => {}
other => panic!("expected Internal error on a closed channel, got {other:?}"),
}
}
}