Optimize Inner::shallow_clone (#217)

- Clones when the kind is INLINE or STATIC are sped up by over double.
- Clones when the kind is ARC are spec up by about 1/3.

benches/bytes.rs

@@ -113,6 +113,39 @@ fn deref_two(b: &mut Bencher) {
     })
 }
 
+#[bench]
+fn clone_inline(b: &mut Bencher) {
+    let bytes = Bytes::from_static(b"hello world");
+
+    b.iter(|| {
+        for _ in 0..1024 {
+            test::black_box(&bytes.clone());
+        }
+    })
+}
+
+#[bench]
+fn clone_static(b: &mut Bencher) {
+    let bytes = Bytes::from_static("hello world 1234567890 and have a good byte 0987654321".as_bytes());
+
+    b.iter(|| {
+        for _ in 0..1024 {
+            test::black_box(&bytes.clone());
+        }
+    })
+}
+
+#[bench]
+fn clone_arc(b: &mut Bencher) {
+    let bytes = Bytes::from("hello world 1234567890 and have a good byte 0987654321".as_bytes());
+
+    b.iter(|| {
+        for _ in 0..1024 {
+            test::black_box(&bytes.clone());
+        }
+    })
+}
+
 #[bench]
 fn alloc_write_split_to_mid(b: &mut Bencher) {
     b.iter(|| {

ci/tsan

@@ -16,6 +16,6 @@ race:test::run_tests_console::*closure
 # Probably more fences in std.
 race:__call_tls_dtors
 
-# `is_inline` is explicitly called concurrently without synchronization. The
+# `is_inline_or_static` is explicitly called concurrently without synchronization.
-# safety explanation can be found in a comment.
+# The safety explanation can be found in a comment.
-race:Inner::is_inline
+race:Inner::is_inline_or_static

src/bytes.rs

@@ -2051,126 +2051,139 @@ impl Inner {
     unsafe fn shallow_clone(&self, mut_self: bool) -> Inner {
         // Always check `inline` first, because if the handle is using inline
         // data storage, all of the `Inner` struct fields will be gibberish.
-        if self.is_inline() {
+        //
-            // In this case, a shallow_clone still involves copying the data.
+        // Additionally, if kind is STATIC, then Arc is *never* changed, making
-            //
+        // it safe and faster to check for it now before an atomic acquire.
-            // TODO: Just copy the fields
-            let mut inner: Inner = mem::uninitialized();
-            let len = self.inline_len();
 
-            inner.arc = AtomicPtr::new(KIND_INLINE as *mut Shared);
+        if self.is_inline_or_static() {
-            inner.set_inline_len(len);
+            // In this case, a shallow_clone still involves copying the data.
-            inner.as_raw()[0..len].copy_from_slice(self.as_ref());
+            let mut inner: Inner = mem::uninitialized();
+            ptr::copy_nonoverlapping(
+                self,
+                &mut inner,
+                1,
+            );
             inner
         } else {
-            // The function requires `&self`, this means that `shallow_clone`
+            self.shallow_clone_sync(mut_self)
-            // could be called concurrently.
-            //
-            // The first step is to load the value of `arc`. This will determine
-            // how to proceed. The `Acquire` ordering synchronizes with the
-            // `compare_and_swap` that comes later in this function. The goal is
-            // to ensure that if `arc` is currently set to point to a `Shared`,
-            // that the current thread acquires the associated memory.
-            let mut arc = self.arc.load(Acquire);
-
-            // If  the buffer is still tracked in a `Vec<u8>`. It is time to
-            // promote the vec to an `Arc`. This could potentially be called
-            // concurrently, so some care must be taken.
-            if arc as usize & KIND_MASK == KIND_VEC {
-                let original_capacity_repr =
-                    (arc as usize & ORIGINAL_CAPACITY_MASK) >> ORIGINAL_CAPACITY_OFFSET;
-
-                // The vec offset cannot be concurrently mutated, so there
-                // should be no danger reading it.
-                let off = (arc as usize) >> VEC_POS_OFFSET;
-
-                // First, allocate a new `Shared` instance containing the
-                // `Vec` fields. It's important to note that `ptr`, `len`,
-                // and `cap` cannot be mutated without having `&mut self`.
-                // This means that these fields will not be concurrently
-                // updated and since the buffer hasn't been promoted to an
-                // `Arc`, those three fields still are the components of the
-                // vector.
-                let shared = Box::new(Shared {
-                    vec: rebuild_vec(self.ptr, self.len, self.cap, off),
-                    original_capacity_repr: original_capacity_repr,
-                    // Initialize refcount to 2. One for this reference, and one
-                    // for the new clone that will be returned from
-                    // `shallow_clone`.
-                    ref_count: AtomicUsize::new(2),
-                });
-
-                let shared = Box::into_raw(shared);
-
-                // The pointer should be aligned, so this assert should
-                // always succeed.
-                debug_assert!(0 == (shared as usize & 0b11));
-
-                // If there are no references to self in other threads,
-                // expensive atomic operations can be avoided.
-                if mut_self {
-                    self.arc.store(shared, Relaxed);
-                    return Inner {
-                        arc: AtomicPtr::new(shared),
-                        .. *self
-                    };
-                }
-
-                // Try compare & swapping the pointer into the `arc` field.
-                // `Release` is used synchronize with other threads that
-                // will load the `arc` field.
-                //
-                // If the `compare_and_swap` fails, then the thread lost the
-                // race to promote the buffer to shared. The `Acquire`
-                // ordering will synchronize with the `compare_and_swap`
-                // that happened in the other thread and the `Shared`
-                // pointed to by `actual` will be visible.
-                let actual = self.arc.compare_and_swap(arc, shared, AcqRel);
-
-                if actual == arc {
-                    // The upgrade was successful, the new handle can be
-                    // returned.
-                    return Inner {
-                        arc: AtomicPtr::new(shared),
-                        .. *self
-                    };
-                }
-
-                // The upgrade failed, a concurrent clone happened. Release
-                // the allocation that was made in this thread, it will not
-                // be needed.
-                let shared = Box::from_raw(shared);
-                mem::forget(*shared);
-
-                // Update the `arc` local variable and fall through to a ref
-                // count update
-                arc = actual;
-            } else if arc as usize & KIND_MASK == KIND_STATIC {
-                // Static buffer
-                return Inner {
-                    arc: AtomicPtr::new(arc),
-                    .. *self
-                };
-            }
-
-            // Buffer already promoted to shared storage, so increment ref
-            // count.
-            //
-            // Relaxed ordering is acceptable as the memory has already been
-            // acquired via the `Acquire` load above.
-            let old_size = (*arc).ref_count.fetch_add(1, Relaxed);
-
-            if old_size == usize::MAX {
-                panic!(); // TODO: abort
-            }
-
-            Inner {
-                arc: AtomicPtr::new(arc),
-                .. *self
-            }
         }
     }
 
+
+    #[cold]
+    unsafe fn shallow_clone_sync(&self, mut_self: bool) -> Inner {
+        // The function requires `&self`, this means that `shallow_clone`
+        // could be called concurrently.
+        //
+        // The first step is to load the value of `arc`. This will determine
+        // how to proceed. The `Acquire` ordering synchronizes with the
+        // `compare_and_swap` that comes later in this function. The goal is
+        // to ensure that if `arc` is currently set to point to a `Shared`,
+        // that the current thread acquires the associated memory.
+        let arc = self.arc.load(Acquire);
+        let kind = arc as usize & KIND_MASK;
+
+        if kind == KIND_ARC {
+            self.shallow_clone_arc(arc)
+        } else {
+            assert!(kind == KIND_VEC);
+            self.shallow_clone_vec(arc as usize, mut_self)
+        }
+    }
+
+    unsafe fn shallow_clone_arc(&self, arc: *mut Shared) -> Inner {
+        debug_assert!(arc as usize & KIND_MASK == KIND_ARC);
+
+        let old_size = (*arc).ref_count.fetch_add(1, Relaxed);
+
+        if old_size == usize::MAX {
+            abort();
+        }
+
+        Inner {
+            arc: AtomicPtr::new(arc),
+            .. *self
+        }
+    }
+
+    #[cold]
+    unsafe fn shallow_clone_vec(&self, arc: usize, mut_self: bool) -> Inner {
+        // If  the buffer is still tracked in a `Vec<u8>`. It is time to
+        // promote the vec to an `Arc`. This could potentially be called
+        // concurrently, so some care must be taken.
+
+        debug_assert!(arc & KIND_MASK == KIND_VEC);
+
+        let original_capacity_repr =
+            (arc as usize & ORIGINAL_CAPACITY_MASK) >> ORIGINAL_CAPACITY_OFFSET;
+
+        // The vec offset cannot be concurrently mutated, so there
+        // should be no danger reading it.
+        let off = (arc as usize) >> VEC_POS_OFFSET;
+
+        // First, allocate a new `Shared` instance containing the
+        // `Vec` fields. It's important to note that `ptr`, `len`,
+        // and `cap` cannot be mutated without having `&mut self`.
+        // This means that these fields will not be concurrently
+        // updated and since the buffer hasn't been promoted to an
+        // `Arc`, those three fields still are the components of the
+        // vector.
+        let shared = Box::new(Shared {
+            vec: rebuild_vec(self.ptr, self.len, self.cap, off),
+            original_capacity_repr: original_capacity_repr,
+            // Initialize refcount to 2. One for this reference, and one
+            // for the new clone that will be returned from
+            // `shallow_clone`.
+            ref_count: AtomicUsize::new(2),
+        });
+
+        let shared = Box::into_raw(shared);
+
+        // The pointer should be aligned, so this assert should
+        // always succeed.
+        debug_assert!(0 == (shared as usize & 0b11));
+
+        // If there are no references to self in other threads,
+        // expensive atomic operations can be avoided.
+        if mut_self {
+            self.arc.store(shared, Relaxed);
+            return Inner {
+                arc: AtomicPtr::new(shared),
+                .. *self
+            };
+        }
+
+        // Try compare & swapping the pointer into the `arc` field.
+        // `Release` is used synchronize with other threads that
+        // will load the `arc` field.
+        //
+        // If the `compare_and_swap` fails, then the thread lost the
+        // race to promote the buffer to shared. The `Acquire`
+        // ordering will synchronize with the `compare_and_swap`
+        // that happened in the other thread and the `Shared`
+        // pointed to by `actual` will be visible.
+        let actual = self.arc.compare_and_swap(arc as *mut Shared, shared, AcqRel);
+
+        if actual as usize == arc {
+            // The upgrade was successful, the new handle can be
+            // returned.
+            return Inner {
+                arc: AtomicPtr::new(shared),
+                .. *self
+            };
+        }
+
+        // The upgrade failed, a concurrent clone happened. Release
+        // the allocation that was made in this thread, it will not
+        // be needed.
+        let shared = Box::from_raw(shared);
+        mem::forget(*shared);
+
+        // Buffer already promoted to shared storage, so increment ref
+        // count.
+        self.shallow_clone_arc(actual)
+    }
+
     #[inline]
     fn reserve(&mut self, additional: usize) {
         let len = self.len();
@@ -2327,6 +2340,18 @@ impl Inner {
         self.kind() == KIND_INLINE
     }
 
+    #[inline]
+    fn is_inline_or_static(&self) -> bool {
+        // The value returned by `kind` isn't itself safe, but the value could
+        // inform what operations to take, and unsafely do something without
+        // synchronization.
+        //
+        // KIND_INLINE and KIND_STATIC will *never* change, so branches on that
+        // information is safe.
+        let kind = self.kind();
+        kind == KIND_INLINE || kind == KIND_STATIC
+    }
+
     /// Used for `debug_assert` statements. &mut is used to guarantee that it is
     /// safe to check VEC_KIND
     #[inline]
@@ -2827,3 +2852,21 @@ impl PartialEq<Bytes> for BytesMut
         &other[..] == &self[..]
     }
 }
+
+// While there is `std::process:abort`, it's only available in Rust 1.17, and
+// our minimum supported version is currently 1.15. So, this acts as an abort
+// by triggering a double panic, which always aborts in Rust.
+struct Abort;
+
+impl Drop for Abort {
+    fn drop(&mut self) {
+        panic!();
+    }
+}
+
+#[inline(never)]
+#[cold]
+fn abort() {
+    let _a = Abort;
+    panic!();
+}