storage_device/
splittable_buffer.rs1use crate::buffer::{BufferAllocator, OwnedBuffer};
6use std::ops::Range;
7use std::ptr::slice_from_raw_parts_mut;
8use std::sync::Arc;
9use storage_ptr_slice::MutPtrByteSlice;
10
11#[derive(Debug)]
12pub(crate) struct SplittableBufferInner {
13 parent_buffer: OwnedBuffer,
14}
15
16impl BufferAllocator for SplittableBufferInner {
17 fn free_buffer(&self, _range: Range<usize>) {
18 }
21}
22
23#[derive(Clone, Debug)]
28pub struct SplittableBufferHandle {
29 inner: Arc<SplittableBufferInner>,
30}
31
32impl SplittableBufferHandle {
33 pub fn into_buffer(self) -> Option<OwnedBuffer> {
38 Arc::into_inner(self.inner).map(|inner| inner.parent_buffer)
39 }
40}
41
42#[derive(Debug)]
46pub struct SplittableBuffer {
47 inner: Arc<SplittableBufferInner>,
48 current_ptr: *mut u8,
49 remaining_range: Range<usize>,
50}
51
52unsafe impl Send for SplittableBuffer {}
55unsafe impl Sync for SplittableBuffer {}
56
57impl SplittableBuffer {
58 pub fn new(mut buffer: OwnedBuffer) -> (Self, SplittableBufferHandle) {
61 let remaining_range = buffer.range();
62 let current_ptr = buffer.as_mut_ptr();
63 let inner = Arc::new(SplittableBufferInner { parent_buffer: buffer });
64 let handle = SplittableBufferHandle { inner: inner.clone() };
65 let splittable = Self { inner, current_ptr, remaining_range };
66 (splittable, handle)
67 }
68
69 pub fn remaining_range(&self) -> Range<usize> {
71 self.remaining_range.clone()
72 }
73
74 pub fn take_prefix(&mut self, len: usize) -> OwnedBuffer {
80 assert!(len <= self.remaining_range.len());
81 let child_range = self.remaining_range.start..self.remaining_range.start + len;
82 self.remaining_range.start += len;
83 let ptr = self.current_ptr;
84 self.current_ptr = self.current_ptr.wrapping_add(len);
85
86 let slice = unsafe { MutPtrByteSlice::new(slice_from_raw_parts_mut(ptr, len)) };
90 OwnedBuffer::new(slice, child_range, self.inner.clone() as Arc<dyn BufferAllocator>)
91 }
92}
93
94#[cfg(test)]
95mod tests {
96 use super::*;
97 use crate::buffer_allocator::{BufferAllocator as PoolBufferAllocator, BufferSource};
98
99 #[fuchsia::test]
100 async fn test_splittable_buffer_handle_into_buffer() {
101 let source = BufferSource::new(4096);
102 let pool = Arc::new(PoolBufferAllocator::new(512, source));
103 let owned = pool.allocate_buffer_sync_owned(2048);
104
105 let (mut splittable, handle) = SplittableBuffer::new(owned);
106 let mut child1 = splittable.take_prefix(1024);
107 let mut child2 = splittable.take_prefix(1024);
108 child1.as_mut_slice().fill(0x33);
109 child2.as_mut_slice().fill(0x44);
110
111 drop(splittable);
113 drop(child1);
114
115 let handle_clone = handle.clone();
116 assert!(handle_clone.into_buffer().is_none());
117
118 drop(child2);
120 let merged = handle.into_buffer().expect("into_buffer must succeed when sole reference");
121 assert_eq!(merged.len(), 2048);
122 assert!(merged.as_slice()[..1024].iter().all(|&b| b == 0x33));
123 assert!(merged.as_slice()[1024..].iter().all(|&b| b == 0x44));
124 }
125}