use std::sync::Arc;
use std::time::Duration;
use crate::pool::ConnectionPool;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CheckpointTick {
Skipped,
Observed(u64),
}
#[derive(Clone, Debug)]
pub struct CheckpointConfig {
pub interval: Duration,
pub warn_pages: u64,
pub high_water_pages: u64,
}
impl Default for CheckpointConfig {
fn default() -> Self {
Self {
interval: Duration::from_millis(500),
warn_pages: 2000,
high_water_pages: 6000,
}
}
}
impl CheckpointConfig {
pub fn from_env() -> Self {
let mut cfg = Self::default();
if let Ok(ms) = std::env::var("KHIVE_CHECKPOINT_INTERVAL_MS") {
if let Ok(v) = ms.parse::<u64>() {
if v > 0 {
cfg.interval = Duration::from_millis(v);
}
}
}
if let Ok(v) = std::env::var("KHIVE_WAL_WARN_PAGES") {
if let Ok(n) = v.parse::<u64>() {
if n > 0 {
cfg.warn_pages = n;
}
}
}
if let Ok(v) = std::env::var("KHIVE_WAL_HIGH_WATER_PAGES") {
if let Ok(n) = v.parse::<u64>() {
if n > 0 {
cfg.high_water_pages = n;
}
}
}
cfg
}
}
pub async fn run_checkpoint_task(pool: Arc<ConnectionPool>, config: CheckpointConfig) {
let mut interval = tokio::time::interval(config.interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut was_above_warn = false;
let mut was_above_high_water = false;
loop {
interval.tick().await;
if Arc::strong_count(&pool) <= 1 {
break;
}
let tick = checkpoint_once(&pool);
let wal_pages = match tick {
CheckpointTick::Skipped => continue,
CheckpointTick::Observed(n) => n,
};
let in_warn_band = wal_pages >= config.warn_pages && wal_pages < config.high_water_pages;
if crossing_warn(in_warn_band, &mut was_above_warn) {
tracing::warn!(
wal_pages,
warn_threshold = config.warn_pages,
"WAL page count approaching checkpoint threshold"
);
}
let above_high_water = wal_pages >= config.high_water_pages;
if crossing_warn(above_high_water, &mut was_above_high_water) {
tracing::warn!(
wal_pages,
high_water = config.high_water_pages,
"WAL high-water mark exceeded; sustained WAL pressure — \
a long-lived reader may be pinning an old snapshot that PASSIVE cannot reclaim"
);
}
}
}
pub fn checkpoint_once(pool: &ConnectionPool) -> CheckpointTick {
let writer = match pool.try_writer_nowait() {
Ok(w) => w,
Err(_) => return CheckpointTick::Skipped,
};
let wal_pages = query_wal_pages(writer.conn());
if let Err(e) = writer
.conn()
.execute_batch("PRAGMA wal_checkpoint(PASSIVE)")
{
tracing::warn!(error = %e, "WAL checkpoint failed");
} else {
tracing::debug!(wal_pages, "WAL checkpoint issued");
}
CheckpointTick::Observed(wal_pages)
}
fn crossing_warn(now_above: bool, was_above: &mut bool) -> bool {
let fire = now_above && !*was_above;
*was_above = now_above;
fire
}
fn query_wal_pages(conn: &rusqlite::Connection) -> u64 {
conn.query_row("PRAGMA wal_checkpoint", [], |row| row.get::<_, i64>(1))
.unwrap_or(0)
.max(0) as u64
}
#[cfg(test)]
mod tests {
use super::*;
use crate::pool::PoolConfig;
use serial_test::serial;
fn file_pool(path: &std::path::Path) -> Arc<ConnectionPool> {
let cfg = PoolConfig {
path: Some(path.to_path_buf()),
..PoolConfig::default()
};
Arc::new(ConnectionPool::new(cfg).expect("pool open"))
}
#[test]
fn checkpoint_once_succeeds_on_file_backed_pool() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("wal_test.db");
let pool = file_pool(&path);
{
let writer = pool.try_writer().unwrap();
writer
.conn()
.execute_batch("CREATE TABLE IF NOT EXISTS t (x INTEGER);")
.unwrap();
writer
.conn()
.execute_batch("INSERT INTO t VALUES (1);")
.unwrap();
}
checkpoint_once(&pool);
}
#[test]
fn checkpoint_once_is_noop_on_in_memory_pool() {
let cfg = PoolConfig {
path: None,
..PoolConfig::default()
};
let pool = Arc::new(ConnectionPool::new(cfg).expect("in-memory pool"));
checkpoint_once(&pool);
}
#[tokio::test]
async fn checkpoint_task_exits_when_pool_dropped() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("wal_task_drop.db");
let pool = file_pool(&path);
let cfg = CheckpointConfig {
interval: Duration::from_millis(10),
..Default::default()
};
let weak = Arc::downgrade(&pool);
let task_pool = Arc::clone(&pool);
let handle = tokio::spawn(run_checkpoint_task(task_pool, cfg));
drop(pool);
tokio::time::timeout(Duration::from_secs(1), handle)
.await
.expect("checkpoint task should exit within 1s")
.expect("checkpoint task panicked");
assert!(weak.upgrade().is_none(), "pool should be fully dropped");
}
#[test]
#[serial]
fn checkpoint_config_env_override() {
std::env::set_var("KHIVE_CHECKPOINT_INTERVAL_MS", "250");
std::env::set_var("KHIVE_WAL_WARN_PAGES", "1500");
std::env::set_var("KHIVE_WAL_HIGH_WATER_PAGES", "8000");
let cfg = CheckpointConfig::from_env();
std::env::remove_var("KHIVE_CHECKPOINT_INTERVAL_MS");
std::env::remove_var("KHIVE_WAL_WARN_PAGES");
std::env::remove_var("KHIVE_WAL_HIGH_WATER_PAGES");
assert_eq!(cfg.interval, Duration::from_millis(250));
assert_eq!(cfg.warn_pages, 1500);
assert_eq!(cfg.high_water_pages, 8000);
}
#[test]
#[serial]
fn checkpoint_config_defaults_on_invalid_env() {
let default = CheckpointConfig::default();
std::env::set_var("KHIVE_CHECKPOINT_INTERVAL_MS", "not_a_number");
std::env::set_var("KHIVE_WAL_WARN_PAGES", "");
std::env::set_var("KHIVE_WAL_HIGH_WATER_PAGES", "0");
let cfg = CheckpointConfig::from_env();
std::env::remove_var("KHIVE_CHECKPOINT_INTERVAL_MS");
std::env::remove_var("KHIVE_WAL_WARN_PAGES");
std::env::remove_var("KHIVE_WAL_HIGH_WATER_PAGES");
assert_eq!(cfg.interval, default.interval);
assert_eq!(cfg.warn_pages, default.warn_pages);
assert_eq!(cfg.high_water_pages, default.high_water_pages);
}
#[test]
fn checkpoint_high_water_does_not_block_behind_reader() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("high_water_test.db");
let pool = Arc::new(
ConnectionPool::new(PoolConfig {
path: Some(path.clone()),
busy_timeout: Duration::from_millis(2000),
..PoolConfig::default()
})
.expect("pool open"),
);
{
let writer = pool.try_writer().unwrap();
writer
.conn()
.execute_batch(
"CREATE TABLE IF NOT EXISTS t (x INTEGER); INSERT INTO t VALUES (1);",
)
.unwrap();
}
let reader = pool.reader().expect("reader");
reader
.execute_batch("BEGIN DEFERRED; SELECT * FROM t;")
.expect("begin read tx");
{
let writer = pool.try_writer().unwrap();
writer
.conn()
.execute_batch("INSERT INTO t VALUES (2);")
.unwrap();
}
let start = std::time::Instant::now();
checkpoint_once(&pool);
let elapsed = start.elapsed();
reader.execute_batch("COMMIT;").ok();
drop(reader);
assert!(
elapsed < std::time::Duration::from_millis(500),
"checkpoint_once with active reader snapshot took {:?}; \
expected <500ms (PASSIVE must not block on readers; \
a TRUNCATE regression would block ~2000ms)",
elapsed
);
}
#[test]
#[serial]
fn checkpoint_config_rejects_zero_for_all_fields() {
let default = CheckpointConfig::default();
std::env::set_var("KHIVE_CHECKPOINT_INTERVAL_MS", "0");
std::env::set_var("KHIVE_WAL_WARN_PAGES", "0");
std::env::set_var("KHIVE_WAL_HIGH_WATER_PAGES", "0");
let cfg = CheckpointConfig::from_env();
std::env::remove_var("KHIVE_CHECKPOINT_INTERVAL_MS");
std::env::remove_var("KHIVE_WAL_WARN_PAGES");
std::env::remove_var("KHIVE_WAL_HIGH_WATER_PAGES");
assert_eq!(
cfg.interval, default.interval,
"zero interval must fall back to default"
);
assert_eq!(
cfg.warn_pages, default.warn_pages,
"zero warn_pages must fall back to default"
);
assert_eq!(
cfg.high_water_pages, default.high_water_pages,
"zero high_water_pages must fall back to default"
);
}
#[test]
fn skipped_tick_does_not_reset_high_water_crossing_state() {
let mut was_above = false;
assert!(
crossing_warn(true, &mut was_above),
"should fire on first crossing"
);
assert!(was_above);
assert!(was_above, "was_above must stay true across skipped ticks");
let fired = crossing_warn(true, &mut was_above);
assert!(!fired, "WARN must not re-fire while still above threshold");
let fired = crossing_warn(false, &mut was_above);
assert!(!fired);
assert!(!was_above);
let fired = crossing_warn(true, &mut was_above);
assert!(fired, "WARN must fire again on a new below→above crossing");
}
#[test]
fn warn_pages_fires_once_on_crossing_not_every_tick() {
let mut was_above_warn = false;
let fired_1 = crossing_warn(true, &mut was_above_warn);
let fired_2 = crossing_warn(true, &mut was_above_warn);
let fired_3 = crossing_warn(true, &mut was_above_warn);
assert!(fired_1, "WARN must fire on the first in-band tick");
assert!(
!fired_2,
"WARN must not fire on the second consecutive in-band tick"
);
assert!(
!fired_3,
"WARN must not fire on the third consecutive in-band tick"
);
crossing_warn(false, &mut was_above_warn);
assert!(!was_above_warn);
let fired_reentry = crossing_warn(true, &mut was_above_warn);
assert!(
fired_reentry,
"WARN must fire again on re-entry into warn band"
);
}
}