harmos_storage/
file_batch.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10// Native File Batches Retain Their Allocation Charge Through Every Shared View.
11
12/* === Modules === */
13
14/* === Re-Exports === */
15
16/* === Imports === */
17
18// Standard Library Imports
19use std::borrow::Borrow;
20use std::num::NonZeroUsize;
21use std::ops::Range;
22use std::sync::Arc;
23
24// Third Party Imports
25
26// Local Imports
27use crate::arrow::RetainedMemory;
28use crate::{Error, FileBatch, MemorySize, ReadBatch, Result};
29use harmos_domain::model::storage::files::FileEvent;
30
31/* === Constants === */
32
33/* === Macros === */
34
35/* === Type Aliases === */
36
37/* === Structs === */
38
39#[derive(Clone)]
40pub struct FileReadBatch {
41    owner: ReadBatch<FileBatch>,
42}
43
44#[derive(Clone)]
45pub struct FileReadSlice {
46    owner: FileReadBatch,
47    range: Range<usize>,
48}
49
50pub type FileReadEvent = ReadBatch<FileEvent>;
51
52/* === Enums === */
53
54/* === Traits === */
55
56/* === Implementations === */
57
58impl FileReadEvent {
59    pub fn event(&self) -> &FileEvent {
60        self.as_ref()
61    }
62
63    pub fn data(&self) -> Option<&[u8]> {
64        match self.event() {
65            FileEvent::Data(batch) => Some(&batch.bytes),
66            _ => None,
67        }
68    }
69
70    pub fn offset(&self) -> Option<u64> {
71        match self.event() {
72            FileEvent::Data(batch) => Some(batch.offset),
73            _ => None,
74        }
75    }
76}
77
78impl FileReadBatch {
79    pub(crate) fn retained(batch: FileBatch, memory: RetainedMemory) -> Self {
80        Self { owner: ReadBatch::retained(batch, memory) }
81    }
82
83    /// Moves a unique file batch through an admitted projection.
84    /// Shared batches or slices refuse; allocation and escape rules match `ReadBatch::try_map`.
85    #[tracing::instrument(name = "storage.file_batch.try_map", level = "trace", skip_all, fields(additional_peak = additional_peak, outcome = "cancelled"))]
86    pub fn try_map<U, E>(
87        self, additional_peak: usize, map: impl FnOnce(FileBatch) -> std::result::Result<U, E>,
88    ) -> std::result::Result<ReadBatch<U>, E>
89    where
90        E: From<Error>,
91    {
92        let result: std::result::Result<ReadBatch<U>, E> = self.owner.try_map(additional_peak, map);
93        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
94        result
95    }
96
97    pub fn offset(&self) -> u64 {
98        self.owner.offset
99    }
100
101    pub fn bytes(&self) -> &[u8] {
102        &self.owner.bytes
103    }
104
105    pub fn owner(&self) -> Arc<impl Borrow<FileBatch> + Send + Sync + 'static> {
106        self.owner.owner()
107    }
108
109    pub fn slice(&self, range: Range<usize>) -> FileReadSlice {
110        assert!(range.start <= range.end && range.end <= self.bytes().len(), "file slice is outside its batch");
111        FileReadSlice { owner: self.clone(), range }
112    }
113}
114
115impl Borrow<FileBatch> for FileReadBatch {
116    fn borrow(&self) -> &FileBatch {
117        self.owner.as_ref()
118    }
119}
120
121impl FileReadSlice {
122    pub fn offset(&self) -> u64 {
123        self.owner.offset() + self.range.start as u64
124    }
125
126    pub fn bytes(&self) -> &[u8] {
127        &self.owner.bytes()[self.range.clone()]
128    }
129}
130
131impl MemorySize for FileReadBatch {
132    fn memory_size(&self) -> usize {
133        self.owner.memory_size().saturating_add(std::mem::size_of::<Self>())
134    }
135}
136
137impl std::fmt::Debug for FileReadBatch {
138    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
139        self.owner.fmt(formatter)
140    }
141}
142
143/* === Public Functions === */
144
145pub(crate) fn batch_limit(ceiling: usize, offset: u64, remaining: u64, event: bool) -> Result<NonZeroUsize> {
146    let mut count = ceiling.min(usize::try_from(remaining).unwrap_or(usize::MAX));
147    loop {
148        let batch = 2usize.saturating_add(varint(offset)).saturating_add(varint(count as u64)).saturating_add(count);
149        let framed = if event { batch.saturating_add(1 + varint(batch as u64)) } else { batch };
150        if framed <= ceiling {
151            return NonZeroUsize::new(count).ok_or(Error::ResourceExhausted { resource: "file response batch bytes" });
152        }
153        count = count
154            .checked_sub(framed - ceiling)
155            .filter(|count| *count != 0)
156            .ok_or(Error::ResourceExhausted { resource: "file response batch bytes" })?;
157    }
158}
159
160/* === Private Functions === */
161
162fn varint(mut value: u64) -> usize {
163    let mut bytes = 1;
164    while value >= 128 {
165        value >>= 7;
166        bytes += 1;
167    }
168    bytes
169}
170
171/* === Tests === */
harmos_storage/file_batch.rs

harmos_storage/
file_batch.rs
