use rusqlite::Connection;
use std::collections::HashMap;
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};
use tokio::sync::{mpsc, oneshot};
use khive_storage::error::{StorageError, WriterTaskRequestState};
use crate::error::SqliteError;
use crate::pool::{ConnectionPool, WriterAcquisitionCounters};
#[path = "writer_task_labels.rs"]
mod labels;
pub(crate) use labels::{collect_statement_labels, MAX_WRITE_LABELS};
use labels::{dedup_labels, SubmitLabels};
const WRITER_BEGIN_RETRY_DELAYS: [Duration; 2] =
[Duration::from_millis(5), Duration::from_millis(10)];
type WriteOp<R> = Box<dyn FnOnce(&Connection) -> Result<R, StorageError> + Send>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WriterStageObservation {
pub queue_wait_micros: u64,
pub transaction_acquire_micros: u64,
pub body_micros: u64,
pub commit_micros: u64,
pub total_micros: u64,
pub queue_depth_at_entry: u64,
pub observed_at_unix_ms: u64,
pub labels: Vec<String>,
pub labels_truncated: Option<u64>,
}
static WRITER_STAGE_OBSERVATIONS: OnceLock<
Mutex<HashMap<Option<PathBuf>, WriterStageObservation>>,
> = OnceLock::new();
fn writer_stage_observations() -> &'static Mutex<HashMap<Option<PathBuf>, WriterStageObservation>> {
WRITER_STAGE_OBSERVATIONS.get_or_init(|| Mutex::new(HashMap::new()))
}
fn writer_db_key_from_path(path: Option<&Path>) -> Option<PathBuf> {
path.map(Path::to_path_buf)
}
fn writer_db_key(pool: &ConnectionPool) -> Option<PathBuf> {
writer_db_key_from_path(pool.canonical_path())
}
fn duration_micros(duration: Duration) -> u64 {
duration.as_micros().min(u128::from(u64::MAX)) as u64
}
fn observed_at_unix_ms() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|duration| duration.as_millis() as u64)
.unwrap_or(0)
}
pub fn last_writer_stage_observation(pool: &ConnectionPool) -> Option<WriterStageObservation> {
writer_stage_observations()
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(&writer_db_key(pool))
.cloned()
}
struct WriteTelemetry {
backend_key: Option<PathBuf>,
db: String,
submitted_at: Instant,
queue_depth_at_entry: usize,
slow_write_threshold: Option<Duration>,
labels: SubmitLabels,
}
impl WriteTelemetry {
fn new(
backend_key: Option<PathBuf>,
db: String,
queue_depth_at_entry: usize,
slow_write_threshold: Option<Duration>,
labels: SubmitLabels,
) -> Self {
Self {
backend_key,
db,
submitted_at: Instant::now(),
queue_depth_at_entry,
slow_write_threshold,
labels,
}
}
fn queue_wait(&self) -> Duration {
self.submitted_at.elapsed()
}
fn resolve_labels(&self) -> (Vec<String>, Option<u64>) {
match &self.labels {
SubmitLabels::None => (Vec::new(), None),
SubmitLabels::Site(site) => (vec![(*site).to_owned()], None),
SubmitLabels::Statements(labels) => {
if labels.len() <= MAX_WRITE_LABELS {
(labels.clone(), None)
} else {
let dropped = (labels.len() - MAX_WRITE_LABELS) as u64;
(labels[..MAX_WRITE_LABELS].to_vec(), Some(dropped))
}
}
}
}
fn finish(
self,
queue_wait: Duration,
transaction_acquire: Duration,
body: Duration,
commit: Duration,
) {
let total = self.submitted_at.elapsed();
let slow = self
.slow_write_threshold
.is_some_and(|threshold| total >= threshold);
let (labels, labels_truncated) = if slow {
self.resolve_labels()
} else {
(Vec::new(), None)
};
let observation = WriterStageObservation {
queue_wait_micros: duration_micros(queue_wait),
transaction_acquire_micros: duration_micros(transaction_acquire),
body_micros: duration_micros(body),
commit_micros: duration_micros(commit),
total_micros: duration_micros(total),
queue_depth_at_entry: self.queue_depth_at_entry as u64,
observed_at_unix_ms: observed_at_unix_ms(),
labels,
labels_truncated,
};
writer_stage_observations()
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(self.backend_key, observation.clone());
if slow {
crate::timeout_sink::emit_slow_write(&self.db, &observation);
}
}
}
pub struct WriteRequest<R: Send + 'static> {
op: WriteOp<R>,
reply: oneshot::Sender<Result<R, StorageError>>,
top_level: bool,
checkpoint_bypass: bool,
vacuum_copy_headroom: bool,
telemetry: WriteTelemetry,
}
mod sealed {
pub trait Sealed {
fn execute_and_reply_reporting_terminal(
self: Box<Self>,
conn: &rusqlite::Connection,
tx_span: Option<khive_storage::tx_registry::TxHandle>,
queue_wait: std::time::Duration,
transaction_acquire: std::time::Duration,
) -> Option<khive_storage::error::WriterTaskRequestState>;
fn execute_and_reply_top_level_reporting_terminal(
self: Box<Self>,
conn: &rusqlite::Connection,
queue_wait: std::time::Duration,
) -> Option<khive_storage::error::WriterTaskRequestState>;
fn reply_error_after_begin(
self: Box<Self>,
err: khive_storage::error::StorageError,
queue_wait: std::time::Duration,
transaction_acquire: std::time::Duration,
);
}
}
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;
fn is_checkpoint_bypass(&self) -> bool;
fn needs_vacuum_headroom(&self) -> bool;
fn queue_wait(&self) -> Duration;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RollbackDisposition {
RolledBack,
SideEffectsUnknown,
}
fn rollback_after_failure(conn: &Connection, failure_context: &'static str) -> RollbackDisposition {
match conn.execute_batch("ROLLBACK") {
Ok(()) if conn.is_autocommit() => RollbackDisposition::RolledBack,
Ok(()) => {
tracing::error!(
failure_context,
"writer transaction: ROLLBACK returned success but the connection is still in a \
transaction; request side effects are unknown"
);
RollbackDisposition::SideEffectsUnknown
}
Err(rollback_error) => {
tracing::error!(
error = %rollback_error,
failure_context,
"writer transaction: rollback after request failure failed; request side effects are \
unknown"
);
RollbackDisposition::SideEffectsUnknown
}
}
}
pub(crate) fn execute_wrapped_transaction<R, F>(
conn: &Connection,
commit_operation: &'static str,
operation: F,
) -> (Result<R, StorageError>, Option<WriterTaskRequestState>)
where
F: FnOnce(&Connection) -> Result<R, StorageError>,
{
let profiled = execute_wrapped_transaction_profiled(conn, commit_operation, None, operation);
(profiled.result, profiled.terminal_state)
}
struct ProfiledWrappedTransaction<R> {
result: Result<R, StorageError>,
terminal_state: Option<WriterTaskRequestState>,
body: Duration,
commit: Duration,
}
fn execute_wrapped_transaction_profiled<R, F>(
conn: &Connection,
commit_operation: &'static str,
db: Option<&str>,
operation: F,
) -> ProfiledWrappedTransaction<R>
where
F: FnOnce(&Connection) -> Result<R, StorageError>,
{
let body_started = Instant::now();
let operation_outcome = catch_unwind(AssertUnwindSafe(|| {
operation(conn).map_err(crate::error::statement_failure)
}));
let body = body_started.elapsed();
match operation_outcome {
Ok(Ok(value)) => {
let commit_started = Instant::now();
let commit_outcome = conn.execute_batch("COMMIT");
let commit = commit_started.elapsed();
match commit_outcome {
Ok(()) if conn.is_autocommit() => ProfiledWrappedTransaction {
result: Ok(value),
terminal_state: None,
body,
commit,
},
Ok(()) => {
tracing::error!(
"writer transaction: COMMIT returned success but the connection is still in \
a transaction; request side effects are unknown"
);
let request_state = WriterTaskRequestState::SideEffectsUnknown;
ProfiledWrappedTransaction {
result: Err(writer_task_terminated(request_state)),
terminal_state: Some(request_state),
body,
commit,
}
}
Err(commit_error) => {
if let Some(db) = db {
crate::timeout_sink::maybe_emit_sqlite_full(db, &commit_error);
}
match rollback_after_failure(conn, "commit failure") {
RollbackDisposition::RolledBack => ProfiledWrappedTransaction {
result: Err(StorageError::WriterTaskRequestFailed {
request_state: WriterTaskRequestState::TransactionRolledBack,
source: Box::new({
let failure = crate::error::native_write_failure(
&commit_error,
khive_storage::error::SqliteWriteStage::Commit,
);
let error = StorageError::Pool {
operation: commit_operation.into(),
message: commit_error.to_string(),
};
match failure {
Some(failure) => error.with_sqlite_write_failure(failure),
None => error,
}
}),
}),
terminal_state: None,
body,
commit,
},
RollbackDisposition::SideEffectsUnknown => {
let request_state = WriterTaskRequestState::SideEffectsUnknown;
ProfiledWrappedTransaction {
result: Err(writer_task_terminated_with_cause(
request_state,
&commit_error,
khive_storage::error::SqliteWriteStage::Commit,
)),
terminal_state: Some(request_state),
body,
commit,
}
}
}
}
}
}
Ok(Err(operation_error)) => {
match rollback_after_failure(conn, "request operation failure") {
RollbackDisposition::RolledBack => ProfiledWrappedTransaction {
result: Err(StorageError::WriterTaskRequestFailed {
request_state: WriterTaskRequestState::TransactionRolledBack,
source: Box::new(operation_error),
}),
terminal_state: None,
body,
commit: Duration::ZERO,
},
RollbackDisposition::SideEffectsUnknown => {
if let Some(db) = db {
crate::timeout_sink::maybe_emit_sqlite_full(db, &operation_error);
}
let request_state = WriterTaskRequestState::SideEffectsUnknown;
ProfiledWrappedTransaction {
result: Err(writer_task_terminated_with_cause(
request_state,
&operation_error,
khive_storage::error::SqliteWriteStage::Statement,
)),
terminal_state: Some(request_state),
body,
commit: Duration::ZERO,
}
}
}
}
Err(_panic_payload) => {
let request_state = match rollback_after_failure(conn, "request panic") {
RollbackDisposition::RolledBack => WriterTaskRequestState::TransactionRolledBack,
RollbackDisposition::SideEffectsUnknown => {
WriterTaskRequestState::SideEffectsUnknown
}
};
ProfiledWrappedTransaction {
result: Err(writer_task_terminated(request_state)),
terminal_state: Some(request_state),
body,
commit: Duration::ZERO,
}
}
}
}
impl<R: Send + 'static> sealed::Sealed for WriteRequest<R> {
fn execute_and_reply_reporting_terminal(
self: Box<Self>,
conn: &Connection,
tx_span: Option<khive_storage::tx_registry::TxHandle>,
queue_wait: Duration,
transaction_acquire: Duration,
) -> Option<WriterTaskRequestState> {
let WriteRequest {
op,
reply,
telemetry,
..
} = *self;
let profiled = execute_wrapped_transaction_profiled(
conn,
"writer_task_commit",
Some(&telemetry.db),
op,
);
if let Err(error) = &profiled.result {
crate::timeout_sink::maybe_emit_sqlite_full(&telemetry.db, error);
}
drop(tx_span);
telemetry.finish(
queue_wait,
transaction_acquire,
profiled.body,
profiled.commit,
);
let _ = reply.send(profiled.result);
profiled.terminal_state
}
fn execute_and_reply_top_level_reporting_terminal(
self: Box<Self>,
conn: &Connection,
queue_wait: Duration,
) -> Option<WriterTaskRequestState> {
let WriteRequest {
op,
reply,
telemetry,
..
} = *self;
let body_started = Instant::now();
let outcome = catch_unwind(AssertUnwindSafe(|| {
op(conn).map_err(crate::error::statement_failure)
}));
let body = body_started.elapsed();
let telemetry_db = telemetry.db.clone();
telemetry.finish(queue_wait, Duration::ZERO, body, Duration::ZERO);
match outcome {
Ok(outcome) if conn.is_autocommit() => {
if let Err(error) = &outcome {
crate::timeout_sink::maybe_emit_sqlite_full(&telemetry_db, error);
}
let _ = reply.send(outcome);
None
}
Ok(outcome) => {
tracing::error!(
"writer task: top-level request returned with an open transaction; request \
side effects are unknown"
);
let request_state = WriterTaskRequestState::SideEffectsUnknown;
let error = match outcome {
Err(error) => writer_task_terminated_with_cause(
request_state,
&error,
khive_storage::error::SqliteWriteStage::Statement,
),
Ok(_) => writer_task_terminated(request_state),
};
let _ = reply.send(Err(error));
Some(request_state)
}
Err(_panic_payload) => {
let request_state = WriterTaskRequestState::SideEffectsUnknown;
let _ = reply.send(Err(writer_task_terminated(request_state)));
Some(request_state)
}
}
}
fn reply_error_after_begin(
self: Box<Self>,
err: StorageError,
queue_wait: Duration,
transaction_acquire: Duration,
) {
let WriteRequest {
reply, telemetry, ..
} = *self;
telemetry.finish(
queue_wait,
transaction_acquire,
Duration::ZERO,
Duration::ZERO,
);
let _ = reply.send(Err(err));
}
}
impl<R: Send + 'static> AnyWriteRequest for WriteRequest<R> {
fn execute_and_reply(self: Box<Self>, conn: &Connection) {
let queue_wait = self.queue_wait();
let _ = sealed::Sealed::execute_and_reply_reporting_terminal(
self,
conn,
None,
queue_wait,
Duration::ZERO,
);
}
fn execute_and_reply_top_level(self: Box<Self>, conn: &Connection) {
let queue_wait = self.queue_wait();
let _ =
sealed::Sealed::execute_and_reply_top_level_reporting_terminal(self, conn, queue_wait);
}
fn reply_error(self: Box<Self>, err: StorageError) {
let queue_wait = self.queue_wait();
sealed::Sealed::reply_error_after_begin(self, err, queue_wait, Duration::ZERO);
}
fn is_top_level(&self) -> bool {
self.top_level
}
fn is_checkpoint_bypass(&self) -> bool {
self.checkpoint_bypass
}
fn needs_vacuum_headroom(&self) -> bool {
self.vacuum_copy_headroom
}
fn queue_wait(&self) -> Duration {
self.telemetry.queue_wait()
}
}
fn writer_task_terminated(request_state: WriterTaskRequestState) -> StorageError {
StorageError::writer_task_terminated(request_state)
}
fn writer_task_terminated_with_cause(
request_state: WriterTaskRequestState,
error: &(dyn std::error::Error + 'static),
stage: khive_storage::error::SqliteWriteStage,
) -> StorageError {
let mut source = Some(error);
let mut sqlite_full_codes = None;
while let Some(current) = source {
if let Some(rusqlite::Error::SqliteFailure(code, _)) =
current.downcast_ref::<rusqlite::Error>()
{
if code.code == rusqlite::ErrorCode::DiskFull {
sqlite_full_codes = Some((code.extended_code & 0xff, code.extended_code));
break;
}
}
source = current.source();
}
let terminal = StorageError::WriterTaskTerminated {
request_state,
sqlite_full_codes,
};
match crate::error::native_write_failure(error, stage) {
Some(mut failure) => {
failure.settlement_unknown =
request_state == WriterTaskRequestState::SideEffectsUnknown;
terminal.with_sqlite_write_failure(failure)
}
None => terminal,
}
}
fn writer_task_begin_error(error: rusqlite::Error, busy_timeout: Duration) -> StorageError {
let failure =
crate::error::native_write_failure(&error, khive_storage::error::SqliteWriteStage::Begin);
let mapped = if crate::timeout_sink::is_busy_or_locked(&error) {
StorageError::WriterTaskBusy {
timeout_ms: u64::try_from(busy_timeout.as_millis()).unwrap_or(u64::MAX),
}
} else {
StorageError::Pool {
operation: "writer_task_begin".into(),
message: error.to_string(),
}
};
match failure {
Some(failure) => mapped.with_sqlite_write_failure(failure),
None => mapped,
}
}
#[derive(Clone, Debug)]
pub struct WriterTaskHandle {
tx: mpsc::Sender<Box<dyn AnyWriteRequest + Send>>,
backend_key: Option<PathBuf>,
db: String,
slow_write_threshold: Option<std::time::Duration>,
enqueue_timeout: std::time::Duration,
}
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, false, false, SubmitLabels::None)
.await
}
async fn enqueue_inner<R, F>(
&self,
op: F,
top_level: bool,
checkpoint_bypass: bool,
vacuum_copy_headroom: bool,
labels: SubmitLabels,
) -> 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 labels = match labels {
SubmitLabels::Statements(labels) => {
SubmitLabels::Statements(dedup_labels(labels.iter().map(String::as_str)))
}
other => other,
};
let telemetry = WriteTelemetry::new(
self.backend_key.clone(),
self.db.clone(),
self.queue_depth(),
self.slow_write_threshold,
labels,
);
let request = WriteRequest {
op: Box::new(op),
reply: reply_tx,
top_level,
checkpoint_bypass,
vacuum_copy_headroom,
telemetry,
};
self.tx
.send(Box::new(request))
.await
.map_err(|_| writer_task_terminated(WriterTaskRequestState::NotStarted))?;
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(|_| writer_task_terminated(WriterTaskRequestState::SideEffectsUnknown))?
}
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,
{
self.send_with_timeout_labelled(op, timeout, SubmitLabels::None)
.await
}
async fn send_with_timeout_labelled<R, F>(
&self,
op: F,
timeout: std::time::Duration,
labels: SubmitLabels,
) -> 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_inner(op, false, false, false, labels),
)
.await
{
Ok(Ok(reply_rx)) => reply_rx,
Ok(Err(e)) => return Err(e),
Err(_elapsed) => {
let timeout_ms = timeout.as_millis() as u64;
crate::timeout_sink::emit_queue_saturation(&self.db, timeout_ms);
return Err(StorageError::WriteQueueFull { timeout_ms });
}
};
reply_rx
.await
.map_err(|_| writer_task_terminated(WriterTaskRequestState::SideEffectsUnknown))?
}
pub async fn send_bounded<R, F>(&self, op: F) -> Result<R, StorageError>
where
R: Send + 'static,
F: FnOnce(&Connection) -> Result<R, StorageError> + Send + 'static,
{
self.send_with_timeout(op, self.enqueue_timeout).await
}
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, false, false, SubmitLabels::None)
.await?;
reply_rx
.await
.map_err(|_| writer_task_terminated(WriterTaskRequestState::SideEffectsUnknown))?
}
pub async fn send_top_level_bounded<R, F>(&self, op: F) -> Result<R, StorageError>
where
R: Send + 'static,
F: FnOnce(&Connection) -> Result<R, StorageError> + Send + 'static,
{
self.send_top_level_bounded_inner(op, false, false, SubmitLabels::None)
.await
}
pub(crate) async fn send_checkpoint_bounded<R, F>(&self, op: F) -> Result<R, StorageError>
where
R: Send + 'static,
F: FnOnce(&Connection) -> Result<R, StorageError> + Send + 'static,
{
self.send_top_level_bounded_inner(op, true, false, SubmitLabels::Site("wal_checkpoint"))
.await
}
pub(crate) async fn send_vacuum_bounded<R, F>(&self, op: F) -> Result<R, StorageError>
where
R: Send + 'static,
F: FnOnce(&Connection) -> Result<R, StorageError> + Send + 'static,
{
self.send_top_level_bounded_inner(op, false, true, SubmitLabels::Site("vacuum"))
.await
}
async fn send_top_level_bounded_inner<R, F>(
&self,
op: F,
checkpoint_bypass: bool,
vacuum_copy_headroom: bool,
labels: SubmitLabels,
) -> Result<R, StorageError>
where
R: Send + 'static,
F: FnOnce(&Connection) -> Result<R, StorageError> + Send + 'static,
{
let reply_rx = match tokio::time::timeout(
self.enqueue_timeout,
self.enqueue_inner(op, true, checkpoint_bypass, vacuum_copy_headroom, labels),
)
.await
{
Ok(Ok(reply_rx)) => reply_rx,
Ok(Err(e)) => return Err(e),
Err(_elapsed) => {
let timeout_ms = self.enqueue_timeout.as_millis() as u64;
crate::timeout_sink::emit_queue_saturation(&self.db, timeout_ms);
return Err(StorageError::WriteQueueFull { timeout_ms });
}
};
reply_rx
.await
.map_err(|_| writer_task_terminated(WriterTaskRequestState::SideEffectsUnknown))?
}
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_untracked()?;
let acquisition_counters = pool.writer_acquisition_counters();
let write_admission = pool.write_admission();
let busy_timeout = pool.config().busy_timeout;
let origin = pool.origin();
let backend_key = writer_db_key(pool);
let db = crate::timeout_sink::db_label(pool);
let (tx, rx) = mpsc::channel(capacity.max(1));
let join = tokio::spawn(run_writer_task(
conn,
rx,
origin,
db.clone(),
acquisition_counters,
write_admission,
busy_timeout,
));
pool.set_writer_task_join(join);
Ok(WriterTaskHandle {
tx,
backend_key,
db,
slow_write_threshold: crate::timeout_sink::slow_write_threshold(),
enqueue_timeout: std::time::Duration::from_millis(
pool.config().write_admission_deadline_ms,
),
})
}
async fn close_and_fail_queued_requests(rx: &mut mpsc::Receiver<Box<dyn AnyWriteRequest + Send>>) {
rx.close();
while let Some(request) = rx.recv().await {
request.reply_error(writer_task_terminated(WriterTaskRequestState::NotStarted));
}
}
fn begin_immediate_with_retry(
conn: &Connection,
acquisition_counters: &WriterAcquisitionCounters,
busy_timeout: Duration,
mut set_busy_timeout: impl FnMut(&Connection, Duration) -> rusqlite::Result<()>,
) -> (rusqlite::Result<()>, Duration, u32) {
let transaction_acquire_started = Instant::now();
let mut begin_attempt = 1_u32;
let mut retry_delays = WRITER_BEGIN_RETRY_DELAYS.into_iter();
let mut busy_timeout_lowered = false;
let begin_outcome = loop {
match conn.execute_batch("BEGIN IMMEDIATE") {
Ok(()) => break Ok(()),
Err(error) if crate::timeout_sink::is_busy_or_locked(&error) => {
acquisition_counters.record_writer_task_begin_busy();
let Some(delay) = retry_delays.next() else {
break Err(error);
};
let remaining_budget =
busy_timeout.saturating_sub(transaction_acquire_started.elapsed());
if remaining_budget.is_zero() {
break Err(error);
}
if let Err(set_err) = set_busy_timeout(conn, remaining_budget) {
tracing::warn!(
error = %set_err,
"writer task: failed to lower busy_timeout for BEGIN \
retry; surfacing the original busy refusal"
);
break Err(error);
}
busy_timeout_lowered = true;
acquisition_counters.record_writer_task_begin_busy_absorbed();
tracing::debug!(
attempt = begin_attempt,
backoff_ms = delay.as_millis() as u64,
budget_remaining_ms = remaining_budget.as_millis() as u64,
"writer task: BEGIN IMMEDIATE refused busy; retrying before \
request execution"
);
std::thread::sleep(delay);
begin_attempt = begin_attempt.saturating_add(1);
}
Err(error) => break Err(error),
}
};
let transaction_acquire = transaction_acquire_started.elapsed();
if busy_timeout_lowered {
if let Err(restore_err) = set_busy_timeout(conn, busy_timeout) {
tracing::warn!(
error = %restore_err,
"writer task: failed to restore busy_timeout after a BEGIN retry \
sequence"
);
}
}
(begin_outcome, transaction_acquire, begin_attempt)
}
struct BlockingWriterConnection {
conn: Option<Connection>,
admission: Arc<crate::pool::WriteAdmission>,
volume_lease: Option<crate::disk_guard::VolumeLease>,
}
impl BlockingWriterConnection {
fn new(conn: Connection, admission: Arc<crate::pool::WriteAdmission>) -> Self {
Self {
conn: Some(conn),
admission,
volume_lease: None,
}
}
fn finish(
mut self,
state: Option<WriterTaskRequestState>,
) -> (Option<Connection>, Option<WriterTaskRequestState>) {
if state.is_some() {
self.retire();
(None, state)
} else {
(self.conn.take(), None)
}
}
fn retire(&mut self) {
if let Some(conn) = self.conn.take() {
let _ = self.admission.close_retired_connection(conn);
}
}
}
impl Drop for BlockingWriterConnection {
fn drop(&mut self) {
self.retire();
}
}
impl std::ops::Deref for BlockingWriterConnection {
type Target = Connection;
fn deref(&self) -> &Connection {
self.conn
.as_ref()
.expect("blocking writer owns its connection until settlement")
}
}
async fn run_writer_task(
mut conn: Connection,
mut rx: mpsc::Receiver<Box<dyn AnyWriteRequest + Send>>,
origin: khive_storage::tx_registry::TxOrigin,
db: String,
acquisition_counters: Arc<WriterAcquisitionCounters>,
write_admission: Arc<crate::pool::WriteAdmission>,
busy_timeout: Duration,
) {
while let Some(request) = rx.recv().await {
let queue_wait = request.queue_wait();
let origin = origin.clone();
let blocking_counters = Arc::clone(&acquisition_counters);
let blocking_admission = Arc::clone(&write_admission);
let blocking_db = db.clone();
let outcome = tokio::task::spawn_blocking(move || {
let mut conn = BlockingWriterConnection::new(conn, Arc::clone(&blocking_admission));
let acquisition_counters = blocking_counters;
if !conn.is_autocommit() {
tracing::error!(
"writer task: connection is not in autocommit mode before request dispatch; \
retiring the poisoned writer without running the request"
);
let request_state = WriterTaskRequestState::NotStarted;
request.reply_error(writer_task_terminated(request_state));
return conn.finish(Some(request_state));
}
conn.volume_lease = if request.is_checkpoint_bypass() {
None
} else {
match blocking_admission.acquire() {
Ok(lease) => lease,
Err(error) => {
request.reply_error(error.into_storage_error(
khive_storage::StorageCapability::Sql,
"writer_task_admission",
));
return conn.finish(None);
}
}
};
let terminal_state = if request.is_top_level() {
if !request.is_checkpoint_bypass() {
let admission = if request.needs_vacuum_headroom() {
blocking_admission.check_for_vacuum()
} else {
blocking_admission.check()
};
if let Err(error) = admission {
request.reply_error(error.into_storage_error(
khive_storage::StorageCapability::Sql,
"writer_task_admission",
));
return conn.finish(None);
}
}
acquisition_counters.record_writer_task_acquisition();
sealed::Sealed::execute_and_reply_top_level_reporting_terminal(
request, &conn, queue_wait,
)
} else {
let tx_span = khive_storage::tx_registry::register_scoped(
Some("writer_task_tx".to_string()),
origin,
);
let (begin_outcome, transaction_acquire, begin_attempt) =
begin_immediate_with_retry(
&conn,
&acquisition_counters,
busy_timeout,
Connection::busy_timeout,
);
match begin_outcome {
Ok(()) => {
if let Err(error) = blocking_admission.check() {
let request_state =
match rollback_after_failure(&conn, "capacity admission") {
RollbackDisposition::RolledBack => None,
RollbackDisposition::SideEffectsUnknown => {
Some(WriterTaskRequestState::SideEffectsUnknown)
}
};
drop(tx_span);
let error = if let Some(state) = request_state {
writer_task_terminated(state)
} else {
error.into_storage_error(
khive_storage::StorageCapability::Sql,
"writer_task_admission",
)
};
sealed::Sealed::reply_error_after_begin(
request,
error,
queue_wait,
transaction_acquire,
);
return conn.finish(request_state);
}
acquisition_counters.record_writer_task_acquisition();
sealed::Sealed::execute_and_reply_reporting_terminal(
request,
&conn,
Some(tx_span),
queue_wait,
transaction_acquire,
)
}
Err(e) => {
crate::timeout_sink::maybe_emit_sqlite_full(&blocking_db, &e);
tracing::warn!(
error = %e,
attempts = begin_attempt,
"writer task: BEGIN IMMEDIATE failed; replying an \
error without running the request's operation"
);
drop(tx_span);
let begin_error = writer_task_begin_error(e, busy_timeout);
if !matches!(
begin_error.without_sqlite_write_stage(),
StorageError::WriterTaskBusy { .. }
) {
acquisition_counters.record_writer_task_begin_error();
}
sealed::Sealed::reply_error_after_begin(
request,
begin_error,
queue_wait,
transaction_acquire,
);
None
}
}
};
conn.finish(terminal_state)
})
.await;
match outcome {
Ok((Some(returned_conn), None)) => conn = returned_conn,
Ok((_returned_conn, Some(request_state))) => {
acquisition_counters.record_writer_task_request_failure();
if request_state == WriterTaskRequestState::SideEffectsUnknown {
acquisition_counters.record_writer_task_side_effects_unknown();
}
tracing::error!(
request_state = %request_state,
"writer task reached a terminal request or connection state; closing and \
failing the queue without restarting"
);
crate::timeout_sink::emit_writer_task_retirement(
&db,
&format!("terminal request state: {request_state}"),
);
close_and_fail_queued_requests(&mut rx).await;
return;
}
Ok((None, None)) => unreachable!("a reusable writer returns its connection"),
Err(join_err) => {
acquisition_counters.record_writer_task_request_failure();
tracing::error!(
error = %join_err,
"writer task blocking closure failed outside the request \
panic boundary; closing and failing the queue without restarting"
);
crate::timeout_sink::emit_writer_task_retirement(
&db,
&format!("blocking closure join failure: {join_err}"),
);
close_and_fail_queued_requests(&mut rx).await;
return;
}
}
}
}
#[cfg(test)]
#[path = "writer_task_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "writer_task_lease_close_tests.rs"]
mod volume_lease_close_tests;
#[cfg(test)]
#[path = "writer_task_slow_write_label_tests.rs"]
mod slow_write_label_tests;
#[cfg(all(test, any(unix, windows)))]
#[path = "writer_task_identity_tests.rs"]
mod identity_tests;