Reduce contention in broadcast channel

benches/Cargo.toml

@@ -26,6 +26,11 @@ name = "spawn"
 path = "spawn.rs"
 harness = false
 
+[[bench]]
+name = "sync_broadcast"
+path = "sync_broadcast.rs"
+harness = false
+
 [[bench]]
 name = "sync_mpsc"
 path = "sync_mpsc.rs"

benches/sync_broadcast.rs

@@ -0,0 +1,82 @@
+use rand::{Rng, RngCore, SeedableRng};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use tokio::sync::{broadcast, Notify};
+
+use criterion::measurement::WallTime;
+use criterion::{black_box, criterion_group, criterion_main, BenchmarkGroup, Criterion};
+
+fn rt() -> tokio::runtime::Runtime {
+    tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(6)
+        .build()
+        .unwrap()
+}
+
+fn do_work(rng: &mut impl RngCore) -> u32 {
+    use std::fmt::Write;
+    let mut message = String::new();
+    for i in 1..=10 {
+        let _ = write!(&mut message, " {i}={}", rng.gen::<f64>());
+    }
+    message
+        .as_bytes()
+        .iter()
+        .map(|&c| c as u32)
+        .fold(0, u32::wrapping_add)
+}
+
+fn contention_impl<const N_TASKS: usize>(g: &mut BenchmarkGroup<WallTime>) {
+    let rt = rt();
+
+    let (tx, _rx) = broadcast::channel::<usize>(1000);
+    let wg = Arc::new((AtomicUsize::new(0), Notify::new()));
+
+    for n in 0..N_TASKS {
+        let wg = wg.clone();
+        let mut rx = tx.subscribe();
+        let mut rng = rand::rngs::StdRng::seed_from_u64(n as u64);
+        rt.spawn(async move {
+            while let Ok(_) = rx.recv().await {
+                let r = do_work(&mut rng);
+                let _ = black_box(r);
+                if wg.0.fetch_sub(1, Ordering::Relaxed) == 1 {
+                    wg.1.notify_one();
+                }
+            }
+        });
+    }
+
+    const N_ITERS: usize = 100;
+
+    g.bench_function(N_TASKS.to_string(), |b| {
+        b.iter(|| {
+            rt.block_on({
+                let wg = wg.clone();
+                let tx = tx.clone();
+                async move {
+                    for i in 0..N_ITERS {
+                        assert_eq!(wg.0.fetch_add(N_TASKS, Ordering::Relaxed), 0);
+                        tx.send(i).unwrap();
+                        while wg.0.load(Ordering::Relaxed) > 0 {
+                            wg.1.notified().await;
+                        }
+                    }
+                }
+            })
+        })
+    });
+}
+
+fn bench_contention(c: &mut Criterion) {
+    let mut group = c.benchmark_group("contention");
+    contention_impl::<10>(&mut group);
+    contention_impl::<100>(&mut group);
+    contention_impl::<500>(&mut group);
+    contention_impl::<1000>(&mut group);
+    group.finish();
+}
+
+criterion_group!(contention, bench_contention);
+
+criterion_main!(contention);

tokio/src/loom/mocked.rs

@@ -1,12 +1,25 @@
 pub(crate) use loom::*;
 
 pub(crate) mod sync {
+    use std::{
+        ops::{Deref, DerefMut},
+        sync::{LockResult, PoisonError},
+    };
 
     pub(crate) use loom::sync::MutexGuard;
 
     #[derive(Debug)]
     pub(crate) struct Mutex<T>(loom::sync::Mutex<T>);
 
+    #[derive(Debug)]
+    pub(crate) struct RwLock<T>(loom::sync::RwLock<T>);
+
+    #[derive(Debug)]
+    pub(crate) struct RwLockReadGuard<'a, T>(loom::sync::RwLockReadGuard<'a, T>);
+
+    #[derive(Debug)]
+    pub(crate) struct RwLockWriteGuard<'a, T>(loom::sync::RwLockWriteGuard<'a, T>);
+
     #[allow(dead_code)]
     impl<T> Mutex<T> {
         #[inline]
@@ -25,6 +38,63 @@ pub(crate) mod sync {
             self.0.try_lock().ok()
         }
     }
+
+    #[allow(dead_code)]
+    impl<T> RwLock<T> {
+        #[inline]
+        pub(crate) fn new(t: T) -> RwLock<T> {
+            RwLock(loom::sync::RwLock::new(t))
+        }
+
+        #[inline]
+        pub(crate) fn read(&self) -> LockResult<RwLockReadGuard<'_, T>> {
+            match self.0.read() {
+                Ok(inner) => Ok(RwLockReadGuard(inner)),
+                Err(err) => Err(PoisonError::new(RwLockReadGuard(err.into_inner()))),
+            }
+        }
+
+        #[inline]
+        pub(crate) fn write(&self) -> LockResult<RwLockWriteGuard<'_, T>> {
+            match self.0.write() {
+                Ok(inner) => Ok(RwLockWriteGuard(inner)),
+                Err(err) => Err(PoisonError::new(RwLockWriteGuard(err.into_inner()))),
+            }
+        }
+    }
+
+    impl<'a, T> Deref for RwLockReadGuard<'a, T> {
+        type Target = T;
+        fn deref(&self) -> &T {
+            self.0.deref()
+        }
+    }
+
+    #[allow(dead_code)]
+    impl<'a, T> RwLockWriteGuard<'a, T> {
+        pub(crate) fn downgrade(
+            s: Self,
+            rwlock: &'a RwLock<T>,
+        ) -> LockResult<RwLockReadGuard<'a, T>> {
+            // Std rwlock does not support downgrading.
+            drop(s);
+            rwlock.read()
+        }
+    }
+
+    impl<'a, T> Deref for RwLockWriteGuard<'a, T> {
+        type Target = T;
+        fn deref(&self) -> &T {
+            self.0.deref()
+        }
+    }
+
+    impl<'a, T> DerefMut for RwLockWriteGuard<'a, T> {
+        fn deref_mut(&mut self) -> &mut T {
+            self.0.deref_mut()
+        }
+    }
+
     pub(crate) use loom::sync::*;
 
     pub(crate) mod atomic {

tokio/src/loom/std/mod.rs

@@ -8,6 +8,7 @@ mod barrier;
 mod mutex;
 #[cfg(all(feature = "parking_lot", not(miri)))]
 mod parking_lot;
+mod rwlock;
 mod unsafe_cell;
 
 pub(crate) mod cell {
@@ -59,15 +60,18 @@ pub(crate) mod sync {
     #[cfg(all(feature = "parking_lot", not(miri)))]
     #[allow(unused_imports)]
     pub(crate) use crate::loom::std::parking_lot::{
-        Condvar, Mutex, MutexGuard, RwLock, RwLockReadGuard, WaitTimeoutResult,
+        Condvar, Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard, WaitTimeoutResult,
     };
 
     #[cfg(not(all(feature = "parking_lot", not(miri))))]
     #[allow(unused_imports)]
-    pub(crate) use std::sync::{Condvar, MutexGuard, RwLock, RwLockReadGuard, WaitTimeoutResult};
+    pub(crate) use std::sync::{Condvar, MutexGuard, WaitTimeoutResult};
 
     #[cfg(not(all(feature = "parking_lot", not(miri))))]
-    pub(crate) use crate::loom::std::mutex::Mutex;
+    pub(crate) use crate::loom::std::{
+        mutex::Mutex,
+        rwlock::{RwLock, RwLockReadGuard, RwLockWriteGuard},
+    };
 
     pub(crate) mod atomic {
         pub(crate) use crate::loom::std::atomic_u16::AtomicU16;

tokio/src/loom/std/parking_lot.rs

@@ -112,6 +112,16 @@ impl<'a, T: ?Sized> Deref for RwLockReadGuard<'a, T> {
     }
 }
 
+impl<'a, T> RwLockWriteGuard<'a, T> {
+    // The corresponding std method requires the rwlock.
+    pub(crate) fn downgrade(s: Self, _rwlock: &'a RwLock<T>) -> LockResult<RwLockReadGuard<'a, T>> {
+        Ok(RwLockReadGuard(
+            PhantomData,
+            parking_lot::RwLockWriteGuard::downgrade(s.1),
+        ))
+    }
+}
+
 impl<'a, T: ?Sized> Deref for RwLockWriteGuard<'a, T> {
     type Target = T;
     fn deref(&self) -> &T {

tokio/src/loom/std/rwlock.rs

@@ -0,0 +1,67 @@
+use std::{
+    ops::{Deref, DerefMut},
+    sync::{self, LockResult, PoisonError},
+};
+
+/// Adapter for std::RwLock that adds `downgrade` method.
+#[derive(Debug)]
+pub(crate) struct RwLock<T>(sync::RwLock<T>);
+
+#[derive(Debug)]
+pub(crate) struct RwLockReadGuard<'a, T>(sync::RwLockReadGuard<'a, T>);
+
+#[derive(Debug)]
+pub(crate) struct RwLockWriteGuard<'a, T>(sync::RwLockWriteGuard<'a, T>);
+
+#[allow(dead_code)]
+impl<T> RwLock<T> {
+    #[inline]
+    pub(crate) fn new(t: T) -> RwLock<T> {
+        RwLock(sync::RwLock::new(t))
+    }
+
+    #[inline]
+    pub(crate) fn read(&self) -> LockResult<RwLockReadGuard<'_, T>> {
+        match self.0.read() {
+            Ok(inner) => Ok(RwLockReadGuard(inner)),
+            Err(err) => Err(PoisonError::new(RwLockReadGuard(err.into_inner()))),
+        }
+    }
+
+    #[inline]
+    pub(crate) fn write(&self) -> LockResult<RwLockWriteGuard<'_, T>> {
+        match self.0.write() {
+            Ok(inner) => Ok(RwLockWriteGuard(inner)),
+            Err(err) => Err(PoisonError::new(RwLockWriteGuard(err.into_inner()))),
+        }
+    }
+}
+
+impl<'a, T> Deref for RwLockReadGuard<'a, T> {
+    type Target = T;
+    fn deref(&self) -> &T {
+        self.0.deref()
+    }
+}
+
+#[allow(dead_code)]
+impl<'a, T> RwLockWriteGuard<'a, T> {
+    pub(crate) fn downgrade(s: Self, rwlock: &'a RwLock<T>) -> LockResult<RwLockReadGuard<'a, T>> {
+        // Std rwlock does not support downgrading.
+        drop(s);
+        rwlock.read()
+    }
+}
+
+impl<'a, T> Deref for RwLockWriteGuard<'a, T> {
+    type Target = T;
+    fn deref(&self) -> &T {
+        self.0.deref()
+    }
+}
+
+impl<'a, T> DerefMut for RwLockWriteGuard<'a, T> {
+    fn deref_mut(&mut self) -> &mut T {
+        self.0.deref_mut()
+    }
+}