Performance: register atomic waker lazily in AsyncRingBuf

[Congyuwang]

For the Async interface, registering atomic waker seems to affect performance quite a bit.

One way to reduce this cost is to register the waker only lazily: registering waker when the future is about to return Pending. So in a lot of cases when the buffer is readily available no registering is needed.

There is one caveat though: the caller must check again after register_waker, that the ring-buffer's expected state has not changed (i.e., the producer after buffer full and not closed -> about to return Pending -> register waker must check again that buffer is still full and not closed -> Pending, and the consumer after buffer empty and not closed -> about to return Pending -> register waker must check again that buffer is still empty and not closed). If the state has changed before the waker is registered, we might miss new notification of close or write / read.

As long as actual Pending return cases is relatively just a small portion of all polling, this might result in a better performance, since it prevents wasted synchronization of unnecessary waker registrations.

I've written a simple benchmark with Tokio runtime. And here is one GitHub Action run:

master:

running 1 test
test bytes_stream ... bench: 102,902,861 ns/iter (+/- 269,243)
test result: ok. 0 passed; 0 failed; 0 ignored; 1 measured; 0 filtered out; finished in 30.99s

lazy register:

running 1 test
test bytes_stream ... bench:  69,654,620 ns/iter (+/- 202,799)
test result: ok. 0 passed; 0 failed; 0 ignored; 1 measured; 0 filtered out; finished in 21.05s

This particular task runs about 30% faster:

#![feature(test)]
extern crate test;
use test::Bencher;
use async_ringbuf::traits::Split;
use futures::{AsyncReadExt, AsyncWriteExt};

const BUF_SIZE: usize = 8 * 1024;
const MSG_SIZE: usize = 128;
const MSG_COUNT: usize = 1024 * 1024;

#[bench]
fn bytes_stream(b: &mut Bencher) {
    let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();
    b.iter(|| {
        rt.block_on(send_receive());
    });
}

async fn send_receive() {

    let (mut prd, mut cons) = async_ringbuf::AsyncHeapRb::<u8>::new(BUF_SIZE).split();

    // prd task
    let prd_task = tokio::spawn(async move {
        let msg = [0u8; MSG_SIZE];
        for _ in 0..MSG_COUNT {
            prd.write_all(&msg).await.unwrap();
        }
    });

    // cons task
    let cons_task = tokio::spawn(async move {
        let mut buf = [0u8; MSG_SIZE];
        for _ in 0..MSG_COUNT {
            cons.read_exact(&mut buf).await.unwrap();
        }
    });

    // await finish
    prd_task.await.unwrap();
    cons_task.await.unwrap();
}

[agerasev]

Nice solution, checking ring buffer in a loop really eliminates the need to register waker in advance.

[agerasev]

30% is really nice performance improvement! It's interesting that registering an AtomicWaker is so expensive.

I'm merging the request.

Also I've added your benchmark to the crate if you don't mind.

Thank you for your interest in this project and for such an effort to make it better!