perf: sharded OnceTable

asyncband/src/internal/mod.rs

@@ -40,12 +40,6 @@ pub(crate) mod arena;
 #[allow(dead_code)]
 pub(crate) mod countdown;
 
-#[cfg(any(feature = "once-map", feature = "singleflight"))]
-// `OnceMap` and `singleflight` use different subsets of `OnceTable`, so single-feature builds
-// leave some operations in the shared implementation unused.
-#[allow(dead_code)]
-pub(crate) mod once_table;
-
 #[cfg(any(feature = "lazy-cell", feature = "once-cell"))]
 // `LazyCell` and `OnceCell` use different subsets of `ValueCell`, so single-feature builds leave
 // some operations in the shared implementation unused.
@@ -65,6 +59,9 @@ pub(crate) mod value_cell;
 ))]
 pub(crate) mod mutex;
 
+#[cfg(feature = "once-map")]
+pub(crate) mod rwlock;
+
 #[cfg(any(
     feature = "mpsc",
     feature = "mutex",
@@ -97,3 +94,13 @@ pub(crate) mod waitlist;
 // `new`. One constructor is therefore unused in every single-primitive build.
 #[allow(dead_code)]
 pub(crate) mod waitset;
+
+#[cfg(any(feature = "once-map", feature = "singleflight"))]
+pub fn default_shard_count() -> usize {
+    // Tested on a 32-core machine, the optimal shard count for `OnceMap` and `Singleflight` is 256.
+    // So I use 8 as the coefficient, which is 256 / 32.
+    // Need to test on other machines to see if this coefficient is optimal.
+    // Dashmap use 4.
+    (std::thread::available_parallelism().map_or(1, |parallelism| parallelism.get()) * 8)
+        .next_power_of_two()
+}

asyncband/src/internal/mutex.rs

@@ -36,6 +36,13 @@ impl<T> Mutex<T> {
     }
 }
 
+#[cfg(any(feature = "once-map", feature = "singleflight"))]
+/// Alignment uses 60% more memory (64/40) but improves write performance by 25% at 32 threads. (On
+/// Zen 5 CPUs)
+/// Need to test on other architectures.
+#[repr(align(64))]
+pub struct CachePaddedMutex<T>(pub Mutex<T>);
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

asyncband/src/internal/rwlock.rs

@@ -0,0 +1,41 @@
+// 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)
+    }
+}