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 fn is_trusted(&self) -> bool {
23 self.parent_buffer.try_as_slice().is_some()
24 }
25}
26
27#[derive(Clone, Debug)]
32pub struct SplittableBufferHandle {
33 inner: Arc<SplittableBufferInner>,
34}
35
36impl SplittableBufferHandle {
37 pub fn into_buffer(self) -> Option<OwnedBuffer> {
42 Arc::into_inner(self.inner).map(|inner| inner.parent_buffer)
43 }
44}
45
46#[derive(Debug)]
50pub struct SplittableBuffer {
51 inner: Arc<SplittableBufferInner>,
52 current_ptr: *mut u8,
53 remaining_range: Range<usize>,
54}
55
56unsafe impl Send for SplittableBuffer {}
59unsafe impl Sync for SplittableBuffer {}
60
61impl SplittableBuffer {
62 pub fn new(mut buffer: OwnedBuffer) -> (Self, SplittableBufferHandle) {
65 let remaining_range = buffer.range();
66 let current_ptr = buffer.as_mut_ptr();
67 let inner = Arc::new(SplittableBufferInner { parent_buffer: buffer });
68 let handle = SplittableBufferHandle { inner: inner.clone() };
69 let splittable = Self { inner, current_ptr, remaining_range };
70 (splittable, handle)
71 }
72
73 pub fn remaining_range(&self) -> Range<usize> {
75 self.remaining_range.clone()
76 }
77
78 pub fn take_prefix(&mut self, len: usize) -> OwnedBuffer {
84 assert!(len <= self.remaining_range.len());
85 let child_range = self.remaining_range.start..self.remaining_range.start + len;
86 self.remaining_range.start += len;
87 let ptr = self.current_ptr;
88 self.current_ptr = self.current_ptr.wrapping_add(len);
89
90 let slice = unsafe { MutPtrByteSlice::new(slice_from_raw_parts_mut(ptr, len)) };
94 OwnedBuffer::new(slice, child_range, self.inner.clone() as Arc<dyn BufferAllocator>)
95 }
96}
97
98#[cfg(test)]
99mod tests {
100 use super::*;
101 use crate::buffer_allocator::{BufferAllocator as PoolBufferAllocator, BufferSource};
102
103 #[fuchsia::test]
104 async fn test_splittable_buffer_handle_into_buffer() {
105 let source = BufferSource::new(4096);
106 let pool = Arc::new(PoolBufferAllocator::new(512, source));
107 let owned = pool.allocate_buffer_sync_owned(2048);
108
109 let (mut splittable, handle) = SplittableBuffer::new(owned);
110 let mut child1 = splittable.take_prefix(1024);
111 let mut child2 = splittable.take_prefix(1024);
112 child1.as_mut_ptr_slice().fill(0x33);
113 child2.as_mut_ptr_slice().fill(0x44);
114
115 drop(splittable);
117 drop(child1);
118
119 let handle_clone = handle.clone();
120 assert!(handle_clone.into_buffer().is_none());
121
122 drop(child2);
124 let merged = handle.into_buffer().expect("into_buffer must succeed when sole reference");
125 assert_eq!(merged.len(), 2048);
126 assert!(merged.as_ptr_slice().subslice(0..1024).iter_as::<u8>().all(|b| b.read() == 0x33));
127 assert!(
128 merged.as_ptr_slice().subslice(1024..2048).iter_as::<u8>().all(|b| b.read() == 0x44)
129 );
130 }
131}