Skip to main content

storage_device/
splittable_buffer.rs

1// Copyright 2026 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use 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        // No-op: Dropping the child `OwnedBuffer` drops its `Arc<dyn BufferAllocator>`,
19        // which automatically decrements the `Arc` reference count of `SplittableBufferInner`.
20    }
21
22    fn is_trusted(&self) -> bool {
23        self.parent_buffer.try_as_slice().is_some()
24    }
25}
26
27/// A clonable handle referencing the underlying `SplittableBufferInner`.
28///
29/// Can be captured inside background sub-read callbacks and consumed via `into_buffer(self)`
30/// once all child `OwnedBuffer`s and other handles drop to recover the merged `OwnedBuffer`.
31#[derive(Clone, Debug)]
32pub struct SplittableBufferHandle {
33    inner: Arc<SplittableBufferInner>,
34}
35
36impl SplittableBufferHandle {
37    /// Consumes this handle and attempts to unwrap the underlying `Arc<SplittableBufferInner>`.
38    ///
39    /// Returns `Some(parent_buffer)` if this handle and all child `OwnedBuffer`s carved from
40    /// `SplittableBuffer` have been dropped (`Arc::strong_count == 1`). Otherwise returns `None`.
41    pub fn into_buffer(self) -> Option<OwnedBuffer> {
42        Arc::into_inner(self.inner).map(|inner| inner.parent_buffer)
43    }
44}
45
46/// A wrapper around `OwnedBuffer` that allows carving out independent child `OwnedBuffer`s
47/// and recovering the original `OwnedBuffer` via a `SplittableBufferHandle` once all child
48/// buffers have been dropped.
49#[derive(Debug)]
50pub struct SplittableBuffer {
51    inner: Arc<SplittableBufferInner>,
52    current_ptr: *mut u8,
53    remaining_range: Range<usize>,
54}
55
56// SAFETY: `current_ptr` points into `inner.parent_buffer`'s VMO / memory region, which can be
57// sent across threads.
58unsafe impl Send for SplittableBuffer {}
59unsafe impl Sync for SplittableBuffer {}
60
61impl SplittableBuffer {
62    /// Creates a new `SplittableBuffer` along with a clonable `SplittableBufferHandle` that can be
63    /// used to recover the merged `buffer` once all child `OwnedBuffer`s and other handles drop.
64    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    /// Returns the remaining unallocated range available for splitting.
74    pub fn remaining_range(&self) -> Range<usize> {
75        self.remaining_range.clone()
76    }
77
78    /// Carves out the first `len` bytes of the remaining unsplit buffer as a new `OwnedBuffer`.
79    ///
80    /// # Panics
81    ///
82    /// Panics if `len` exceeds `remaining_range.len()`.
83    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        // SAFETY: `child_range` is strictly within the original parent buffer bounds and
91        // never overlaps with any other prefix taken from `remaining_range`. The
92        // `Arc<SplittableBufferInner>` keeps the parent `OwnedBuffer` alive for `'static`.
93        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 and child1 first while child2 is still active.
116        drop(splittable);
117        drop(child1);
118
119        let handle_clone = handle.clone();
120        assert!(handle_clone.into_buffer().is_none());
121
122        // Drop child2. Now only handle remains (strong count == 1), so into_buffer succeeds!
123        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}