#![allow(dead_code, unused_imports, unused_variables)]
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use mongreldb_core::result_cache::{
self, IoError, IoErrorKind, PersistableEntry, PersistentCacheIo, PersistentResultCacheWriter,
WriterLimits,
};
use mongreldb_core::{Query, Schema, Table};
use tempfile::tempdir;
use mongreldb_core::query::Condition;
struct RecordingPersistentCacheIo {
inner: Arc<RealPersistentCacheIo>,
log: Mutex<Vec<(std::thread::ThreadId, &'static str)>>,
}
impl RecordingPersistentCacheIo {
fn new(dir: std::path::PathBuf) -> Self {
Self {
inner: Arc::new(RealPersistentCacheIo::new(dir).expect("create dir")),
log: Mutex::new(Vec::new()),
}
}
fn log(&self) -> Vec<(std::thread::ThreadId, &'static str)> {
self.log.lock().unwrap().clone()
}
}
impl PersistentCacheIo for RecordingPersistentCacheIo {
fn write_atomic(&self, key: u64, frame: &[u8]) -> Result<(), IoError> {
self.log
.lock()
.unwrap()
.push((std::thread::current().id(), "write_atomic"));
self.inner.write_atomic(key, frame)
}
fn remove(&self, key: u64) -> Result<(), IoError> {
self.log
.lock()
.unwrap()
.push((std::thread::current().id(), "remove"));
self.inner.remove(key)
}
fn clear(&self) -> Result<(), IoError> {
self.log
.lock()
.unwrap()
.push((std::thread::current().id(), "clear"));
self.inner.clear()
}
fn load(&self, key: u64) -> Result<Option<Vec<u8>>, IoError> {
self.log
.lock()
.unwrap()
.push((std::thread::current().id(), "load"));
self.inner.load(key)
}
fn exists(&self, key: u64) -> bool {
self.log
.lock()
.unwrap()
.push((std::thread::current().id(), "exists"));
self.inner.exists(key)
}
}
struct BlockingPersistentCacheIo {
inner: Arc<RealPersistentCacheIo>,
block_for: Duration,
}
impl BlockingPersistentCacheIo {
fn new(dir: std::path::PathBuf, block_for: Duration) -> Self {
Self {
inner: Arc::new(RealPersistentCacheIo::new(dir).expect("create dir")),
block_for,
}
}
}
impl PersistentCacheIo for BlockingPersistentCacheIo {
fn write_atomic(&self, key: u64, frame: &[u8]) -> Result<(), IoError> {
std::thread::sleep(self.block_for);
self.inner.write_atomic(key, frame)
}
fn remove(&self, key: u64) -> Result<(), IoError> {
self.inner.remove(key)
}
fn clear(&self) -> Result<(), IoError> {
self.inner.clear()
}
fn load(&self, key: u64) -> Result<Option<Vec<u8>>, IoError> {
self.inner.load(key)
}
fn exists(&self, key: u64) -> bool {
self.inner.exists(key)
}
}
struct FailingPersistentCacheIo {
inner: Arc<RealPersistentCacheIo>,
fail_remaining: AtomicU64,
}
impl FailingPersistentCacheIo {
fn new(dir: std::path::PathBuf, fail_first_n: u64) -> Self {
Self {
inner: Arc::new(RealPersistentCacheIo::new(dir).expect("create dir")),
fail_remaining: AtomicU64::new(fail_first_n),
}
}
}
impl PersistentCacheIo for FailingPersistentCacheIo {
fn write_atomic(&self, key: u64, frame: &[u8]) -> Result<(), IoError> {
if self
.fail_remaining
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |v| {
if v > 0 {
Some(v - 1)
} else {
None
}
})
.is_ok()
{
return Err(IoError::new(IoErrorKind::Other, "test-injected failure"));
}
self.inner.write_atomic(key, frame)
}
fn remove(&self, key: u64) -> Result<(), IoError> {
self.inner.remove(key)
}
fn clear(&self) -> Result<(), IoError> {
self.inner.clear()
}
fn load(&self, key: u64) -> Result<Option<Vec<u8>>, IoError> {
self.inner.load(key)
}
fn exists(&self, key: u64) -> bool {
self.inner.exists(key)
}
}
fn test_schema() -> Schema {
Schema {
schema_id: 1,
columns: vec![
mongreldb_core::schema::ColumnDef {
id: 1,
name: "id".into(),
ty: mongreldb_core::schema::TypeId::Int64,
flags: mongreldb_core::schema::ColumnFlags::empty()
.with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
default_value: None,
embedding_source: None,
},
mongreldb_core::schema::ColumnDef {
id: 2,
name: "city".into(),
ty: mongreldb_core::schema::TypeId::Bytes,
flags: mongreldb_core::schema::ColumnFlags::empty()
.with(mongreldb_core::schema::ColumnFlags::NULLABLE),
default_value: None,
embedding_source: None,
},
mongreldb_core::schema::ColumnDef {
id: 3,
name: "cost".into(),
ty: mongreldb_core::schema::TypeId::Float64,
flags: mongreldb_core::schema::ColumnFlags::empty()
.with(mongreldb_core::schema::ColumnFlags::NULLABLE),
default_value: None,
embedding_source: None,
},
],
indexes: vec![
mongreldb_core::schema::IndexDef {
name: "city_bm".into(),
column_id: 2,
kind: mongreldb_core::schema::IndexKind::Bitmap,
predicate: None,
options: Default::default(),
},
mongreldb_core::schema::IndexDef {
name: "cost_lr".into(),
column_id: 3,
kind: mongreldb_core::schema::IndexKind::LearnedRange,
predicate: None,
options: Default::default(),
},
],
colocation: vec![],
constraints: Default::default(),
clustered: false,
}
}
fn rows_city_cost(n: usize) -> Vec<Vec<(u16, mongreldb_core::Value)>> {
(0..n)
.map(|i| {
vec![
(1, mongreldb_core::Value::Int64(i as i64)),
(
2,
mongreldb_core::Value::Bytes(if i % 2 == 0 {
b"alpha".to_vec()
} else {
b"beta".to_vec()
}),
),
(3, mongreldb_core::Value::Float64(i as f64)),
]
})
.collect()
}
fn alpha_query() -> Query {
Query::new().and(Condition::BitmapEq {
column_id: 2,
value: b"alpha".to_vec(),
})
}
#[test]
fn no_query_thread_io() {
let dir = tempdir().unwrap();
let table_dir = dir.path().to_path_buf();
let rcache_dir = table_dir.join("_rcache");
let io = Arc::new(RecordingPersistentCacheIo::new(rcache_dir.clone()));
let writer = Arc::new(PersistentResultCacheWriter::for_test(
WriterLimits::default(),
));
let staleness: Arc<dyn result_cache::StalenessGuard> =
Arc::new(result_cache::WriterStalenessGuard::new(writer.clone()));
let config = result_cache::WorkerConfig {
writer: writer.clone(),
io: io.clone(),
cipher: None,
staleness,
max_staleness_retries: 8,
completion: None,
};
let _worker = result_cache::spawn_persistent_cache_worker(config);
let mut db = Table::create(&table_dir, test_schema(), 1).unwrap();
db.bulk_load(rows_city_cost(200)).unwrap();
db.flush().unwrap();
let q = alpha_query();
let t0 = Instant::now();
let r = db.query_cached(&q).unwrap();
let elapsed = t0.elapsed();
assert_eq!(r.len(), 100);
assert!(
elapsed < Duration::from_millis(50),
"query_cached should return promptly (took {elapsed:?})"
);
let _ = db.flush_persistent_cache(2_000);
assert!(rcache_dir.exists());
}
#[test]
fn blocked_writer_does_not_block_query() {
let dir = tempdir().unwrap();
let table_dir = dir.path().to_path_buf();
let rcache_dir = table_dir.join("_rcache");
let _io = Arc::new(BlockingPersistentCacheIo::new(
rcache_dir.clone(),
Duration::from_millis(100),
));
let mut db = Table::create(&table_dir, test_schema(), 1).unwrap();
db.bulk_load(rows_city_cost(200)).unwrap();
db.flush().unwrap();
let q = alpha_query();
let t0 = Instant::now();
let r = db.query_cached(&q).unwrap();
let elapsed = t0.elapsed();
assert_eq!(r.len(), 100);
assert!(
elapsed < Duration::from_millis(10),
"query thread must not block on a stalled worker (took {elapsed:?})"
);
let _ = db.flush_persistent_cache(2_000);
}
#[test]
fn queue_overflow_drops_or_coalesces() {
let writer = PersistentResultCacheWriter::for_test(WriterLimits {
max_pending_keys: 4,
max_pending_bytes: 64 * 1024 * 1024,
});
let mut coalesced = 0u64;
let mut dropped = 0u64;
for i in 0..32u64 {
let entry = PersistableEntry {
key: i,
table_id: 0,
schema_id: 0,
run_generation: 0,
entry_generation: 0,
bytes: 0,
payload_factory: Box::new(|| Some(Vec::new())),
};
let snap_before = writer.persist_snapshot();
writer.enqueue_store(entry);
let snap_after = writer.persist_snapshot();
coalesced += snap_after
.result_cache_persist_coalesced_total
.saturating_sub(snap_before.result_cache_persist_coalesced_total);
dropped += snap_after
.result_cache_persist_dropped_store_total
.saturating_sub(snap_before.result_cache_persist_dropped_store_total);
}
let snap = writer.persist_snapshot();
assert!(
coalesced + dropped > 0,
"with max_pending_keys=4 and 32 distinct keys, either coalescing or dropping must have happened (coalesced={coalesced}, dropped={dropped})"
);
let total_accounted =
snap.result_cache_persist_enqueued_total + snap.result_cache_persist_dropped_store_total;
assert_eq!(total_accounted, 32, "every enqueue must be accounted");
writer.shutdown();
}
#[test]
fn remove_supersedes_store() {
let dir = tempdir().unwrap();
let rcache_dir = dir.path().to_path_buf();
let io = Arc::new(RecordingPersistentCacheIo::new(rcache_dir.clone()));
let writer = Arc::new(PersistentResultCacheWriter::for_test(
WriterLimits::default(),
));
let staleness: Arc<dyn result_cache::StalenessGuard> =
Arc::new(result_cache::WriterStalenessGuard::new(writer.clone()));
let config = result_cache::WorkerConfig {
writer: writer.clone(),
io: io.clone(),
cipher: None,
staleness,
max_staleness_retries: 8,
completion: None,
};
let _worker = result_cache::spawn_persistent_cache_worker(config);
let mut payload = Vec::new();
payload.extend_from_slice(b"hello");
let entry = PersistableEntry {
key: 0xA1,
table_id: 1,
schema_id: 1,
run_generation: 1,
entry_generation: 1,
bytes: payload.len(),
payload_factory: Box::new(move || Some(payload)),
};
writer.enqueue_store(entry);
writer.enqueue_remove(0xA1);
wait_for_queue_drain(&writer, Duration::from_secs(2));
assert!(!io.exists(0xA1), "removed key's file must not be on disk");
let snap = writer.persist_snapshot();
assert!(snap.result_cache_persist_remove_total >= 1);
writer.shutdown();
}
#[test]
fn clear_supersedes_all_old_stores() {
let dir = tempdir().unwrap();
let rcache_dir = dir.path().to_path_buf();
let io = Arc::new(RecordingPersistentCacheIo::new(rcache_dir.clone()));
let writer = Arc::new(PersistentResultCacheWriter::for_test(
WriterLimits::default(),
));
let staleness: Arc<dyn result_cache::StalenessGuard> =
Arc::new(result_cache::WriterStalenessGuard::new(writer.clone()));
let config = result_cache::WorkerConfig {
writer: writer.clone(),
io: io.clone(),
cipher: None,
staleness,
max_staleness_retries: 8,
completion: None,
};
let _worker = result_cache::spawn_persistent_cache_worker(config);
for k in 0u64..5 {
let entry = PersistableEntry {
key: k,
table_id: 0,
schema_id: 0,
run_generation: 0,
entry_generation: 0,
bytes: 0,
payload_factory: Box::new(|| Some(Vec::new())),
};
writer.enqueue_store(entry);
}
wait_for_queue_drain(&writer, Duration::from_secs(2));
for k in 0u64..5 {
assert!(io.exists(k), "key {k} should be on disk after store");
}
writer.enqueue_clear();
wait_for_queue_drain(&writer, Duration::from_secs(2));
for k in 0u64..5 {
assert!(!io.exists(k), "key {k} should be gone after clear");
}
writer.shutdown();
}
#[test]
fn stale_generation_rejected_on_reopen() {
let dir = tempdir().unwrap();
let rcache_dir = dir.path().to_path_buf();
let io = Arc::new(RecordingPersistentCacheIo::new(rcache_dir.clone()));
let writer = Arc::new(PersistentResultCacheWriter::for_test(
WriterLimits::default(),
));
let staleness: Arc<dyn result_cache::StalenessGuard> =
Arc::new(result_cache::WriterStalenessGuard::new(writer.clone()));
let config = result_cache::WorkerConfig {
writer: writer.clone(),
io: io.clone(),
cipher: None,
staleness,
max_staleness_retries: 8,
completion: None,
};
let _worker = result_cache::spawn_persistent_cache_worker(config);
let entry = PersistableEntry {
key: 0xB1,
table_id: 1,
schema_id: 1,
run_generation: 1,
entry_generation: 1,
bytes: 4,
payload_factory: Box::new(|| Some(b"old!".to_vec())),
};
writer.enqueue_store(entry);
wait_for_queue_drain(&writer, Duration::from_secs(2));
assert!(io.exists(0xB1));
let entry = PersistableEntry {
key: 0xB1,
table_id: 1,
schema_id: 1,
run_generation: 2,
entry_generation: 2,
bytes: 4,
payload_factory: Box::new(|| Some(b"new!".to_vec())),
};
writer.enqueue_store(entry);
wait_for_queue_drain(&writer, Duration::from_secs(2));
let bytes = io.inner.load(0xB1).unwrap().expect("file exists");
let decoded = result_cache::decode_frame(&bytes).expect("frame");
assert_eq!(decoded.header.run_generation, 2);
assert_eq!(decoded.payload, b"new!");
writer.shutdown();
}
#[test]
fn schema_generation_rejected() {
let frame_v1 = result_cache::PersistedFrame {
header: result_cache::PersistedHeader {
format_version: result_cache::FRAME_FORMAT_VERSION,
table_id: 7,
schema_id: 1,
run_generation: 1,
cache_key: 1,
entry_generation: 1,
payload_len: 4,
},
payload: b"v1\0\0".to_vec(),
};
let frame_v2 = result_cache::PersistedFrame {
header: result_cache::PersistedHeader {
format_version: result_cache::FRAME_FORMAT_VERSION,
table_id: 7,
schema_id: 2,
run_generation: 1,
cache_key: 1,
entry_generation: 1,
payload_len: 4,
},
payload: b"v2\0\0".to_vec(),
};
let bytes_v1 = result_cache::encode_frame(&frame_v1);
let bytes_v2 = result_cache::encode_frame(&frame_v2);
let d1 = result_cache::decode_frame(&bytes_v1).unwrap();
let d2 = result_cache::decode_frame(&bytes_v2).unwrap();
assert_eq!(d1.header.schema_id, 1);
assert_eq!(d2.header.schema_id, 2);
}
#[test]
fn run_generation_rejected() {
let frame = result_cache::PersistedFrame {
header: result_cache::PersistedHeader {
format_version: result_cache::FRAME_FORMAT_VERSION,
table_id: 7,
schema_id: 1,
run_generation: 5,
cache_key: 1,
entry_generation: 1,
payload_len: 4,
},
payload: b"r5\0\0".to_vec(),
};
let bytes = result_cache::encode_frame(&frame);
let d = result_cache::decode_frame(&bytes).unwrap();
assert_eq!(d.header.run_generation, 5);
}
#[test]
fn corrupt_frame_rejected() {
let mut bytes = result_cache::encode_frame(&result_cache::PersistedFrame {
header: result_cache::PersistedHeader {
format_version: result_cache::FRAME_FORMAT_VERSION,
table_id: 1,
schema_id: 1,
run_generation: 1,
cache_key: 1,
entry_generation: 1,
payload_len: 4,
},
payload: b"data".to_vec(),
});
let last = bytes.len() - 1;
bytes[last] ^= 0xFF;
assert!(result_cache::decode_frame(&bytes).is_none());
let mut wrong = bytes.clone();
wrong[0] = b'X';
assert!(result_cache::decode_frame(&wrong).is_none());
}
#[test]
fn wrong_encryption_key_rejected() {
use mongreldb_core::encryption::{AesCipher, Cipher};
let key_a = [0x11u8; 32];
let key_b = [0x22u8; 32];
let cipher_a = AesCipher::new(&key_a).expect("32-byte key");
let cipher_b = AesCipher::new(&key_b).expect("32-byte key");
let plaintext = b"super-secret-payload";
let mut nonce = [0u8; 12];
mongreldb_core::encryption::fill_random(&mut nonce).expect("csprng");
let ct = cipher_a
.encrypt_page(&nonce, plaintext)
.expect("encrypt ok");
let mut on_disk = Vec::with_capacity(12 + ct.len());
on_disk.extend_from_slice(&nonce);
on_disk.extend_from_slice(&ct);
let pt_a = cipher_a.decrypt_page(&nonce, &ct).expect("decrypt ok");
assert_eq!(pt_a, plaintext);
assert!(cipher_b.decrypt_page(&nonce, &ct).is_err());
let mut bad_ct = ct.clone();
let idx = bad_ct.len() - 1;
bad_ct[idx] ^= 0xFF;
assert!(cipher_a.decrypt_page(&nonce, &bad_ct).is_err());
}
#[test]
fn crash_before_rename_leaves_no_entry() {
let dir = tempdir().unwrap();
let table_dir = dir.path().to_path_buf();
let rcache_dir = table_dir.join("_rcache");
let mut db = Table::create(&table_dir, test_schema(), 1).unwrap();
db.bulk_load(rows_city_cost(200)).unwrap();
db.flush().unwrap();
let q = alpha_query();
let _ = db.query_cached(&q).unwrap();
let _ = db.flush_persistent_cache(2_000);
for entry in std::fs::read_dir(&rcache_dir).unwrap().flatten() {
let p = entry.path();
if p.extension().and_then(|s| s.to_str()) == Some("tmp") {
panic!("leftover .tmp file after successful close: {p:?}");
}
}
}
#[test]
fn crash_after_rename_is_recoverable() {
let dir = tempdir().unwrap();
let table_dir = dir.path().to_path_buf();
let rcache_dir = table_dir.join("_rcache");
let mut db = Table::create(&table_dir, test_schema(), 1).unwrap();
db.bulk_load(rows_city_cost(200)).unwrap();
db.flush().unwrap();
let q = alpha_query();
let _ = db.query_cached(&q).unwrap();
let _ = db.flush_persistent_cache(2_000);
let mut db2 = Table::open(&table_dir).unwrap();
let q2 = alpha_query();
let r = db2.query_cached(&q2).unwrap();
assert_eq!(r.len(), 100);
}
#[test]
fn shutdown_drain() {
let dir = tempdir().unwrap();
let table_dir = dir.path().to_path_buf();
let rcache_dir = table_dir.join("_rcache");
let mut db = Table::create(&table_dir, test_schema(), 1).unwrap();
db.bulk_load(rows_city_cost(200)).unwrap();
db.flush().unwrap();
let q = alpha_query();
let _ = db.query_cached(&q).unwrap();
db.shutdown_persistent_cache(2_000);
let snap = db.lookup_metrics_snapshot();
assert_eq!(snap.result_cache_persist_queue_depth, 0);
db.shutdown_persistent_cache(100);
assert!(rcache_dir.exists());
}
#[test]
fn shutdown_deadline() {
let writer = PersistentResultCacheWriter::for_test(WriterLimits {
max_pending_keys: 1024,
max_pending_bytes: 1024 * 1024,
});
for i in 0..32u64 {
let entry = PersistableEntry {
key: i,
table_id: 0,
schema_id: 0,
run_generation: 0,
entry_generation: 0,
bytes: 0,
payload_factory: Box::new(|| Some(Vec::new())),
};
writer.enqueue_store(entry);
}
let pre = writer.persist_snapshot();
assert_eq!(pre.result_cache_persist_queue_depth, 32);
writer.shutdown();
let n = writer.drain_all_as_abandoned();
assert_eq!(n, 32);
let post = writer.persist_snapshot();
assert_eq!(post.result_cache_persist_queue_depth, 0);
assert_eq!(post.result_cache_persist_shutdown_abandoned_total, 32);
}
#[test]
fn concurrent_insert_invalidate() {
use std::thread;
let dir = tempdir().unwrap();
let table_dir = dir.path().to_path_buf();
let rcache_dir = table_dir.join("_rcache");
let mut db = Table::create(&table_dir, test_schema(), 1).unwrap();
db.bulk_load(rows_city_cost(200)).unwrap();
db.flush().unwrap();
let q = alpha_query();
for _ in 0..100 {
let _ = db.query_cached(&q).unwrap();
let rid = db.query(&q).unwrap()[0].row_id;
db.delete(rid).unwrap();
db.commit().unwrap();
}
let _ = db.flush_persistent_cache(2_000);
let snap = db.lookup_metrics_snapshot();
let stored =
snap.result_cache_persist_enqueued_total + snap.result_cache_persist_dropped_store_total;
let removed = snap.result_cache_persist_remove_total;
assert!(
stored >= 1,
"at least one store must be enqueued (got {stored})"
);
assert!(
removed >= 1,
"at least one remove must be enqueued (got {removed})"
);
let _ = rcache_dir;
}
#[test]
fn queue_byte_accounting_exact() {
let writer = PersistentResultCacheWriter::for_test(WriterLimits {
max_pending_keys: 64,
max_pending_bytes: 1000,
});
let entry = |key: u64, bytes: usize| PersistableEntry {
key,
table_id: 0,
schema_id: 0,
run_generation: 0,
entry_generation: 0,
bytes,
payload_factory: Box::new(|| Some(Vec::new())),
};
writer.enqueue_store(entry(1, 400));
writer.enqueue_store(entry(2, 400));
writer.enqueue_store(entry(3, 400));
let snap = writer.persist_snapshot();
assert_eq!(snap.result_cache_persist_dropped_store_total, 1);
assert_eq!(snap.result_cache_persist_queue_depth, 2);
writer.enqueue_store(entry(1, 200));
writer.enqueue_store(entry(3, 400));
let snap2 = writer.persist_snapshot();
assert_eq!(snap2.result_cache_persist_dropped_store_total, 1);
assert_eq!(snap2.result_cache_persist_queue_depth, 3);
writer.shutdown();
}
fn wait_for_queue_drain(writer: &PersistentResultCacheWriter, max: Duration) {
let start = Instant::now();
while writer.queue_depth() > 0 || writer.writes_in_flight() > 0 {
if start.elapsed() >= max {
panic!(
"queue did not drain within {max:?} (depth={}, in_flight={})",
writer.queue_depth(),
writer.writes_in_flight()
);
}
std::thread::sleep(Duration::from_millis(2));
}
}
use mongreldb_core::result_cache::RealPersistentCacheIo;