commonware_storage/ordinal/
storage.rs

1use super::{Config, Error};
2use crate::rmap::RMap;
3use bytes::{Buf, BufMut};
4use commonware_codec::{Encode, FixedSize, Read, ReadExt, Write as CodecWrite};
5use commonware_runtime::{
6    buffer::{Read as ReadBuffer, Write},
7    Blob, Clock, Error as RError, Metrics, Storage,
8};
9use commonware_utils::{hex, Array, BitVec};
10use futures::future::try_join_all;
11use prometheus_client::metrics::counter::Counter;
12use std::{
13    collections::{btree_map::Entry, BTreeMap, BTreeSet},
14    marker::PhantomData,
15    mem::take,
16};
17use tracing::{debug, warn};
18
19/// Value stored in the index file.
20#[derive(Debug, Clone)]
21struct Record<V: Array> {
22    value: V,
23    crc: u32,
24}
25
26impl<V: Array> Record<V> {
27    fn new(value: V) -> Self {
28        let crc = crc32fast::hash(value.as_ref());
29        Self { value, crc }
30    }
31
32    fn is_valid(&self) -> bool {
33        self.crc == crc32fast::hash(self.value.as_ref())
34    }
35}
36
37impl<V: Array> FixedSize for Record<V> {
38    const SIZE: usize = V::SIZE + u32::SIZE;
39}
40
41impl<V: Array> CodecWrite for Record<V> {
42    fn write(&self, buf: &mut impl BufMut) {
43        self.value.write(buf);
44        self.crc.write(buf);
45    }
46}
47
48impl<V: Array> Read for Record<V> {
49    type Cfg = ();
50
51    fn read_cfg(buf: &mut impl Buf, _: &Self::Cfg) -> Result<Self, commonware_codec::Error> {
52        let value = V::read(buf)?;
53        let crc = u32::read(buf)?;
54
55        Ok(Self { value, crc })
56    }
57}
58
59/// Implementation of [Ordinal].
60pub struct Ordinal<E: Storage + Metrics + Clock, V: Array> {
61    // Configuration and context
62    context: E,
63    config: Config,
64
65    // Index blobs for storing key records
66    blobs: BTreeMap<u64, Write<E::Blob>>,
67
68    // RMap for interval tracking
69    intervals: RMap,
70
71    // Pending index entries to be synced, grouped by section
72    pending: BTreeSet<u64>,
73
74    // Metrics
75    puts: Counter,
76    gets: Counter,
77    has: Counter,
78    syncs: Counter,
79    pruned: Counter,
80
81    _phantom: PhantomData<V>,
82}
83
84impl<E: Storage + Metrics + Clock, V: Array> Ordinal<E, V> {
85    /// Initialize a new [Ordinal] instance.
86    pub async fn init(context: E, config: Config) -> Result<Self, Error> {
87        Self::init_with_bits(context, config, None).await
88    }
89
90    /// Initialize a new [Ordinal] instance with a collection of [BitVec]s (indicating which
91    /// records should be considered available).
92    ///
93    /// If a section is not provided in the [BTreeMap], all records in that section are considered
94    /// unavailable. If a [BitVec] is provided for a section, all records in that section are
95    /// considered available if and only if the [BitVec] is set for the record. If a section is provided
96    /// but no [BitVec] is populated, all records in that section are considered available.
97    // TODO(#1227): Hide this complexity from the caller.
98    pub async fn init_with_bits(
99        context: E,
100        config: Config,
101        bits: Option<BTreeMap<u64, &Option<BitVec>>>,
102    ) -> Result<Self, Error> {
103        // Scan for all blobs in the partition
104        let mut blobs = BTreeMap::new();
105        let stored_blobs = match context.scan(&config.partition).await {
106            Ok(blobs) => blobs,
107            Err(commonware_runtime::Error::PartitionMissing(_)) => Vec::new(),
108            Err(err) => return Err(Error::Runtime(err)),
109        };
110
111        // Open all blobs and check for partial records
112        for name in stored_blobs {
113            let (blob, mut len) = context.open(&config.partition, &name).await?;
114            let index = match name.try_into() {
115                Ok(index) => u64::from_be_bytes(index),
116                Err(nm) => Err(Error::InvalidBlobName(hex(&nm)))?,
117            };
118
119            // Check if blob size is aligned to record size
120            let record_size = Record::<V>::SIZE as u64;
121            if len % record_size != 0 {
122                warn!(
123                    blob = index,
124                    invalid_size = len,
125                    record_size,
126                    "blob size is not a multiple of record size, truncating"
127                );
128                len -= len % record_size;
129                blob.resize(len).await?;
130                blob.sync().await?;
131            }
132
133            debug!(blob = index, len, "found index blob");
134            let wrapped_blob = Write::new(blob, len, config.write_buffer);
135            blobs.insert(index, wrapped_blob);
136        }
137
138        // Initialize intervals by scanning existing records
139        debug!(
140            blobs = blobs.len(),
141            "rebuilding intervals from existing index"
142        );
143        let start = context.current();
144        let mut items = 0;
145        let mut intervals = RMap::new();
146        for (section, blob) in &blobs {
147            // Skip if bits are provided and the section is not in the bits
148            if let Some(bits) = &bits {
149                if !bits.contains_key(section) {
150                    warn!(section, "skipping section without bits");
151                    continue;
152                }
153            }
154
155            // Initialize read buffer
156            let size = blob.size().await;
157            let mut replay_blob = ReadBuffer::new(blob.clone(), size, config.replay_buffer);
158
159            // Iterate over all records in the blob
160            let mut offset = 0;
161            let items_per_blob = config.items_per_blob.get();
162            while offset < size {
163                // Calculate index for this record
164                let index = section * items_per_blob + (offset / Record::<V>::SIZE as u64);
165
166                // If bits are provided, skip if not set
167                let mut must_exist = false;
168                if let Some(bits) = &bits {
169                    // If bits are provided, check if the record exists
170                    let bits = bits.get(section).unwrap();
171                    if let Some(bits) = bits {
172                        let bit_index = offset as usize / Record::<V>::SIZE;
173                        if !bits.get(bit_index).expect("invalid index") {
174                            offset += Record::<V>::SIZE as u64;
175                            continue;
176                        }
177                    }
178
179                    // If bit section exists but it is empty, we must have all records
180                    must_exist = true;
181                }
182
183                // Attempt to read record at offset
184                replay_blob.seek_to(offset)?;
185                let mut record_buf = vec![0u8; Record::<V>::SIZE];
186                replay_blob
187                    .read_exact(&mut record_buf, Record::<V>::SIZE)
188                    .await?;
189                let record = Record::<V>::read(&mut record_buf.as_slice())?;
190                offset += Record::<V>::SIZE as u64;
191
192                // If record is valid, add to intervals
193                if record.is_valid() {
194                    items += 1;
195                    intervals.insert(index);
196                    continue;
197                }
198
199                // If record is invalid, it may either be empty or corrupted. We only care
200                // which is which if the provided bits indicate that the record must exist.
201                if must_exist {
202                    return Err(Error::MissingRecord(index));
203                }
204            }
205        }
206        debug!(
207            items,
208            elapsed = ?context.current().duration_since(start).unwrap_or_default(),
209            "rebuilt intervals"
210        );
211
212        // Initialize metrics
213        let puts = Counter::default();
214        let gets = Counter::default();
215        let has = Counter::default();
216        let syncs = Counter::default();
217        let pruned = Counter::default();
218        context.register("puts", "Number of put calls", puts.clone());
219        context.register("gets", "Number of get calls", gets.clone());
220        context.register("has", "Number of has calls", has.clone());
221        context.register("syncs", "Number of sync calls", syncs.clone());
222        context.register("pruned", "Number of pruned blobs", pruned.clone());
223
224        Ok(Self {
225            context,
226            config,
227            blobs,
228            intervals,
229            pending: BTreeSet::new(),
230            puts,
231            gets,
232            has,
233            syncs,
234            pruned,
235            _phantom: PhantomData,
236        })
237    }
238
239    /// Add a value at the specified index (pending until sync).
240    pub async fn put(&mut self, index: u64, value: V) -> Result<(), Error> {
241        self.puts.inc();
242
243        // Check if blob exists
244        let items_per_blob = self.config.items_per_blob.get();
245        let section = index / items_per_blob;
246        if let Entry::Vacant(entry) = self.blobs.entry(section) {
247            let (blob, len) = self
248                .context
249                .open(&self.config.partition, &section.to_be_bytes())
250                .await?;
251            entry.insert(Write::new(blob, len, self.config.write_buffer));
252            debug!(section, "created blob");
253        }
254
255        // Write the value to the blob
256        let blob = self.blobs.get(&section).unwrap();
257        let offset = (index % items_per_blob) * Record::<V>::SIZE as u64;
258        let record = Record::new(value);
259        blob.write_at(record.encode(), offset).await?;
260        self.pending.insert(section);
261
262        // Add to intervals
263        self.intervals.insert(index);
264
265        Ok(())
266    }
267
268    /// Get the value for a given index.
269    pub async fn get(&self, index: u64) -> Result<Option<V>, Error> {
270        self.gets.inc();
271
272        // If get isn't in an interval, it doesn't exist and we don't need to access disk
273        if self.intervals.get(&index).is_none() {
274            return Ok(None);
275        }
276
277        // Read from disk
278        let items_per_blob = self.config.items_per_blob.get();
279        let section = index / items_per_blob;
280        let blob = self.blobs.get(&section).unwrap();
281        let offset = (index % items_per_blob) * Record::<V>::SIZE as u64;
282        let read_buf = vec![0u8; Record::<V>::SIZE];
283        let read_buf = blob.read_at(read_buf, offset).await?;
284        let record = Record::<V>::read(&mut read_buf.as_ref())?;
285
286        // If record is valid, return it
287        if record.is_valid() {
288            Ok(Some(record.value))
289        } else {
290            Err(Error::InvalidRecord(index))
291        }
292    }
293
294    /// Check if an index exists.
295    pub fn has(&self, index: u64) -> bool {
296        self.has.inc();
297
298        self.intervals.get(&index).is_some()
299    }
300
301    /// Get the next gap information for backfill operations.
302    pub fn next_gap(&self, index: u64) -> (Option<u64>, Option<u64>) {
303        self.intervals.next_gap(index)
304    }
305
306    /// Get up to the next `max` missing items after `start`.
307    pub fn missing_items(&self, start: u64, max: usize) -> Vec<u64> {
308        self.intervals.missing_items(start, max)
309    }
310
311    /// Prune indices older than `min` by removing entire blobs.
312    ///
313    /// Pruning is done at blob boundaries to avoid partial deletions. A blob is pruned only if
314    /// all possible indices in that blob are less than `min`.
315    pub async fn prune(&mut self, min: u64) -> Result<(), Error> {
316        // Collect sections to remove
317        let items_per_blob = self.config.items_per_blob.get();
318        let min_section = min / items_per_blob;
319        let sections_to_remove: Vec<u64> = self
320            .blobs
321            .keys()
322            .filter(|&&section| section < min_section)
323            .copied()
324            .collect();
325
326        // Remove the collected sections
327        for section in sections_to_remove {
328            if let Some(blob) = self.blobs.remove(&section) {
329                drop(blob);
330                self.context
331                    .remove(&self.config.partition, Some(&section.to_be_bytes()))
332                    .await?;
333
334                // Remove the corresponding index range from intervals
335                let start_index = section * items_per_blob;
336                let end_index = (section + 1) * items_per_blob - 1;
337                self.intervals.remove(start_index, end_index);
338                debug!(section, start_index, end_index, "pruned blob");
339            }
340
341            // Update metrics
342            self.pruned.inc();
343        }
344
345        // Clean pending entries that fall into pruned sections.
346        self.pending.retain(|&section| section >= min_section);
347
348        Ok(())
349    }
350
351    /// Write all pending entries and sync all modified [Blob]s.
352    pub async fn sync(&mut self) -> Result<(), Error> {
353        self.syncs.inc();
354
355        // Sync all modified blobs
356        let mut futures = Vec::with_capacity(self.pending.len());
357        for &section in &self.pending {
358            futures.push(self.blobs.get(&section).unwrap().sync());
359        }
360        try_join_all(futures).await?;
361
362        // Clear pending sections
363        self.pending.clear();
364
365        Ok(())
366    }
367
368    /// Sync all pending entries and [Blob]s.
369    pub async fn close(mut self) -> Result<(), Error> {
370        self.sync().await?;
371        for (_, blob) in take(&mut self.blobs) {
372            blob.sync().await?;
373        }
374        Ok(())
375    }
376
377    /// Destroy [Ordinal] and remove all data.
378    pub async fn destroy(self) -> Result<(), Error> {
379        for (i, blob) in self.blobs.into_iter() {
380            drop(blob);
381            self.context
382                .remove(&self.config.partition, Some(&i.to_be_bytes()))
383                .await?;
384            debug!(section = i, "destroyed blob");
385        }
386        match self.context.remove(&self.config.partition, None).await {
387            Ok(()) => {}
388            Err(RError::PartitionMissing(_)) => {
389                // Partition already removed or never existed.
390            }
391            Err(err) => return Err(Error::Runtime(err)),
392        }
393        Ok(())
394    }
395}
commonware_storage/ordinal/storage.rs

commonware_storage/ordinal/
storage.rs
