armdb/queue/

wait.rs

use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::task::{Context, Poll, Waker};
use std::thread::JoinHandle;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use crate::error::{DbError, DbResult};
use crate::shutdown::ShutdownSignal;
use crate::sync::{Condvar, Mutex, cv_wait_timeout, lock};

use super::{LIVE_CLASS, Queue, QueueHandle, QueueKey, QueueSpec};

/// How long the wake worker parks when nothing is registered.
///
/// Only a bound on how long a *lost* notification could cost; every real wake
/// comes through the condvar.
const WORKER_IDLE_PARK: Duration = Duration::from_secs(60);

/// Debug-only synchronization point inside the wake worker.
///
/// It is absent from release builds and is never configured by production
/// code. The worker stops here after its top-of-loop shutdown check and before
/// it parks, which is exactly the window in which a `shutdown` can issue its
/// notification with nobody waiting on the condvar.
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub mod test_support {
    use std::sync::{Arc, Barrier, LazyLock, Mutex};

    /// `(arrived, release)` — the worker signals the first and waits on the second.
    type ParkPause = (Arc<Barrier>, Arc<Barrier>);

    static PARK_PAUSE: LazyLock<Mutex<Option<ParkPause>>> = LazyLock::new(|| Mutex::new(None));

    /// Arm the pause; returns `(arrived, release)`.
    pub fn pause_wake_worker_before_park() -> ParkPause {
        let pause = (Arc::new(Barrier::new(2)), Arc::new(Barrier::new(2)));
        *PARK_PAUSE.lock().expect("queue wake pause lock") = Some(pause.clone());
        pause
    }

    pub fn disarm_wake_worker_pause() {
        *PARK_PAUSE.lock().expect("queue wake pause lock") = None;
    }

    pub(crate) fn wait_before_park() {
        let pause = PARK_PAUSE
            .lock()
            .expect("queue wake pause lock")
            .as_ref()
            .cloned();
        if let Some((arrived, release)) = pause {
            arrived.wait();
            release.wait();
        }
    }
}

pub(crate) trait QueueHeaderRow {
    fn lease_deadline_ms(&self) -> u64;
}

/// Fresh result of scanning queue visibility and lease state.
#[doc(hidden)]
pub struct QueueScan {
    pub(crate) due: bool,
    pub(crate) delay: Option<Duration>,
}

impl QueueScan {
    pub(crate) fn rows<I, R>(rows: I, wakeup_recheck: Duration) -> DbResult<Self>
    where
        I: IntoIterator<Item = (QueueKey, R)>,
        R: QueueHeaderRow,
    {
        Self::rows_at(rows, wakeup_recheck, now_ms()?)
    }

    pub(crate) fn rows_at<I, R>(rows: I, wakeup_recheck: Duration, now_ms: u64) -> DbResult<Self>
    where
        I: IntoIterator<Item = (QueueKey, R)>,
        R: QueueHeaderRow,
    {
        let mut next = None;
        for (key, row) in rows {
            if key.class() != LIVE_CLASS {
                break;
            }
            if key.visible_at_ms() > now_ms {
                next = min_target(next, key.visible_at_ms());
                break;
            }
            let deadline = row.lease_deadline_ms();
            if deadline > now_ms {
                next = min_target(next, deadline);
                continue;
            }
            return Ok(Self {
                due: true,
                delay: Some(Duration::ZERO),
            });
        }
        Ok(Self {
            due: false,
            delay: next.map(|target| {
                Duration::from_millis(target.saturating_sub(now_ms)).min(wakeup_recheck)
            }),
        })
    }
}

fn min_target(current: Option<u64>, candidate: u64) -> Option<u64> {
    Some(current.map_or(candidate, |current| current.min(candidate)))
}

fn now_ms() -> DbResult<u64> {
    let millis = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_err(|_| DbError::Client("queue time is before unix epoch"))?
        .as_millis();
    u64::try_from(millis).map_err(|_| DbError::Client("queue time exceeds u64 milliseconds"))
}

#[doc(hidden)]
pub struct QueueWait {
    state: Mutex<WaitState>,
    changed: Condvar,
    pub(crate) service: Arc<QueueWakeService>,
    recheck: Duration,
}

struct WaitState {
    generation: u64,
    blockers: HashMap<u64, u64>,
    next_blocker: u64,
    registrations: usize,
}

enum RegisterResult {
    Registered(u64),
    Changed,
    Shutdown,
}

impl QueueWait {
    pub(crate) fn new(service: Arc<QueueWakeService>, recheck: Duration) -> Arc<Self> {
        let wait = Arc::new(Self {
            state: Mutex::new(WaitState {
                generation: 0,
                blockers: HashMap::new(),
                next_blocker: 1,
                registrations: 0,
            }),
            changed: Condvar::new(),
            service,
            recheck,
        });
        // `shutdown_wake` only reaches async registrations, which live on the
        // service. A thread parked in `wait_due_timeout` waits on *this*
        // queue's condvar, so it needs its own listener or it sleeps out its
        // full timeout on a database that is being torn down.
        let weak = Arc::downgrade(&wait);
        wait.service.shutdown.listen(Arc::new(move || {
            if let Some(wait) = weak.upgrade() {
                wait.notify_all_blockers();
            }
        }));
        wait
    }

    /// Publish a visibility change before waking this queue's async waiters.
    /// The queue mutex is never held while taking the service mutex, so worker
    /// expiry (service-only) and direct producers cannot form a lock cycle.
    pub(crate) fn notify_visible_at(&self, visible_at_ms: u64) {
        {
            let mut state = lock(&self.state);
            state.generation = state.generation.wrapping_add(1);
            if visible_at_ms < blocker_threshold(&state) {
                self.changed.notify_all();
            }
        }
        self.service.wake_queue(self, visible_at_ms);
    }

    /// Wake every blocking waiter, whatever it announced.
    ///
    /// Used for events that are not a visibility change — today, shutdown.
    fn notify_all_blockers(&self) {
        let mut state = lock(&self.state);
        state.generation = state.generation.wrapping_add(1);
        self.changed.notify_all();
    }

    fn generation(&self) -> u64 {
        lock(&self.state).generation
    }

    fn wait_for_change(&self, delay: Duration, announced_until_ms: u64, generation: u64) -> bool {
        let mut state = lock(&self.state);
        if state.generation != generation {
            return false;
        }
        let id = state.next_blocker;
        state.next_blocker = state.next_blocker.wrapping_add(1);
        state.blockers.insert(id, announced_until_ms);
        let mut state = cv_wait_timeout(&self.changed, &self.state, state, delay);
        state.blockers.remove(&id);
        true
    }

    fn register(
        &self,
        id: Option<u64>,
        waker: Waker,
        delay: Duration,
        generation: u64,
    ) -> RegisterResult {
        let mut state = lock(&self.state);
        if state.generation != generation {
            return RegisterResult::Changed;
        }
        if self.service.shutdown.is_shutdown() {
            return RegisterResult::Shutdown;
        }
        let is_new = id.is_none();
        let announced_until_ms = now_ms()
            .unwrap_or(u64::MAX)
            .saturating_add(u64::try_from(delay.as_millis()).unwrap_or(u64::MAX));
        let Some(id) = self
            .service
            .register_locked(id, waker, delay, announced_until_ms, self)
        else {
            return RegisterResult::Shutdown;
        };
        if is_new {
            state.registrations += 1;
        }
        RegisterResult::Registered(id)
    }

    fn deregistered(&self) {
        let mut state = lock(&self.state);
        state.registrations = state.registrations.saturating_sub(1);
    }

    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    pub(crate) fn registration_count(&self) -> usize {
        lock(&self.state).registrations
    }
}

/// The visibility below which a `notify_all` is worth issuing.
///
/// `notify_visible_at` wakes blocking waiters all at once, so the aggregate has
/// to be the *latest* announced deadline, not the earliest: taking the minimum
/// would let one waiter with a short deadline suppress the notification for
/// every waiter with a longer one, which then sleeps out its whole recheck
/// instead of re-deriving a nearer delay. Waking a waiter early is free — it
/// rescans and parks again.
fn blocker_threshold(state: &WaitState) -> u64 {
    state.blockers.values().copied().max().unwrap_or(u64::MAX)
}

pub(crate) struct QueueWakeService {
    state: Mutex<WakeState>,
    changed: Condvar,
    pub(crate) shutdown: ShutdownSignal,
    worker: Mutex<Option<JoinHandle<()>>>,
    next_id: AtomicU64,
}

struct WakeState {
    registrations: HashMap<u64, Registration>,
}

struct Registration {
    queue: usize,
    announced_until_ms: u64,
    deadline: Instant,
    waker: Waker,
}

impl QueueWakeService {
    pub(crate) fn new(shutdown: ShutdownSignal) -> Arc<Self> {
        let service = Arc::new(Self {
            state: Mutex::new(WakeState {
                registrations: HashMap::new(),
            }),
            changed: Condvar::new(),
            shutdown: shutdown.clone(),
            worker: Mutex::new(None),
            next_id: AtomicU64::new(1),
        });
        let weak = Arc::downgrade(&service);
        shutdown.listen(Arc::new(move || {
            if let Some(service) = weak.upgrade() {
                service.shutdown_wake();
            }
        }));
        service
    }

    fn register_locked(
        &self,
        id: Option<u64>,
        waker: Waker,
        delay: Duration,
        announced_until_ms: u64,
        queue: &QueueWait,
    ) -> Option<u64> {
        let id = id.unwrap_or_else(|| self.next_id.fetch_add(1, Ordering::Relaxed));
        let mut state = lock(&self.state);
        if self.shutdown.is_shutdown() {
            return None;
        }
        state.registrations.insert(
            id,
            Registration {
                queue: queue as *const QueueWait as usize,
                announced_until_ms,
                deadline: Instant::now() + delay,
                waker,
            },
        );
        self.changed.notify_all();
        Some(id)
    }

    fn deregister(&self, id: u64) {
        lock(&self.state).registrations.remove(&id);
        self.changed.notify_all();
    }

    fn wake_queue(&self, queue: &QueueWait, visible_at_ms: u64) {
        let queue = queue as *const QueueWait as usize;
        let wakers = {
            let mut state = lock(&self.state);
            let ids: Vec<_> = state
                .registrations
                .iter()
                .filter_map(|(id, registration)| {
                    (registration.queue == queue && visible_at_ms < registration.announced_until_ms)
                        .then_some(*id)
                })
                .collect();
            ids.into_iter()
                .filter_map(|id| {
                    state
                        .registrations
                        .remove(&id)
                        .map(|registration| registration.waker)
                })
                .collect::<Vec<_>>()
        };
        for waker in wakers {
            waker.wake();
        }
    }

    fn shutdown_wake(&self) {
        let wakers = {
            let mut state = lock(&self.state);
            state
                .registrations
                .drain()
                .map(|(_, registration)| registration.waker)
                .collect::<Vec<_>>()
        };
        self.changed.notify_all();
        for waker in wakers {
            waker.wake();
        }
    }

    fn start(self: &Arc<Self>) {
        let mut worker = lock(&self.worker);
        if worker.is_some() {
            return;
        }
        let service = Arc::clone(self);
        *worker = Some(std::thread::spawn(move || service.run()));
    }

    fn run(self: Arc<Self>) {
        loop {
            if self.shutdown.is_shutdown() {
                self.shutdown_wake();
                return;
            }
            #[cfg(all(debug_assertions, feature = "queue-test-support"))]
            test_support::wait_before_park();
            let wakers = {
                let mut state = lock(&self.state);
                let now = Instant::now();
                let ids: Vec<_> = state
                    .registrations
                    .iter()
                    .filter_map(|(id, registration)| (registration.deadline <= now).then_some(*id))
                    .collect();
                if !ids.is_empty() {
                    ids.into_iter()
                        .filter_map(|id| {
                            state
                                .registrations
                                .remove(&id)
                                .map(|registration| registration.waker)
                        })
                        .collect()
                } else {
                    let delay = state
                        .registrations
                        .values()
                        .map(|registration| registration.deadline)
                        .min()
                        .map_or(WORKER_IDLE_PARK, |deadline| {
                            deadline.saturating_duration_since(now)
                        });
                    // The shutdown re-check must happen under `state`, not only
                    // at the top of the loop: `shutdown_wake` notifies after
                    // releasing this same mutex, so a worker preempted between
                    // the top-of-loop check and this park would miss the
                    // notification entirely and sleep out `WORKER_IDLE_PARK`
                    // while `Db::drop` blocks in `join`.
                    if self.shutdown.is_shutdown() {
                        drop(state);
                        self.shutdown_wake();
                        return;
                    }
                    drop(cv_wait_timeout(&self.changed, &self.state, state, delay));
                    Vec::new()
                }
            };
            for waker in wakers {
                waker.wake();
            }
        }
    }

    pub(crate) fn join(&self) {
        self.shutdown_wake();
        if let Some(worker) = lock(&self.worker).take() {
            let _ = worker.join();
        }
    }
}

/// Runtime-agnostic future that resolves once queue work can be leased.
pub struct QueueDue<'a, N: QueueSpec>
where
    N::Mode: QueueHandle<N>,
{
    queue: &'a Queue<N>,
    registration: Option<u64>,
}

impl<'a, N> QueueDue<'a, N>
where
    N: QueueSpec,
    N::Mode: QueueHandle<N>,
{
    pub(crate) fn new(queue: &'a Queue<N>) -> Self {
        Self {
            queue,
            registration: None,
        }
    }

    fn unregister(&mut self) {
        if let Some(id) = self.registration.take() {
            let wait = <N::Mode as QueueHandle<N>>::wait(self.queue);
            wait.service.deregister(id);
            wait.deregistered();
        }
    }
}

impl<N> Future for QueueDue<'_, N>
where
    N: QueueSpec,
    N::Mode: QueueHandle<N>,
{
    type Output = DbResult<()>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        loop {
            let wait = <N::Mode as QueueHandle<N>>::wait(self.queue);
            let generation = wait.generation();
            let scan = match <N::Mode as QueueHandle<N>>::scan_due(self.queue) {
                Ok(scan) => scan,
                Err(err) => {
                    self.unregister();
                    return Poll::Ready(Err(err));
                }
            };
            if wait.service.shutdown.is_shutdown() {
                self.unregister();
                return Poll::Ready(Err(DbError::Client("queue wait service is shut down")));
            }
            if scan.due {
                self.unregister();
                return Poll::Ready(Ok(()));
            }
            match wait.register(
                self.registration,
                cx.waker().clone(),
                scan.delay.unwrap_or(wait.recheck),
                generation,
            ) {
                RegisterResult::Registered(id) => {
                    self.registration = Some(id);
                    wait.service.start();
                    return Poll::Pending;
                }
                RegisterResult::Changed => continue,
                RegisterResult::Shutdown => {
                    self.unregister();
                    return Poll::Ready(Err(DbError::Client("queue wait service is shut down")));
                }
            }
        }
    }
}

impl<N> Drop for QueueDue<'_, N>
where
    N: QueueSpec,
    N::Mode: QueueHandle<N>,
{
    fn drop(&mut self) {
        self.unregister();
    }
}

pub(crate) fn wait_due_timeout<N>(queue: &Queue<N>, timeout: Duration) -> DbResult<bool>
where
    N: QueueSpec,
    N::Mode: QueueHandle<N>,
{
    let deadline = Instant::now() + timeout;
    loop {
        let wait = <N::Mode as QueueHandle<N>>::wait(queue);
        let generation = wait.generation();
        let scan = <N::Mode as QueueHandle<N>>::scan_due(queue)?;
        if scan.due {
            return Ok(true);
        }
        // Reported the same way the async twin reports it, so a
        // `while !stop { q.wait_due_timeout(..) }` loop terminates instead of
        // sleeping out its timeout against a shutting-down database.
        if wait.service.shutdown.is_shutdown() {
            return Err(DbError::Client("queue wait service is shut down"));
        }
        let remaining = deadline.saturating_duration_since(Instant::now());
        if remaining.is_zero() {
            return Ok(false);
        }
        let delay = scan
            .delay
            .unwrap_or(wait.recheck)
            .min(wait.recheck)
            .min(remaining);
        let announced_until_ms =
            now_ms()?.saturating_add(u64::try_from(delay.as_millis()).unwrap_or(u64::MAX));
        let _ = wait.wait_for_change(delay, announced_until_ms, generation);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[derive(Clone, Copy)]
    struct Header(u64);

    impl QueueHeaderRow for Header {
        fn lease_deadline_ms(&self) -> u64 {
            self.0
        }
    }

    #[test]
    fn backward_clock_jump_is_capped_by_recheck_from_a_fresh_scan() {
        let rows = [(QueueKey::live(110, 1), Header(0))];
        let recheck = Duration::from_millis(5);
        let before_jump = QueueScan::rows_at(rows, recheck, 100).unwrap();
        let after_jump = QueueScan::rows_at(rows, recheck, 90).unwrap();

        assert_eq!(before_jump.delay, Some(Duration::from_millis(5)));
        assert_eq!(after_jump.delay, Some(Duration::from_millis(5)));
    }
}