modrpc 0.0.3

Runtime for modrpc - a modular RPC framework
Documentation
use std::{
    cell::Cell,
    future::Future,
    pin::Pin,
    rc::Rc,
    task::{Context, Poll, Waker},
    time::{Duration, Instant},
};

pub enum FlushBatcherStatus {
    Snooze { duration: Duration },
    FlushNow,
    DoNotFlush,
}

#[derive(Copy, Clone)]
enum NextFlushTime {
    Instant(Instant),
    Immediate,
}

#[derive(Clone)]
pub struct FlushBatcher {
    inner: Rc<Inner>,
}

struct Inner {
    max_flush_delay: Duration,
    next_flush_time: Cell<Option<NextFlushTime>>,
    waker: Cell<Option<Waker>>,
}

impl FlushBatcher {
    pub fn new(max_flush_delay: Duration) -> Self {
        Self {
            inner: Rc::new(Inner {
                max_flush_delay,
                next_flush_time: Cell::new(None),
                waker: Cell::new(None),
            }),
        }
    }

    pub fn schedule_flush(&self) {
        if self.inner.next_flush_time.get().is_some() {
            return;
        }

        if self.inner.max_flush_delay.is_zero() {
            self.inner
                .next_flush_time
                .set(Some(NextFlushTime::Immediate));
        } else {
            #[cfg(not(target_arch = "wasm32"))]
            self.inner.next_flush_time.set(
                Instant::now()
                    .checked_add(self.inner.max_flush_delay)
                    .map(NextFlushTime::Instant),
            );
            #[cfg(target_arch = "wasm32")]
            self.inner
                .next_flush_time
                .set(Some(NextFlushTime::Immediate));
        }

        if let Some(waker) = self.inner.waker.take() {
            waker.wake();
        }
    }

    pub fn cancel_flush(&self) {
        self.inner.next_flush_time.set(None);
    }

    pub fn handle_flush(&self) -> FlushBatcherStatus {
        if let Some(next_flush_time) = self.inner.next_flush_time.get() {
            match next_flush_time {
                NextFlushTime::Immediate => {
                    self.inner.next_flush_time.set(None);
                    FlushBatcherStatus::FlushNow
                }
                NextFlushTime::Instant(next_flush_time) => {
                    let now = Instant::now();
                    if now >= next_flush_time {
                        self.inner.next_flush_time.set(None);
                        FlushBatcherStatus::FlushNow
                    } else {
                        FlushBatcherStatus::Snooze {
                            duration: next_flush_time - now,
                        }
                    }
                }
            }
        } else {
            FlushBatcherStatus::DoNotFlush
        }
    }

    pub fn wait(&self) -> FlushBatcherWait {
        FlushBatcherWait {
            inner: self.inner.clone(),
        }
    }
}

pub struct FlushBatcherWait {
    inner: Rc<Inner>,
}

impl Future for FlushBatcherWait {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        if self.inner.next_flush_time.get().is_some() {
            return Poll::Ready(());
        }

        self.inner.waker.set(Some(cx.waker().clone()));

        Poll::Pending
    }
}