diff --git a/CHANGELOG.md b/CHANGELOG.md index 17785dc3..2c7222cc 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/asyncband/src/once/lazy_cell/mod.rs b/asyncband/src/once/lazy_cell/mod.rs index a0418765..7867fde6 100644 --- a/asyncband/src/once/lazy_cell/mod.rs +++ b/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 Fut> { value: ValueCell, - state: Mutex>, + // Shared access to `state` requires the driver permits. + semaphore: Semaphore, + state: UnsafeCell>, 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 Sync for LazyCell {} + // 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 { @@ -162,7 +174,8 @@ impl LazyCell { { 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 LazyCell> { 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 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<'_> { diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index fffebeee..adcc653f 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -32,6 +32,7 @@ asyncband = { workspace = true, features = [ "condvar", "event", "latch", + "lazy-cell", "mpmc", "mpsc", "mutex", diff --git a/benchmarks/benches/primitives/once/force.rs b/benchmarks/benches/primitives/once/force.rs new file mode 100644 index 00000000..36ee7529 --- /dev/null +++ b/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); +} diff --git a/benchmarks/benches/primitives/once/mod.rs b/benchmarks/benches/primitives/once/mod.rs index 522a543b..fc154b22 100644 --- a/benchmarks/benches/primitives/once/mod.rs +++ b/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); } diff --git a/tests-integration/tests/lazy_cell_test.rs b/tests-integration/tests/lazy_cell_test.rs index d08f4bef..f500a7bb 100644 --- a/tests-integration/tests/lazy_cell_test.rs +++ b/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]); + }); + }); +} diff --git a/tests-integration/tests/unsafe_paths_test.rs b/tests-integration/tests/unsafe_paths_test.rs index d19b266f..f9cb58a8 100644 --- a/tests-integration/tests/unsafe_paths_test.rs +++ b/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)); +}