yantrikdb 0.7.1

Cognitive memory engine for persistent AI systems
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
//! Decoupled write path RFC, Phase 3.5 — background materializer threads.
//!
//! Spawns N worker threads that continuously drain pending oplog entries
//! (applied=0) and dispatch them via `apply_pending_ops_once`. Returns
//! [`MaterializerGuard`] handles whose `Drop` impl signals shutdown and
//! joins the threads.
//!
//! The workers hold a `Weak<YantrikDB>` reference, so dropping the engine
//! also lets the workers exit cleanly even if guards are leaked.
//!
//! Phase 4 will add explicit notify/wakeup so foreground `record()` can
//! signal the workers immediately rather than waiting for the 100ms timer.

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Weak};
use std::thread::JoinHandle;
use std::time::Duration;

use super::YantrikDB;

/// How many ops a worker drains in a single pass before yielding.
/// Tuned to amortize per-pass overhead while keeping individual workers
/// responsive to shutdown.
const DRAIN_BATCH_SIZE: usize = 64;

/// Sleep duration when no pending work is found.
/// Phase 4 replaces this with a condvar wake on log_op_pending.
const IDLE_POLL_INTERVAL: Duration = Duration::from_millis(100);

/// Sleep duration after a drain error, before retry.
/// Prevents tight loops on persistent failure (e.g. disk full).
const ERROR_BACKOFF_INTERVAL: Duration = Duration::from_millis(100);

/// Owns one materializer thread + its shutdown flag.
///
/// Drop semantics:
/// - sets shutdown=true so the worker exits its next loop iteration
/// - joins the JoinHandle so the thread fully terminates before this guard goes out of scope
///
/// Joining can take up to `IDLE_POLL_INTERVAL` (100ms) in the worst case
/// since the worker only checks shutdown between drain passes. If you need
/// faster shutdown, drop the engine's `Arc<YantrikDB>` first (Weak::upgrade
/// will fail next iteration).
pub struct MaterializerGuard {
    shutdown: Arc<AtomicBool>,
    handle: Option<JoinHandle<()>>,
}

impl Drop for MaterializerGuard {
    fn drop(&mut self) {
        self.shutdown.store(true, Ordering::Relaxed);
        if let Some(handle) = self.handle.take() {
            let _ = handle.join();
        }
    }
}

/// Spawn `count` materializer threads against a shared `YantrikDB`.
///
/// Each worker holds a `Weak<YantrikDB>`. When the last `Arc<YantrikDB>` is
/// dropped, workers exit on their next iteration. Drop the returned guards
/// (or let them go out of scope) for an explicit shutdown signal.
///
/// **Recommendation:** `count = std::thread::available_parallelism().map(|n| n.get() / 2).unwrap_or(2).clamp(2, 16)`.
/// The RFC pins this default; tests may pass smaller values.
pub fn spawn_materializers(db: &Arc<YantrikDB>, count: usize) -> Vec<MaterializerGuard> {
    let mut guards = Vec::with_capacity(count);

    for worker_id in 0..count {
        let shutdown = Arc::new(AtomicBool::new(false));
        let shutdown_clone = Arc::clone(&shutdown);
        let weak = Arc::downgrade(db);

        let handle = std::thread::Builder::new()
            .name(format!("yantrikdb-materializer-{worker_id}"))
            .spawn(move || worker_loop(weak, shutdown_clone, worker_id))
            .expect("spawn materializer thread");

        guards.push(MaterializerGuard {
            shutdown,
            handle: Some(handle),
        });
    }

    guards
}

/// Worker main loop.
///
/// Each iteration:
/// 1. Upgrade the weak ref to YantrikDB. Exit if engine has been dropped.
/// 2. Check shutdown flag. Exit if set.
/// 3. Call apply_pending_ops_once(DRAIN_BATCH_SIZE).
/// 4. If applied 0 ops → sleep IDLE_POLL_INTERVAL.
/// 5. If applied N > 0 ops → loop immediately to drain more.
/// 6. On error → log + sleep ERROR_BACKOFF_INTERVAL.
fn worker_loop(weak: Weak<YantrikDB>, shutdown: Arc<AtomicBool>, worker_id: usize) {
    tracing::debug!(worker_id, "materializer worker started");

    while !shutdown.load(Ordering::Relaxed) {
        let Some(db) = weak.upgrade() else {
            tracing::debug!(worker_id, "engine dropped — materializer exiting");
            break;
        };

        match db.apply_pending_ops_once(DRAIN_BATCH_SIZE) {
            Ok(0) => {
                // No work — release the strong ref before sleeping so the
                // engine can be dropped during our sleep without blocking.
                drop(db);
                std::thread::sleep(IDLE_POLL_INTERVAL);
            }
            Ok(n) => {
                tracing::trace!(worker_id, applied = n, "drained batch");
                drop(db);
                // Loop immediately — drain more if available.
            }
            Err(e) => {
                tracing::warn!(worker_id, error = %e, "drain failed; backing off");
                drop(db);
                std::thread::sleep(ERROR_BACKOFF_INTERVAL);
            }
        }
    }

    tracing::debug!(worker_id, "materializer worker exited");
}

/// Recommended worker count for the current host.
///
/// `cores / 2`, clamped to [2, 16] per the RFC. Stays modest because each
/// worker does both SQLite reads (oplog SELECT) and writes (UPDATE applied=1)
/// — too many workers fight for the conn mutex.
pub fn recommended_worker_count() -> usize {
    std::thread::available_parallelism()
        .map(|n| n.get() / 2)
        .unwrap_or(2)
        .clamp(2, 16)
}

/// **Decoupled write path RFC, Phase 5 — compactor.**
///
/// Background thread that periodically calls `db.vec_index.compact()` to
/// drain the delta tier into the cold tier, bounding read latency growth.
///
/// Polls every `COMPACTOR_INTERVAL` (250ms by default in v0.6.7+,
/// was 1s in v0.6.6). On each tick:
///   1. Check `should_compact()` — fires when delta is past half-capacity
///      OR the oldest dirty entry has aged past `max_dirty_age` (default
///      60s; epic 5 task 13). The age trigger closes the gap where a
///      low-write namespace's delta sits forever and reads pay the
///      linear scan indefinitely.
///   2. Run `compact()` — clone-rebuild cold from old cold + sealed delta,
///      then ArcSwap the new cold in.
///
/// Drop the returned [`CompactorGuard`] (or let the engine `Arc<YantrikDB>`
/// drop) for clean shutdown.

pub struct CompactorGuard {
    shutdown: Arc<AtomicBool>,
    handle: Option<JoinHandle<()>>,
}

impl Drop for CompactorGuard {
    fn drop(&mut self) {
        self.shutdown.store(true, Ordering::Relaxed);
        if let Some(handle) = self.handle.take() {
            let _ = handle.join();
        }
    }
}

/// How often the compactor wakes to check `should_compact()`.
///
/// **v0.6.7+ default 250ms (was 1s in v0.6.6).** Tuned in response to
/// the 2026-05-08 32-writer empirical run (saga task 18) showing read
/// p99 regression under sustained writer pressure: `compact()`'s
/// per-cycle `(*cold.load_full()).clone()` dominates cold-grow cost,
/// and at 1s intervals the compactor cycle (~5-20ms at cold=7600
/// entries) couldn't keep pace with delta fill rate, so writes
/// back-pressured and reads stalled behind delta.read() during seal.
///
/// At 250ms the compactor wakes 4× more often, drains smaller batches,
/// produces shorter pauses. Total clone work per second is unchanged
/// (same number of entries cycle through delta-then-cold), but the
/// p99 tail spreads thinner.
///
/// CPU cost trade-off: under sustained writer pressure the compactor
/// is now near-constantly busy. That's intentional — it is a P3
/// background worker (see CONCURRENCY.md Rule 3) and giving it a
/// CPU is what the priority hierarchy says to do when the engine is
/// under load. Idle deployments still pay near-zero CPU because
/// `should_compact()` short-circuits when delta is empty.
const COMPACTOR_INTERVAL: Duration = Duration::from_millis(250);

pub fn spawn_compactor(db: &Arc<YantrikDB>) -> CompactorGuard {
    let shutdown = Arc::new(AtomicBool::new(false));
    let shutdown_clone = Arc::clone(&shutdown);
    let weak = Arc::downgrade(db);

    let handle = std::thread::Builder::new()
        .name("yantrikdb-compactor".to_string())
        .spawn(move || compactor_loop(weak, shutdown_clone))
        .expect("spawn compactor thread");

    CompactorGuard {
        shutdown,
        handle: Some(handle),
    }
}

fn compactor_loop(weak: Weak<YantrikDB>, shutdown: Arc<AtomicBool>) {
    tracing::debug!("compactor started");

    while !shutdown.load(Ordering::Relaxed) {
        let Some(db) = weak.upgrade() else {
            tracing::debug!("engine dropped — compactor exiting");
            break;
        };

        if db.vec_index.should_compact() {
            match db.vec_index.compact() {
                Ok(0) => {}
                Ok(n) => {
                    tracing::debug!(applied = n, "compaction drained delta into cold");
                }
                Err(e) => {
                    tracing::warn!(error = %e, "compaction failed; retrying next tick");
                }
            }
        }

        drop(db);
        std::thread::sleep(COMPACTOR_INTERVAL);
    }

    tracing::debug!("compactor exited");
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::YantrikDB;

    fn open_test_db() -> Arc<YantrikDB> {
        Arc::new(YantrikDB::new(":memory:", 64).expect("open"))
    }

    #[test]
    fn worker_drains_pending_ops() {
        let db = open_test_db();

        // Push 5 pending ops.
        for i in 0..5 {
            db.log_op_pending(
                "record",
                Some(&format!("rid_{i}")),
                &serde_json::json!({}),
                None,
                None,
            )
            .unwrap();
        }
        assert_eq!(db.count_pending_ops().unwrap(), 5);

        // Spawn one worker.
        let _guards = spawn_materializers(&db, 1);

        // Worker should drain within ~150ms (one IDLE_POLL_INTERVAL).
        let mut tries = 0;
        while db.count_pending_ops().unwrap() > 0 && tries < 20 {
            std::thread::sleep(Duration::from_millis(50));
            tries += 1;
        }
        assert_eq!(
            db.count_pending_ops().unwrap(),
            0,
            "worker should drain all 5 pending ops within 1s"
        );
    }

    #[test]
    fn guard_drop_shuts_down_worker() {
        let db = open_test_db();
        let guards = spawn_materializers(&db, 2);
        // Push some work.
        for i in 0..3 {
            db.log_op_pending(
                "record",
                Some(&format!("rid_{i}")),
                &serde_json::json!({}),
                None,
                None,
            )
            .unwrap();
        }
        std::thread::sleep(Duration::from_millis(200));

        // Drop the guards — workers should exit cleanly within ~200ms.
        let start = std::time::Instant::now();
        drop(guards);
        let elapsed = start.elapsed();
        assert!(
            elapsed < Duration::from_secs(1),
            "guard drop should join workers within 1s, took {elapsed:?}"
        );
    }

    #[test]
    fn engine_drop_lets_workers_exit_via_weak_upgrade_fail() {
        let db = open_test_db();
        let guards = spawn_materializers(&db, 1);
        // Drop the engine while the worker is presumably sleeping.
        drop(db);

        // Now drop the guards. They should join quickly because the worker
        // saw weak.upgrade() fail and exited.
        let start = std::time::Instant::now();
        drop(guards);
        let elapsed = start.elapsed();
        assert!(
            elapsed < Duration::from_secs(1),
            "engine drop should let worker exit within 1s, took {elapsed:?}"
        );
    }

    #[test]
    fn multiple_workers_dont_double_apply() {
        // Two workers racing on the same pending ops should still result
        // in each op applied exactly once. apply_pending_ops_once is
        // idempotent on op_id (mark_op_applied + the SELECT WHERE applied=0
        // filter), so even if both workers SELECT the same op, only one
        // will see it through to applied=1; the other's UPDATE is a no-op.
        let db = open_test_db();
        for i in 0..50 {
            db.log_op_pending(
                "record",
                Some(&format!("rid_{i}")),
                &serde_json::json!({}),
                None,
                None,
            )
            .unwrap();
        }

        let _guards = spawn_materializers(&db, 4);

        // Wait for drain.
        let mut tries = 0;
        while db.count_pending_ops().unwrap() > 0 && tries < 40 {
            std::thread::sleep(Duration::from_millis(50));
            tries += 1;
        }
        assert_eq!(db.count_pending_ops().unwrap(), 0);

        // Verify all 50 ops are now applied=1 (no duplicates / loss).
        let conn = db.read_conn();
        let total: i64 = conn
            .query_row(
                "SELECT COUNT(*) FROM oplog WHERE applied = 1 AND op_type = 'record'",
                [],
                |row| row.get(0),
            )
            .unwrap();
        assert_eq!(total, 50);
    }

    #[test]
    fn recommended_worker_count_in_range() {
        let n = recommended_worker_count();
        assert!((2..=16).contains(&n), "expected [2,16], got {n}");
    }

    #[test]
    fn compactor_drains_delta_periodically() {
        // Test scaled to DEFAULT_DELTA_MAX = 256 (v0.6.7+):
        //   half-cap (compaction threshold) = 128
        //   backpressure trigger             = 256
        //
        // Push 250 entries — past half-cap so the compactor's
        // `should_compact` returns true, but under the backpressure
        // ceiling so the burst fits in the delta. Compactor wakes every
        // COMPACTOR_INTERVAL (250ms in v0.6.7+); we wait up to 4s for at least 200 to
        // land in cold.
        let db = open_test_db();
        let _guard = spawn_compactor(&db);

        for i in 0..250 {
            let emb: Vec<f32> = (0..64).map(|j| (i + j) as f32 * 0.001).collect();
            let norm: f32 = emb.iter().map(|x| x * x).sum::<f32>().sqrt();
            let normalized: Vec<f32> = emb.iter().map(|x| x / norm).collect();
            let seq = db.vec_seq.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
            db.vec_index.append(format!("rid_{i}"), normalized, seq).unwrap();
        }
        assert!(
            db.vec_index.delta_len() >= 128,
            "delta should have crossed compact threshold (half of 256), got {}",
            db.vec_index.delta_len()
        );

        let mut tries = 0;
        while db.vec_index.cold_len() < 200 && tries < 40 {
            std::thread::sleep(Duration::from_millis(100));
            tries += 1;
        }

        assert!(
            db.vec_index.cold_len() >= 200,
            "compactor should have moved >=200 entries to cold within 4s, got cold={} delta={}",
            db.vec_index.cold_len(),
            db.vec_index.delta_len()
        );
        assert!(
            db.vec_index.delta_len() < 128,
            "delta should be drained below half-cap, got {}",
            db.vec_index.delta_len()
        );
    }
    #[test]
    fn compactor_guard_drop_shuts_down_clean() {
        let db = open_test_db();
        let guard = spawn_compactor(&db);
        std::thread::sleep(Duration::from_millis(100));
        let start = std::time::Instant::now();
        drop(guard);
        assert!(
            start.elapsed() < Duration::from_secs(2),
            "compactor guard drop must join within 2s"
        );
    }
}