kevy_store/keyspace_load.rs
1//! Snapshot / replication / reshard load hooks on [`Store`] — the
2//! `load_value` re-home dispatch and the stream loader. Split out of
3//! `keyspace.rs` (500-LOC house rule); behaviour unchanged.
4
5#[cfg(not(feature = "std"))]
6use crate::nostd_prelude::*;
7
8use alloc::sync::Arc;
9
10use crate::value::Value;
11use crate::{SmallBytes, Store};
12
13impl Store {
14 /// Insert one already-typed `(key, value, ttl)` triple, e.g. straight out
15 /// of another store's [`Self::snapshot_each`] — the redistribution step
16 /// both reshard paths (embedded `shards` bring-up, server routing
17 /// migration) use to re-home keys after a layout change.
18 // LOC-WAIVER: pure per-Value-variant dispatch table — one arm per
19 // stored type routing it to its typed loader; no control flow.
20 pub fn load_value(&mut self, key: &[u8], value: &Value, ttl_ms: Option<u64>) {
21 let k = key.to_vec();
22 match value {
23 Value::Str(v) => self.load_str(k, v.to_vec(), ttl_ms),
24 // L2: snapshot/replication load keeps the encoding — store as
25 // Int directly to preserve the in-memory shape (and avoid the
26 // SET-detect parse on the load path).
27 Value::Int(n) => self.insert_loaded(k, Value::Int(*n), ttl_ms),
28 // L1: preserve the Arc-backed encoding on snapshot/replication
29 // load. Arc::clone is cheap; avoids re-copying the bytes.
30 Value::ArcBulk(a) => self.insert_loaded(k, Value::ArcBulk(a.clone()), ttl_ms),
31 Value::Hash(h) => {
32 self.load_hash(k, h.iter().map(|(f, v)| (f.to_vec(), v.to_vec())).collect(), ttl_ms)
33 }
34 // A.8: same shape as A.7 — re-materialise to the heap-backed
35 // variant on snapshot/replication load. First mutation that
36 // targets the key will go through the encoding-switch path
37 // and (if size still fits) re-promote to the inline variant.
38 Value::SmallHashInline(h) => {
39 self.load_hash(k, h.iter().map(|(f, v)| (f.to_vec(), v.to_vec())).collect(), ttl_ms)
40 }
41 Value::SegHash(h) => {
42 self.load_hash(k, h.iter().map(|(f, v)| (f.to_vec(), v.to_vec())).collect(), ttl_ms)
43 }
44 Value::SegSet(s) => {
45 self.load_set(k, s.keys().map(|m| m.to_vec()).collect(), ttl_ms)
46 }
47 Value::List(l) => self.load_list(k, l.iter().cloned().collect(), ttl_ms),
48 Value::SegList(l) => self.load_list(k, l.iter().cloned().collect(), ttl_ms),
49 Value::SmallListInline(l) => {
50 self.load_list(k, l.iter().map(<[u8]>::to_vec).collect(), ttl_ms)
51 }
52 Value::Set(s) => {
53 self.load_set(k, s.iter().map(kevy_bytes::SmallBytes::to_vec).collect(), ttl_ms)
54 }
55 // A.7 O5: snapshot/replication load — re-materialise the
56 // inline-encoded set as a `Value::Set`-backed `KevySet`. We
57 // don't preserve the SmallSetInline encoding on reload
58 // because (a) the upgrade path will naturally rebuild it on
59 // the first SADD that targets the key, and (b) the snapshot
60 // wire format already uses the OP_SET length-prefixed
61 // payload — losing the inline encoding bit costs nothing
62 // beyond one re-promotion on the first mutation.
63 Value::SmallSetInline(s) => {
64 self.load_set(k, s.iter_slices().map(<[u8]>::to_vec).collect(), ttl_ms)
65 }
66 Value::ZSet(z) => {
67 self.load_zset(k, z.ordered().map(|(m, sc)| (m.to_vec(), sc)).collect(), ttl_ms)
68 }
69 Value::SegZSet(z) => {
70 self.load_zset(k, z.ordered().map(|(m, sc)| (m.to_vec(), sc)).collect(), ttl_ms)
71 }
72 Value::SmallZSetInline(z) => {
73 self.load_zset(k, z.iter().map(|(m, sc)| (m.to_vec(), sc)).collect(), ttl_ms)
74 }
75 Value::Stream(st) => self.load_stream_value(k, st, ttl_ms),
76 // Unreachable by construction: every producer of a shipped
77 // value (take_with_ttl / clone_with_ttl / snapshot_each
78 // consumers) materializes cold values on ITS side —
79 // a ColdRef names the source shard's vlog, which this store
80 // cannot read. Skip rather than alias a foreign record;
81 // loud in debug builds.
82 Value::Cold(_) => debug_assert!(
83 false,
84 "load_value received a cold stub — source must materialize before shipping"
85 ),
86 }
87 }
88
89 /// [`Self::load_value`]'s stream arm: decode the live `StreamData`
90 /// into the primitive tuples [`Self::load_stream`] takes.
91 fn load_stream_value(&mut self, k: Vec<u8>, st: &crate::StreamData, ttl_ms: Option<u64>) {
92 let entries: Vec<crate::stream::LoadedStreamEntry> = st
93 .iter_entries()
94 .map(|(id, fv)| {
95 let fvv = fv
96 .iter()
97 .map(|(f, v)| (f.as_slice().to_vec(), v.as_slice().to_vec()))
98 .collect();
99 (id.ms, id.seq, fvv)
100 })
101 .collect();
102 let last = st.last_id();
103 let mxd = st.max_deleted_id();
104 self.load_stream(
105 k,
106 entries,
107 (last.ms, last.seq),
108 (mxd.ms, mxd.seq),
109 st.entries_added(),
110 st.export_groups(),
111 ttl_ms,
112 );
113 }
114
115 /// Snapshot-load a stream: every entry plus the per-stream scalar
116 /// state (last_id, max_deleted_id, entries_added) and the consumer
117 /// groups are restored verbatim. Caller passes already-decoded
118 /// primitive tuples; this fn does the [`SmallBytes`] /
119 /// [`crate::StreamData`] conversion.
120 #[allow(clippy::too_many_arguments)]
121 pub fn load_stream(
122 &mut self,
123 key: Vec<u8>,
124 entries: Vec<crate::stream::LoadedStreamEntry>,
125 last_id: (u64, u64),
126 max_deleted_id: (u64, u64),
127 entries_added: u64,
128 groups: Vec<crate::stream::LoadedGroup>,
129 ttl_ms: Option<u64>,
130 ) {
131 let mut s = crate::stream::StreamData::default();
132 for (ms, seq, fv) in entries {
133 let id = crate::stream::StreamId { ms, seq };
134 let fv_small: Vec<(SmallBytes, SmallBytes)> = fv
135 .into_iter()
136 .map(|(f, v)| (SmallBytes::from_vec(f), SmallBytes::from_vec(v)))
137 .collect();
138 s.load_entry(id, fv_small);
139 }
140 s.set_loaded_state(
141 crate::stream::StreamId { ms: last_id.0, seq: last_id.1 },
142 crate::stream::StreamId { ms: max_deleted_id.0, seq: max_deleted_id.1 },
143 entries_added,
144 );
145 s.import_groups(groups);
146 self.insert_loaded(key, Value::Stream(Arc::new(s)), ttl_ms);
147 }
148}