perf(waitgroup): streamline completion for participants and observers

CHANGELOG.md

@@ -23,6 +23,10 @@ All notable changes to this project will be documented in this file.
 
 ## Unreleased
 
+### Improvements
+
+* Reduce `WaitGroup` completion overhead by sharing allocation ownership across participants and storing the first pending observer's waker inline.
+
 ## v0.8.0
 
 ### Breaking changes

Cargo.toml

@@ -40,10 +40,12 @@ async-channel = { version = "2.5.0" }
 cargo_metadata = { version = "0.23.1" }
 clap = { version = "4.6.5" }
 flume = { version = "0.12.0", default-features = false }
+futures-util = { version = "0.3.34" }
 semver = { version = "1.0.28" }
 serde = { version = "1.0.229" }
 tokio = { version = "1.53.1" }
 tokio-test = { version = "0.4.5" }
+tokio-util = { version = "0.7.20" }
 ureq = { version = "3.3.0", default-features = false }
 waitgroup = { version = "0.1.2" }
 which = { version = "8.0.5" }

asyncband/src/barrier/mod.rs

@@ -33,17 +33,16 @@
 //! use asyncband::barrier::Barrier;
 //!
 //! let barrier = Arc::new(Barrier::new(3));
-//! let stage = Arc::new([const { AtomicUsize::new(0) }; 3]);
+//! let arrived = Arc::new(AtomicUsize::new(0));
 //! let mut workers = Vec::with_capacity(3);
 //!
-//! for slot in 0..3 {
+//! for _ in 0..3 {
 //!     let barrier = barrier.clone();
-//!     let stage = stage.clone();
+//!     let arrived = arrived.clone();
 //!     workers.push(tokio::spawn(async move {
-//!         stage[slot].store(slot + 1, Ordering::Release);
+//!         arrived.fetch_add(1, Ordering::Relaxed);
 //!         let result = barrier.wait().await;
-//!         let sum: usize = stage.iter().map(|v| v.load(Ordering::Acquire)).sum();
-//!         assert_eq!(sum, 6);
+//!         assert_eq!(arrived.load(Ordering::Relaxed), 3);
 //!         result.is_leader()
 //!     }));
 //! }

asyncband/src/internal/mod.rs

@@ -76,13 +76,17 @@ pub(crate) mod value_cell;
     feature = "phaser",
     feature = "rwlock",
     feature = "semaphore",
-    feature = "waitgroup",
     feature = "watch",
 ))]
 // Some primitives use only shared access, leaving `Mutex::get_mut` unused.
 #[allow(dead_code)]
 pub(crate) mod mutex;
 
+#[cfg(feature = "waitgroup")]
+// Some primitives use only exclusive access, leaving `read` and `get_mut` unused.
+#[allow(dead_code)]
+pub(crate) mod rwlock;
+
 #[cfg(any(
     feature = "broadcast",
     feature = "mutex",

asyncband/src/internal/rwlock.rs

@@ -0,0 +1,45 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::fmt;
+use std::sync::PoisonError;
+
+pub struct RwLock<T: ?Sized>(std::sync::RwLock<T>);
+
+impl<T: ?Sized + fmt::Debug> fmt::Debug for RwLock<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl<T> RwLock<T> {
+    pub const fn new(t: T) -> Self {
+        Self(std::sync::RwLock::new(t))
+    }
+
+    pub fn read(&self) -> std::sync::RwLockReadGuard<'_, T> {
+        self.0.read().unwrap_or_else(PoisonError::into_inner)
+    }
+
+    pub fn write(&self) -> std::sync::RwLockWriteGuard<'_, T> {
+        self.0.write().unwrap_or_else(PoisonError::into_inner)
+    }
+
+    pub fn get_mut(&mut self) -> &mut T {
+        self.0.get_mut().unwrap_or_else(PoisonError::into_inner)
+    }
+}

asyncband/src/internal/wakerset.rs

@@ -80,32 +80,54 @@ impl WakerSet {
         self.wakers.into_iter().for_each(Waker::wake);
     }
 
+    /// Inserts a new waker slot and returns its exclusive token.
+    #[inline]
+    pub fn insert(&mut self, waker: &Waker) -> WakerToken {
+        WakerToken(self.wakers.insert(waker.clone()))
+    }
+
+    /// Replaces the waker at `token` unless it already wakes the same task.
+    #[inline]
+    #[must_use = "drop the returned waker after releasing the waker set's state lock"]
+    pub fn update(&mut self, token: &WakerToken, waker: &Waker) -> Option<Waker> {
+        let current = self
+            .wakers
+            .get_mut(token.0)
+            .expect("waker token must refer to an occupied slot");
+        if current.will_wake(waker) {
+            None
+        } else {
+            Some(mem::replace(current, waker.clone()))
+        }
+    }
+
+    /// Removes and returns the waker identified by `token`.
+    #[inline]
+    #[must_use = "drop the returned waker after releasing the waker set's state lock"]
+    pub fn remove(&mut self, token: WakerToken) -> Waker {
+        self.wakers.remove(token.0)
+    }
+
     /// Registers or updates a waker.
     ///
     /// Returns the previous waker only if it was replaced.
     #[inline]
     #[must_use = "drop the returned waker after releasing the waker set's state lock"]
     pub fn register(&mut self, token: &mut Option<WakerToken>, waker: &Waker) -> Option<Waker> {
-        if let Some(token) = token {
-            let current = self
-                .wakers
-                .get_mut(token.0)
-                .expect("waker token must refer to an occupied slot");
-            if current.will_wake(waker) {
-                return None;
+        match token {
+            Some(token) => self.update(token, waker),
+            None => {
+                *token = Some(self.insert(waker));
+                None
             }
-            return Some(mem::replace(current, waker.clone()));
         }
-
-        *token = Some(WakerToken(self.wakers.insert(waker.clone())));
-        None
     }
 
     /// Removes and returns the waker identified by `token`, clearing the token.
     #[inline]
     #[must_use = "drop the returned waker after releasing the waker set's state lock"]
     pub fn unregister(&mut self, token: &mut Option<WakerToken>) -> Option<Waker> {
-        token.take().map(|token| self.wakers.remove(token.0))
+        token.take().map(|token| self.remove(token))
     }
 }