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 // Column order comes from the declaration, not from the shard, so
32 // a re-home keeps the packed encoding as it stands.
33 Value::PackedRow(r) => self.insert_loaded(k, Value::PackedRow(r.clone()), ttl_ms),
34 Value::Hash(h) => {
35 self.load_hash(k, h.iter().map(|(f, v)| (f.to_vec(), v.to_vec())).collect(), ttl_ms)
36 }
37 // A.8: same shape as A.7 — re-materialise to the heap-backed
38 // variant on snapshot/replication load. First mutation that
39 // targets the key will go through the encoding-switch path
40 // and (if size still fits) re-promote to the inline variant.
41 Value::SmallHashInline(h) => {
42 self.load_hash(k, h.iter().map(|(f, v)| (f.to_vec(), v.to_vec())).collect(), ttl_ms)
43 }
44 Value::SegHash(h) => {
45 self.load_hash(k, h.iter().map(|(f, v)| (f.to_vec(), v.to_vec())).collect(), ttl_ms)
46 }
47 Value::SegSet(s) => self.load_set(k, s.keys().map(|m| m.to_vec()).collect(), ttl_ms),
48 Value::List(l) => self.load_list(k, l.iter().cloned().collect(), ttl_ms),
49 Value::SegList(l) => self.load_list(k, l.iter().cloned().collect(), ttl_ms),
50 Value::SmallListInline(l) => {
51 self.load_list(k, l.iter().map(<[u8]>::to_vec).collect(), ttl_ms)
52 }
53 Value::Set(s) => {
54 self.load_set(k, s.iter().map(kevy_bytes::SmallBytes::to_vec).collect(), ttl_ms)
55 }
56 // A.7 O5: snapshot/replication load — re-materialise the
57 // inline-encoded set as a `Value::Set`-backed `KevySet`. We
58 // don't preserve the SmallSetInline encoding on reload
59 // because (a) the upgrade path will naturally rebuild it on
60 // the first SADD that targets the key, and (b) the snapshot
61 // wire format already uses the OP_SET length-prefixed
62 // payload — losing the inline encoding bit costs nothing
63 // beyond one re-promotion on the first mutation.
64 Value::SmallSetInline(s) => {
65 self.load_set(k, s.iter_slices().map(<[u8]>::to_vec).collect(), ttl_ms)
66 }
67 Value::ZSet(z) => {
68 self.load_zset(k, z.ordered().map(|(m, sc)| (m.to_vec(), sc)).collect(), ttl_ms)
69 }
70 Value::SegZSet(z) => {
71 self.load_zset(k, z.ordered().map(|(m, sc)| (m.to_vec(), sc)).collect(), ttl_ms)
72 }
73 Value::SmallZSetInline(z) => {
74 self.load_zset(k, z.iter().map(|(m, sc)| (m.to_vec(), sc)).collect(), ttl_ms)
75 }
76 Value::Stream(st) => self.load_stream_value(k, st, ttl_ms),
77 // Unreachable by construction: every producer of a shipped
78 // value (take_with_ttl / clone_with_ttl / snapshot_each
79 // consumers) materializes cold values on ITS side —
80 // a ColdRef names the source shard's vlog, which this store
81 // cannot read. Skip rather than alias a foreign record;
82 // loud in debug builds.
83 Value::Cold(_) => debug_assert!(
84 false,
85 "load_value received a cold stub — source must materialize before shipping"
86 ),
87 }
88 }
89
90 /// [`Self::load_value`]'s stream arm: decode the live `StreamData`
91 /// into the primitive tuples [`Self::load_stream`] takes.
92 fn load_stream_value(&mut self, k: Vec<u8>, st: &crate::StreamData, ttl_ms: Option<u64>) {
93 let entries: Vec<crate::stream::LoadedStreamEntry> = st
94 .iter_entries()
95 .map(|(id, fv)| {
96 let fvv = fv
97 .iter()
98 .map(|(f, v)| (f.as_slice().to_vec(), v.as_slice().to_vec()))
99 .collect();
100 (id.ms, id.seq, fvv)
101 })
102 .collect();
103 let last = st.last_id();
104 let mxd = st.max_deleted_id();
105 self.load_stream(
106 k,
107 entries,
108 (last.ms, last.seq),
109 (mxd.ms, mxd.seq),
110 st.entries_added(),
111 st.export_groups(),
112 ttl_ms,
113 );
114 }
115
116 /// Snapshot-load a stream: every entry plus the per-stream scalar
117 /// state (last_id, max_deleted_id, entries_added) and the consumer
118 /// groups are restored verbatim. Caller passes already-decoded
119 /// primitive tuples; this fn does the [`SmallBytes`] /
120 /// [`crate::StreamData`] conversion.
121 #[allow(clippy::too_many_arguments)]
122 pub fn load_stream(
123 &mut self,
124 key: Vec<u8>,
125 entries: Vec<crate::stream::LoadedStreamEntry>,
126 last_id: (u64, u64),
127 max_deleted_id: (u64, u64),
128 entries_added: u64,
129 groups: Vec<crate::stream::LoadedGroup>,
130 ttl_ms: Option<u64>,
131 ) {
132 let mut s = crate::stream::StreamData::default();
133 for (ms, seq, fv) in entries {
134 let id = crate::stream::StreamId { ms, seq };
135 let fv_small: Vec<(SmallBytes, SmallBytes)> = fv
136 .into_iter()
137 .map(|(f, v)| (SmallBytes::from_vec(f), SmallBytes::from_vec(v)))
138 .collect();
139 s.load_entry(id, fv_small);
140 }
141 s.set_loaded_state(
142 crate::stream::StreamId { ms: last_id.0, seq: last_id.1 },
143 crate::stream::StreamId { ms: max_deleted_id.0, seq: max_deleted_id.1 },
144 entries_added,
145 );
146 s.import_groups(groups);
147 self.insert_loaded(key, Value::Stream(Arc::new(s)), ttl_ms);
148 }
149}