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
23/// A clonable handle referencing the underlying `SplittableBufferInner`.
24///
25/// Can be captured inside background sub-read callbacks and consumed via `into_buffer(self)`
26/// once all child `OwnedBuffer`s and other handles drop to recover the merged `OwnedBuffer`.
27#[derive(Clone, Debug)]
28pub struct SplittableBufferHandle {
29    inner: Arc<SplittableBufferInner>,
30}
31
32impl SplittableBufferHandle {
33    /// Consumes this handle and attempts to unwrap the underlying `Arc<SplittableBufferInner>`.
34    ///
35    /// Returns `Some(parent_buffer)` if this handle and all child `OwnedBuffer`s carved from
36    /// `SplittableBuffer` have been dropped (`Arc::strong_count == 1`). Otherwise returns `None`.
37    pub fn into_buffer(self) -> Option<OwnedBuffer> {
38        Arc::into_inner(self.inner).map(|inner| inner.parent_buffer)
39    }
40}
41
42/// A wrapper around `OwnedBuffer` that allows carving out independent child `OwnedBuffer`s
43/// and recovering the original `OwnedBuffer` via a `SplittableBufferHandle` once all child
44/// buffers have been dropped.
45#[derive(Debug)]
46pub struct SplittableBuffer {
47    inner: Arc<SplittableBufferInner>,
48    current_ptr: *mut u8,
49    remaining_range: Range<usize>,
50}
51
52// SAFETY: `current_ptr` points into `inner.parent_buffer`'s VMO / memory region, which can be
53// sent across threads.
54unsafe impl Send for SplittableBuffer {}
55unsafe impl Sync for SplittableBuffer {}
56
57impl SplittableBuffer {
58    /// Creates a new `SplittableBuffer` along with a clonable `SplittableBufferHandle` that can be
59    /// used to recover the merged `buffer` once all child `OwnedBuffer`s and other handles drop.
60    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    /// Returns the remaining unallocated range available for splitting.
70    pub fn remaining_range(&self) -> Range<usize> {
71        self.remaining_range.clone()
72    }
73
74    /// Carves out the first `len` bytes of the remaining unsplit buffer as a new `OwnedBuffer`.
75    ///
76    /// # Panics
77    ///
78    /// Panics if `len` exceeds `remaining_range.len()`.
79    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        // SAFETY: `child_range` is strictly within the original parent buffer bounds and
87        // never overlaps with any other prefix taken from `remaining_range`. The
88        // `Arc<SplittableBufferInner>` keeps the parent `OwnedBuffer` alive for `'static`.
89        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 and child1 first while child2 is still active.
112        drop(splittable);
113        drop(child1);
114
115        let handle_clone = handle.clone();
116        assert!(handle_clone.into_buffer().is_none());
117
118        // Drop child2. Now only handle remains (strong count == 1), so into_buffer succeeds!
119        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}