sync: reduce contention in Notify

[satakuma]

This PR replaces the notification field used by waiters in Notify with an atomic variable to reduce lock contention when multiple pending Notified futures are completing.

Motivation

The typical lifecycle of a Notified future looks like this:

1. It gets polled for the first time, and it is inserted into the waiting list.
2. It receives a notification from a notify_one or notify_waiters call and is removed from the waiting list.
3. The notifying function calls wake() on the future's waker.
4. The future gets polled again by the executor, consumes the notification and returns Ready.

Currently, step 4 requires acquiring the lock to read the notification field and the lock is released immediately after. Using an atomic variable here should significantly reduce lock contention, especially when multiple futures get notified at once by a notify_waiters call.

Benchmarks

Simple benchmarks on my machine (Linux 6.1.11, 4-core x86_64 developer laptop) show improved performance by 15-35%.

Click to expand results

master (c894069)

test notify_one<10>      ... bench:     128,130 ns/iter (+/- 18,168)
test notify_one<50>      ... bench:     226,002 ns/iter (+/- 20,095)
test notify_one<100>     ... bench:     199,244 ns/iter (+/- 18,873)
test notify_one<200>     ... bench:     210,647 ns/iter (+/- 28,222)
test notify_one<500>     ... bench:     218,932 ns/iter (+/- 29,464)

test notify_waiters<10>  ... bench:     289,161 ns/iter (+/- 25,546)
test notify_waiters<50>  ... bench:     248,969 ns/iter (+/- 47,217)
test notify_waiters<100> ... bench:     246,124 ns/iter (+/- 51,925)
test notify_waiters<200> ... bench:     250,637 ns/iter (+/- 44,959)
test notify_waiters<500> ... bench:     245,787 ns/iter (+/- 62,522)

with atomic notification:

test notify_one<10>      ... bench:     112,035 ns/iter (+/- 18,930)
test notify_one<50>      ... bench:     153,866 ns/iter (+/- 21,207)
test notify_one<100>     ... bench:     150,171 ns/iter (+/- 36,634)
test notify_one<200>     ... bench:     154,960 ns/iter (+/- 12,386)
test notify_one<500>     ... bench:     152,904 ns/iter (+/- 17,885)

test notify_waiters<10>  ... bench:     210,752 ns/iter (+/- 23,719)
test notify_waiters<50>  ... bench:     161,876 ns/iter (+/- 7,037)
test notify_waiters<100> ... bench:     157,448 ns/iter (+/- 25,909)
test notify_waiters<200> ... bench:     167,677 ns/iter (+/- 22,640)
test notify_waiters<500> ... bench:     290,953 ns/iter (+/- 39,285)

With many contending waiters (>500) benchmarks show degraded performance, but I did some profiling and the bottleneck seems to be somewhere in the scheduler.

[Darksonn]

This PR is rather large. Will take me a bit until I can review it.

[satakuma]

Sure, no problem. I've pushed an update which slightly reduces the size of this change. I can write down some key points to make the review easier.

Adding the atomic variable is quite straightforward, it boils down to adding atomic stores to notify_{one,waiters} functions and an atomic load to poll_notified. The reason why the diff ended up being so large is that I had to change the existing code to be correct aliasing-wise regarding the atomic field.

Previously, a Waiter lived inside an UnsafeCell and could be mutably borrowed after acquiring the lock which guards the waitlist. This PR avoids mutable borrows of Waiter and instead wraps the only mutable field (the waker) in an UnsafeCell to allow mutation. Technically, there is one, indirect mutable borrow of Waiter coming from poll's Pin<&mut Notified>, but that's a known issue of self-referential structs and this is safe because Waiter is !Unpin.

Waiters for the batch semaphore also use an atomic field for tracking the number of permits, and the code there is organized similarly to this PR. If this PR gets merged, I would like to also investigate if similar optimization can be applied in the IO driver.

[Darksonn]

I think the code is ok, but the safety invariants invovled with this atomic should be documented more clearly.

[satakuma]

I've added some comments to better document the safety invariants. I've also moved the atomic variable into a wrapper type to use an enum instead of constants.

[Darksonn]

I reviewed this again, and it LGTM.