net: replace RwLock<Slab> with a lock free slab

azure-pipelines.yml

@@ -83,6 +83,7 @@ jobs:
     rust: beta
     crates:
       - tokio-executor
+      - tokio
 
 # Try cross compiling
 - template: ci/azure-cross-compile.yml

tokio/Cargo.toml

@@ -39,7 +39,7 @@ fs = ["tokio-executor/blocking"]
 io = ["tokio-io", "bytes", "iovec"]
 macros = ["tokio-macros"]
 net-full = ["tcp", "udp", "uds"]
-net-driver = ["mio", "tokio-executor/blocking"]
+net-driver = ["mio", "tokio-executor/blocking", "lazy_static"]
 rt-current-thread = [
   "timer",
   "tokio-executor/current-thread",
@@ -110,6 +110,10 @@ version = "0.3.8"
 default-features = false
 optional = true
 
+[target.'cfg(loom)'.dependencies]
+# play nice with loom tests in other crates.
+loom = "0.2.11"
+
 [dev-dependencies]
 tokio-test = { version = "=0.2.0-alpha.6", path = "../tokio-test" }
 tokio-util = { version = "=0.2.0-alpha.6", path = "../tokio-util" }
@@ -127,5 +131,9 @@ serde_json = "1.0"
 tempfile = "3.1.0"
 time = "0.1"
 
+# sharded slab tests
+loom = "0.2.11"
+proptest = "0.9.4"
+
 [package.metadata.docs.rs]
 all-features = true

tokio/src/lib.rs

@@ -69,7 +69,6 @@
 //!     }
 //! }
 //! ```
-
 macro_rules! if_runtime {
     ($($i:item)*) => ($(
         #[cfg(any(
@@ -97,6 +96,9 @@ pub mod io;
 #[cfg(feature = "net-driver")]
 pub mod net;
 
+#[cfg(feature = "net-driver")]
+mod loom;
+
 pub mod prelude;
 
 #[cfg(feature = "process")]

tokio/src/loom.rs

@@ -0,0 +1,45 @@
+//! This module abstracts over `loom` and `std::sync` depending on whether we
+//! are running tests or not.
+pub(crate) use self::inner::*;
+
+#[cfg(all(test, loom))]
+mod inner {
+    pub(crate) use loom::sync::CausalCell;
+    pub(crate) use loom::sync::Mutex;
+    pub(crate) mod atomic {
+        pub(crate) use loom::sync::atomic::*;
+        pub(crate) use std::sync::atomic::Ordering;
+    }
+}
+
+#[cfg(not(all(test, loom)))]
+mod inner {
+    use std::cell::UnsafeCell;
+    pub(crate) use std::sync::atomic;
+    pub(crate) use std::sync::Mutex;
+
+    #[derive(Debug)]
+    pub(crate) struct CausalCell<T>(UnsafeCell<T>);
+
+    impl<T> CausalCell<T> {
+        pub(crate) fn new(data: T) -> CausalCell<T> {
+            CausalCell(UnsafeCell::new(data))
+        }
+
+        #[inline(always)]
+        pub(crate) fn with<F, R>(&self, f: F) -> R
+        where
+            F: FnOnce(*const T) -> R,
+        {
+            f(self.0.get())
+        }
+
+        #[inline(always)]
+        pub(crate) fn with_mut<F, R>(&self, f: F) -> R
+        where
+            F: FnOnce(*mut T) -> R,
+        {
+            f(self.0.get())
+        }
+    }
+}

tokio/src/net/driver/mod.rs

@@ -124,7 +124,15 @@
 //! [`PollEvented`]: struct.PollEvented.html
 //! [`std::io::Read`]: https://doc.rust-lang.org/std/io/trait.Read.html
 //! [`std::io::Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
+#[cfg(loom)]
+macro_rules! loom_thread_local {
+    ($($tts:tt)+) => { loom::thread_local!{ $($tts)+ } }
+}
 
+#[cfg(not(loom))]
+macro_rules! loom_thread_local {
+    ($($tts:tt)+) => { std::thread_local!{ $($tts)+ } }
+}
 pub(crate) mod platform;
 mod reactor;
 mod registration;

tokio/src/net/driver/reactor/dispatch/iter.rs

@@ -0,0 +1,53 @@
+use super::{
+    page::{self, ScheduledIo},
+    Shard,
+};
+use std::slice;
+
+pub(in crate::net::driver::reactor) struct UniqueIter<'a> {
+    pub(super) shards: slice::IterMut<'a, Shard>,
+    pub(super) pages: slice::Iter<'a, page::Shared>,
+    pub(super) slots: Option<page::Iter<'a>>,
+}
+
+impl<'a> Iterator for UniqueIter<'a> {
+    type Item = &'a ScheduledIo;
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let Some(item) = self.slots.as_mut().and_then(|slots| slots.next()) {
+                return Some(item);
+            }
+
+            if let Some(page) = self.pages.next() {
+                self.slots = page.iter();
+            }
+
+            if let Some(shard) = self.shards.next() {
+                self.pages = shard.iter();
+            } else {
+                return None;
+            }
+        }
+    }
+}
+
+pub(in crate::net::driver::reactor) struct ShardIter<'a> {
+    pub(super) pages: slice::IterMut<'a, page::Shared>,
+    pub(super) slots: Option<page::Iter<'a>>,
+}
+
+impl<'a> Iterator for ShardIter<'a> {
+    type Item = &'a ScheduledIo;
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let Some(item) = self.slots.as_mut().and_then(|slots| slots.next()) {
+                return Some(item);
+            }
+            if let Some(page) = self.pages.next() {
+                self.slots = page.iter();
+            } else {
+                return None;
+            }
+        }
+    }
+}

tokio/src/net/driver/reactor/dispatch/mod.rs

@@ -0,0 +1,36 @@
+//! A lock-free concurrent slab.
+
+#[cfg(all(test, loom))]
+macro_rules! test_println {
+    ($($arg:tt)*) => {
+        println!("{:?} {}", crate::net::driver::reactor::dispatch::Tid::current(), format_args!($($arg)*))
+    }
+}
+
+mod iter;
+mod pack;
+mod page;
+mod sharded_slab;
+mod tid;
+
+#[cfg(all(test, loom))]
+// this is used by sub-modules
+use self::tests::test_util;
+use pack::{Pack, WIDTH};
+use sharded_slab::Shard;
+#[cfg(all(test, loom))]
+pub(crate) use sharded_slab::Slab;
+pub(crate) use sharded_slab::{SingleShard, MAX_SOURCES};
+use tid::Tid;
+
+#[cfg(target_pointer_width = "64")]
+const MAX_THREADS: usize = 4096;
+#[cfg(target_pointer_width = "32")]
+const MAX_THREADS: usize = 2048;
+const INITIAL_PAGE_SIZE: usize = 32;
+const MAX_PAGES: usize = WIDTH / 4;
+// Chosen arbitrarily.
+const RESERVED_BITS: usize = 5;
+
+#[cfg(test)]
+mod tests;

tokio/src/net/driver/reactor/dispatch/pack.rs

@@ -0,0 +1,89 @@
+pub(super) const WIDTH: usize = std::mem::size_of::<usize>() * 8;
+
+/// Trait encapsulating the calculations required for bit-packing slab indices.
+///
+/// This allows us to avoid manually repeating some calculations when packing
+/// and unpacking indices.
+pub(crate) trait Pack: Sized {
+    // ====== provided by each implementation =================================
+
+    /// The number of bits occupied by this type when packed into a usize.
+    ///
+    /// This must be provided to determine the number of bits into which to pack
+    /// the type.
+    const LEN: usize;
+    /// The type packed on the less significant side of this type.
+    ///
+    /// If this type is packed into the least significant bit of a usize, this
+    /// should be `()`, which occupies no bytes.
+    ///
+    /// This is used to calculate the shift amount for packing this value.
+    type Prev: Pack;
+
+    // ====== calculated automatically ========================================
+
+    /// A number consisting of `Self::LEN` 1 bits, starting at the least
+    /// significant bit.
+    ///
+    /// This is the higest value this type can represent. This number is shifted
+    /// left by `Self::SHIFT` bits to calculate this type's `MASK`.
+    ///
+    /// This is computed automatically based on `Self::LEN`.
+    const BITS: usize = {
+        let shift = 1 << (Self::LEN - 1);
+        shift | (shift - 1)
+    };
+    /// The number of bits to shift a number to pack it into a usize with other
+    /// values.
+    ///
+    /// This is caculated automatically based on the `LEN` and `SHIFT` constants
+    /// of the previous value.
+    const SHIFT: usize = Self::Prev::SHIFT + Self::Prev::LEN;
+
+    /// The mask to extract only this type from a packed `usize`.
+    ///
+    /// This is calculated by shifting `Self::BITS` left by `Self::SHIFT`.
+    const MASK: usize = Self::BITS << Self::SHIFT;
+
+    fn as_usize(&self) -> usize;
+    fn from_usize(val: usize) -> Self;
+
+    #[inline(always)]
+    fn pack(&self, to: usize) -> usize {
+        let value = self.as_usize();
+        debug_assert!(value <= Self::BITS);
+
+        (to & !Self::MASK) | (value << Self::SHIFT)
+    }
+
+    #[inline(always)]
+    fn from_packed(from: usize) -> Self {
+        let value = (from & Self::MASK) >> Self::SHIFT;
+        debug_assert!(value <= Self::BITS);
+        Self::from_usize(value)
+    }
+}
+
+impl Pack for () {
+    const BITS: usize = 0;
+    const LEN: usize = 0;
+    const SHIFT: usize = 0;
+    const MASK: usize = 0;
+
+    type Prev = ();
+
+    fn as_usize(&self) -> usize {
+        unreachable!()
+    }
+    fn from_usize(_val: usize) -> Self {
+        unreachable!()
+    }
+
+    fn pack(&self, _to: usize) -> usize {
+        unreachable!()
+    }
+
+    fn from_packed(_from: usize) -> Self {
+        unreachable!()
+    }
+}