Skip to main content

timely_bytes/
lib.rs

1//! A simplified implementation of the `bytes` crate, with different features, less safety.
2//!
3//! The crate is currently minimalist rather than maximalist, and for example does not support
4//! methods on `BytesMut` that seem like they should be safe, because they are not yet needed.
5//! For example, `BytesMut` should be able to implement `Send`, and `BytesMut::extract_to` could
6//! return a `BytesMut` rather than a `Bytes`.
7//!
8//! # Examples
9//!
10//! ```
11//! use timely_bytes::arc::BytesMut;
12//!
13//! let bytes = vec![0u8; 1024];
14//! let mut shared1 = BytesMut::from(bytes);
15//! let mut shared2 = shared1.extract_to(100);
16//! let mut shared3 = shared1.extract_to(100);
17//! let mut shared4 = shared2.extract_to(60);
18//!
19//! assert_eq!(shared1.len(), 824);
20//! assert_eq!(shared2.len(), 40);
21//! assert_eq!(shared3.len(), 100);
22//! assert_eq!(shared4.len(), 60);
23//!
24//! for byte in shared1.iter_mut() { *byte = 1u8; }
25//!
26//! // memory in slabs [4, 2, 3, 1]: merge back in arbitrary order.
27//! shared2.try_merge(shared3).ok().expect("Failed to merge 2 and 3");
28//! shared2.try_merge(shared1.freeze()).ok().expect("Failed to merge 23 and 1");
29//! shared4.try_merge(shared2).ok().expect("Failed to merge 4 and 231");
30//!
31//! assert_eq!(shared4.len(), 1024);
32//! ```
33#![forbid(missing_docs)]
34
35/// An `Arc`-backed mutable byte slice backed by a common allocation.
36pub mod arc {
37
38    use std::ops::{Deref, DerefMut};
39    use std::sync::Arc;
40    use std::any::Any;
41    use std::cell::UnsafeCell;
42
43    /// A mutable byte slice backed by a shared allocation.
44    ///
45    /// This type is neither `Send` nor `Sync`.
46    /// It can produce immutable [`Bytes`] views that are both `Send` and `Sync`.
47    ///
48    /// An instance of this type contends that `ptr` is valid for `len` bytes,
49    /// and that no other reference to these bytes exists, other than through
50    /// the type currently held in `sequestered`.
51    pub struct BytesMut {
52        /// Pointer to the start of this slice (not the allocation).
53        ptr: *mut u8,
54        /// Length of this slice.
55        len: usize,
56        /// Shared access to underlying resources.
57        ///
58        /// The backing object is accessed only during construction or unique regeneration.
59        /// Regeneration clears the old view before accessing the cell's contents.
60        /// See the safety argument above `Bytes`'s `Sync` implementation.
61        sequestered: Arc<UnsafeCell<dyn Any + Send>>,
62    }
63
64    impl BytesMut {
65
66        /// Create a new instance from a byte allocation.
67        pub fn from<B>(bytes: B) -> BytesMut where B : DerefMut<Target=[u8]>+Send+'static {
68
69            // Sequester allocation behind an `Arc`, which *should* keep the address
70            // stable for the lifetime of `sequestered`. The `Arc` also serves as our
71            // source of truth for the allocation, which we use to re-connect slices
72            // of the same allocation.
73            // `Arc` bookkeeping can form shared references covering the backing object.
74            // `UnsafeCell` lets those references coexist with writes to bytes stored inline in that object.
75            // The mutable slice must still remain disjoint from every published `Bytes` slice.
76            let mut sequestered = Arc::new(UnsafeCell::new(bytes)) as Arc<UnsafeCell<dyn Any + Send>>;
77            let (ptr, len) =
78            Arc::get_mut(&mut sequestered)
79                .unwrap()
80                .get_mut()
81                .downcast_mut::<B>()
82                .map(|a| {
83                    // Acquire one slice, so that the two next calls must agree.
84                    // Otherwise, adversarial implementations could lie to us.
85                    let slice = a.deref_mut();
86                    (slice.as_mut_ptr(), slice.len())
87                })
88                .unwrap();
89
90            BytesMut {
91                ptr,
92                len,
93                sequestered,
94            }
95        }
96
97        /// Extracts [0, index) into a new `Bytes` which is returned, updating `self`.
98        ///
99        /// # Implementation
100        ///
101        /// This method first tests `index` against `self.len`, which should ensure that both
102        /// the returned `Bytes` contains valid memory, and that `self` can no longer access it.
103        pub fn extract_to(&mut self, index: usize) -> Bytes {
104
105            assert!(index <= self.len);
106
107            let result = BytesMut {
108                ptr: self.ptr,
109                len: index,
110                sequestered: Arc::clone(&self.sequestered),
111            };
112
113            self.ptr = self.ptr.wrapping_add(index);
114            self.len -= index;
115
116            result.freeze()
117        }
118
119        /// Regenerates the BytesMut if it is uniquely held.
120        ///
121        /// If uniquely held, this method obtains a fresh mutable slice from the backing object and re-initializes the BytesMut.
122        /// `Some(true)` indicates success.
123        /// `Some(false)` indicates that the allocation is shared and leaves the view unchanged.
124        /// `None` indicates that the downcast to `B` failed and leaves the view empty.
125        /// The backing object remains owned by `self`, so regeneration with the correct type can be retried.
126        ///
127        /// # Panics
128        ///
129        /// If the backing object's `deref_mut` panics, `self` is left empty.
130        /// The backing object remains owned by `self`.
131        ///
132        /// # Examples
133        ///
134        /// ```
135        /// use timely_bytes::arc::BytesMut;
136        ///
137        /// let bytes = vec![0u8; 1024];
138        /// let mut shared1 = BytesMut::from(bytes);
139        /// let mut shared2 = shared1.extract_to(100);
140        /// let mut shared3 = shared1.extract_to(100);
141        /// let mut shared4 = shared2.extract_to(60);
142        ///
143        /// drop(shared3);
144        /// drop(shared2);
145        /// drop(shared4);
146        /// assert_eq!(shared1.try_regenerate::<Vec<u8>>(), Some(true));
147        /// assert!(shared1.len() == 1024);
148        /// ```
149        pub fn try_regenerate<B>(&mut self) -> Option<bool> where B: DerefMut<Target=[u8]>+'static {
150            // Only possible if this is the only reference to the sequestered allocation.
151            if let Some(boxed) = Arc::get_mut(&mut self.sequestered) {
152                // Clear the view before accessing the cell's contents.
153                // References formed during downcasting can invalidate the old pointer for inline storage.
154                // The backing object's `deref_mut` may also invalidate the old slice and then panic.
155                // Neither a failed downcast nor a caught panic may leave the old pointer accessible.
156                self.ptr = std::ptr::dangling_mut::<u8>();
157                self.len = 0;
158                let downcast = boxed.get_mut().downcast_mut::<B>()?;
159
160                // Acquire one slice, so that the two next calls must agree.
161                // Otherwise, adversarial implementations could lie to us.
162                let slice = downcast.deref_mut();
163                self.ptr = slice.as_mut_ptr();
164                self.len = slice.len();
165                Some(true)
166            }
167            else {
168                Some(false)
169            }
170        }
171
172        /// Converts a writeable byte slice to a shareable byte slice.
173        #[inline(always)]
174        pub fn freeze(self) -> Bytes {
175            Bytes {
176                ptr: self.ptr,
177                len: self.len,
178                sequestered: self.sequestered,
179            }
180        }
181    }
182
183    impl Deref for BytesMut {
184        type Target = [u8];
185        #[inline(always)]
186        fn deref(&self) -> &[u8] {
187            unsafe { ::std::slice::from_raw_parts(self.ptr, self.len) }
188        }
189    }
190
191    impl DerefMut for BytesMut {
192        #[inline(always)]
193        fn deref_mut(&mut self) -> &mut [u8] {
194            unsafe { ::std::slice::from_raw_parts_mut(self.ptr, self.len) }
195        }
196    }
197
198
199    /// A thread-safe shared byte buffer backed by a shared allocation.
200    ///
201    /// An instance of this type contends that `ptr` is valid for `len` bytes,
202    /// and that no other mutable reference to these bytes exists, other than
203    /// through the type currently held in `sequestered`.
204    #[derive(Clone)]
205    pub struct Bytes {
206        /// Pointer to the start of this slice (not the allocation).
207        ptr: *const u8,
208        /// Length of this slice.
209        len: usize,
210        /// Shared access to underlying resources.
211        ///
212        /// The backing object is accessed only during construction or unique regeneration.
213        /// Regeneration clears the old view before accessing the cell's contents.
214        /// See the safety argument above `Bytes`'s `Sync` implementation.
215        sequestered: Arc<UnsafeCell<dyn Any + Send>>,
216    }
217
218    // Synchronization happens through `self.sequestered`, which means to ensure that even
219    // across multiple threads the referenced range of bytes remains valid.
220    unsafe impl Send for Bytes { }
221
222    // `Sync` holds because accesses through `&Bytes` are read-only or atomic.
223    // `Deref` yields `&[u8]` (and `u8: Sync`), the mutating methods take `&mut self`, and cloning only changes the atomic `Arc` refcount.
224    // The backing object is inside an `UnsafeCell`, but `&Bytes` exposes only immutable byte slices and no mutable references.
225    // A `BytesMut` sharing the allocation can only write to a disjoint byte range.
226    // Regeneration accesses the backing object only after `Arc::get_mut` establishes uniqueness and the old view is cleared.
227    // Only a fresh slice from `DerefMut` can reinstall the view; a failed downcast or panic leaves it empty.
228    //
229    // Note this requires only that the sequestered payload `B` be `Send` (enforced by `BytesMut::from`), not `Sync`.
230    // `B` is never exposed by reference, so it is never shared across threads.
231    // The only cross-thread use of `B` is its destructor, which may run on whichever thread drops the last `Arc` clone -- and that needs `Send`, not `Sync`.
232    unsafe impl Sync for Bytes { }
233
234    impl Bytes {
235
236        /// Extracts [0, index) into a new `Bytes` which is returned, updating `self`.
237        ///
238        /// # Implementation
239        ///
240        /// This method first tests `index` against `self.len`, which should ensure that both
241        /// the returned `Bytes` contains valid memory, and that `self` can no longer access it.
242        pub fn extract_to(&mut self, index: usize) -> Bytes {
243
244            assert!(index <= self.len);
245
246            let result = Bytes {
247                ptr: self.ptr,
248                len: index,
249                sequestered: Arc::clone(&self.sequestered),
250            };
251
252            self.ptr = self.ptr.wrapping_add(index);
253            self.len -= index;
254
255            result
256        }
257
258        /// Attempts to merge adjacent slices from the same allocation.
259        ///
260        /// If the merge succeeds then `other.len` is added to `self` and the result is `Ok(())`.
261        /// If the merge fails self is unmodified and the result is `Err(other)`, returning the
262        /// bytes supplied as input.
263        ///
264        /// # Examples
265        ///
266        /// ```
267        /// use timely_bytes::arc::BytesMut;
268        ///
269        /// let bytes = vec![0u8; 1024];
270        /// let mut shared1 = BytesMut::from(bytes).freeze();
271        /// let mut shared2 = shared1.extract_to(100);
272        /// let mut shared3 = shared1.extract_to(100);
273        /// let mut shared4 = shared2.extract_to(60);
274        ///
275        /// // memory in slabs [4, 2, 3, 1]: merge back in arbitrary order.
276        /// shared2.try_merge(shared3).ok().expect("Failed to merge 2 and 3");
277        /// shared2.try_merge(shared1).ok().expect("Failed to merge 23 and 1");
278        /// shared4.try_merge(shared2).ok().expect("Failed to merge 4 and 231");
279        /// ```
280        pub fn try_merge(&mut self, other: Bytes) -> Result<(), Bytes> {
281            if Arc::ptr_eq(&self.sequestered, &other.sequestered) && ::std::ptr::eq(self.ptr.wrapping_add(self.len), other.ptr) {
282                self.len += other.len;
283                Ok(())
284            }
285            else {
286                Err(other)
287            }
288        }
289    }
290
291    impl Deref for Bytes {
292        type Target = [u8];
293        #[inline(always)]
294        fn deref(&self) -> &[u8] {
295            unsafe { ::std::slice::from_raw_parts(self.ptr, self.len) }
296        }
297    }
298}