Skip to main content

nodedb_graph/csr/rebuild/
seed.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Starting a rebuild on the owning thread, and the build that runs off it.
4
5use std::sync::atomic::{AtomicU64, Ordering};
6
7use nodedb_mem::ScopedMemory;
8
9use super::journal::CsrJournal;
10use crate::GraphError;
11use crate::csr::index::CsrIndex;
12
13static NEXT_REBUILD_TOKEN: AtomicU64 = AtomicU64::new(1);
14
15/// Node labels as the index holds them. The checkpoint format omits them,
16/// so a rebuild carries them beside it.
17#[derive(Debug, Clone)]
18pub(crate) struct NodeLabelState {
19    /// Label names in id order.
20    pub(crate) names: Vec<String>,
21    /// Label bitset per node id.
22    pub(crate) bits: Vec<u64>,
23}
24
25/// A rebuild's input, taken on the owning thread. Every field is `Send`.
26#[derive(Debug)]
27pub struct CsrRebuildSeed {
28    token: u64,
29    checkpoint: Vec<u8>,
30    node_labels: NodeLabelState,
31}
32
33/// A compacted copy on its way back to the owning thread. Every field is
34/// `Send`.
35#[derive(Debug)]
36pub struct CsrRebuilt {
37    pub(crate) token: u64,
38    pub(crate) checkpoint: Vec<u8>,
39    pub(crate) node_labels: NodeLabelState,
40}
41
42impl CsrRebuildSeed {
43    /// The token that ties this rebuild to the journal on the live index.
44    pub fn token(&self) -> u64 {
45        self.token
46    }
47
48    /// Restore the snapshot, compact it and serialize the result. Runs on
49    /// any thread. `memory` bounds the copy's allocations.
50    pub fn build(self, memory: ScopedMemory) -> Result<CsrRebuilt, GraphError> {
51        let mut copy = restore_checkpoint(&self.checkpoint, memory)?;
52        copy.compact()?;
53        let checkpoint = copy.checkpoint_to_bytes()?;
54        Ok(CsrRebuilt {
55            token: self.token,
56            checkpoint,
57            node_labels: self.node_labels,
58        })
59    }
60}
61
62impl CsrRebuilt {
63    /// The token of the rebuild that produced this copy.
64    pub fn token(&self) -> u64 {
65        self.token
66    }
67}
68
69impl CsrIndex {
70    /// Snapshot this index and open its write journal.
71    ///
72    /// Every later mutation records itself until `finish_rebuild` or
73    /// `abort_rebuild` closes the journal. `max_journal_bytes` bounds the
74    /// journal. Returns [`GraphError::RebuildInProgress`] when a journal is
75    /// already open.
76    pub fn begin_rebuild(
77        &mut self,
78        max_journal_bytes: usize,
79    ) -> Result<CsrRebuildSeed, GraphError> {
80        if self.rebuild_journal.is_some() {
81            return Err(GraphError::RebuildInProgress);
82        }
83        let checkpoint = self.checkpoint_to_bytes()?;
84        let token = NEXT_REBUILD_TOKEN.fetch_add(1, Ordering::Relaxed);
85        self.rebuild_journal = Some(CsrJournal::new(token, max_journal_bytes));
86        Ok(CsrRebuildSeed {
87            token,
88            checkpoint,
89            node_labels: NodeLabelState {
90                names: self.node_label_names.clone(),
91                bits: self.node_label_bits.clone(),
92            },
93        })
94    }
95
96    /// Whether a rebuild journal is open on this index.
97    pub fn rebuild_in_progress(&self) -> bool {
98        self.rebuild_journal.is_some()
99    }
100
101    /// Close the journal of rebuild `token`. The index itself is unchanged.
102    /// A journal of another rebuild stays open.
103    pub fn abort_rebuild(&mut self, token: u64) {
104        if self
105            .rebuild_journal
106            .as_ref()
107            .is_some_and(|journal| journal.token == token)
108        {
109            self.rebuild_journal = None;
110        }
111    }
112}
113
114/// Decode a checkpoint a rebuild produced.
115pub(crate) fn restore_checkpoint(
116    bytes: &[u8],
117    memory: ScopedMemory,
118) -> Result<CsrIndex, GraphError> {
119    CsrIndex::from_checkpoint(bytes, memory)
120        .map_err(|e| GraphError::RebuildSnapshotInvalid {
121            detail: e.to_string(),
122        })?
123        .ok_or_else(|| GraphError::RebuildSnapshotInvalid {
124            detail: "checkpoint bytes carry no CSR header".to_string(),
125        })
126}