timerwheel 0.1.0

// Copyright © 2026-present The Timerwheel Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::{
    Arc, Barrier, Mutex,
    atomic::{AtomicUsize, Ordering},
    mpsc,
};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

use timerwheel::executor::{BoxTask, Executor, ExecutorMetrics, Pool, RejectedTask};
use timerwheel::{BackpressurePolicy, Error, ExpiredTaskPolicy, Timer};

struct CountingTimerMetricSink {
    observations: Arc<AtomicUsize>,
}

impl timerwheel::timer::MetricSink for CountingTimerMetricSink {
    fn observe_metrics(&self, metrics: timerwheel::TimerMetrics) {
        if metrics.scheduled_timeouts > 0 {
            self.observations.fetch_add(1, Ordering::SeqCst);
        }
    }
}

struct RejectingExecutor {
    calls: AtomicUsize,
}

impl RejectingExecutor {
    fn new() -> Arc<Self> {
        Arc::new(Self {
            calls: AtomicUsize::new(0),
        })
    }
}

impl Executor for RejectingExecutor {
    fn try_execute(&self, task: BoxTask) -> std::result::Result<(), RejectedTask> {
        self.calls.fetch_add(1, Ordering::SeqCst);
        Err(RejectedTask::new(Error::Saturated, task))
    }

    fn shutdown(&self) -> timerwheel::Result<()> {
        Ok(())
    }

    fn metrics(&self) -> ExecutorMetrics {
        ExecutorMetrics::default()
    }
}

struct ScriptedExecutor {
    rejections: AtomicUsize,
    accepted: AtomicUsize,
    sender: Mutex<Option<mpsc::Sender<BoxTask>>>,
    worker: Mutex<Option<JoinHandle<()>>>,
}

impl ScriptedExecutor {
    fn new(rejections: usize) -> Arc<Self> {
        let (sender, receiver) = mpsc::channel::<BoxTask>();
        let worker = thread::spawn(move || {
            while let Ok(task) = receiver.recv() {
                task.run();
            }
        });

        Arc::new(Self {
            rejections: AtomicUsize::new(rejections),
            accepted: AtomicUsize::new(0),
            sender: Mutex::new(Some(sender)),
            worker: Mutex::new(Some(worker)),
        })
    }
}

impl Executor for ScriptedExecutor {
    fn try_execute(&self, task: BoxTask) -> std::result::Result<(), RejectedTask> {
        let mut remaining = self.rejections.load(Ordering::SeqCst);
        while remaining > 0 {
            match self.rejections.compare_exchange(
                remaining,
                remaining - 1,
                Ordering::SeqCst,
                Ordering::SeqCst,
            ) {
                Ok(_) => return Err(RejectedTask::new(Error::Saturated, task)),
                Err(actual) => remaining = actual,
            }
        }

        let Some(sender) = self.sender.lock().expect("sender lock").as_ref().cloned() else {
            return Err(RejectedTask::new(Error::Closed, task));
        };
        sender
            .send(task)
            .map_err(|error| RejectedTask::new(Error::Closed, error.0))?;
        self.accepted.fetch_add(1, Ordering::SeqCst);
        Ok(())
    }

    fn shutdown(&self) -> timerwheel::Result<()> {
        self.sender.lock().expect("sender lock").take();
        if let Some(worker) = self.worker.lock().expect("worker lock").take() {
            let _ = worker.join();
        }
        Ok(())
    }

    fn metrics(&self) -> ExecutorMetrics {
        ExecutorMetrics::default()
    }
}

#[test]
fn timer_executes_delayed_task() {
    let timer = Timer::builder()
        .tick(Duration::from_millis(1))
        .bucket_count(64)
        .build()
        .expect("timer builds");
    let (tx, rx) = mpsc::channel();

    timer
        .schedule(Duration::from_millis(10), move || {
            tx.send("fired").expect("send succeeds");
        })
        .expect("schedule succeeds");

    assert_eq!(
        rx.recv_timeout(Duration::from_secs(1))
            .expect("task should fire"),
        "fired"
    );
    assert!(timer.shutdown().is_ok());
}

#[test]
fn cancelled_timeout_does_not_run_task() {
    let timer = Timer::builder()
        .tick(Duration::from_millis(1))
        .bucket_count(64)
        .build()
        .expect("timer builds");
    let (tx, rx) = mpsc::channel();

    let timeout = timer
        .schedule(Duration::from_millis(40), move || {
            tx.send("fired").expect("send succeeds");
        })
        .expect("schedule succeeds");

    assert!(timeout.cancel());
    assert!(timeout.is_cancelled());
    assert!(rx.recv_timeout(Duration::from_millis(120)).is_err());
    assert!(timer.shutdown().is_ok());
}

#[test]
fn timer_rejects_when_max_pending_is_reached() {
    let timer = Timer::builder()
        .tick(Duration::from_millis(1))
        .bucket_count(64)
        .max_pending(1)
        .backpressure_policy(BackpressurePolicy::Reject)
        .build()
        .expect("timer builds");

    let first = timer
        .schedule(Duration::from_secs(1), || {})
        .expect("first schedule succeeds");
    let second = timer.schedule(Duration::from_secs(1), || {});

    assert_eq!(
        second.expect_err("second schedule rejects"),
        Error::Saturated
    );
    assert!(first.cancel());
    assert!(timer.shutdown().is_ok());
}

#[test]
fn timer_rejects_delay_that_is_too_large() {
    let timer = Timer::builder().build().expect("timer builds");

    let result = timer.schedule(Duration::MAX, || {});

    assert_eq!(
        result.expect_err("too large delay should reject"),
        Error::InvalidConfig("delay is too large")
    );
    assert_eq!(timer.len(), 0);
    assert!(timer.shutdown().is_ok());
}

#[test]
fn timer_releases_capacity_after_too_large_delay_rejection() {
    let timer = Timer::builder()
        .max_pending(1)
        .build()
        .expect("timer builds");

    let result = timer.schedule(Duration::MAX, || {});

    assert_eq!(
        result.expect_err("too large delay should reject"),
        Error::InvalidConfig("delay is too large")
    );
    timer
        .schedule(Duration::from_millis(1), || {})
        .expect("capacity should be reusable");
    assert!(timer.shutdown().is_ok());
}

#[test]
fn timer_shutdown_is_idempotent_and_closes_schedule_path() {
    let timer = Timer::builder().build().expect("timer builds");

    assert!(timer.shutdown().is_ok());
    assert!(timer.shutdown().is_ok());
    assert_eq!(
        timer
            .schedule(Duration::from_millis(1), || {})
            .expect_err("schedule after shutdown rejects"),
        Error::Closed
    );
}

#[test]
fn timer_updates_metrics_for_expire_and_cancel() {
    let timer = Timer::builder()
        .tick(Duration::from_millis(1))
        .bucket_count(64)
        .build()
        .expect("timer builds");
    let (tx, rx) = mpsc::channel();

    let cancelled = timer
        .schedule(Duration::from_millis(50), || {})
        .expect("schedule succeeds");
    assert!(cancelled.cancel());

    timer
        .schedule(Duration::from_millis(5), move || {
            tx.send(()).expect("send succeeds");
        })
        .expect("schedule succeeds");

    rx.recv_timeout(Duration::from_secs(1))
        .expect("task should fire");
    let metrics = timer.metrics();
    assert_eq!(metrics.scheduled_timeouts, 2);
    assert_eq!(metrics.cancelled_timeouts, 1);
    assert_eq!(metrics.expired_timeouts, 1);
    assert_eq!(timer.len(), 0);
    assert!(timer.is_empty());
    assert!(timer.shutdown().is_ok());
}

#[test]
fn timer_reports_metrics_to_metric_sink() {
    let observations = Arc::new(AtomicUsize::new(0));
    let timer = Timer::builder()
        .tick(Duration::from_millis(1))
        .bucket_count(64)
        .metric_sink(CountingTimerMetricSink {
            observations: Arc::clone(&observations),
        })
        .build()
        .expect("timer builds");
    let (tx, rx) = mpsc::channel();

    timer
        .schedule(Duration::from_millis(5), move || {
            tx.send(()).expect("send succeeds");
        })
        .expect("schedule succeeds");

    rx.recv_timeout(Duration::from_secs(1))
        .expect("task should run");
    assert!(timer.shutdown().is_ok());
    assert!(observations.load(Ordering::SeqCst) > 0);
}

#[test]
fn timer_cascades_far_deadline() {
    let timer = Timer::builder()
        .tick(Duration::from_millis(5))
        .bucket_count(8)
        .build()
        .expect("timer builds");
    let (tx, rx) = mpsc::channel();

    timer
        .schedule(Duration::from_millis(90), move || {
            tx.send("far").expect("send succeeds");
        })
        .expect("schedule succeeds");

    assert_eq!(
        rx.recv_timeout(Duration::from_secs(1))
            .expect("far task should fire"),
        "far"
    );
    assert!(timer.shutdown().is_ok());
}

#[test]
fn timer_releases_pending_timeout_on_task_panic() {
    let pool = Arc::new(
        Pool::builder()
            .workers(1)
            .queue_capacity(8)
            .build()
            .expect("pool builds"),
    );
    let timer = Timer::builder()
        .executor(pool.clone())
        .build()
        .expect("timer builds");
    let (tx, rx) = mpsc::channel();

    timer
        .schedule(Duration::from_millis(1), || panic!("expected panic"))
        .expect("panic task schedule succeeds");
    timer
        .schedule(Duration::from_millis(5), move || {
            tx.send("after panic").expect("send succeeds");
        })
        .expect("second schedule succeeds");

    assert_eq!(
        rx.recv_timeout(Duration::from_secs(1))
            .expect("worker should continue"),
        "after panic"
    );
    thread::sleep(Duration::from_millis(20));

    let metrics = timer.metrics();
    assert_eq!(metrics.panicked_timeouts, 1);
    assert_eq!(metrics.pending_timeouts, 0);
    assert_eq!(timer.len(), 0);
    assert!(timer.shutdown().is_ok());
    assert!(pool.shutdown().is_ok());
}

#[test]
fn timer_shutdown_releases_pending_timeouts() {
    let timer = Timer::builder().build().expect("timer builds");

    timer
        .schedule(Duration::from_secs(10), || {})
        .expect("schedule succeeds");
    assert_eq!(timer.len(), 1);

    assert!(timer.shutdown().is_ok());
    assert_eq!(timer.len(), 0);
    assert_eq!(timer.metrics().pending_timeouts, 0);
}

#[test]
fn timer_retry_dispatch_keeps_task_until_executor_accepts_it() {
    let pool = Arc::new(
        Pool::builder()
            .workers(1)
            .queue_capacity(1)
            .build()
            .expect("pool builds"),
    );
    let started = Arc::new(Barrier::new(2));
    let release = Arc::new(Barrier::new(2));
    let started_task = Arc::clone(&started);
    let release_task = Arc::clone(&release);
    pool.execute(move || {
        started_task.wait();
        release_task.wait();
    })
    .expect("blocking task accepted");
    started.wait();
    pool.try_execute(|| {}).expect("queued task accepted");

    let timer = Timer::builder()
        .executor(pool.clone())
        .expired_task_policy(ExpiredTaskPolicy::Retry)
        .expired_task_retry(Duration::from_millis(5))
        .build()
        .expect("timer builds");
    let (tx, rx) = mpsc::channel();

    timer
        .schedule(Duration::from_millis(1), move || {
            tx.send("retried").expect("send succeeds");
        })
        .expect("schedule succeeds");
    thread::sleep(Duration::from_millis(30));
    assert!(rx.try_recv().is_err());

    release.wait();
    assert_eq!(
        rx.recv_timeout(Duration::from_secs(1))
            .expect("retried task should run"),
        "retried"
    );
    assert!(timer.shutdown().is_ok());
    assert!(pool.shutdown().is_ok());
}

#[test]
fn scheduler_retries_rejected_dispatch_without_losing_task() {
    let executor = RejectingExecutor::new();
    let timer_executor: Arc<dyn Executor> = executor.clone();
    let timer = Timer::builder()
        .tick(Duration::from_millis(1))
        .expired_task_policy(ExpiredTaskPolicy::Retry)
        .expired_task_retry(Duration::from_millis(1))
        .executor(timer_executor)
        .build()
        .expect("timer builds");

    timer
        .schedule(Duration::from_millis(1), || {})
        .expect("schedule succeeds");
    thread::sleep(Duration::from_millis(20));

    assert!(executor.calls.load(Ordering::SeqCst) > 1);
    timer.shutdown().expect("timer shuts down");
}

#[test]
fn scheduler_runs_retried_task_once_after_acceptance() {
    let executor = ScriptedExecutor::new(3);
    let timer_executor: Arc<dyn Executor> = executor.clone();
    let timer = Timer::builder()
        .tick(Duration::from_millis(1))
        .expired_task_policy(ExpiredTaskPolicy::Retry)
        .expired_task_retry(Duration::from_millis(1))
        .executor(timer_executor)
        .build()
        .expect("timer builds");
    let runs = Arc::new(AtomicUsize::new(0));
    let runs_for_task = Arc::clone(&runs);
    let (tx, rx) = mpsc::channel();

    timer
        .schedule(Duration::from_millis(1), move || {
            runs_for_task.fetch_add(1, Ordering::SeqCst);
            tx.send(()).expect("send succeeds");
        })
        .expect("schedule succeeds");

    rx.recv_timeout(Duration::from_secs(1))
        .expect("retried task should run");
    thread::sleep(Duration::from_millis(20));

    assert_eq!(runs.load(Ordering::SeqCst), 1);
    assert_eq!(executor.accepted.load(Ordering::SeqCst), 1);
    timer.shutdown().expect("timer shuts down");
    executor.shutdown().expect("executor shuts down");
}

#[test]
fn scheduler_keeps_draining_commands_while_executor_rejects() {
    let executor = RejectingExecutor::new();
    let timer_executor: Arc<dyn Executor> = executor.clone();
    let timer = Timer::builder()
        .tick(Duration::from_millis(1))
        .bucket_count(64)
        .expired_task_policy(ExpiredTaskPolicy::Retry)
        .expired_task_retry(Duration::from_millis(1))
        .executor(timer_executor)
        .build()
        .expect("timer builds");

    let _first = timer
        .schedule(Duration::from_millis(1), || {})
        .expect("first schedule succeeds");
    let second = timer
        .schedule(Duration::from_millis(20), || {})
        .expect("second schedule succeeds");

    let drain_deadline = Instant::now() + Duration::from_secs(1);
    while timer.metrics().command_queue_depth != 0 && Instant::now() < drain_deadline {
        thread::sleep(Duration::from_millis(1));
    }

    assert_eq!(timer.metrics().command_queue_depth, 0);
    let retry_deadline = Instant::now() + Duration::from_secs(1);
    while executor.calls.load(Ordering::SeqCst) <= 1 && Instant::now() < retry_deadline {
        thread::sleep(Duration::from_millis(1));
    }

    assert!(second.cancel());
    assert!(timer.shutdown().is_ok());
    assert!(executor.calls.load(Ordering::SeqCst) > 1);
}

#[test]
fn long_running_worker_task_does_not_block_schedule_and_cancel() {
    let pool = Arc::new(
        Pool::builder()
            .workers(1)
            .queue_capacity(1)
            .build()
            .expect("pool builds"),
    );
    let timer_pool: Arc<dyn Executor> = pool.clone();
    let timer = Timer::builder()
        .tick(Duration::from_millis(1))
        .executor(timer_pool)
        .build()
        .expect("timer builds");
    let (release_tx, release_rx) = mpsc::channel();
    let (started_tx, started_rx) = mpsc::channel();

    timer
        .schedule(Duration::from_millis(1), move || {
            started_tx.send(()).expect("send succeeds");
            release_rx.recv().expect("release received");
        })
        .expect("first schedule succeeds");
    timer
        .schedule(Duration::from_millis(1), || {})
        .expect("second schedule succeeds");

    started_rx
        .recv_timeout(Duration::from_secs(1))
        .expect("first task starts");
    let queue_deadline = Instant::now() + Duration::from_secs(1);
    while pool.metrics().queue_depth == 0 && Instant::now() < queue_deadline {
        thread::sleep(Duration::from_millis(1));
    }
    assert!(pool.metrics().queue_depth >= 1);

    let start = Instant::now();
    let timeout = timer
        .schedule(Duration::from_millis(50), || {})
        .expect("late schedule succeeds");
    assert!(timeout.cancel());
    assert!(start.elapsed() < Duration::from_millis(20));

    release_tx.send(()).expect("release succeeds");
    timer.shutdown().expect("timer shuts down");
    pool.shutdown().expect("pool shuts down");
}

#[test]
fn timer_releases_pending_capacity_after_dispatch_not_task_completion() {
    let timer = Timer::builder()
        .max_pending(1)
        .build()
        .expect("timer builds");
    let (started_tx, started_rx) = mpsc::channel();
    let (release_tx, release_rx) = mpsc::channel();

    timer
        .schedule(Duration::from_millis(1), move || {
            started_tx.send(()).expect("send succeeds");
            release_rx.recv().expect("release received");
        })
        .expect("first schedule succeeds");
    started_rx
        .recv_timeout(Duration::from_secs(1))
        .expect("first task starts");

    assert_eq!(timer.len(), 0);
    let schedule_result = timer.schedule(Duration::from_millis(1), || {}).map(|_| ());

    release_tx.send(()).expect("release succeeds");
    schedule_result.expect("capacity should be released after dispatch");
    assert!(timer.shutdown().is_ok());
}