harmos_storage/
batch.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12/* === Re-Exports === */
13
14/* === Imports === */
15
16// Standard Library Imports
17use crate::manager::Inner;
18use std::ops::Range;
19use std::sync::{
20    Arc,
21    atomic::{AtomicBool, AtomicUsize, Ordering},
22};
23// Third Party Imports
24// Local Imports
25use crate::MemorySize;
26
27/* === Constants === */
28
29/* === Macros === */
30
31/* === Type Aliases === */
32
33pub type FileChunk = StorageBatch<u8>;
34
35/* === Structs === */
36
37#[derive(Debug)]
38pub struct StorageBatch<T> {
39    values: Arc<Vec<T>>,
40    range: Range<usize>,
41    charges: Arc<Vec<Arc<WriteCharge>>>,
42    memory: Arc<Vec<crate::arrow::RetainedMemory>>,
43}
44
45#[derive(Debug)]
46pub(crate) struct WriteCharge {
47    bytes: usize,
48    retired: AtomicBool,
49    held: Arc<AtomicUsize>,
50    changed: tokio::sync::watch::Sender<u64>,
51}
52
53/* === Enums === */
54
55/* === Traits === */
56
57/* === Implementations === */
58
59impl<T> StorageBatch<T> {
60    pub fn new(values: Vec<T>) -> Self {
61        let range = 0..values.len();
62        Self { values: Arc::new(values), range, charges: Arc::new(Vec::new()), memory: Arc::new(Vec::new()) }
63    }
64
65    pub(crate) fn retained(mut self, charge: Arc<WriteCharge>) -> Self {
66        let mut charges = self.charges.as_ref().clone();
67        charges.push(charge);
68        self.charges = Arc::new(charges);
69        self
70    }
71
72    pub(crate) fn retain_memory(mut self, memory: crate::arrow::RetainedMemory) -> Self {
73        Arc::make_mut(&mut self.memory).push(memory);
74        self
75    }
76
77    pub(crate) fn charged(&self) -> bool {
78        self.charges.iter().any(|charge| charge.retired.load(Ordering::SeqCst))
79    }
80
81    pub(crate) fn retire(&self) {
82        for charge in self.charges.iter() {
83            charge.retire();
84        }
85    }
86
87    pub(crate) fn merge(batches: &mut [&mut Self]) -> Option<Self> {
88        let first = batches.first()?;
89        let mut range = first.range.clone();
90        let adjoining = batches.iter().skip(1).all(|batch| {
91            let joined = Arc::ptr_eq(&first.values, &batch.values) && range.end == batch.range.start;
92            range.end = batch.range.end;
93            joined
94        });
95        if adjoining {
96            return Some(Self {
97                values: Arc::clone(&first.values),
98                range,
99                charges: Arc::clone(&first.charges),
100                memory: Arc::clone(&first.memory),
101            });
102        }
103        if batches.iter().any(|batch| Arc::strong_count(&batch.values) != 1 || batch.range != (0..batch.values.len())) {
104            return None;
105        }
106        let count = batches.iter().map(|batch| batch.len()).sum();
107        let lengths: Vec<_> = batches.iter().map(|batch| batch.len()).collect();
108        let charges = batches.iter().flat_map(|batch| batch.charges.iter().cloned()).collect();
109        let memory = batches.iter().flat_map(|batch| batch.memory.iter().cloned()).collect();
110        let mut values = Vec::with_capacity(count);
111        for batch in batches.iter_mut() {
112            let source = std::mem::take(Arc::get_mut(&mut batch.values).expect("unique queued backing"));
113            values.extend(source);
114        }
115        let merged = Self { values: Arc::new(values), range: 0..count, charges: Arc::new(charges), memory: Arc::new(memory) };
116        let mut start = 0;
117        for (batch, length) in batches.iter_mut().zip(lengths) {
118            **batch = merged.slice(start..start + length);
119            start += length;
120        }
121        Some(merged)
122    }
123
124    pub fn as_slice(&self) -> &[T] {
125        &self.values[self.range.clone()]
126    }
127
128    pub fn slice(&self, range: Range<usize>) -> Self {
129        assert!(range.start <= range.end && range.end <= self.len(), "batch slice is outside its span");
130        let range = self.range.start + range.start..self.range.start + range.end;
131        Self { values: Arc::clone(&self.values), range, charges: Arc::clone(&self.charges), memory: Arc::clone(&self.memory) }
132    }
133
134    pub fn len(&self) -> usize {
135        self.range.len()
136    }
137
138    pub fn is_empty(&self) -> bool {
139        self.range.is_empty()
140    }
141
142    pub fn capacity(&self) -> usize {
143        self.values.capacity()
144    }
145
146    pub fn allocation_bytes(&self) -> usize {
147        self.values.capacity().saturating_mul(std::mem::size_of::<T>())
148    }
149}
150
151impl WriteCharge {
152    pub(crate) fn new(inner: &Arc<Inner>, bytes: usize) -> Arc<Self> {
153        Arc::new(Self {
154            bytes,
155            retired: AtomicBool::new(false),
156            held: Arc::clone(&inner.retained_write_bytes),
157            changed: inner.changed.clone(),
158        })
159    }
160
161    pub(crate) fn read(inner: &Arc<Inner>, bytes: usize) -> Arc<Self> {
162        Arc::new(Self {
163            bytes,
164            retired: AtomicBool::new(false),
165            held: Arc::clone(&inner.retained_read_bytes),
166            changed: inner.changed.clone(),
167        })
168    }
169
170    pub(crate) fn retire(&self) {
171        if !self.retired.swap(true, Ordering::SeqCst) {
172            self.held.fetch_add(self.bytes, Ordering::SeqCst);
173        }
174    }
175}
176
177impl Drop for WriteCharge {
178    fn drop(&mut self) {
179        if self.retired.load(Ordering::SeqCst) {
180            self.held.fetch_sub(self.bytes, Ordering::SeqCst);
181            self.changed.send_modify(|revision| *revision = revision.wrapping_add(1));
182        }
183    }
184}
185
186impl From<Vec<u8>> for StorageBatch<u8> {
187    fn from(values: Vec<u8>) -> Self {
188        Self::new(values)
189    }
190}
191
192impl<T> Clone for StorageBatch<T> {
193    fn clone(&self) -> Self {
194        Self {
195            values: Arc::clone(&self.values),
196            range: self.range.clone(),
197            charges: Arc::clone(&self.charges),
198            memory: Arc::clone(&self.memory),
199        }
200    }
201}
202
203impl<T: MemorySize> MemorySize for StorageBatch<T> {
204    fn memory_size(&self) -> usize {
205        std::mem::size_of::<Self>().saturating_add(self.values.as_ref().memory_size()).saturating_add(2 * std::mem::size_of::<usize>())
206    }
207}
208
209/* === Public Functions === */
210
211/* === Private Functions === */
212
213/* === Tests === */
214
215#[cfg(test)]
216mod tests {
217    use super::{FileChunk, StorageBatch};
218    use crate::MemorySize;
219
220    #[test]
221    fn chunk_allocation() {
222        let mut bytes = Vec::with_capacity(1024);
223        bytes.extend_from_slice(&[1, 2, 3, 4]);
224        let pointer = bytes.as_ptr();
225        let capacity = bytes.capacity();
226        let batch = FileChunk::from(bytes);
227        assert_eq!(batch.as_slice().as_ptr(), pointer);
228        let retained = batch.slice(1..3);
229        let empty = batch.slice(4..4);
230        assert_eq!(retained.allocation_bytes(), capacity);
231        assert_eq!(empty.memory_size(), batch.memory_size());
232        drop(batch);
233        assert_eq!(retained.as_slice(), &[2, 3]);
234        assert_eq!(retained.as_slice().as_ptr(), pointer.wrapping_add(1));
235    }
236
237    #[test]
238    fn slice_accounting() {
239        let mut payload = Vec::with_capacity(4096);
240        payload.push(7u8);
241        let batch = StorageBatch::new(vec![payload, vec![8u8]]);
242        let slice = batch.slice(1..2);
243        assert_eq!(slice.memory_size(), batch.memory_size());
244        assert!(slice.memory_size() >= 4096);
245    }
246
247    fn storage() -> crate::StorageManager {
248        use crate::test_backend::{Data, Recovery, history, storage_settings};
249        use std::sync::Arc;
250        crate::StorageManager::new(
251            crate::StorageBackends { primary: Arc::new(Data::<()>::new(history())), state: Arc::new(Recovery::new(history())) },
252            storage_settings(),
253        )
254        .unwrap()
255    }
256
257    fn pending(storage: &crate::StorageManager) -> std::sync::Arc<crate::manager::Pending> {
258        use crate::{Identifier, PayloadEncoding, SchemaKey, StoredRecord};
259        let record = StoredRecord {
260            position: storage.position(1),
261            schema: SchemaKey {
262                id: Identifier::new("record").unwrap(),
263                version: std::num::NonZeroU32::MIN,
264                payload_encoding: PayloadEncoding::PostcardV1,
265            },
266            scope: None,
267            item: Identifier::new("item").unwrap(),
268            source: None,
269            source_time: None,
270            payload: vec![7; 32],
271        };
272        std::sync::Arc::new(crate::manager::Pending::new(
273            &storage.inner,
274            None,
275            1,
276            4096,
277            crate::manager::Operation::Records(StorageBatch::new(vec![record])),
278        ))
279    }
280
281    #[test]
282    fn pending_retention() {
283        let storage = storage();
284        let pending = pending(&storage);
285        let crate::manager::Operation::Records(values) = &pending.operation else { unreachable!() };
286        let detached = values.clone();
287        let empty = detached.slice(1..1);
288        {
289            let mut state = storage.inner.state.lock().unwrap();
290            state.queue.push_back(std::sync::Arc::clone(&pending));
291            state.collect();
292            assert_eq!(state.status.pending_bytes, 4096);
293            assert_eq!(state.retained_write_bytes.load(std::sync::atomic::Ordering::SeqCst), 0);
294            state.queue.retire_data(1, 1);
295            state.collect();
296            assert_eq!(state.status.pending_bytes, 4096);
297        }
298        drop(pending);
299        drop(detached);
300        assert_eq!(storage.status().pending_bytes, 4096);
301        drop(empty);
302        assert_eq!(storage.status().pending_bytes, 0);
303    }
304
305    #[test]
306    fn failure_retention() {
307        let storage = storage();
308        let pending = pending(&storage);
309        let crate::manager::Operation::Records(values) = &pending.operation else { unreachable!() };
310        let detached = values.clone();
311        {
312            let mut state = storage.inner.state.lock().unwrap();
313            state.queue.push_back(std::sync::Arc::clone(&pending));
314            state.status.failure = Some(std::sync::Arc::new(crate::Error::NotReady { domain: "injected backend" }));
315            state.status.closing = true;
316            state.collect();
317            assert_eq!(state.status.pending_bytes, 4096);
318            assert_eq!(state.retained_write_bytes.load(std::sync::atomic::Ordering::SeqCst), 0);
319            state.queue.retire_data(1, 1);
320            state.status.closed = true;
321        }
322        drop(pending);
323        assert_eq!(storage.status().pending_bytes, 4096);
324        drop(detached);
325        assert_eq!(storage.status().pending_bytes, 0);
326    }
327
328    #[test]
329    fn payload_merge() {
330        let left = vec![7u8; 32];
331        let right = vec![8u8; 32];
332        let pointers = [left.as_ptr(), right.as_ptr()];
333        let mut first = StorageBatch::new(vec![left]);
334        let mut second = StorageBatch::new(vec![right]);
335        let merged = StorageBatch::merge(&mut [&mut first, &mut second]).unwrap();
336        assert_eq!(merged.as_slice()[0].as_ptr(), pointers[0]);
337        assert_eq!(merged.as_slice()[1].as_ptr(), pointers[1]);
338        let detached = merged.clone();
339        let retry = StorageBatch::merge(&mut [&mut first, &mut second]).unwrap();
340        assert_eq!(retry.len(), 2);
341        assert_eq!(retry.as_slice()[1].as_ptr(), pointers[1]);
342        drop(detached);
343    }
344
345    #[test]
346    fn shared_merge_refused() {
347        let mut first = StorageBatch::new(vec![vec![7u8; 32]]);
348        let mut second = StorageBatch::new(vec![vec![8u8; 32]]);
349        let detached = first.clone();
350        assert!(StorageBatch::merge(&mut [&mut first, &mut second]).is_none());
351        assert_eq!(first.as_slice(), detached.as_slice());
352        assert_eq!(second.as_slice()[0], vec![8; 32]);
353    }
354}
harmos_storage/batch.rs

harmos_storage/
batch.rs
