perf(lazy-cell): avoid locking on uncontended initialization

CHANGELOG.md

@@ -41,7 +41,7 @@ All notable changes to this project will be documented in this file.
 
 ### Improvements
 
-* Reduce the cost of initializing an uncontended `OnceCell` by avoiding its internal wait-queue lock.
+* Reduce the cost of initializing an uncontended `OnceCell` or `LazyCell` by avoiding its internal wait-queue lock.
 * Reduce the cost of polling a pending oneshot receive again from the same task, which now keeps its registered waker instead of cloning and dropping it.
 * Reduce the state kept by futures waiting for a `Mutex`, `RwLock`, `Semaphore`, `Once`, `OnceCell`, or bounded broadcast capacity, making lock, permit acquisition, and bounded broadcast send futures smaller.
 

asyncband/src/once/lazy_cell/mod.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::cell::UnsafeCell;
 use std::fmt;
 use std::future::Future;
 use std::future::Pending;
@@ -28,8 +29,13 @@ use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::thread;
 
+use crate::internal::semaphore::Semaphore;
 use crate::internal::value_cell::ValueCell;
-use crate::mutex::Mutex;
+
+// Each forcing caller takes two of three permits, so only one can drive the attempt at a time. The
+// spare permit lets an uncontended release avoid locking the semaphore's wait queue.
+const PERMITS: usize = 3;
+const DRIVER_PERMITS: usize = 2;
 
 /// A thread-safe value initialized by a stored asynchronous function on first access.
 ///
@@ -111,10 +117,16 @@ use crate::mutex::Mutex;
 /// ```
 pub struct LazyCell<T, Fut, F = fn() -> Fut> {
     value: ValueCell<T>,
-    state: Mutex<State<F, Fut>>,
+    // Shared access to `state` requires the driver permits.
+    semaphore: Semaphore,
+    state: UnsafeCell<State<F, Fut>>,
     poisoned: AtomicBool,
 }
 
+// SAFETY: Shared access to the initializer state is serialized by the driver permits, so the state
+// only needs to be sendable, as with a mutex. The value is shared once initialized.
+unsafe impl<T: Send + Sync, Fut: Send, F: Send> Sync for LazyCell<T, Fut, F> {}
+
 // Keep mutually exclusive phases in an enum so an inline future shares storage with its
 // initializer instead of making the cell large enough to hold both at once.
 enum State<F, Fut> {
@@ -162,7 +174,8 @@ impl<T, Fut, F> LazyCell<T, Fut, F> {
     {
         Self {
             value: ValueCell::new(),
-            state: Mutex::new(State::Initializer(Some(initializer))),
+            semaphore: Semaphore::new(PERMITS),
+            state: UnsafeCell::new(State::Initializer(Some(initializer))),
             poisoned: AtomicBool::new(false),
         }
     }
@@ -286,18 +299,20 @@ where
         }
         this_ref.assert_unpoisoned();
 
-        let mut state = this_ref.state.lock().await;
+        this_ref.semaphore.acquire(DRIVER_PERMITS).await;
+        let _permits = DriverPermits(&this_ref.semaphore);
         if let Some(value) = Self::get(this_ref) {
             return value;
         }
         this_ref.assert_unpoisoned();
 
-        // SAFETY: `this` remains pinned for the duration of the returned future. The state is
-        // stored in that cell, and a mutex guard never relocates its protected value.
-        let state = unsafe { Pin::new_unchecked(&mut *state) };
+        // SAFETY: The driver permits give this caller exclusive access to the state until they are
+        // released. `this` remains pinned for the duration of the returned future, and the state is
+        // stored inline in that cell.
+        let state = unsafe { Pin::new_unchecked(&mut *this_ref.state.get()) };
         let value = drive_pinned_attempt(state, &this_ref.poisoned).await;
 
-        // SAFETY: The state mutex serializes initialization, and the double check above verified
+        // SAFETY: The driver permits serialize initialization, and the double check above verified
         // that the value was not initialized by another caller.
         unsafe { this_ref.value.set(value) }
     }
@@ -426,7 +441,8 @@ where
     pub fn from_future(future: Fut) -> Self {
         Self {
             value: ValueCell::new(),
-            state: Mutex::new(State::Attempt(future)),
+            semaphore: Semaphore::new(PERMITS),
+            state: UnsafeCell::new(State::Attempt(future)),
             poisoned: AtomicBool::new(false),
         }
     }
@@ -439,7 +455,8 @@ impl<T> LazyCell<T, Pending<T>> {
     pub const fn from_value(value: T) -> Self {
         Self {
             value: ValueCell::from_value(value),
-            state: Mutex::new(State::Complete),
+            semaphore: Semaphore::new(PERMITS),
+            state: UnsafeCell::new(State::Complete),
             poisoned: AtomicBool::new(false),
         }
     }
@@ -465,6 +482,15 @@ impl<T: RefUnwindSafe + UnwindSafe, Fut: UnwindSafe, F: UnwindSafe> RefUnwindSaf
 {
 }
 
+// Returns the driver permits when forcing completes, is cancelled, or unwinds.
+struct DriverPermits<'a>(&'a Semaphore);
+
+impl Drop for DriverPermits<'_> {
+    fn drop(&mut self) {
+        self.0.release(DRIVER_PERMITS);
+    }
+}
+
 struct PoisonOnPanic<'a>(&'a AtomicBool);
 
 impl Drop for PoisonOnPanic<'_> {

benchmarks/Cargo.toml

@@ -32,6 +32,7 @@ asyncband = { workspace = true, features = [
   "condvar",
   "event",
   "latch",
+  "lazy-cell",
   "mpmc",
   "mpsc",
   "mutex",

benchmarks/benches/primitives/once/force.rs

@@ -0,0 +1,63 @@
+// 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::future::ready;
+use std::hint::black_box;
+use std::pin::pin;
+
+use asyncband::once::LazyCell;
+use benchmarks::harness::Bencher;
+use benchmarks::harness::Runner;
+use benchmarks::primitive;
+use benchmarks::support::bench_context;
+use benchmarks::support::poll_ready;
+
+fn initialize_cell(bencher: Bencher) {
+    let mut context = bench_context();
+
+    bencher.bench_local(|| {
+        let cell = pin!(LazyCell::new(|| ready(black_box(1usize))));
+        let value = poll_ready(LazyCell::force_pin(cell.as_ref()), &mut context);
+        black_box(*value)
+    });
+}
+
+fn initialize_cell_from_future(bencher: Bencher) {
+    let mut context = bench_context();
+
+    bencher.bench_local(|| {
+        let cell = pin!(LazyCell::from_future(ready(black_box(1usize))));
+        let value = poll_ready(LazyCell::force_pin(cell.as_ref()), &mut context);
+        black_box(*value)
+    });
+}
+
+fn initialized_cell(bencher: Bencher) {
+    let mut context = bench_context();
+    let cell = LazyCell::from_value(1usize);
+
+    bencher.bench_local(|| {
+        let value = poll_ready(LazyCell::force(black_box(&cell)), &mut context);
+        black_box(*value)
+    });
+}
+
+pub fn register(runner: &mut Runner) {
+    primitive!(runner, initialize_cell, 7_000_000);
+    primitive!(runner, initialize_cell_from_future, 7_000_000);
+    primitive!(runner, initialized_cell, 85_000_000);
+}

benchmarks/benches/primitives/once/mod.rs

@@ -15,12 +15,14 @@
 // specific language governing permissions and limitations
 // under the License.
 
+mod force;
 mod get_or_init;
 mod wait;
 
 use benchmarks::harness::Runner;
 
 pub fn register(runner: &mut Runner) {
+    force::register(runner);
     get_or_init::register(runner);
     wait::register(runner);
 }

tests-integration/tests/lazy_cell_test.rs

@@ -20,14 +20,24 @@ use std::future::Ready;
 use std::future::pending;
 use std::future::poll_fn;
 use std::future::ready;
+use std::panic::AssertUnwindSafe;
+use std::panic::catch_unwind;
 use std::pin::pin;
 use std::rc::Rc;
 use std::sync::Arc;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
+use std::sync::mpsc;
 use std::task::Poll;
+use std::thread;
 
+use asyncband::blocking::FutureExt;
 use asyncband::once::LazyCell;
+use tests_integration::WakeCounter;
+use tests_integration::assert_completes_without_deadlock;
+use tests_integration::expect_ready;
+use tests_integration::poll_once;
+use tests_integration::poll_with;
 use tokio::sync::Notify;
 
 struct DropPanic;
@@ -448,3 +458,136 @@ async fn cancellation_preserves_inline_initialization_future() {
     assert_eq!(*LazyCell::force_pin(lazy.as_ref()).await, 42);
     assert_eq!(attempts.load(Ordering::SeqCst), 1);
 }
+
+#[test]
+fn queued_caller_resumes_the_same_future_after_the_driver_is_cancelled() {
+    let starts = Cell::new(0);
+    let (finish, ready) = tokio::sync::oneshot::channel();
+    let lazy = pin!(LazyCell::new({
+        let starts = &starts;
+        async move || {
+            starts.set(starts.get() + 1);
+            ready.await.unwrap();
+            42
+        }
+    }));
+    let mut driver = Box::pin(LazyCell::force_pin(lazy.as_ref()));
+    let mut queued = pin!(LazyCell::force_pin(lazy.as_ref()));
+    let (waker, wakes) = WakeCounter::new();
+
+    assert!(poll_once(driver.as_mut()).is_pending());
+    assert!(poll_with(queued.as_mut(), &waker).is_pending());
+    drop(driver);
+    assert_eq!(wakes.take(), 1);
+
+    assert!(poll_with(queued.as_mut(), &waker).is_pending());
+    finish.send(()).unwrap();
+    assert_eq!(wakes.take(), 1);
+    assert_eq!(expect_ready(poll_with(queued.as_mut(), &waker)), &42);
+    assert_eq!(starts.get(), 1);
+}
+
+#[test]
+fn queued_caller_observes_poison_after_the_driver_panics() {
+    let (finish, ready) = tokio::sync::oneshot::channel::<()>();
+    let lazy = pin!(LazyCell::new(async move || -> u8 {
+        ready.await.unwrap();
+        panic!("initializer panic");
+    }));
+    let mut driver = pin!(LazyCell::force_pin(lazy.as_ref()));
+    let mut queued = pin!(LazyCell::force_pin(lazy.as_ref()));
+    let (waker, wakes) = WakeCounter::new();
+
+    assert!(poll_once(driver.as_mut()).is_pending());
+    assert!(poll_with(queued.as_mut(), &waker).is_pending());
+    finish.send(()).unwrap();
+    let panic = catch_unwind(AssertUnwindSafe(|| poll_once(driver.as_mut()))).unwrap_err();
+    assert_eq!(panic.downcast_ref::<&str>(), Some(&"initializer panic"));
+    assert_eq!(wakes.count(), 1);
+
+    let panic = catch_unwind(AssertUnwindSafe(|| poll_with(queued.as_mut(), &waker))).unwrap_err();
+    assert_eq!(
+        panic.downcast_ref::<&str>(),
+        Some(&"LazyCell instance has previously been poisoned")
+    );
+    assert_eq!(LazyCell::get(&lazy), None);
+}
+
+#[test]
+fn cancelling_a_partial_waiter_does_not_start_another_driver() {
+    let (finish, ready) = tokio::sync::oneshot::channel();
+    let lazy = pin!(LazyCell::from_future(async move {
+        ready.await.unwrap();
+        2
+    }));
+    let mut driver = Box::pin(LazyCell::force_pin(lazy.as_ref()));
+    let mut cancelled = Box::pin(LazyCell::force_pin(lazy.as_ref()));
+    let mut next = pin!(LazyCell::force_pin(lazy.as_ref()));
+    let (waker, wakes) = WakeCounter::new();
+
+    assert!(poll_once(driver.as_mut()).is_pending());
+    // This waiter holds the spare permit, but still cannot drive the attempt.
+    assert!(poll_once(cancelled.as_mut()).is_pending());
+    assert!(poll_with(next.as_mut(), &waker).is_pending());
+    drop(cancelled);
+    assert_eq!(wakes.count(), 0);
+    assert!(poll_with(next.as_mut(), &waker).is_pending());
+
+    drop(driver);
+    assert_eq!(wakes.take(), 1);
+    finish.send(()).unwrap();
+    assert_eq!(expect_ready(poll_with(next.as_mut(), &waker)), &2);
+}
+
+#[test]
+fn cancelling_a_granted_waiter_passes_the_attempt_to_the_next() {
+    let (finish, ready) = tokio::sync::oneshot::channel();
+    let lazy = pin!(LazyCell::from_future(async move {
+        ready.await.unwrap();
+        3
+    }));
+    let mut driver = Box::pin(LazyCell::force_pin(lazy.as_ref()));
+    let mut cancelled = Box::pin(LazyCell::force_pin(lazy.as_ref()));
+    let mut next = pin!(LazyCell::force_pin(lazy.as_ref()));
+    let (cancelled_waker, cancelled_wakes) = WakeCounter::new();
+    let (next_waker, next_wakes) = WakeCounter::new();
+
+    assert!(poll_once(driver.as_mut()).is_pending());
+    assert!(poll_with(cancelled.as_mut(), &cancelled_waker).is_pending());
+    assert!(poll_with(next.as_mut(), &next_waker).is_pending());
+    drop(driver);
+    assert_eq!(cancelled_wakes.count(), 1);
+    assert_eq!(next_wakes.count(), 0);
+
+    drop(cancelled);
+    assert_eq!(next_wakes.take(), 1);
+    finish.send(()).unwrap();
+    assert_eq!(expect_ready(poll_with(next.as_mut(), &next_waker)), &3);
+}
+
+#[test]
+fn queued_caller_observes_completion_from_another_thread() {
+    assert_completes_without_deadlock(|| {
+        let (finish, ready) = tokio::sync::oneshot::channel();
+        let lazy = pin!(LazyCell::new(async move || {
+            ready.await.unwrap();
+            [42; 32]
+        }));
+        let lazy = lazy.as_ref();
+        let mut driver = pin!(LazyCell::force_pin(lazy));
+        assert!(poll_once(driver.as_mut()).is_pending());
+
+        let (registered, registration) = mpsc::channel();
+        thread::scope(|scope| {
+            scope.spawn(move || {
+                let mut waiter = pin!(LazyCell::force_pin(lazy));
+                assert!(poll_once(waiter.as_mut()).is_pending());
+                registered.send(()).unwrap();
+                assert_eq!(waiter.block_on(), &[42; 32]);
+            });
+            registration.recv().unwrap();
+            finish.send(()).unwrap();
+            assert_eq!(expect_ready(poll_once(driver.as_mut())), &[42; 32]);
+        });
+    });
+}

tests-integration/tests/unsafe_paths_test.rs

@@ -213,3 +213,16 @@ fn lazy_cell_resumes_a_pinned_attempt_in_place() {
     let mut force = pin!(LazyCell::force_pin(lazy.as_ref()));
     assert_eq!(poll_once(force.as_mut()), Poll::Ready(&42));
 }
+
+#[test]
+fn lazy_cell_hands_a_pinned_attempt_to_a_queued_caller() {
+    let lazy = LazyCell::from_future(AddressSensitiveFuture::new());
+    let lazy = pin!(lazy);
+    let mut driver = Box::pin(LazyCell::force_pin(lazy.as_ref()));
+    let mut queued = pin!(LazyCell::force_pin(lazy.as_ref()));
+
+    assert!(poll_once(driver.as_mut()).is_pending());
+    assert!(poll_once(queued.as_mut()).is_pending());
+    drop(driver);
+    assert_eq!(poll_once(queued.as_mut()), Poll::Ready(&42));
+}