iroh_blobs/store/
mutable_mem_storage.rs

1use bao_tree::{
2    io::{fsm::BaoContentItem, sync::WriteAt},
3    BaoTree,
4};
5use bytes::Bytes;
6
7use crate::{
8    util::{compute_outboard, copy_limited_slice, SparseMemFile},
9    IROH_BLOCK_SIZE,
10};
11
12/// Mutable in memory storage for a bao file.
13///
14/// This is used for incomplete files if they are not big enough to warrant
15/// writing to disk. We must keep track of ranges in both data and outboard
16/// that have been written to, and track the most precise known size.
17#[derive(Debug, Default)]
18pub struct MutableMemStorage {
19    /// Data file, can be any size.
20    pub data: SparseMemFile,
21    /// Outboard file, must be a multiple of 64 bytes.
22    pub outboard: SparseMemFile,
23    /// Size that was announced as we wrote that chunk
24    pub sizes: SizeInfo,
25}
26
27/// Keep track of the most precise size we know of.
28///
29/// When in memory, we don't have to write the size for every chunk to a separate
30/// slot, but can just keep the best one.
31#[derive(Debug, Default)]
32pub struct SizeInfo {
33    pub offset: u64,
34    pub size: u64,
35}
36
37impl SizeInfo {
38    /// Create a new size info for a complete file of size `size`.
39    pub(crate) fn complete(size: u64) -> Self {
40        let mask = (1 << IROH_BLOCK_SIZE.chunk_log()) - 1;
41        // offset of the last bao chunk in a file of size `size`
42        let last_chunk_offset = size & mask;
43        Self {
44            offset: last_chunk_offset,
45            size,
46        }
47    }
48
49    /// Write a size at the given offset. The size at the highest offset is going to be kept.
50    fn write(&mut self, offset: u64, size: u64) {
51        // >= instead of > because we want to be able to update size 0, the initial value.
52        if offset >= self.offset {
53            self.offset = offset;
54            self.size = size;
55        }
56    }
57
58    /// The current size, representing the most correct size we know.
59    pub fn current_size(&self) -> u64 {
60        self.size
61    }
62}
63
64impl MutableMemStorage {
65    /// Create a new mutable mem storage from the given data
66    pub fn complete(bytes: Bytes, cb: impl Fn(u64) + Send + Sync + 'static) -> (Self, crate::Hash) {
67        let (hash, outboard) = compute_outboard(&bytes[..], bytes.len() as u64, move |offset| {
68            cb(offset);
69            Ok(())
70        })
71        .unwrap();
72        let outboard = outboard.unwrap_or_default();
73        let res = Self {
74            data: bytes.to_vec().into(),
75            outboard: outboard.into(),
76            sizes: SizeInfo::complete(bytes.len() as u64),
77        };
78        (res, hash)
79    }
80
81    pub(super) fn current_size(&self) -> u64 {
82        self.sizes.current_size()
83    }
84
85    pub(super) fn read_data_at(&self, offset: u64, len: usize) -> Bytes {
86        copy_limited_slice(&self.data, offset, len)
87    }
88
89    pub(super) fn data_len(&self) -> u64 {
90        self.data.len() as u64
91    }
92
93    pub(super) fn read_outboard_at(&self, offset: u64, len: usize) -> Bytes {
94        copy_limited_slice(&self.outboard, offset, len)
95    }
96
97    pub(super) fn outboard_len(&self) -> u64 {
98        self.outboard.len() as u64
99    }
100
101    pub(super) fn write_batch(
102        &mut self,
103        size: u64,
104        batch: &[BaoContentItem],
105    ) -> std::io::Result<()> {
106        let tree = BaoTree::new(size, IROH_BLOCK_SIZE);
107        for item in batch {
108            match item {
109                BaoContentItem::Parent(parent) => {
110                    if let Some(offset) = tree.pre_order_offset(parent.node) {
111                        let o0 = offset
112                            .checked_mul(64)
113                            .expect("u64 overflow multiplying to hash pair offset");
114                        let o1 = o0.checked_add(32).expect("u64 overflow");
115                        let outboard = &mut self.outboard;
116                        outboard.write_all_at(o0, parent.pair.0.as_bytes().as_slice())?;
117                        outboard.write_all_at(o1, parent.pair.1.as_bytes().as_slice())?;
118                    }
119                }
120                BaoContentItem::Leaf(leaf) => {
121                    self.sizes.write(leaf.offset, size);
122                    self.data.write_all_at(leaf.offset, leaf.data.as_ref())?;
123                }
124            }
125        }
126        Ok(())
127    }
128}
iroh_blobs/store/mutable_mem_storage.rs

iroh_blobs/store/
mutable_mem_storage.rs
