armdb 0.11.2

use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
use std::time::Duration;

use armdb::{QueueConfig, queue};
use armour_core::GetType;
use rapira::Rapira;

#[derive(
    Clone,
    Copy,
    GetType,
    zerocopy::FromBytes,
    zerocopy::IntoBytes,
    zerocopy::Immutable,
    zerocopy::KnownLayout,
)]
#[repr(transparent)]
pub struct Job(u64);

queue!(WaitJobs, payload = Job, name = "queue_wait_jobs");

#[derive(Clone, Debug, PartialEq, Eq, GetType, Rapira)]
pub struct VarWaitJob {
    body: String,
}

queue!(
    VarWaitJobs,
    payload = VarWaitJob,
    name = "queue_wait_var_jobs",
    var
);

fn open_wait_queue(
    recheck: Duration,
) -> (
    std::sync::Arc<armdb::Queue<WaitJobs>>,
    std::sync::Arc<armdb::armour::Db>,
    tempfile::TempDir,
) {
    open_wait_queue_config(
        QueueConfig::balanced()
            .wakeup_recheck(recheck)
            .build()
            .unwrap(),
    )
}

fn open_wait_queue_config(
    config: QueueConfig,
) -> (
    std::sync::Arc<armdb::Queue<WaitJobs>>,
    std::sync::Arc<armdb::armour::Db>,
    tempfile::TempDir,
) {
    // The directory is returned, never `keep()`-ed: these tests run under
    // `just test`, whose TMPDIR is a RAM disk with a bounded free-space guard.
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(armdb::armour::Db::open_test(dir.path()).unwrap());
    let q = db.open_queue::<WaitJobs>(config).unwrap();
    (q, db, dir)
}

fn open_var_wait_queue(
    config: QueueConfig,
) -> (
    std::sync::Arc<armdb::Queue<VarWaitJobs>>,
    std::sync::Arc<armdb::armour::Db>,
    tempfile::TempDir,
) {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(armdb::armour::Db::open_test(dir.path()).unwrap());
    let q = db.open_queue::<VarWaitJobs>(config).unwrap();
    (q, db, dir)
}

#[test]
fn push_wakes_a_blocking_waiter_before_its_timeout() {
    let (q, _db, _dir) = open_wait_queue(Duration::from_secs(5));
    let wait_q = q.clone();
    let waiter =
        std::thread::spawn(move || wait_q.wait_due_timeout(Duration::from_secs(2)).unwrap());
    std::thread::sleep(Duration::from_millis(10));
    q.push(Job(1), None).unwrap();
    assert!(waiter.join().unwrap());
}

#[test]
fn next_delay_considers_an_expiring_lease_and_is_capped_by_recheck() {
    let (q, _db, _dir) = open_wait_queue(Duration::from_millis(5));
    q.push(Job(1), None).unwrap();
    let _leased = q.lease(1, Some(Duration::from_secs(30))).unwrap();
    assert!(q.next_delay().unwrap().unwrap() <= Duration::from_millis(5));
}

#[test]
fn dropping_pending_due_future_removes_its_waker_registration() {
    let (q, _db, _dir) = open_wait_queue(Duration::from_secs(5));
    let mut due = Box::pin(q.due());
    let waker = noop_waker();
    let mut cx = Context::from_waker(&waker);
    assert!(matches!(Pin::as_mut(&mut due).poll(&mut cx), Poll::Pending));
    assert_eq!(q.debug_wait_registration_count(), 1);
    drop(due);
    assert_eq!(q.debug_wait_registration_count(), 0);
}

#[test]
fn direct_push_wakes_a_pending_due_future_before_its_recheck() {
    let (q, _db, _dir) = open_wait_queue(Duration::from_secs(5));
    let counter = Arc::new(WakeCounter::default());
    let mut due = Box::pin(q.due());
    poll_pending(&mut due, &counter);

    q.push(Job(1), None).unwrap();
    counter.wait_for(1);
}

#[test]
fn transactional_fixed_push_wakes_a_pending_due_future_before_its_recheck() {
    let (q, db, _dir) = open_wait_queue(Duration::from_secs(5));
    let counter = Arc::new(WakeCounter::default());
    let mut due = Box::pin(q.due());
    poll_pending(&mut due, &counter);

    let producer_q = Arc::clone(&q);
    let producer_db = Arc::clone(&db);
    let producer = std::thread::spawn(move || {
        producer_db
            .tx()
            .push(&producer_q, Job(1), None)
            .unwrap()
            .commit()
            .unwrap();
    });

    counter.wait_for(1);
    producer.join().unwrap();
}

#[test]
fn transactional_var_push_wakes_a_pending_due_future_before_its_recheck() {
    let (q, db, _dir) = open_var_wait_queue(
        QueueConfig::balanced()
            .wakeup_recheck(Duration::from_secs(5))
            .build()
            .unwrap(),
    );
    let counter = Arc::new(WakeCounter::default());
    let mut due = Box::pin(q.due());
    let waker = counting_waker(Arc::clone(&counter));
    let mut cx = Context::from_waker(&waker);
    assert!(matches!(due.as_mut().poll(&mut cx), Poll::Pending));

    let producer_q = Arc::clone(&q);
    let producer_db = Arc::clone(&db);
    let producer = std::thread::spawn(move || {
        producer_db
            .tx()
            .push(
                &producer_q,
                VarWaitJob {
                    body: "wake".into(),
                },
                None,
            )
            .unwrap()
            .commit()
            .unwrap();
    });

    counter.wait_for(1);
    producer.join().unwrap();
}

#[test]
fn transactional_fixed_nack_wakes_a_pending_due_future_before_its_recheck() {
    let (q, db, _dir) = open_wait_queue_config(
        QueueConfig::balanced()
            .max_attempts(2)
            .backoff_secs(&[0])
            .wakeup_recheck(Duration::from_secs(5))
            .build()
            .unwrap(),
    );
    q.push(Job(1), None).unwrap();
    let token = q.lease(1, Some(Duration::from_secs(30))).unwrap()[0].token();
    let counter = Arc::new(WakeCounter::default());
    let mut due = Box::pin(q.due());
    poll_pending(&mut due, &counter);

    let producer_q = Arc::clone(&q);
    let producer_db = Arc::clone(&db);
    let producer = std::thread::spawn(move || {
        producer_db
            .tx()
            .nack(&producer_q, token)
            .unwrap()
            .commit()
            .unwrap();
    });

    counter.wait_for(1);
    producer.join().unwrap();
}

#[test]
fn transactional_var_nack_wakes_a_pending_due_future_before_its_recheck() {
    let (q, db, _dir) = open_var_wait_queue(
        QueueConfig::balanced()
            .max_attempts(2)
            .backoff_secs(&[0])
            .wakeup_recheck(Duration::from_secs(5))
            .build()
            .unwrap(),
    );
    q.push(
        VarWaitJob {
            body: "retry".into(),
        },
        None,
    )
    .unwrap();
    let token = q.lease(1, Some(Duration::from_secs(30))).unwrap()[0].token();
    let counter = Arc::new(WakeCounter::default());
    let mut due = Box::pin(q.due());
    let waker = counting_waker(Arc::clone(&counter));
    let mut cx = Context::from_waker(&waker);
    assert!(matches!(due.as_mut().poll(&mut cx), Poll::Pending));

    let producer_q = Arc::clone(&q);
    let producer_db = Arc::clone(&db);
    let producer = std::thread::spawn(move || {
        producer_db
            .tx()
            .nack(&producer_q, token)
            .unwrap()
            .commit()
            .unwrap();
    });

    counter.wait_for(1);
    producer.join().unwrap();
}

#[test]
fn direct_nack_wakes_a_pending_due_future_before_its_recheck() {
    let (q, _db, _dir) = open_wait_queue_config(
        QueueConfig::balanced()
            .max_attempts(2)
            .backoff_secs(&[0])
            .wakeup_recheck(Duration::from_secs(5))
            .build()
            .unwrap(),
    );
    q.push(Job(1), None).unwrap();
    let leased = q.lease(1, Some(Duration::from_secs(30))).unwrap();
    let token = leased[0].token();
    let counter = Arc::new(WakeCounter::default());
    let mut due = Box::pin(q.due());
    poll_pending(&mut due, &counter);

    q.nack(token).unwrap();
    counter.wait_for(1);
}

#[test]
fn direct_var_nack_wakes_a_pending_due_future_before_its_recheck() {
    let (q, _db, _dir) = open_var_wait_queue(
        QueueConfig::balanced()
            .max_attempts(2)
            .backoff_secs(&[0])
            .wakeup_recheck(Duration::from_secs(5))
            .build()
            .unwrap(),
    );
    q.push(
        VarWaitJob {
            body: "retry".into(),
        },
        None,
    )
    .unwrap();
    let token = q.lease(1, Some(Duration::from_secs(30))).unwrap()[0].token();
    let counter = Arc::new(WakeCounter::default());
    let mut due = Box::pin(q.due());
    let waker = counting_waker(Arc::clone(&counter));
    let mut cx = Context::from_waker(&waker);
    assert!(matches!(due.as_mut().poll(&mut cx), Poll::Pending));

    q.nack(token).unwrap();
    counter.wait_for(1);
}

#[test]
fn expired_due_registration_wakes_once_until_repolled() {
    let (q, _db, _dir) = open_wait_queue(Duration::from_millis(5));
    let counter = Arc::new(WakeCounter::default());
    let mut due = Box::pin(q.due());
    poll_pending(&mut due, &counter);

    counter.wait_for(1);
    std::thread::sleep(Duration::from_millis(20));
    assert_eq!(counter.wakes.load(Ordering::Acquire), 1);
    poll_pending(&mut due, &counter);
}

#[test]
fn repolled_expired_due_rearms_its_threshold_for_a_direct_push() {
    let (q, _db, _dir) = open_wait_queue(Duration::from_millis(5));
    let counter = Arc::new(WakeCounter::default());
    let mut due = Box::pin(q.due());
    poll_pending(&mut due, &counter);

    counter.wait_for(1);
    poll_pending(&mut due, &counter);
    q.push(Job(1), None).unwrap();
    counter.wait_for(2);
}

#[test]
fn shutdown_wakes_pending_due_and_repoll_returns_error() {
    let (q, db, _dir) = open_wait_queue(Duration::from_secs(5));
    let counter = Arc::new(WakeCounter::default());
    let mut due = Box::pin(q.due());
    poll_pending(&mut due, &counter);

    db.shutdown_signal().shutdown();
    counter.wait_for(1);
    let waker = counting_waker(Arc::clone(&counter));
    let mut cx = Context::from_waker(&waker);
    assert!(matches!(
        Pin::as_mut(&mut due).poll(&mut cx),
        Poll::Ready(Err(_))
    ));
}

#[test]
fn a_short_deadline_waiter_does_not_suppress_the_wake_of_a_longer_one() {
    // `notify_visible_at` wakes blocking waiters with a single `notify_all`, so
    // the announced deadlines must aggregate with `max`. Both waiters below are
    // parked when a row is pushed with a visibility *between* their announced
    // deadlines: under a `min` aggregate the 300 ms waiter pins the threshold,
    // no notification is issued, and the 4 s waiter sleeps out its entire wait
    // instead of re-deriving the row's much nearer delay.
    const VISIBLE_IN: Duration = Duration::from_millis(900);
    let (q, _db, _dir) = open_wait_queue(Duration::from_secs(30));

    let short_q = Arc::clone(&q);
    let short = std::thread::spawn(move || short_q.wait_due_timeout(Duration::from_millis(300)));
    let long_q = Arc::clone(&q);
    let long = std::thread::spawn(move || {
        let started = std::time::Instant::now();
        let due = long_q.wait_due_timeout(Duration::from_secs(4)).unwrap();
        (due, started.elapsed())
    });

    std::thread::sleep(Duration::from_millis(100));
    q.push(Job(1), Some(std::time::SystemTime::now() + VISIBLE_IN))
        .unwrap();

    assert!(
        !short.join().unwrap().unwrap(),
        "the short waiter times out before the row becomes visible"
    );
    let (due, elapsed) = long.join().unwrap();
    assert!(due, "the long waiter must observe the row");
    assert!(
        elapsed < Duration::from_secs(3),
        "the long waiter was never notified and slept out its own timeout: {elapsed:?}"
    );
}

#[test]
fn shutdown_ends_a_blocking_wait_instead_of_sleeping_out_its_timeout() {
    let (q, db, _dir) = open_wait_queue(Duration::from_secs(30));
    let waiter_q = Arc::clone(&q);
    let waiter = std::thread::spawn(move || {
        let started = std::time::Instant::now();
        (
            waiter_q.wait_due_timeout(Duration::from_secs(30)),
            started.elapsed(),
        )
    });
    std::thread::sleep(Duration::from_millis(50));
    db.shutdown_signal().shutdown();

    let (result, elapsed) = waiter.join().unwrap();
    assert!(result.is_err(), "a shut-down wait reports an error");
    assert!(
        elapsed < Duration::from_secs(5),
        "the waiter slept past shutdown: {elapsed:?}"
    );
}

#[test]
fn shutdown_racing_a_wake_worker_about_to_park_does_not_stall_the_join() {
    // The worker checks the shutdown flag at the top of its loop; the shutdown
    // listener notifies after releasing the state mutex. A worker paused
    // between those two points therefore misses the notification, and without
    // a re-check under the mutex it parks for its whole idle interval while
    // `Db::drop` blocks in `join`. The debug pause makes that window exact
    // instead of hoping a loop hits it.
    let (arrived, release) = armdb::queue_wake_test_support::pause_wake_worker_before_park();
    let (q, db, _dir) = open_wait_queue(Duration::from_secs(30));
    let counter = Arc::new(WakeCounter::default());
    let mut due = Box::pin(q.due());
    poll_pending(&mut due, &counter); // starts the worker

    arrived.wait(); // the worker is inside the window
    drop(due);
    db.shutdown_signal().shutdown(); // this notification finds nobody parked

    // Drop from another thread so the join is already in progress — and so its
    // own wake has also been issued and lost — before the worker is allowed to
    // reach the park.
    let joiner = std::thread::spawn(move || {
        let started = std::time::Instant::now();
        drop(q);
        drop(db); // `Drop for Db` runs shutdown, which joins the wake worker.
        started.elapsed()
    });
    std::thread::sleep(Duration::from_millis(200));
    release.wait();
    armdb::queue_wake_test_support::disarm_wake_worker_pause();

    let elapsed = joiner.join().unwrap();
    assert!(
        elapsed < Duration::from_secs(5),
        "dropping the database stalled on the wake worker: {elapsed:?}"
    );
}

fn poll_pending(due: &mut Pin<Box<armdb::QueueDue<'_, WaitJobs>>>, counter: &Arc<WakeCounter>) {
    let waker = counting_waker(Arc::clone(counter));
    let mut cx = Context::from_waker(&waker);
    assert!(matches!(due.as_mut().poll(&mut cx), Poll::Pending));
}

#[derive(Default)]
struct WakeCounter {
    wakes: AtomicUsize,
    observed: Mutex<usize>,
    changed: Condvar,
}

impl WakeCounter {
    fn wake(&self) {
        let wakes = self.wakes.fetch_add(1, Ordering::AcqRel) + 1;
        *self.observed.lock().unwrap() = wakes;
        self.changed.notify_all();
    }

    fn wait_for(&self, expected: usize) {
        let mut observed = self.observed.lock().unwrap();
        while *observed < expected {
            let (next, result) = self
                .changed
                .wait_timeout(observed, Duration::from_secs(1))
                .unwrap();
            assert!(!result.timed_out(), "waker was not called");
            observed = next;
        }
    }
}

fn noop_waker() -> Waker {
    unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &NOOP_WAKER_VTABLE)) }
}

fn counting_waker(counter: Arc<WakeCounter>) -> Waker {
    unsafe {
        Waker::from_raw(RawWaker::new(
            Arc::into_raw(counter).cast(),
            &COUNTING_WAKER_VTABLE,
        ))
    }
}

unsafe fn counting_clone(raw: *const ()) -> RawWaker {
    unsafe { Arc::increment_strong_count(raw.cast::<WakeCounter>()) };
    RawWaker::new(raw, &COUNTING_WAKER_VTABLE)
}

unsafe fn counting_wake(raw: *const ()) {
    let counter = unsafe { Arc::from_raw(raw.cast::<WakeCounter>()) };
    counter.wake();
}

unsafe fn counting_wake_by_ref(raw: *const ()) {
    let counter = unsafe { Arc::from_raw(raw.cast::<WakeCounter>()) };
    counter.wake();
    let _ = Arc::into_raw(counter);
}

unsafe fn counting_drop(raw: *const ()) {
    drop(unsafe { Arc::from_raw(raw.cast::<WakeCounter>()) });
}

static COUNTING_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
    counting_clone,
    counting_wake,
    counting_wake_by_ref,
    counting_drop,
);

unsafe fn noop_clone(_: *const ()) -> RawWaker {
    RawWaker::new(std::ptr::null(), &NOOP_WAKER_VTABLE)
}

unsafe fn noop(_: *const ()) {}

static NOOP_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(noop_clone, noop, noop, noop);