manual_share/
shared_box.rs

1//! # Manually shared box
2//! ```
3//! use std::thread;
4//! use manual_share::SharedBox;
5//!
6//! let b = Box::new(13);
7//! let mut b = SharedBox::from_box(b);
8//!
9//! let br = b.borrow();
10//! let j1 = thread::spawn(move || {
11//!     println!("{}", br.get());
12//!     br
13//! });
14//!
15//! let br = b.borrow();
16//! let j2 = thread::spawn(move || {
17//!     println!("{}", br.get() + 1);
18//!     br
19//! });
20//!
21//! b.try_return(j1.join().unwrap()).unwrap();
22//! b.try_return(j2.join().unwrap()).unwrap();
23//!
24//! let b = b.try_into_box().unwrap();
25//! println!("{:?}", b);
26//! ```
27
28/// A structure owning the original `Box`
29/// which can be used to create multiple `SharedBoxRef` to send to other thread.
30/// It uses a counter to record the number of `SharedBoxRef` that has been created and not given back.
31///
32/// Dropping `SharedBox` without returning all `SharedBoxRef` that it has created leaks its heap memory.
33/// When **`panic-on-drop`** feature is enabled, it will panic:
34/// ```should_panic
35/// let r = {
36///     let mut b = manual_share::SharedBox::new(0);
37///     b.borrow()
38///
39///     // dropping SharedBox, panic here
40/// };
41/// println!("{}", r.get());
42/// ```
43///
44/// When all `SharedBoxRef` have been returned back to `SharedBox`,
45/// it will free and properly release its heap memory when `SharedBox` is dropped.
46/// ```
47/// let mut b = manual_share::SharedBox::new(0);
48/// let r = b.borrow();
49/// b.try_return(r).unwrap();
50///
51/// // Dropping SharedBox here will free the heap memory it owns. No panic will occur.
52/// ```
53#[derive(Debug)]
54pub struct SharedBox<T: ?Sized> {
55    borrow_count: usize,
56    ptr: *mut T,
57}
58
59impl<T> SharedBox<T> {
60    /// Create a `SharedBox` by creating a `Box` first,
61    /// then convert it to a `SharedBox`.
62    pub fn new(value: T) -> Self {
63        let b = Box::new(value);
64        Self::from_box(b)
65    }
66}
67
68impl<T: ?Sized> SharedBox<T> {
69    /// Create a `SharedBox` by consuming a `Box`.
70    pub fn from_box(unique: Box<T>) -> Self {
71        Self {
72            borrow_count: 0,
73            ptr: Box::into_raw(unique),
74        }
75    }
76
77    /// Create a `SharedBoxRef` and increase the borrow count.
78    /// # panics
79    /// Panics when borrow count overflows `usize`.
80    pub fn borrow(&mut self) -> SharedBoxRef<T> {
81        self.borrow_count = self.borrow_count.checked_add(1).unwrap();
82        SharedBoxRef { ptr: self.ptr }
83    }
84
85    /// Try to return back the `SharedBoxRef`.
86    /// Returns `Err` if the `SharedBoxRef` does not originate from the same `SharedBox`.
87    ///
88    /// Decrease the borrow count if not error occurs.
89    ///
90    /// ```
91    /// use manual_share::SharedBox;
92    ///
93    /// let mut b1 = SharedBox::from_box(Box::new(8));
94    /// let r1 = b1.borrow();
95    /// b1.try_return(r1).unwrap();
96    ///
97    /// let mut b2 = SharedBox::from_box(Box::new(9));
98    /// let r2 = b2.borrow();
99    ///
100    /// // Giving SharedBoxRef to the wrong SharedBox returns Err.
101    /// let r2 = b1.try_return(r2).unwrap_err();
102    ///
103    /// b2.try_return(r2).unwrap();
104    /// ```
105    pub fn try_return(&mut self, reference: SharedBoxRef<T>) -> Result<(), SharedBoxRef<T>> {
106        if !core::ptr::eq(self.ptr, reference.ptr) {
107            return Err(reference);
108        }
109
110        if size_of_val(unsafe { &*self.ptr }) == 0 {
111            // ZST types can have multiple allocations to the same address, so we need to check for overflow.
112            if let Some(new_count) = self.borrow_count.checked_sub(1) {
113                self.borrow_count = new_count;
114                let _ = core::mem::ManuallyDrop::new(reference);
115                Ok(())
116            } else {
117                Err(reference)
118            }
119        } else {
120            self.borrow_count -= 1;
121            let _ = core::mem::ManuallyDrop::new(reference);
122            Ok(())
123        }
124    }
125
126    /// Try to convert `Self` into a `Box` if all borrowed `SharedBoxRef` has been given back.
127    ///
128    /// ```
129    /// use manual_share::SharedBox;
130    ///
131    /// let b = Box::new(0);
132    /// let mut b = SharedBox::from_box(b);
133    ///
134    /// let r = b.borrow();
135    ///
136    /// // Try to convert to Box without returning all SharedBoxRef returns Err.
137    /// let mut b = b.try_into_box().unwrap_err();
138    ///
139    /// b.try_return(r).unwrap();
140    ///
141    /// let b = b.try_into_box().unwrap();
142    /// assert_eq!(b, Box::new(0));
143    /// ```
144    pub fn try_into_box(self) -> Result<Box<T>, Self> {
145        if self.borrow_count > 0 {
146            Err(self)
147        } else {
148            let r = core::mem::ManuallyDrop::new(self);
149            Ok(unsafe { Box::from_raw(r.ptr) })
150        }
151    }
152    /// Directly get a reference to the value inside the `SharedBox`.
153    /// This use rust built-in lifetime check to ensure the reference is valid as long as the `SharedBox` is alive,
154    /// and has no runtime overhead.
155    pub fn get(&self) -> &T {
156        // SAFETY:
157        // The pointer is valid as long as the SharedBox is alive.
158        // All other references can only get immutable reference.
159        unsafe { &*self.ptr }
160    }
161}
162
163/// See https://users.rust-lang.org/t/built-a-crate-to-safely-share-box-and-vec-manually/141138/25?u=newdino for why it require `Sync`.
164unsafe impl<T: Send + Sync> Send for SharedBox<T> {}
165
166unsafe impl<T: Sync> Sync for SharedBox<T> {}
167
168impl<T: ?Sized> Drop for SharedBox<T> {
169    fn drop(&mut self) {
170        #[cfg(feature = "panic-on-drop")]
171        {
172            // Let user deal with other panics.
173            #[cfg(feature = "do-not-panic-when-panicking")]
174            if std::thread::panicking() {
175                return;
176            }
177
178            if self.borrow_count > 0 {
179                panic!("Dropping a SharedBox without giving back all SharedBoxRef")
180            }
181        }
182        // Only drops when there are no outstanding SharedBoxRef values to prevent use-after-free.
183        if self.borrow_count == 0 {
184            unsafe {
185                drop(Box::from_raw(self.ptr));
186            }
187        }
188    }
189}
190
191/// A Reference to `SharedBox` that can be sent to other threads.
192///
193/// Dropping a `SharedBoxRef` leaks the heap memory it points to.
194/// When **`panic-on-drop`** feature is enabled, dropping it will panic:
195/// ```should_panic
196/// let mut b = manual_share::SharedBox::new(1);
197/// b.borrow();
198///
199/// // forget SharedBox to make sure the panic is not caused by dropping it first.
200/// std::mem::forget(b);
201///
202/// // panic here due to dropping SharedBoxRef
203/// ```
204///
205/// Use `SharedBox::try_return` to consume it without causing panic.
206/// ```
207/// let mut b = manual_share::SharedBox::new(1);
208/// let r = b.borrow();
209/// b.try_return(r).unwrap();
210/// ```
211#[derive(Debug)]
212pub struct SharedBoxRef<T: ?Sized> {
213    ptr: *const T,
214}
215
216/// `SharedBoxRef` is like `&T`, which only requires `T: Sync` to implement `Send`.
217unsafe impl<T: Sync> Send for SharedBoxRef<T> {}
218unsafe impl<T: Sync> Sync for SharedBoxRef<T> {}
219
220impl<T: ?Sized> SharedBoxRef<T> {
221    /// Example usage:
222    /// ```
223    /// let mut b = manual_share::SharedBox::new(42);
224    /// let r = b.borrow();
225    /// let value = *r.get();
226    /// assert_eq!(value, 42);
227    ///
228    /// b.try_return(r).unwrap();
229    /// ```
230    ///
231    /// The reference got from this method has the same lifetime of the `SharedBoxRef`,
232    /// which means it will be invalidated after `SharedBoxRef` is given back to `SharedBox`:
233    /// ```compile_fail
234    /// let mut b = manual_share::SharedBox::new(42);
235    /// let br = b.borrow();
236    /// let r = br.get();
237    ///
238    /// b.try_return(br).unwrap();
239    ///
240    /// // r is no longer valid here.
241    /// println!("{}", r);
242    /// ```
243    pub fn get(&self) -> &T {
244        unsafe { &*self.ptr }
245    }
246}
247
248impl<T: ?Sized> Drop for SharedBoxRef<T> {
249    fn drop(&mut self) {
250        #[cfg(feature = "panic-on-drop")]
251        {
252            #[cfg(feature = "do-not-panic-when-panicking")]
253            // Let user deal with other panics.
254            if std::thread::panicking() {
255                return;
256            }
257
258            panic!("SharedBoxRef should not be dropped. Use SharedBox::try_return to consume it.");
259        }
260    }
261}
262
263#[cfg(test)]
264mod test {
265    use super::*;
266
267    #[test]
268    fn zst() {
269        let mut b1 = SharedBox::new(());
270        let mut b2 = SharedBox::new(());
271
272        let r11 = b1.borrow();
273        let r12 = b1.borrow();
274
275        let r2 = b2.borrow();
276
277        b1.try_return(r2).unwrap();
278
279        b1.try_return(r11).unwrap();
280        let r12 = b1.try_return(r12).unwrap_err();
281
282        b2.try_return(r12).unwrap();
283    }
284}
manual_share/shared_box.rs

manual_share/
shared_box.rs
