sui_core/
rpc_store_restore_source.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! [`RestoreSource`] backed by a validator's
5//! [`AuthorityPerpetualTables`].
6//!
7//! Streams every `LiveObject::Normal` in the perpetual store
8//! into the `sui-consistent-store` restore driver, sharded by
9//! `ObjectID` prefix so multiple shards can iterate in parallel.
10//!
11//! # Sharding
12//!
13//! The `ObjectID` space is split into 32 shards by the top
14//! `SHARD_BITS = 5` bits of the first byte.
15//! Each shard yields chunks of [`CHUNK_SIZE`] objects; the
16//! `RestoreChunk::cursor` is the 32-byte ObjectID of the last
17//! object in that chunk, so resuming with `Some(c)` starts the
18//! next iteration immediately after that id.
19//!
20//! # Snapshot consistency
21//!
22//! Each shard's stream opens exactly one RocksDB iterator and
23//! drives it to completion from a single `spawn_blocking` task,
24//! pushing chunks back over a tokio mpsc. RocksDB iterators
25//! created without an explicit snapshot implicitly pin one at
26//! construction time, so a shard sees a single point-in-time
27//! view for its full run — including the merge-based `balance`
28//! pipeline, which is safe against concurrent execution.
29//!
30//! Different shards take their snapshots at the moments their
31//! `spawn_blocking` tasks start, so cross-shard skew can still
32//! exist if the validator commits between shard launches. This
33//! does not affect any of the `sui-rpc-store` pipelines because
34//! every object lives in exactly one shard.
35//!
36//! A side-effect of holding open one iterator per shard for the
37//! full restore is that the SSTs it references stay pinned and
38//! cannot compact away for the duration. That is acceptable for
39//! a one-shot bootstrap.
40
41use std::sync::Arc;
42
43use async_trait::async_trait;
44use bytes::Bytes;
45use futures::StreamExt;
46use futures::stream;
47use futures::stream::BoxStream;
48use sui_consistent_store::ChainId;
49use sui_consistent_store::restore::RestoreChunk;
50use sui_consistent_store::restore::RestoreSource;
51use sui_types::base_types::ObjectID;
52use sui_types::object::Object;
53use tokio::sync::mpsc;
54use tokio_stream::wrappers::ReceiverStream;
55
56use crate::authority::authority_store_tables::AuthorityPerpetualTables;
57use crate::authority::authority_store_tables::LiveObject;
58
59/// Bits of the first `ObjectID` byte used to choose a shard.
60/// `1 << SHARD_BITS` shards.
61const SHARD_BITS: u32 = 5;
62
63/// Total number of shards (`1 << SHARD_BITS`).
64const SHARDS: u32 = 1 << SHARD_BITS;
65
66/// Bit shift placing the shard id in the high bits of the first
67/// `ObjectID` byte.
68const SHARD_PREFIX_SHIFT: u32 = 8 - SHARD_BITS;
69
70/// Default objects per [`RestoreChunk`]. Tuned to keep the
71/// per-pipeline batch comfortably under a few MB of writes
72/// while still amortising the per-chunk commit overhead.
73pub const CHUNK_SIZE: usize = 50_000;
74
75/// [`RestoreSource`] over an
76/// [`AuthorityPerpetualTables`]. Construct via
77/// [`PerpetualStoreRestoreSource::new`].
78pub struct PerpetualStoreRestoreSource {
79    perpetual: Arc<AuthorityPerpetualTables>,
80    target_checkpoint: u64,
81    chain_id: ChainId,
82    chunk_size: usize,
83}
84
85impl PerpetualStoreRestoreSource {
86    /// Build a source rooted at `perpetual`, anchored at
87    /// `target_checkpoint` and `chain_id`. Tip indexing will
88    /// resume at `target_checkpoint + 1` once restore finishes
89    /// — pick the highest executed checkpoint the perpetual
90    /// store has seen at restore time. `chain_id` is pinned
91    /// into `__chain_id` on finalise so subsequent tip
92    /// indexing refuses checkpoints from the wrong chain.
93    pub fn new(
94        perpetual: Arc<AuthorityPerpetualTables>,
95        target_checkpoint: u64,
96        chain_id: ChainId,
97    ) -> Self {
98        Self {
99            perpetual,
100            target_checkpoint,
101            chain_id,
102            chunk_size: CHUNK_SIZE,
103        }
104    }
105
106    /// Override the per-chunk object count. Useful for tests
107    /// that want to exercise multi-chunk shards without
108    /// materialising 50k objects.
109    pub fn with_chunk_size(mut self, chunk_size: usize) -> Self {
110        assert!(chunk_size > 0, "chunk_size must be > 0");
111        self.chunk_size = chunk_size;
112        self
113    }
114}
115
116/// Inclusive `[start, end]` `ObjectID` range covered by `shard_id`.
117fn shard_range(shard_id: u32) -> (ObjectID, ObjectID) {
118    let prefix = (shard_id as u8) << SHARD_PREFIX_SHIFT;
119    let mut start = [0u8; ObjectID::LENGTH];
120    start[0] = prefix;
121    let mut end = [0xffu8; ObjectID::LENGTH];
122    end[0] = prefix | ((1 << SHARD_PREFIX_SHIFT) - 1);
123    (ObjectID::new(start), ObjectID::new(end))
124}
125
126/// Increment `id` as a 256-bit big-endian integer, returning
127/// `None` on overflow.
128fn next_id(id: ObjectID) -> Option<ObjectID> {
129    let mut bytes = id.into_bytes();
130    for byte in bytes.iter_mut().rev() {
131        if *byte == 0xff {
132            *byte = 0;
133        } else {
134            *byte += 1;
135            return Some(ObjectID::new(bytes));
136        }
137    }
138    None
139}
140
141#[async_trait]
142impl RestoreSource for PerpetualStoreRestoreSource {
143    fn target_checkpoint(&self) -> u64 {
144        self.target_checkpoint
145    }
146
147    fn target_chain_id(&self) -> ChainId {
148        self.chain_id
149    }
150
151    fn shards(&self) -> u32 {
152        SHARDS
153    }
154
155    fn stream(
156        &self,
157        shard_id: u32,
158        cursor: Option<Bytes>,
159    ) -> BoxStream<'_, anyhow::Result<RestoreChunk>> {
160        let (shard_start, shard_end) = shard_range(shard_id);
161
162        let start_id = match cursor {
163            None => Some(shard_start),
164            Some(bytes) => match ObjectID::from_bytes(&bytes[..]) {
165                Ok(id) => next_id(id).filter(|n| *n <= shard_end),
166                Err(e) => {
167                    return stream::once(async move {
168                        Err(anyhow::anyhow!("invalid perpetual-store cursor: {e}"))
169                    })
170                    .boxed();
171                }
172            },
173        };
174
175        let Some(start_id) = start_id else {
176            return stream::empty().boxed();
177        };
178
179        // Bounded mpsc applies backpressure on the iterator
180        // task so it pauses when the driver hasn't committed
181        // the previous chunk yet.
182        let (tx, rx) = mpsc::channel::<anyhow::Result<RestoreChunk>>(2);
183        let perpetual = self.perpetual.clone();
184        let chunk_size = self.chunk_size;
185
186        tokio::task::spawn_blocking(move || {
187            iterate_shard(perpetual, start_id, shard_end, chunk_size, tx);
188        });
189
190        ReceiverStream::new(rx).boxed()
191    }
192}
193
194/// Drive one shard's iteration end-to-end in a single
195/// `spawn_blocking` task.
196///
197/// Opens exactly one `range_iter_live_object_set` and pushes
198/// chunks of up to `chunk_size` `LiveObject::Normal` rows over
199/// `tx`. The iterator's implicit RocksDB snapshot is held for
200/// the lifetime of this function, so the whole shard observes
201/// a single point-in-time view of the perpetual store.
202///
203/// Returns early without sending anything if the receiver is
204/// dropped (e.g. the driver was cancelled).
205fn iterate_shard(
206    perpetual: Arc<AuthorityPerpetualTables>,
207    start_id: ObjectID,
208    shard_end: ObjectID,
209    chunk_size: usize,
210    tx: mpsc::Sender<anyhow::Result<RestoreChunk>>,
211) {
212    let iter = perpetual.range_iter_live_object_set(Some(start_id), Some(shard_end), false);
213    let mut buffer: Vec<Object> = Vec::with_capacity(chunk_size.min(1024));
214
215    for live in iter {
216        let LiveObject::Normal(obj) = live else {
217            continue;
218        };
219        buffer.push(obj);
220        if buffer.len() >= chunk_size {
221            let chunk = std::mem::replace(&mut buffer, Vec::with_capacity(chunk_size.min(1024)));
222            if send_chunk(&tx, chunk).is_err() {
223                return;
224            }
225        }
226    }
227
228    if !buffer.is_empty() {
229        let _ = send_chunk(&tx, buffer);
230    }
231}
232
233/// Wrap `objects` in a [`RestoreChunk`] (cursor = last object's
234/// id) and blocking-send it. Returns `Err(())` if the receiver
235/// is closed so the caller can stop iterating.
236fn send_chunk(
237    tx: &mpsc::Sender<anyhow::Result<RestoreChunk>>,
238    objects: Vec<Object>,
239) -> Result<(), ()> {
240    let last_id = objects.last().expect("non-empty chunk").id();
241    let chunk = RestoreChunk {
242        objects,
243        cursor: Bytes::copy_from_slice(&last_id.into_bytes()),
244    };
245    tx.blocking_send(Ok(chunk)).map_err(|_| ())
246}
247
248#[cfg(test)]
249mod tests {
250    use std::collections::BTreeSet;
251
252    use tempfile::TempDir;
253
254    use super::*;
255
256    fn open_perpetual() -> (TempDir, Arc<AuthorityPerpetualTables>) {
257        let dir = TempDir::new().unwrap();
258        let perpetual = Arc::new(AuthorityPerpetualTables::open(dir.path(), None, None));
259        (dir, perpetual)
260    }
261
262    fn obj_with_first_byte(first: u8, last: u8) -> Object {
263        let mut bytes = [0u8; ObjectID::LENGTH];
264        bytes[0] = first;
265        bytes[ObjectID::LENGTH - 1] = last;
266        Object::immutable_with_id_for_testing(ObjectID::new(bytes))
267    }
268
269    /// Hand-pick a representative shard and verify the shard
270    /// range covers the right ObjectID prefixes.
271    #[test]
272    fn shard_range_covers_correct_prefixes() {
273        let (s0, e0) = shard_range(0);
274        assert_eq!(s0.into_bytes()[0], 0x00);
275        assert_eq!(e0.into_bytes()[0], 0x07);
276
277        let (s1, e1) = shard_range(1);
278        assert_eq!(s1.into_bytes()[0], 0x08);
279        assert_eq!(e1.into_bytes()[0], 0x0F);
280
281        let (s31, e31) = shard_range(31);
282        assert_eq!(s31.into_bytes()[0], 0xF8);
283        assert_eq!(e31.into_bytes()[0], 0xFF);
284        // Last byte of the upper bound is 0xFF.
285        assert_eq!(e31.into_bytes()[ObjectID::LENGTH - 1], 0xFF);
286    }
287
288    #[test]
289    fn next_id_increments_with_carry() {
290        let mut bytes = [0u8; ObjectID::LENGTH];
291        bytes[ObjectID::LENGTH - 1] = 0xff;
292        bytes[ObjectID::LENGTH - 2] = 0x01;
293        let inc = next_id(ObjectID::new(bytes)).unwrap().into_bytes();
294        let mut expected = [0u8; ObjectID::LENGTH];
295        expected[ObjectID::LENGTH - 1] = 0x00;
296        expected[ObjectID::LENGTH - 2] = 0x02;
297        assert_eq!(inc, expected);
298    }
299
300    #[test]
301    fn next_id_overflow_returns_none() {
302        let max = ObjectID::new([0xff; ObjectID::LENGTH]);
303        assert_eq!(next_id(max), None);
304    }
305
306    /// End-to-end smoke: seed objects across two shards, drain
307    /// every shard's stream, confirm every object lands exactly
308    /// once and shard boundaries are respected.
309    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
310    async fn streams_objects_across_shards() {
311        let (_dir, perpetual) = open_perpetual();
312
313        // Insert four objects across shard 0 (first byte in
314        // 0x00..=0x07) and shard 1 (0x08..=0x0F).
315        let inserted: Vec<Object> = [(0x01, 0xaa), (0x05, 0xbb), (0x0a, 0xcc), (0x0f, 0xdd)]
316            .into_iter()
317            .map(|(first, last)| obj_with_first_byte(first, last))
318            .collect();
319        for o in &inserted {
320            perpetual.insert_object_test_only(o.clone()).unwrap();
321        }
322
323        let source = PerpetualStoreRestoreSource::new(perpetual.clone(), 7, ChainId([9u8; 32]))
324            .with_chunk_size(1);
325        assert_eq!(source.target_checkpoint(), 7);
326        assert_eq!(source.shards(), SHARDS);
327
328        // Drain shard 0 and shard 1; assert every other shard is empty.
329        let mut got = BTreeSet::new();
330        for shard in 0..SHARDS {
331            let mut stream = source.stream(shard, None);
332            while let Some(chunk) = stream.next().await {
333                let chunk = chunk.unwrap();
334                for o in chunk.objects {
335                    got.insert(o.id());
336                }
337            }
338        }
339        let want: BTreeSet<_> = inserted.iter().map(|o| o.id()).collect();
340        assert_eq!(got, want);
341    }
342
343    /// Resume from a cursor that points at the first object in
344    /// a shard and confirm the second object (and only the
345    /// second) is yielded.
346    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
347    async fn resume_from_cursor_skips_already_yielded() {
348        let (_dir, perpetual) = open_perpetual();
349
350        let a = obj_with_first_byte(0x01, 0x10);
351        let b = obj_with_first_byte(0x01, 0x20);
352        perpetual.insert_object_test_only(a.clone()).unwrap();
353        perpetual.insert_object_test_only(b.clone()).unwrap();
354
355        // Shard 0 covers first byte 0x00..=0x07, so both
356        // objects live there.
357        let source = PerpetualStoreRestoreSource::new(perpetual.clone(), 0, ChainId([0u8; 32]));
358        let cursor = Bytes::copy_from_slice(&a.id().into_bytes());
359        let mut stream = source.stream(0, Some(cursor));
360        let mut yielded = Vec::new();
361        while let Some(chunk) = stream.next().await {
362            for o in chunk.unwrap().objects {
363                yielded.push(o.id());
364            }
365        }
366        assert_eq!(yielded, vec![b.id()]);
367    }
368}
sui_core/rpc_store_restore_source.rs

sui_core/
rpc_store_restore_source.rs
