Skip to main content

rushdb_storage/
derive.rs

1//! Apply-time derivation.
2//!
3//! Derived rows (presence, value dictionary, counters, declared-index
4//! entries) are computed by one deterministic [`Deriver`] from a commit's
5//! *base* ops plus the state before that commit, and written in the same
6//! backend batch at the same commit index. They are never logged in the WAL
7//! and never enter a transaction's write set, so they cannot cause write
8//! conflicts. Publication, command replay, and WAL replay (crash recovery,
9//! PITR, replicas) all call [`Store::derive_ops`], so every node derives the
10//! same rows from the same base log.
11
12use std::collections::BTreeMap;
13use std::ops::Bound;
14use std::sync::Arc;
15
16use crate::backend::StorageReadView;
17use crate::read_view::{ReadLimits, ScanProjection};
18use crate::{Op, StorageError, Store};
19
20/// Read access to the state *before* a commit, for a [`Deriver`].
21pub trait DeriveRead: Sync {
22    /// Current value of `key`, or `None` when absent.
23    ///
24    /// # Errors
25    /// Returns [`StorageError`] if the underlying read fails.
26    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError>;
27
28    /// Every live `(key, value)` whose key starts with `prefix`, in key order.
29    ///
30    /// # Errors
31    /// Returns [`StorageError`] if the underlying read fails.
32    fn range_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError>;
33
34    /// Every live `(key, value)` with `start <= key < end`, in key order. An
35    /// empty `end` is unbounded above.
36    ///
37    /// # Errors
38    /// Returns [`StorageError`] if the underlying read fails.
39    fn range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError>;
40}
41
42/// Rows of `map` in `[start, end)` (empty `end` = unbounded).
43fn map_range<'a, V: 'a>(
44    map: &'a BTreeMap<Vec<u8>, V>,
45    start: &[u8],
46    end: &[u8],
47) -> impl Iterator<Item = (&'a Vec<u8>, &'a V)> {
48    let upper = if end.is_empty() {
49        Bound::Unbounded
50    } else {
51        Bound::Excluded(end.to_vec())
52    };
53    map.range::<Vec<u8>, _>((Bound::Included(start.to_vec()), upper))
54}
55
56/// Deterministic derivation of derived rows from a commit's base ops.
57///
58/// Implementations must be pure functions of `(ci, base_ops, before)`: no
59/// clock, randomness, or global state, because replay and replicas call the
60/// same function and must reach byte-identical state.
61pub trait Deriver: Send + Sync {
62    /// Materialized derived ops (final values, not deltas) for one commit.
63    ///
64    /// # Errors
65    /// Returns [`StorageError`] when the base ops or the prior state are
66    /// inconsistent; the commit (or replay) then fails.
67    fn derive(
68        &self,
69        ci: u64,
70        base_ops: &[Op],
71        before: &dyn DeriveRead,
72    ) -> Result<Vec<Op>, StorageError>;
73
74    /// Whether `key` belongs to a derived family (so an overlay must derive
75    /// its staged writes before answering a read of it).
76    fn is_derived_key(&self, key: &[u8]) -> bool;
77
78    /// Whether a scan of `prefix` can reach derived keys.
79    fn prefix_reaches_derived(&self, prefix: &[u8]) -> bool;
80}
81
82impl DeriveRead for crate::ReadView {
83    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
84        crate::ReadView::get(self, key)
85    }
86
87    fn range_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
88        crate::ReadView::range_prefix(self, prefix)
89    }
90
91    fn range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
92        crate::ReadView::range(self, start, end)
93    }
94}
95
96/// A read base plus the writes of earlier commits of the same publish group
97/// not yet visible in it (`None` = tombstone): the state before a commit
98/// inside the group. The buffer borrows the group's ops instead of copying
99/// them.
100pub(crate) struct Layered<'a, 'b> {
101    pub(crate) base: &'a dyn DeriveRead,
102    pub(crate) buffer: &'a BTreeMap<&'b [u8], Option<&'b [u8]>>,
103}
104
105impl Layered<'_, '_> {
106    fn overlay(
107        &self,
108        mut merged: BTreeMap<Vec<u8>, Vec<u8>>,
109        start: &[u8],
110        end: Bound<&[u8]>,
111    ) -> Vec<(Vec<u8>, Vec<u8>)> {
112        for (key, value) in self.buffer.range::<[u8], _>((Bound::Included(start), end)) {
113            match value {
114                Some(value) => {
115                    merged.insert(key.to_vec(), value.to_vec());
116                }
117                None => {
118                    merged.remove(*key);
119                }
120            }
121        }
122        merged.into_iter().collect()
123    }
124}
125
126impl DeriveRead for Layered<'_, '_> {
127    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
128        match self.buffer.get(key) {
129            Some(value) => Ok(value.map(<[u8]>::to_vec)),
130            None => self.base.get(key),
131        }
132    }
133
134    fn range_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
135        let merged = self.base.range_prefix(prefix)?.into_iter().collect();
136        let upper = crate::prefix_range(prefix).1;
137        let end = if upper.is_empty() {
138            Bound::Unbounded
139        } else {
140            Bound::Excluded(upper.as_slice())
141        };
142        Ok(self.overlay(merged, prefix, end))
143    }
144
145    fn range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
146        let merged = self.base.range(start, end)?.into_iter().collect();
147        let end = if end.is_empty() {
148            Bound::Unbounded
149        } else {
150            Bound::Excluded(end)
151        };
152        Ok(self.overlay(merged, start, end))
153    }
154}
155
156/// Fold ops into a borrowed layered buffer (history is irrelevant to
157/// derivation).
158pub(crate) fn fold<'b>(
159    buffer: &mut BTreeMap<&'b [u8], Option<&'b [u8]>>,
160    ops: impl IntoIterator<Item = &'b Op>,
161) {
162    for op in ops {
163        match op {
164            Op::Put(key, value) | Op::Set(key, value) => {
165                buffer.insert(key, Some(value));
166            }
167            Op::Del(key) | Op::Purge(key) => {
168                buffer.insert(key, None);
169            }
170        }
171    }
172}
173
174/// [`DeriveRead`] over one backend's current state.
175pub(crate) struct BackendRead(pub(crate) Box<dyn StorageReadView>);
176
177impl DeriveRead for BackendRead {
178    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
179        self.0.read_value(
180            key,
181            u64::MAX,
182            true,
183            ReadLimits::default(),
184            ScanProjection::KeyValue,
185        )
186    }
187
188    fn range_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
189        let end = crate::prefix_range(prefix).1;
190        self.0
191            .current_range(
192                prefix,
193                &end,
194                ReadLimits::default(),
195                ScanProjection::KeyValue,
196            )?
197            .collect()
198    }
199
200    fn range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
201        self.0
202            .current_range(start, end, ReadLimits::default(), ScanProjection::KeyValue)?
203            .collect()
204    }
205}
206
207/// [`DeriveRead`] over an in-memory key map (command-replay shadow state).
208pub(crate) struct MapRead<'a>(pub(crate) &'a BTreeMap<Vec<u8>, Vec<u8>>);
209
210impl DeriveRead for MapRead<'_> {
211    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
212        Ok(self.0.get(key).cloned())
213    }
214
215    fn range_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
216        Ok(self
217            .0
218            .range::<[u8], _>((Bound::Included(prefix), Bound::Unbounded))
219            .take_while(|(key, _)| key.starts_with(prefix))
220            .map(|(key, value)| (key.clone(), value.clone()))
221            .collect())
222    }
223
224    fn range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
225        Ok(map_range(self.0, start, end)
226            .map(|(key, value)| (key.clone(), value.clone()))
227            .collect())
228    }
229}
230
231/// [`DeriveRead`] over nothing (an empty prior state).
232pub(crate) struct Empty;
233
234impl DeriveRead for Empty {
235    fn get(&self, _key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
236        Ok(None)
237    }
238
239    fn range_prefix(&self, _prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
240        Ok(Vec::new())
241    }
242
243    fn range(&self, _start: &[u8], _end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
244        Ok(Vec::new())
245    }
246}
247
248impl Store {
249    /// Install the apply-time deriver. Install it before any commit or replay
250    /// that must produce derived rows (see [`Store::open_with_deriver`]).
251    ///
252    /// Lock poisoning fences the store until reopen.
253    pub fn set_deriver(&self, deriver: Arc<dyn Deriver>) {
254        // WAL records replayed at open without a deriver applied base rows
255        // only; their derived rows are missing. Serving reads from that state
256        // would be silently wrong, so fence the store (reopen with the
257        // deriver installed, e.g. `rushdb_lmpg::open_store`).
258        if self
259            .underived_replay
260            .load(std::sync::atomic::Ordering::SeqCst)
261        {
262            self.health.fence();
263        }
264        *self.health.cleanup(self.deriver.write()) = Some(deriver);
265    }
266
267    /// **Fault injection for tests only.** Run `f` (typically one commit of
268    /// malformed base rows) with apply-time derivation suspended, so a test
269    /// can plant corruption the deriver would refuse and prove that readers
270    /// fail on it. The deriver is restored afterwards; derived rows for the
271    /// injected commit are intentionally absent.
272    #[doc(hidden)]
273    pub fn inject_without_derivation<T>(&self, f: impl FnOnce() -> T) -> T {
274        let saved = self.health.cleanup(self.deriver.write()).take();
275        let out = f();
276        *self.health.cleanup(self.deriver.write()) = saved;
277        out
278    }
279
280    pub(crate) fn deriver(&self) -> Option<Arc<dyn Deriver>> {
281        self.health.cleanup(self.deriver.read()).clone()
282    }
283
284    /// The ops derived from `base_ops` (applied after them, at the same
285    /// commit index), or none when no deriver is installed. The base ops are
286    /// never copied: callers apply the two slices together
287    /// ([`crate::backend::CommitOps`]).
288    pub(crate) fn derive_ops(
289        deriver: Option<&Arc<dyn Deriver>>,
290        ci: u64,
291        base_ops: &[Op],
292        before: &dyn DeriveRead,
293    ) -> Result<Vec<Op>, StorageError> {
294        match deriver {
295            Some(deriver) => deriver.derive(ci, base_ops, before),
296            None => Ok(Vec::new()),
297        }
298    }
299}