Skip to main content

memory_mapped_vmo/
memory_mapped_arc_vmo.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 std::marker::PhantomData;
6use std::mem::{ManuallyDrop, size_of};
7use std::ops::{Deref, DerefMut};
8use std::sync::atomic::{AtomicUsize, Ordering};
9
10use crate::ReadWriteMemoryMappedVmo;
11use crate::memory_mapped_vmo::MemoryMappedVmo;
12
13struct ControlBlock<M> {
14    metadata: M,
15    ref_count: AtomicUsize,
16}
17
18/// Reference-counted MemoryMappedVmo backed by an auxiliary control VMO to store the reference
19/// count and, optionally, extra shared metadata.
20pub struct MemoryMappedArcVmo<const WRITABLE: bool, M: Send + Sync = ()> {
21    control: ManuallyDrop<ReadWriteMemoryMappedVmo>,
22    storage: ManuallyDrop<MemoryMappedVmo<WRITABLE>>,
23    _marker: PhantomData<M>,
24}
25
26impl<const WRITABLE: bool, M: Send + Sync> MemoryMappedArcVmo<WRITABLE, M> {
27    /// Wraps a MemoryMappedVmo into a MemoryMappedArcVmo with the given extra metadata.
28    ///
29    /// # Safety
30    /// The caller must either guarantee that the VMO is not modified by others while the returned
31    /// instance is alive or that accesses are synchronized in application-specific ways.
32    pub unsafe fn wrap(
33        storage: MemoryMappedVmo<WRITABLE>,
34        metadata: M,
35    ) -> Result<Self, zx::Status> {
36        unsafe { Self::wrap_internal(storage, metadata, None) }
37    }
38
39    /// Like wrap, but also sets the control VMO name.
40    ///
41    /// # Safety
42    /// The caller must either guarantee that the VMO is not modified by others while the returned
43    /// instance is alive or that accesses are synchronized in application-specific ways.
44    pub unsafe fn wrap_with_name(
45        storage: MemoryMappedVmo<WRITABLE>,
46        metadata: M,
47        control_vmo_name: &zx::Name,
48    ) -> Result<Self, zx::Status> {
49        unsafe { Self::wrap_internal(storage, metadata, Some(control_vmo_name)) }
50    }
51
52    unsafe fn wrap_internal(
53        storage: MemoryMappedVmo<WRITABLE>,
54        metadata: M,
55        control_vmo_name: Option<&zx::Name>,
56    ) -> Result<Self, zx::Status> {
57        // Create the control VMO.
58        let control_vmo = zx::Vmo::create(size_of::<ControlBlock<M>>() as u64)?;
59        if let Some(name) = control_vmo_name {
60            control_vmo.set_name(name)?;
61        }
62
63        // Fill it with initial data.
64        let mut control = unsafe { ReadWriteMemoryMappedVmo::new(&control_vmo)? };
65        let control_ptr = control.validate_and_get_mut_ptr::<ControlBlock<M>>(0, 1).unwrap();
66        unsafe {
67            control_ptr.write(ControlBlock { metadata, ref_count: AtomicUsize::new(1) });
68        }
69
70        Ok(MemoryMappedArcVmo {
71            control: ManuallyDrop::new(control),
72            storage: ManuallyDrop::new(storage),
73            _marker: PhantomData,
74        })
75    }
76
77    fn control_block(&self) -> &ControlBlock<M> {
78        let ptr = self.control.validate_and_get_ptr::<ControlBlock<M>>(0, 1).unwrap();
79        unsafe { &*ptr }
80    }
81
82    /// Returns a shared reference to the metadata stored in the control block.
83    pub fn metadata(&self) -> &M {
84        &self.control_block().metadata
85    }
86}
87
88impl<M: Send + Sync> MemoryMappedArcVmo<true, M> {
89    /// Downgrades a read-write mapping handle into a read-only mapping handle without unmapping.
90    pub fn into_readonly(self) -> MemoryMappedArcVmo<false, M> {
91        let mut this = ManuallyDrop::new(self);
92        // SAFETY: self is now wrapped in ManuallyDrop. Therefore, control and storage are moved
93        // once without running self's Drop implementation.
94        unsafe {
95            MemoryMappedArcVmo {
96                control: ManuallyDrop::new(ManuallyDrop::take(&mut this.control)),
97                storage: ManuallyDrop::new(ManuallyDrop::take(&mut this.storage).into_readonly()),
98                _marker: PhantomData,
99            }
100        }
101    }
102}
103
104impl<const WRITABLE: bool, M: Send + Sync> Deref for MemoryMappedArcVmo<WRITABLE, M> {
105    type Target = MemoryMappedVmo<WRITABLE>;
106
107    fn deref(&self) -> &Self::Target {
108        &self.storage
109    }
110}
111
112impl<M: Send + Sync> DerefMut for MemoryMappedArcVmo<true, M> {
113    fn deref_mut(&mut self) -> &mut Self::Target {
114        &mut self.storage
115    }
116}
117
118impl<const WRITABLE: bool, M: Send + Sync> Clone for MemoryMappedArcVmo<WRITABLE, M> {
119    fn clone(&self) -> Self {
120        // Increment the reference count.
121        self.control_block().ref_count.fetch_add(1, Ordering::Relaxed);
122
123        // SAFETY: ref_count was just incremented, ensuring Drop will only unmap control and storage
124        // when the last clone is dropped.
125        unsafe {
126            MemoryMappedArcVmo {
127                control: self.control.create_alias_with_manually_drop(),
128                storage: self.storage.create_alias_with_manually_drop(),
129                _marker: PhantomData,
130            }
131        }
132    }
133}
134
135impl<const WRITABLE: bool, M: Send + Sync> Drop for MemoryMappedArcVmo<WRITABLE, M> {
136    fn drop(&mut self) {
137        if self.control_block().ref_count.fetch_sub(1, Ordering::AcqRel) != 1 {
138            return; // This was not the last reference.
139        }
140
141        let control_ptr = self.control.validate_and_get_mut_ptr::<ControlBlock<M>>(0, 1).unwrap();
142        // SAFETY: We are the last reference to control and storage. No others exist anymore.
143        unsafe {
144            ManuallyDrop::drop(&mut self.storage);
145
146            std::ptr::drop_in_place(control_ptr);
147            ManuallyDrop::drop(&mut self.control);
148        }
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use crate::{ReadWriteMemoryMappedArcVmo, ReadWriteMemoryMappedVmo};
155    use std::mem::size_of;
156    use std::sync::atomic::{AtomicBool, Ordering};
157
158    const TEST_DATA: u64 = 42;
159    const TEST_DATA_SIZE: usize = size_of::<u64>();
160
161    struct DropOnceChecker<'a>(&'a AtomicBool);
162
163    impl Drop for DropOnceChecker<'_> {
164        fn drop(&mut self) {
165            assert!(!self.0.swap(true, Ordering::SeqCst));
166        }
167    }
168
169    #[test]
170    fn test_clone_and_drop_metadata() {
171        let vmo = zx::Vmo::create(TEST_DATA_SIZE as u64).unwrap();
172
173        // Create an initial MemoryMappedVmo and write data into it.
174        let mut rw = unsafe { ReadWriteMemoryMappedVmo::new(&vmo) }.unwrap();
175        *rw.get_object_mut::<u64>(0).unwrap() = TEST_DATA;
176
177        // Wrap with reference-counting and with sentinel metadata that reports when the mapping
178        // is dropped.
179        let dropped = AtomicBool::new(false);
180        let m1 =
181            unsafe { ReadWriteMemoryMappedArcVmo::wrap(rw, DropOnceChecker(&dropped)) }.unwrap();
182        let m2 = m1.clone();
183
184        // Drop one reference.
185        assert_eq!(m2.control_block().ref_count.load(Ordering::SeqCst), 2);
186        drop(m1);
187
188        // Verify that the other reference is still valid.
189        assert_eq!(m2.control_block().ref_count.load(Ordering::SeqCst), 1);
190        assert_eq!(dropped.load(Ordering::SeqCst), false);
191        assert_eq!(*m2.get_object::<u64>(0).unwrap(), TEST_DATA);
192
193        // Verify downgrading into a read-only mapping.
194        let m2 = m2.into_readonly();
195        assert_eq!(m2.control_block().ref_count.load(Ordering::SeqCst), 1);
196        assert_eq!(dropped.load(Ordering::SeqCst), false);
197        assert_eq!(*m2.get_object::<u64>(0).unwrap(), TEST_DATA);
198
199        drop(m2);
200        assert_eq!(dropped.load(Ordering::SeqCst), true);
201    }
202}