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::collections::HashMap;
use std::fmt;
use std::panic::{self, AssertUnwindSafe};
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, Mutex, mpsc};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

use crate::error::{Error, Result};
use crate::executor::{self, BoxTask};
use crate::policy::{BackpressurePolicy, ExpiredTaskPolicy};
use crate::timer::builder::TimerBuilder;
use crate::timer::command::{Command, ScheduleCommand};
use crate::timer::metrics::{TimerMetrics, TimerMetricsInner};
use crate::timer::timeout::Timeout;
use crate::timing_wheel::TimingWheel;
use crate::timing_wheel::entry::Entry;
use crate::timing_wheel::state::TimeoutState;

/// Hierarchical timer-wheel scheduler.
pub struct Timer {
    sender: mpsc::SyncSender<Command>,
    scheduler: Mutex<Option<JoinHandle<()>>>,
    wake: Arc<(Mutex<bool>, Condvar)>,
    metrics: Arc<TimerMetricsInner>,
    executor: Arc<dyn executor::Executor>,
    owns_executor: bool,
    closed: AtomicBool,
    permits: Arc<AtomicUsize>,
    max_pending: usize,
    backpressure_policy: BackpressurePolicy,
    next_id: AtomicU64,
}

impl Timer {
    /// Returns a builder for a timer.
    pub fn builder() -> TimerBuilder {
        TimerBuilder::default()
    }

    pub(crate) fn new(builder: TimerBuilder) -> Result<Self> {
        let executor: Arc<dyn executor::Executor> = match builder.executor {
            Some(executor) => executor,
            None => Arc::new(executor::Pool::builder().build()?),
        };
        let (sender, receiver) = mpsc::sync_channel(builder.command_capacity);
        let metrics = Arc::new(TimerMetricsInner::new(Arc::clone(&builder.metric_sink)));
        let permits = Arc::new(AtomicUsize::new(0));
        let wake = Arc::new((Mutex::new(false), Condvar::new()));
        let scheduler_metrics = Arc::clone(&metrics);
        let scheduler_wake = Arc::clone(&wake);
        let scheduler_executor = Arc::clone(&executor);
        let scheduler_permits = Arc::clone(&permits);
        let tick_ms =
            duration_to_millis(builder.tick).ok_or(Error::InvalidConfig("tick is too large"))?;
        let retry_ms = duration_to_millis(builder.expired_task_retry)
            .ok_or(Error::InvalidConfig("expired_task_retry is too large"))?;
        let bucket_count = builder.bucket_count;
        let expired_task_policy = builder.expired_task_policy;

        let scheduler = thread::Builder::new()
            .name("timerwheel-scheduler".to_owned())
            .spawn(move || {
                let mut scheduler = Scheduler {
                    receiver,
                    wake: scheduler_wake,
                    metrics: scheduler_metrics,
                    executor: scheduler_executor,
                    permits: scheduler_permits,
                    wheel: TimingWheel::new(tick_ms, bucket_count, 0),
                    start: Instant::now(),
                    tick_ms,
                    retry_ms,
                    expired_task_policy,
                    shutting_down: false,
                    tasks: HashMap::new(),
                };
                scheduler.run();
            })
            .map_err(|_| Error::InvalidConfig("failed to start scheduler"))?;

        Ok(Self {
            sender,
            scheduler: Mutex::new(Some(scheduler)),
            wake,
            metrics,
            executor,
            owns_executor: builder.owns_executor,
            closed: AtomicBool::new(false),
            permits,
            max_pending: builder.max_pending,
            backpressure_policy: builder.backpressure_policy,
            next_id: AtomicU64::new(1),
        })
    }

    /// Schedules a closure to run after the given delay.
    pub fn schedule<F>(&self, delay: Duration, task: F) -> Result<Timeout>
    where
        F: FnOnce() + Send + 'static,
    {
        if self.closed.load(Ordering::Acquire) {
            return Err(Error::Closed);
        }
        let delay_ms =
            duration_to_millis(delay).ok_or(Error::InvalidConfig("delay is too large"))?;

        // Reserve scheduler capacity before allocating an entry. The permit is
        // released when the timeout is cancelled, rejected, or handed to the
        // executor; it is not held while user code runs on a worker thread.
        let mut current = self.permits.load(Ordering::Acquire);
        loop {
            if current >= self.max_pending {
                self.metrics.rejected_schedule();
                return Err(Error::Saturated);
            }
            match self.permits.compare_exchange(
                current,
                current + 1,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => break,
                Err(actual) => current = actual,
            }
        }

        let state = Arc::new(TimeoutState::scheduled());
        let timeout = Timeout::new(
            Arc::clone(&state),
            Arc::clone(&self.metrics),
            Arc::clone(&self.permits),
        );
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        let entry = Arc::new(Entry::new(id, delay_ms, state));
        let command = Command::Schedule(ScheduleCommand {
            delay_ms,
            entry,
            task: Box::new(task),
        });

        // Scheduling is a command submission. The scheduler thread remains the
        // only owner of the timing wheel and wakes after a successful enqueue.
        let result = match self.backpressure_policy {
            BackpressurePolicy::Reject => {
                self.sender.try_send(command).map_err(|error| match error {
                    mpsc::TrySendError::Full(_) => Error::Saturated,
                    mpsc::TrySendError::Disconnected(_) => Error::Closed,
                })
            }
            BackpressurePolicy::Block => self.sender.send(command).map_err(|_| Error::Closed),
        };

        match result {
            Ok(()) => {
                self.metrics.scheduled();
                self.metrics.command_queued();
                self.wake_scheduler();
                Ok(timeout)
            }
            Err(error) => {
                self.permits.fetch_sub(1, Ordering::AcqRel);
                self.metrics.rejected_schedule();
                Err(error)
            }
        }
    }

    /// Stops the scheduler and releases timer-owned resources.
    pub fn shutdown(&self) -> Result<()> {
        if !self.closed.swap(true, Ordering::AcqRel) {
            // Shutdown is command-driven so the scheduler drains already
            // accepted work and releases pending timeout state in one place.
            let _ = self.sender.send(Command::Shutdown);
            self.wake_scheduler();
            if let Ok(mut scheduler) = self.scheduler.lock() {
                if let Some(handle) = scheduler.take() {
                    let _ = handle.join();
                }
            }
            if self.owns_executor {
                self.executor.shutdown()?;
            }
        }
        Ok(())
    }

    /// Returns an immutable metrics snapshot.
    pub fn metrics(&self) -> TimerMetrics {
        self.metrics.snapshot()
    }

    /// Returns the number of pending timeouts currently owned by the timer.
    pub fn len(&self) -> usize {
        self.metrics.snapshot().pending_timeouts
    }

    /// Returns true when the timer has no pending timeouts.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn wake_scheduler(&self) {
        let (lock, condvar) = &*self.wake;
        if let Ok(mut signalled) = lock.lock() {
            *signalled = true;
            condvar.notify_one();
        }
    }
}

impl Drop for Timer {
    fn drop(&mut self) {
        let _ = self.shutdown();
    }
}

impl fmt::Debug for Timer {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Timer")
            .field("metrics", &self.metrics())
            .field("owns_executor", &self.owns_executor)
            .field("closed", &self.closed.load(Ordering::Acquire))
            .field("max_pending", &self.max_pending)
            .field("backpressure_policy", &self.backpressure_policy)
            .finish_non_exhaustive()
    }
}

struct Scheduler {
    receiver: mpsc::Receiver<Command>,
    wake: Arc<(Mutex<bool>, Condvar)>,
    metrics: Arc<TimerMetricsInner>,
    executor: Arc<dyn executor::Executor>,
    permits: Arc<AtomicUsize>,
    wheel: TimingWheel,
    start: Instant,
    tick_ms: u64,
    retry_ms: u64,
    expired_task_policy: ExpiredTaskPolicy,
    shutting_down: bool,
    tasks: HashMap<u64, PendingTask>,
}

impl Scheduler {
    fn run(&mut self) {
        // The loop always drains commands before polling buckets so newly
        // scheduled earlier deadlines can affect the next sleep duration.
        while !self.shutting_down {
            self.drain_commands();
            self.dispatch_due();
            self.wait();
        }
        self.drain_commands();
        self.cancel_pending_tasks();
    }

    fn drain_commands(&mut self) {
        loop {
            match self.receiver.try_recv() {
                Ok(Command::Schedule(command)) => {
                    self.metrics.command_drained();
                    self.add_schedule(command);
                }
                Ok(Command::Shutdown) => {
                    self.shutting_down = true;
                    return;
                }
                Err(mpsc::TryRecvError::Empty) => return,
                Err(mpsc::TryRecvError::Disconnected) => {
                    self.shutting_down = true;
                    return;
                }
            }
        }
    }

    fn add_schedule(&mut self, command: ScheduleCommand) {
        let deadline_ms = self.now_ms().saturating_add(command.delay_ms);
        command.entry.set_deadline_ms(deadline_ms);
        if self.wheel.add(Arc::clone(&command.entry)) {
            // Tasks are stored outside the wheel. Buckets only hold entry
            // handles, which keeps cancellation and cascade movement cheap.
            self.metrics.bucket_offered();
            self.tasks.insert(
                command.entry.id(),
                PendingTask::new(command.entry, command.task),
            );
            return;
        }
        self.dispatch(command.entry, command.task);
    }

    fn dispatch_due(&mut self) {
        let now_ms = self.now_ms();
        let entries = self.wheel.poll_expired(now_ms);
        for entry in entries {
            self.metrics.bucket_expired();
            if let Some(task) = self.tasks.remove(&entry.id()) {
                self.dispatch_pending(task);
            }
        }
    }

    fn dispatch_pending(&mut self, pending: PendingTask) {
        match pending.task {
            PendingTaskKind::Timer(task) => self.dispatch(pending.entry, task),
            PendingTaskKind::Executor(task) => self.dispatch_executor_task(pending.entry, task),
        }
    }

    fn dispatch(&mut self, entry: Arc<Entry>, task: BoxTask) {
        let state = Arc::clone(entry.state());
        let metrics = Arc::clone(&self.metrics);
        // Timer task panics are caught before the executor worker boundary sees
        // them. That lets timer metrics record the panic while the worker still
        // remains alive for later tasks.
        let wrapped = move || {
            let result = panic::catch_unwind(AssertUnwindSafe(|| task.run()));
            match result {
                Ok(()) => {
                    let _ = state.expire();
                }
                Err(_) => {
                    metrics.panicked();
                    let _ = state.expire();
                }
            }
        };

        self.dispatch_executor_task(entry, Box::new(wrapped));
    }

    fn dispatch_executor_task(&mut self, entry: Arc<Entry>, task: BoxTask) {
        if !entry.state().dispatch() {
            return;
        }

        match self.executor.try_execute(task) {
            Ok(()) => {
                // A timeout leaves the timer once the executor accepts it. Long
                // user work should not consume timer pending capacity.
                self.metrics.expired();
                self.permits.fetch_sub(1, Ordering::AcqRel);
            }
            Err(rejected)
                if rejected.error() == &Error::Saturated
                    && self.expired_task_policy == ExpiredTaskPolicy::Retry =>
            {
                let (_, task) = rejected.into_parts();
                if entry.state().reschedule() {
                    // The executor returned the boxed task on saturation, so a
                    // retry keeps the exact same task and does not lose work.
                    entry.set_deadline_ms(self.now_ms().saturating_add(self.retry_ms));
                    let id = entry.id();
                    let pending = PendingTask::new_executor(entry, task);
                    let _ = self.wheel.add(Arc::clone(&pending.entry));
                    self.tasks.insert(id, pending);
                    self.metrics.rejected_dispatch();
                }
            }
            Err(rejected) => {
                drop(rejected.into_task());
                if entry.state().reject() {
                    self.metrics.dropped_after_dispatch_reject();
                    self.permits.fetch_sub(1, Ordering::AcqRel);
                }
            }
        }
    }

    fn cancel_pending_tasks(&mut self) {
        // On shutdown, entries still owned by the scheduler are cancelled
        // logically and their permits are released without executing user code.
        for (_, pending) in self.tasks.drain() {
            if pending.entry.state().cancel() {
                self.metrics.cancelled_pending();
                self.permits.fetch_sub(1, Ordering::AcqRel);
            }
        }
    }

    fn wait(&mut self) {
        let now_ms = self.now_ms();
        let wait_duration = self
            .wheel
            .next_expiration_ms()
            .map(|expiration_ms| expiration_ms.saturating_sub(now_ms))
            .unwrap_or(self.tick_ms)
            .max(1);
        // The condition variable handles both new command wakeups and timed
        // sleeps. The scheduler never spins while waiting for the next bucket.
        let (lock, condvar) = &*self.wake;
        let Ok(mut signalled) = lock.lock() else {
            return;
        };
        if !*signalled {
            let timeout = Duration::from_millis(wait_duration);
            match condvar.wait_timeout(signalled, timeout) {
                Ok((guard, _)) => signalled = guard,
                Err(_) => return,
            }
        }
        *signalled = false;
    }

    fn now_ms(&self) -> u64 {
        duration_to_millis(self.start.elapsed()).unwrap_or(u64::MAX)
    }
}

struct PendingTask {
    entry: Arc<Entry>,
    task: PendingTaskKind,
}

enum PendingTaskKind {
    Timer(BoxTask),
    Executor(BoxTask),
}

impl PendingTask {
    fn new(entry: Arc<Entry>, task: BoxTask) -> Self {
        Self {
            entry,
            task: PendingTaskKind::Timer(task),
        }
    }

    fn new_executor(entry: Arc<Entry>, task: BoxTask) -> Self {
        Self {
            entry,
            task: PendingTaskKind::Executor(task),
        }
    }
}

fn duration_to_millis(duration: Duration) -> Option<u64> {
    let millis = duration.as_millis();
    if millis > u128::from(u64::MAX) {
        None
    } else {
        Some(millis as u64)
    }
}