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.

pub(crate) mod bucket;
pub(crate) mod delay_queue;
pub(crate) mod entry;
pub(crate) mod state;

use std::sync::Arc;

use crate::timing_wheel::bucket::Bucket;
use crate::timing_wheel::delay_queue::{BucketOffer, DelayQueue};
use crate::timing_wheel::entry::Entry;

#[derive(Debug)]
pub(crate) struct TimingWheel {
    tick_ms: u64,
    bucket_count: usize,
    interval_ms: u64,
    current_time_ms: u64,
    buckets: Vec<Bucket>,
    delay_queue: DelayQueue,
    overflow: Option<Box<TimingWheel>>,
}

impl TimingWheel {
    pub(crate) fn new(tick_ms: u64, bucket_count: usize, start_ms: u64) -> Self {
        let current_time_ms = start_ms - (start_ms % tick_ms);
        let interval_ms = tick_ms.saturating_mul(bucket_count as u64);
        let buckets = (0..bucket_count).map(|_| Bucket::default()).collect();

        Self {
            tick_ms,
            bucket_count,
            interval_ms,
            current_time_ms,
            buckets,
            delay_queue: DelayQueue::default(),
            overflow: None,
        }
    }

    pub(crate) fn add(&mut self, entry: Arc<Entry>) -> bool {
        self.add_with_deadline(entry, None)
    }

    fn add_with_deadline(&mut self, entry: Arc<Entry>, deadline_override: Option<u64>) -> bool {
        let deadline_ms = deadline_override.unwrap_or_else(|| entry.deadline_ms());
        if deadline_ms < self.current_time_ms.saturating_add(self.tick_ms) {
            // Returning false tells the scheduler that the entry is already due
            // and should be dispatched without touching any bucket.
            return false;
        }

        let level_end_ms = self.current_time_ms.saturating_add(self.interval_ms);
        if deadline_ms < level_end_ms || level_end_ms == u64::MAX {
            // Bucket expiration is rounded down to the current level tick. A
            // stale offer may remain in the delay queue after the bucket moves.
            let virtual_id = deadline_ms / self.tick_ms;
            let bucket_index = (virtual_id as usize) % self.bucket_count;
            let expiration_ms = virtual_id * self.tick_ms;
            let bucket = &mut self.buckets[bucket_index];
            bucket.add(entry);
            if bucket.set_expiration_ms(expiration_ms) {
                self.delay_queue.offer(BucketOffer {
                    expiration_ms,
                    bucket_index,
                });
            }
            return true;
        }

        let overflow = self.overflow.get_or_insert_with(|| {
            // Overflow levels widen the tick to the lower level interval. They
            // keep far deadlines compact until they can cascade downward.
            Box::new(Self::new(
                self.interval_ms,
                self.bucket_count,
                self.current_time_ms,
            ))
        });
        overflow.add_with_deadline(entry, Some(deadline_ms))
    }

    pub(crate) fn advance_clock(&mut self, time_ms: u64) {
        if time_ms >= self.current_time_ms.saturating_add(self.tick_ms) {
            self.current_time_ms = time_ms - (time_ms % self.tick_ms);
            if let Some(overflow) = self.overflow.as_mut() {
                overflow.advance_clock(self.current_time_ms);
            }
        }
    }

    pub(crate) fn next_expiration_ms(&mut self) -> Option<u64> {
        while let Some(offer) = self.delay_queue.peek() {
            // Delay-queue offers are immutable. If a bucket was reused after an
            // older offer, skip the stale offer and keep looking.
            if self.buckets[offer.bucket_index].expiration_ms() == Some(offer.expiration_ms) {
                return Some(offer.expiration_ms);
            }
            self.delay_queue.poll();
        }
        None
    }

    pub(crate) fn poll_expired(&mut self, now_ms: u64) -> Vec<Arc<Entry>> {
        let mut expired = Vec::new();
        self.advance_clock(now_ms);
        if let Some(overflow) = self.overflow.as_mut() {
            // Higher levels are polled first so entries can cascade into this
            // level before the current level is drained.
            for entry in overflow.poll_expired(now_ms) {
                if self.add(Arc::clone(&entry)) {
                    continue;
                }
                expired.push(entry);
            }
        }
        while let Some(offer) = self.delay_queue.peek() {
            if offer.expiration_ms > now_ms {
                break;
            }

            self.delay_queue.poll();
            if self.buckets[offer.bucket_index].expiration_ms() != Some(offer.expiration_ms) {
                continue;
            }

            self.advance_clock(offer.expiration_ms);
            let entries = self.buckets[offer.bucket_index].drain();
            for entry in entries {
                // Re-adding after a drain performs cascade movement. If add
                // returns false, the entry is due at this level.
                if self.add(Arc::clone(&entry)) {
                    continue;
                }
                expired.push(entry);
            }
        }
        expired
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use crate::timing_wheel::entry::Entry;
    use crate::timing_wheel::state::TimeoutState;

    use super::TimingWheel;

    #[test]
    fn immediate_deadline_is_due_without_bucket() {
        let mut wheel = TimingWheel::new(10, 8, 0);
        let entry = Arc::new(Entry::new(1, 5, Arc::new(TimeoutState::scheduled())));

        assert!(!wheel.add(entry));
    }

    #[test]
    fn in_range_deadline_expires_from_bucket() {
        let mut wheel = TimingWheel::new(10, 8, 0);
        let entry = Arc::new(Entry::new(1, 20, Arc::new(TimeoutState::scheduled())));

        assert!(wheel.add(Arc::clone(&entry)));
        assert_eq!(wheel.next_expiration_ms(), Some(20));
        assert!(wheel.poll_expired(19).is_empty());
        assert_eq!(wheel.poll_expired(20)[0].id(), entry.id());
    }

    #[test]
    fn overflow_deadline_cascades_into_lower_level() {
        let mut wheel = TimingWheel::new(10, 8, 0);
        let entry = Arc::new(Entry::new(1, 120, Arc::new(TimeoutState::scheduled())));

        assert!(wheel.add(Arc::clone(&entry)));
        assert!(wheel.poll_expired(79).is_empty());
        assert!(wheel.poll_expired(80).is_empty());
        assert_eq!(wheel.poll_expired(120)[0].id(), entry.id());
    }

    #[test]
    fn stale_delay_queue_offer_is_skipped() {
        let mut wheel = TimingWheel::new(10, 8, 0);
        let old_entry = Arc::new(Entry::new(1, 20, Arc::new(TimeoutState::scheduled())));
        let new_entry = Arc::new(Entry::new(2, 100, Arc::new(TimeoutState::scheduled())));

        assert!(wheel.add(Arc::clone(&old_entry)));
        wheel.advance_clock(80);
        assert!(wheel.add(Arc::clone(&new_entry)));

        assert_eq!(wheel.next_expiration_ms(), Some(100));
    }

    #[test]
    fn very_far_deadline_cascades_across_multiple_levels() {
        let mut wheel = TimingWheel::new(10, 8, 0);
        let entry = Arc::new(Entry::new(1, 800, Arc::new(TimeoutState::scheduled())));

        assert!(wheel.add(Arc::clone(&entry)));
        assert!(wheel.poll_expired(79).is_empty());
        assert!(wheel.poll_expired(640).is_empty());
        assert_eq!(wheel.poll_expired(800)[0].id(), entry.id());
    }
}