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 anyhow::Error;
7use fuchsia_sync::Mutex;
8use std::fmt::Debug;
9use std::ops::Range;
10use std::ptr::slice_from_raw_parts_mut;
11use std::sync::Arc;
12use storage_ptr_slice::MutPtrByteSlice;
13
14type CompletionCallback = Box<dyn FnOnce(Result<OwnedBuffer, Error>) + Send>;
15
16enum State {
17    /// Initial state while `split`'s closure `f` is executing. If an error is reported before
18    /// `f` finishes, `failed` records that error.
19    Pending { failed: Option<Error> },
20    /// `f` succeeded; all sub-operations are in flight.
21    Dispatched { parent_buffer: OwnedBuffer, callback: CompletionCallback },
22    /// An error occurred or the callback has already been invoked.
23    Completed,
24}
25
26struct SplittableBufferInner {
27    state: Mutex<State>,
28    parent_allocator: Arc<dyn BufferAllocator>,
29    is_trusted: bool,
30}
31
32impl Debug for SplittableBufferInner {
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        f.debug_struct("SplittableBufferInner").finish_non_exhaustive()
35    }
36}
37
38impl BufferAllocator for SplittableBufferInner {
39    fn free_buffer(&self, _range: Range<usize>) {
40        // No-op: Dropping the child `OwnedBuffer` drops its `Arc<dyn BufferAllocator>`,
41        // which automatically decrements the `Arc` reference count of `SplittableBufferInner`.
42    }
43
44    fn identifier(&self) -> usize {
45        self.parent_allocator.identifier()
46    }
47
48    fn is_trusted(&self) -> bool {
49        self.is_trusted
50    }
51
52    #[cfg(target_os = "fuchsia")]
53    fn vmo(&self) -> Option<Arc<zx::Vmo>> {
54        self.parent_allocator.vmo()
55    }
56
57    #[cfg(target_os = "fuchsia")]
58    fn paddrs(&self, range: &Range<usize>) -> Option<(&[zx::sys::zx_paddr_t], u64)> {
59        self.parent_allocator.paddrs(range)
60    }
61}
62
63impl Drop for SplittableBufferInner {
64    fn drop(&mut self) {
65        if let State::Dispatched { parent_buffer, callback } =
66            std::mem::replace(self.state.get_mut(), State::Completed)
67        {
68            (callback)(Ok(parent_buffer));
69        }
70    }
71}
72
73/// A handle for an individual in-flight sub-read operation.
74///
75/// Must be consumed via [`merge`] on completion. If dropped before `merge` is called, it
76/// automatically fails the parent `SplittableBuffer` operation.
77#[derive(Debug)]
78pub struct SubHandle {
79    inner: Arc<SplittableBufferInner>,
80    completed: bool,
81}
82
83impl SubHandle {
84    /// Invoked in sub-read callbacks. If `f()` returns `Err(e)`, disarms the callback and delivers
85    /// `Err(e)` immediately.
86    pub fn merge(mut self, f: impl FnOnce() -> Result<(), Error>) {
87        self.completed = true;
88        if let Err(e) = f() {
89            let mut guard = self.inner.state.lock();
90            match &mut *guard {
91                State::Pending { failed } => {
92                    failed.get_or_insert(e);
93                }
94                State::Dispatched { .. } => {
95                    if let State::Dispatched { callback, .. } =
96                        std::mem::replace(&mut *guard, State::Completed)
97                    {
98                        (callback)(Err(e));
99                    }
100                }
101                State::Completed => {}
102            }
103        }
104    }
105}
106
107impl Drop for SubHandle {
108    fn drop(&mut self) {
109        if !self.completed {
110            let mut guard = self.inner.state.lock();
111            match &mut *guard {
112                State::Pending { failed } => {
113                    failed.get_or_insert_with(|| {
114                        anyhow::anyhow!("Read sub-request dropped before completion")
115                    });
116                }
117                State::Dispatched { .. } => {
118                    if let State::Dispatched { callback, .. } =
119                        std::mem::replace(&mut *guard, State::Completed)
120                    {
121                        (callback)(Err(anyhow::anyhow!(
122                            "Read sub-request dropped before completion"
123                        )));
124                    }
125                }
126                State::Completed => {}
127            }
128        }
129    }
130}
131
132/// A wrapper around `OwnedBuffer` that allows carving out independent child `OwnedBuffer`s
133/// and delivering the reconstructed original `OwnedBuffer` to a completion callback once all child
134/// buffers have been dropped.
135#[derive(Debug)]
136pub struct SplittableBuffer {
137    inner: Arc<SplittableBufferInner>,
138    current_ptr: *mut u8,
139    remaining_range: Range<usize>,
140}
141
142// SAFETY: `current_ptr` points into `inner.parent_buffer`'s VMO / memory region, which can be
143// sent across threads.
144unsafe impl Send for SplittableBuffer {}
145unsafe impl Sync for SplittableBuffer {}
146
147impl SplittableBuffer {
148    /// Returns the remaining unallocated range available for splitting.
149    pub fn remaining_range(&self) -> Range<usize> {
150        self.remaining_range.clone()
151    }
152
153    /// Zeroes the next `len` bytes in the buffer and advances the offset.
154    ///
155    /// # Panics
156    ///
157    /// Panics if `len` exceeds `remaining_range.len()`.
158    pub fn fill_zeros(&mut self, len: usize) {
159        assert!(len <= self.remaining_range.len());
160        self.remaining_range.start += len;
161        let ptr = self.current_ptr;
162        self.current_ptr = self.current_ptr.wrapping_add(len);
163
164        // SAFETY: `ptr` points into `inner.parent_buffer`'s allocated memory of at least
165        // `len` bytes.
166        unsafe {
167            std::ptr::write_bytes(ptr, 0, len);
168        }
169    }
170
171    /// Carves out the first `len` bytes of the remaining unsplit buffer as an `OwnedBuffer`
172    /// along with a [`SubHandle`] to track its asynchronous completion.
173    ///
174    /// # Panics
175    ///
176    /// Panics if `len` exceeds `remaining_range.len()`.
177    pub fn take_prefix(&mut self, len: usize) -> (OwnedBuffer, SubHandle) {
178        assert!(len <= self.remaining_range.len());
179        let child_range = self.remaining_range.start..self.remaining_range.start + len;
180        self.remaining_range.start += len;
181        let ptr = self.current_ptr;
182        self.current_ptr = self.current_ptr.wrapping_add(len);
183
184        // SAFETY: `child_range` is strictly within the original parent buffer bounds and
185        // never overlaps with any other prefix taken from `remaining_range`. The
186        // `Arc<SplittableBufferInner>` keeps the parent `OwnedBuffer` alive for `'static`.
187        let slice = unsafe { MutPtrByteSlice::new(slice_from_raw_parts_mut(ptr, len)) };
188        let buffer =
189            OwnedBuffer::new(slice, child_range, self.inner.clone() as Arc<dyn BufferAllocator>);
190        let sub_handle = SubHandle { inner: self.inner.clone(), completed: false };
191        (buffer, sub_handle)
192    }
193}
194
195impl OwnedBuffer {
196    /// Splits the buffer into one or more child `OwnedBuffer`s via `f`.
197    ///
198    /// When all child `OwnedBuffer`s and [`SubHandle`]s drop, the merged parent buffer is
199    /// automatically delivered to `on_complete(Ok(parent_buffer))`. If an error occurs via `f`
200    /// or during sub-operation execution, `on_complete(Err(e))` is delivered.
201    pub fn split<R>(
202        mut self,
203        f: impl FnOnce(&mut SplittableBuffer) -> Result<R, Error>,
204        on_complete: impl FnOnce(Result<OwnedBuffer, Error>) + Send + 'static,
205    ) -> Result<R, Error> {
206        let is_trusted = self.try_as_slice().is_some();
207        let remaining_range = self.range();
208        let current_ptr = self.as_mut_ptr();
209        let inner = Arc::new(SplittableBufferInner {
210            state: Mutex::new(State::Pending { failed: None }),
211            parent_allocator: self.allocator.clone(),
212            is_trusted,
213        });
214        let mut splittable =
215            SplittableBuffer { inner: inner.clone(), current_ptr, remaining_range };
216        let res = f(&mut splittable);
217        match res {
218            Ok(val) => {
219                let mut guard = inner.state.lock();
220                match std::mem::replace(&mut *guard, State::Completed) {
221                    State::Pending { failed: Some(e) } => {
222                        on_complete(Err(e));
223                    }
224                    State::Pending { failed: None } => {
225                        *guard = State::Dispatched {
226                            parent_buffer: self,
227                            callback: Box::new(on_complete),
228                        };
229                    }
230                    State::Dispatched { .. } | State::Completed => {}
231                }
232                Ok(val)
233            }
234            Err(e) => {
235                *inner.state.lock() = State::Completed;
236                Err(e)
237            }
238        }
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245    use crate::buffer_allocator::{BufferAllocator as PoolBufferAllocator, BufferSource};
246    use anyhow::anyhow;
247    use std::sync::atomic::{AtomicBool, Ordering};
248
249    #[fuchsia::test]
250    async fn test_splittable_buffer_with_callback_success() {
251        let source = BufferSource::new(4096);
252        let pool = Arc::new(PoolBufferAllocator::new(512, source));
253        let owned = pool.allocate_buffer_sync_owned(2048);
254
255        let completed = Arc::new(AtomicBool::new(false));
256        let completed_clone = completed.clone();
257
258        let mut sub1_opt = None;
259        let mut sub2_opt = None;
260        let mut child1_opt = None;
261        let mut child2_opt = None;
262
263        owned
264            .split(
265                |splittable| {
266                    let (mut child1, sub1) = splittable.take_prefix(1024);
267                    let (mut child2, sub2) = splittable.take_prefix(1024);
268                    child1.as_mut_ptr_slice().fill(0x11);
269                    child2.as_mut_ptr_slice().fill(0x22);
270                    sub1_opt = Some(sub1);
271                    sub2_opt = Some(sub2);
272                    child1_opt = Some(child1);
273                    child2_opt = Some(child2);
274                    Ok(())
275                },
276                move |res| {
277                    let merged = res.expect("must succeed");
278                    assert_eq!(merged.len(), 2048);
279                    assert!(
280                        merged.as_ptr_slice().subslice(0..1024).iter_as::<u8>().all(|b| b == 0x11)
281                    );
282                    assert!(
283                        merged
284                            .as_ptr_slice()
285                            .subslice(1024..2048)
286                            .iter_as::<u8>()
287                            .all(|b| b == 0x22)
288                    );
289                    completed_clone.store(true, Ordering::Relaxed);
290                },
291            )
292            .unwrap();
293
294        let sub1 = sub1_opt.unwrap();
295        let sub2 = sub2_opt.unwrap();
296        let child1 = child1_opt.unwrap();
297        let child2 = child2_opt.unwrap();
298
299        sub1.merge(|| Ok(()));
300        drop(child1);
301        assert!(!completed.load(Ordering::Relaxed));
302
303        // When child2 drops and sub2 is merged, all references are gone and callback is fired.
304        sub2.merge(|| Ok(()));
305        drop(child2);
306        assert!(completed.load(Ordering::Relaxed));
307    }
308
309    #[fuchsia::test]
310    async fn test_splittable_buffer_with_callback_async_error() {
311        let source = BufferSource::new(4096);
312        let pool = Arc::new(PoolBufferAllocator::new(512, source));
313        let owned = pool.allocate_buffer_sync_owned(2048);
314
315        let err_received = Arc::new(AtomicBool::new(false));
316        let err_clone = err_received.clone();
317
318        let mut sub1_opt = None;
319        let mut sub2_opt = None;
320        let mut child1_opt = None;
321        let mut child2_opt = None;
322
323        owned
324            .split(
325                |splittable| {
326                    let (child1, sub1) = splittable.take_prefix(1024);
327                    let (child2, sub2) = splittable.take_prefix(1024);
328                    sub1_opt = Some(sub1);
329                    sub2_opt = Some(sub2);
330                    child1_opt = Some(child1);
331                    child2_opt = Some(child2);
332                    Ok(())
333                },
334                move |res| {
335                    assert!(res.is_err());
336                    err_clone.store(true, Ordering::Relaxed);
337                },
338            )
339            .unwrap();
340
341        let sub1 = sub1_opt.unwrap();
342        let sub2 = sub2_opt.unwrap();
343        let child1 = child1_opt.unwrap();
344        let child2 = child2_opt.unwrap();
345
346        // Merge error from chunk 1:
347        sub1.merge(|| Err(anyhow!("chunk 1 failure")));
348        assert!(err_received.load(Ordering::Relaxed));
349
350        // Chunk 2 drops later without error, but callback was already consumed:
351        drop(child1);
352        drop(child2);
353        sub2.merge(|| Ok(()));
354    }
355
356    #[fuchsia::test]
357    async fn test_splittable_buffer_sub_handle_dropped_fails_operation() {
358        let source = BufferSource::new(4096);
359        let pool = Arc::new(PoolBufferAllocator::new(512, source));
360        let owned = pool.allocate_buffer_sync_owned(2048);
361
362        let err_received = Arc::new(AtomicBool::new(false));
363        let err_clone = err_received.clone();
364
365        let mut sub1 = None;
366        let mut child1 = None;
367        owned
368            .split(
369                |splittable| {
370                    let (c, s) = splittable.take_prefix(1024);
371                    child1 = Some(c);
372                    sub1 = Some(s);
373                    Ok(())
374                },
375                move |res| {
376                    assert!(res.is_err());
377                    err_clone.store(true, Ordering::Relaxed);
378                },
379            )
380            .unwrap();
381
382        drop(child1);
383
384        // sub1 is dropped without calling merge:
385        drop(sub1);
386        assert!(err_received.load(Ordering::Relaxed));
387    }
388
389    #[fuchsia::test]
390    async fn test_splittable_buffer_sync_error() {
391        let source = BufferSource::new(4096);
392        let pool = Arc::new(PoolBufferAllocator::new(512, source));
393        let owned = pool.allocate_buffer_sync_owned(2048);
394
395        let callback_called = Arc::new(AtomicBool::new(false));
396        let callback_clone = callback_called.clone();
397
398        let mut sub1 = None;
399        let mut child1 = None;
400        let res: Result<(), _> = owned.split(
401            |splittable| {
402                let (c, s) = splittable.take_prefix(1024);
403                child1 = Some(c);
404                sub1 = Some(s);
405                Err(anyhow!("synchronous dispatch failure"))
406            },
407            move |_res| {
408                callback_clone.store(true, Ordering::Relaxed);
409            },
410        );
411
412        assert!(res.is_err());
413        drop(child1);
414        drop(sub1);
415        assert!(!callback_called.load(Ordering::Relaxed));
416    }
417
418    #[fuchsia::test]
419    async fn test_splittable_buffer_fill_zeros() {
420        let source = BufferSource::new(4096);
421        let pool = Arc::new(PoolBufferAllocator::new(512, source));
422        let owned = pool.allocate_buffer_sync_owned(2048);
423
424        owned
425            .split(
426                |splittable| {
427                    splittable.fill_zeros(2048);
428                    Ok(())
429                },
430                |res| {
431                    let buffer = res.unwrap();
432                    assert_eq!(buffer.len(), 2048);
433                    assert!(buffer.as_ptr_slice().iter_as::<u8>().all(|b| b == 0));
434                },
435            )
436            .unwrap();
437    }
438
439    #[cfg(target_os = "fuchsia")]
440    #[fuchsia::test]
441    async fn test_splittable_buffer_preserves_vmo_and_paddrs() {
442        let fake_bti = fake_bti::FakeBti::create().unwrap();
443        fake_bti.set_paddrs(&[4096, 8192]);
444        let source = BufferSource::new(8192);
445        let allocator =
446            Arc::new(crate::pinned_buffer_allocator::PinnedBufferAllocator::with_chunk_size(
447                4096,
448                source,
449                fake_bti.duplicate_handle(zx::Rights::SAME_RIGHTS).unwrap(),
450                4096,
451                8192,
452            ));
453        let owned = allocator.allocate_buffer_sync_owned(8192);
454        let expected_vmo_koid = owned.vmo().unwrap().koid().unwrap();
455        let expected_paddrs: Vec<_> = owned.paddrs().unwrap().to_vec();
456
457        owned
458            .split(
459                |splittable| {
460                    let (child1, sub1) = splittable.take_prefix(4096);
461                    let (child2, sub2) = splittable.take_prefix(4096);
462                    assert_eq!(child1.vmo().unwrap().koid().unwrap(), expected_vmo_koid);
463                    assert_eq!(child2.vmo().unwrap().koid().unwrap(), expected_vmo_koid);
464                    assert_eq!(child1.contiguity(), Some(4096));
465                    assert_eq!(child2.contiguity(), Some(4096));
466                    assert_eq!(child1.paddrs().unwrap(), &expected_paddrs[0..1]);
467                    assert_eq!(child2.paddrs().unwrap(), &expected_paddrs[1..2]);
468                    sub1.merge(|| Ok(()));
469                    sub2.merge(|| Ok(()));
470                    Ok(())
471                },
472                |res| {
473                    assert!(res.is_ok());
474                },
475            )
476            .unwrap();
477    }
478}
storage_device/splittable_buffer.rs

storage_device/
splittable_buffer.rs
