nodedb_lite/nodedb/core/kv_local.rs
1// SPDX-License-Identifier: Apache-2.0
2
3//! KV state the public API keeps in front of storage: the write buffer with
4//! its read overlay, the read cache, and the write-order lock.
5//!
6//! `NodeDbLite` and the query engine share it:
7//! - a SQL-path `TRUNCATE` forgets what the buffer and cache hold for the
8//! collection it cleared, so a buffered put cannot reach storage on the
9//! next flush and a cached value is not served after its row is gone;
10//! - public-API and SQL writes take one write-order lock, so both enqueue
11//! for sync in the order they applied.
12
13use std::collections::HashMap;
14use std::num::NonZeroUsize;
15use std::sync::Mutex;
16
17use crate::error::LiteError;
18use crate::nodedb::lock_ext::LockExt;
19use crate::storage::engine::{StorageEngine, WriteOp};
20
21/// Buffered KV writes for batch commit.
22///
23/// All public KV read and write methods acquire `Mutex<KvWriteBuffer>` before
24/// inspecting or mutating the overlay, so every read-through-overlay access is
25/// serialized against concurrent writes. Reads always lock — there is no
26/// lock-free fast path.
27pub(crate) struct KvWriteBuffer {
28 /// Pending write operations for batch commit.
29 pub ops: Vec<WriteOp>,
30 /// Read overlay: maps composite KV key → value (None = deleted).
31 /// Lets `kv_get` see uncommitted writes without hitting storage.
32 pub overlay: HashMap<Vec<u8>, Option<Vec<u8>>>,
33}
34
35/// The write buffer and read cache of the public KV API.
36pub struct KvLocalState {
37 /// Buffered KV writes awaiting batch commit to storage.
38 /// `kv_flush()`, the 1024-operation threshold, and `flush()` drain this buffer.
39 /// The HashMap overlay lets reads see uncommitted writes.
40 pub(crate) write_buf: Mutex<KvWriteBuffer>,
41 /// In-memory LRU cache for the KV get hot path.
42 ///
43 /// Stores raw encoded bytes (8-byte LE deadline + user value) keyed by the
44 /// composite KV key (`{collection}\0{user_key}`). TTL expiry is re-checked
45 /// on every cache hit so no entry is served past its deadline.
46 ///
47 /// Capacity is controlled by [`crate::config::LiteConfig::kv_cache_capacity`].
48 pub(crate) cache: Mutex<lru::LruCache<Vec<u8>, Vec<u8>>>,
49 /// Serializes every KV write that syncs to Origin, public API and SQL.
50 ///
51 /// A write holds it from its local apply through its outbound enqueue,
52 /// so the outbound queue holds writes in the order they were applied.
53 /// Without it, two writes to one key could apply in one order and
54 /// enqueue in the other, and Origin would keep the value Lite overwrote.
55 /// It also makes a read-modify-write such as `kv_increment` atomic.
56 pub(crate) write_order: tokio::sync::Mutex<()>,
57 /// Secondary indexes every commit of the buffer maintains. Set once by
58 /// the query engine.
59 indexes: std::sync::OnceLock<std::sync::Arc<crate::index::IndexCatalog>>,
60}
61
62impl KvLocalState {
63 pub(crate) fn new(cache_capacity: NonZeroUsize) -> Self {
64 Self {
65 write_buf: Mutex::new(KvWriteBuffer {
66 ops: Vec::with_capacity(1024),
67 overlay: HashMap::new(),
68 }),
69 cache: Mutex::new(lru::LruCache::new(cache_capacity)),
70 write_order: tokio::sync::Mutex::new(()),
71 indexes: std::sync::OnceLock::new(),
72 }
73 }
74
75 /// Maintain `catalog`'s key-value indexes from every buffer commit.
76 pub(crate) fn set_index_catalog(&self, catalog: std::sync::Arc<crate::index::IndexCatalog>) {
77 // Set once, when the query engine is built; a second call keeps the
78 // first catalog.
79 let _ = self.indexes.set(catalog);
80 }
81
82 /// Whether writes to `collection` maintain a key-value index. Such writes
83 /// commit directly rather than wait in the buffer, so a unique index
84 /// refuses the write that breaks it instead of the batch it lands in.
85 pub(crate) fn is_indexed(&self, collection: &str) -> bool {
86 self.indexes
87 .get()
88 .is_some_and(|c| c.has_defs(collection, crate::index::IndexEngine::KeyValue))
89 }
90
91 /// Refuse KV writes `ops` a unique index forbids, writing nothing.
92 pub(crate) async fn check<S: StorageEngine>(
93 &self,
94 storage: &S,
95 ops: &[WriteOp],
96 ) -> Result<(), LiteError> {
97 match self.indexes.get() {
98 Some(catalog) => crate::query::kv_ops::indexes::kv_check(storage, catalog, ops).await,
99 None => Ok(()),
100 }
101 }
102
103 /// Commit KV writes `ops` with the index entries they imply. Every KV
104 /// write reaches storage through here.
105 pub(crate) async fn commit<S: StorageEngine>(
106 &self,
107 storage: &S,
108 ops: Vec<WriteOp>,
109 ) -> Result<(), LiteError> {
110 match self.indexes.get() {
111 Some(catalog) => crate::query::kv_ops::indexes::kv_commit(storage, catalog, ops).await,
112 None if ops.is_empty() => Ok(()),
113 None => storage.batch_write(&ops).await,
114 }
115 }
116
117 /// Commit every buffered write to `storage` as one batch. Returns the
118 /// number of writes committed.
119 ///
120 /// A SQL KV write calls this before it runs, so a buffered public-API
121 /// write to the same key cannot land after it and overwrite it.
122 pub(crate) async fn flush_to<S: StorageEngine>(&self, storage: &S) -> Result<usize, LiteError> {
123 let ops: Vec<WriteOp> = {
124 let mut buf = self.write_buf.lock_or_recover();
125 if buf.ops.is_empty() {
126 return Ok(0);
127 }
128 let ops = std::mem::take(&mut buf.ops);
129 buf.overlay.clear();
130 ops
131 };
132 let count = ops.len();
133 self.commit(storage, ops).await?;
134 Ok(count)
135 }
136
137 /// Drop every buffered write and cached value whose key belongs to
138 /// `collection`. Returns the number of buffered writes dropped.
139 pub(crate) fn forget_collection(&self, collection: &str) -> usize {
140 let mut prefix = collection.as_bytes().to_vec();
141 prefix.push(0);
142 let dropped = {
143 let mut buf = self.write_buf.lock_or_recover();
144 let before = buf.ops.len();
145 buf.ops.retain(|op| {
146 let key = match op {
147 WriteOp::Put { key, .. } | WriteOp::Delete { key, .. } => key,
148 };
149 !key.starts_with(&prefix)
150 });
151 buf.overlay.retain(|key, _| !key.starts_with(&prefix));
152 before - buf.ops.len()
153 };
154 let mut cache = self.cache.lock_or_recover();
155 let stale: Vec<Vec<u8>> = cache
156 .iter()
157 .filter(|(key, _)| key.starts_with(&prefix))
158 .map(|(key, _)| key.clone())
159 .collect();
160 for key in stale {
161 cache.pop(&key);
162 }
163 dropped
164 }
165}
166
167#[cfg(test)]
168mod tests {
169 use nodedb_types::Namespace;
170
171 use super::*;
172
173 fn key(collection: &str, k: &str) -> Vec<u8> {
174 let mut key = collection.as_bytes().to_vec();
175 key.push(0);
176 key.extend_from_slice(k.as_bytes());
177 key
178 }
179
180 #[test]
181 fn forget_collection_drops_only_that_collection() {
182 let state = KvLocalState::new(NonZeroUsize::new(8).expect("cap"));
183 {
184 let mut buf = state.write_buf.lock_or_recover();
185 for (c, k) in [("a", "1"), ("a", "2"), ("b", "1")] {
186 buf.ops.push(WriteOp::Put {
187 ns: Namespace::Kv,
188 key: key(c, k),
189 value: vec![1],
190 });
191 buf.overlay.insert(key(c, k), Some(vec![1]));
192 }
193 buf.ops.push(WriteOp::Delete {
194 ns: Namespace::Kv,
195 key: key("a", "3"),
196 });
197 }
198 {
199 let mut cache = state.cache.lock_or_recover();
200 cache.put(key("a", "1"), vec![1]);
201 cache.put(key("b", "1"), vec![1]);
202 }
203
204 assert_eq!(state.forget_collection("a"), 3);
205 let buf = state.write_buf.lock_or_recover();
206 assert_eq!(buf.ops.len(), 1);
207 assert_eq!(buf.overlay.len(), 1);
208 assert!(buf.overlay.contains_key(&key("b", "1")));
209 let cache = state.cache.lock_or_recover();
210 assert_eq!(cache.len(), 1);
211 assert!(cache.peek(&key("b", "1")).is_some());
212 }
213}