spacetraders-client 0.1.0

mod priority;
mod state;

use std::sync::{Arc, Mutex};

use tokio::sync::Notify;
use tokio::time::{Duration, Instant};
use tracing::warn;

pub use priority::RequestPriority;
use priority::{DEFAULT_AGING_INTERVAL, DEFAULT_BURST, DEFAULT_RATE_PER_SEC};
pub use state::SchedulerPermit;
use state::{State, Waiter};

/// Priority-aware, rate-limited admission control for SpaceTraders API calls.
///
/// A token bucket bounds the request rate (sustained `rate_per_sec`, up to
/// `burst` back-to-back), while a priority queue decides who gets the next
/// token. Cloning shares one underlying queue, so a single scheduler can be
/// threaded through every [`crate::StClient`]. Aging guarantees low-priority
/// work is never starved by a steady stream of higher-priority requests.
#[derive(Clone)]
pub struct RequestScheduler {
    state: Arc<Mutex<State>>,
    notify: Arc<Notify>,
}

/// Removes a waiter from the queue if its `acquire` future is dropped before it
/// receives a permit, so a cancelled request never blocks the queue head.
struct WaiterGuard {
    state: Arc<Mutex<State>>,
    notify: Arc<Notify>,
    seq: u64,
    done: bool,
}

impl Drop for WaiterGuard {
    fn drop(&mut self) {
        if self.done {
            return;
        }
        let mut st = self.state.lock().unwrap();
        st.remove(self.seq);
        drop(st);
        self.notify.notify_waiters();
    }
}

impl RequestScheduler {
    /// Scheduler with the default burst capacity and aging at the given
    /// sustained rate. This is the production constructor.
    ///
    /// ```
    /// use spacetraders_client::RequestScheduler;
    ///
    /// // SpaceTraders enforces ~2 req/s.
    /// let scheduler = RequestScheduler::new(2);
    /// ```
    pub fn new(requests_per_second: u32) -> Self {
        Self::with_config(
            requests_per_second as f64,
            DEFAULT_BURST,
            DEFAULT_AGING_INTERVAL,
        )
    }

    /// Pure token bucket with no aging. Used in tests to exercise burst/priority
    /// behavior in isolation.
    pub fn with_burst(rate_per_sec: f64, burst: f64) -> Self {
        Self::with_config(rate_per_sec, burst, Duration::ZERO)
    }

    fn with_config(rate_per_sec: f64, burst: f64, aging_interval: Duration) -> Self {
        let rate_per_sec = if rate_per_sec > 0.0 {
            rate_per_sec
        } else {
            DEFAULT_RATE_PER_SEC
        };
        let capacity = burst.max(1.0);
        Self {
            state: Arc::new(Mutex::new(State {
                tokens: capacity,
                capacity,
                rate_per_sec,
                aging_interval,
                last_refill: Instant::now(),
                waiters: Vec::new(),
                next_seq: 0,
            })),
            notify: Arc::new(Notify::new()),
        }
    }

    /// Wait for a permit at the given priority before issuing a request.
    ///
    /// Resolves once this caller is at the head of the priority queue and a
    /// token is available; the token is consumed on return.
    ///
    /// ```
    /// use spacetraders_client::{RequestPriority, RequestScheduler};
    ///
    /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
    /// let scheduler = RequestScheduler::new(2);
    /// let _permit = scheduler.acquire(RequestPriority::Normal).await;
    /// // ... a token is now held; issue the request ...
    /// # });
    /// ```
    pub async fn acquire(&self, priority: RequestPriority) -> SchedulerPermit {
        let enqueued_at = Instant::now();
        let seq = {
            let mut st = self.state.lock().unwrap();
            let seq = st.next_seq;
            st.next_seq += 1;
            st.waiters.push(Waiter {
                priority,
                seq,
                enqueued_at,
            });
            seq
        };

        let mut logged_level = priority as u32;

        let mut guard = WaiterGuard {
            state: self.state.clone(),
            notify: self.notify.clone(),
            seq,
            done: false,
        };

        loop {
            let notified = self.notify.notified();
            tokio::pin!(notified);
            notified.as_mut().enable();

            let mut promotion = None;
            let wait = {
                let mut st = self.state.lock().unwrap();
                let now = Instant::now();
                st.refill(now);

                let effective = st.effective_priority(
                    &Waiter {
                        priority,
                        seq,
                        enqueued_at,
                    },
                    now,
                );
                if effective > logged_level {
                    logged_level = effective;
                    promotion = Some((
                        effective,
                        now.saturating_duration_since(enqueued_at),
                        st.waiters.len(),
                    ));
                }

                if st.best_seq(now) == Some(seq) {
                    if st.tokens >= 1.0 {
                        st.tokens -= 1.0;
                        st.remove(seq);
                        let permit = SchedulerPermit {
                            queued: now.saturating_duration_since(enqueued_at),
                            queue_depth: st.waiters.len(),
                            rate_budget_remaining: st.tokens,
                        };
                        guard.done = true;
                        drop(st);
                        self.notify.notify_waiters();
                        return permit;
                    }
                    Some(st.time_until_token())
                } else {
                    None
                }
            };

            if let Some((effective, waited, queue_depth)) = promotion {
                warn!(
                    seq,
                    base_priority = ?priority,
                    effective_priority = ?RequestPriority::from_level(effective),
                    waited_ms = waited.as_millis() as u64,
                    queue_depth,
                    "scheduler queue: request aged up to avoid starvation"
                );
            }

            match wait {
                Some(delay) => {
                    tokio::select! {
                        _ = tokio::time::sleep(delay) => {}
                        _ = &mut notified => {}
                    }
                }
                None => notified.await,
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicBool, Ordering};

    use super::*;

    #[tokio::test(start_paused = true)]
    async fn burst_then_settles_to_rate() {
        let scheduler = RequestScheduler::with_burst(2.0, 10.0);

        let start = Instant::now();
        for _ in 0..10 {
            scheduler.acquire(RequestPriority::Normal).await;
        }
        assert!(start.elapsed() < Duration::from_millis(10));

        let mid = Instant::now();
        for _ in 0..4 {
            scheduler.acquire(RequestPriority::Normal).await;
        }
        let took = mid.elapsed();
        assert!(
            (Duration::from_millis(1900)..=Duration::from_millis(2100)).contains(&took),
            "expected ~2s, got {took:?}"
        );
    }

    #[tokio::test(start_paused = true)]
    async fn critical_jumps_ahead_of_low() {
        let scheduler = RequestScheduler::with_burst(2.0, 1.0);
        scheduler.acquire(RequestPriority::Normal).await;

        let order = Arc::new(Mutex::new(Vec::<String>::new()));
        let mut handles = Vec::new();

        for i in 0..3 {
            let scheduler = scheduler.clone();
            let order = order.clone();
            handles.push(tokio::spawn(async move {
                scheduler.acquire(RequestPriority::Low).await;
                order.lock().unwrap().push(format!("low{i}"));
            }));
        }

        tokio::time::sleep(Duration::from_millis(1)).await;

        let crit = {
            let scheduler = scheduler.clone();
            let order = order.clone();
            tokio::spawn(async move {
                scheduler.acquire(RequestPriority::Critical).await;
                order.lock().unwrap().push("critical".to_string());
            })
        };

        crit.await.unwrap();
        for handle in handles {
            handle.await.unwrap();
        }

        let order = order.lock().unwrap();
        assert_eq!(order[0], "critical", "critical should run first: {order:?}");
    }

    #[tokio::test(start_paused = true)]
    async fn fifo_within_same_priority() {
        let scheduler = RequestScheduler::with_burst(2.0, 1.0);
        scheduler.acquire(RequestPriority::Normal).await;

        let order = Arc::new(Mutex::new(Vec::<u32>::new()));
        let mut handles = Vec::new();

        for i in 0..4 {
            let scheduler = scheduler.clone();
            let order = order.clone();
            handles.push(tokio::spawn(async move {
                scheduler.acquire(RequestPriority::Normal).await;
                order.lock().unwrap().push(i);
            }));
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        for handle in handles {
            handle.await.unwrap();
        }

        let order = order.lock().unwrap();
        assert_eq!(*order, vec![0, 1, 2, 3], "expected FIFO order");
    }

    #[tokio::test(start_paused = true)]
    async fn aging_prevents_low_starvation() {
        let scheduler = RequestScheduler::with_config(2.0, 1.0, Duration::from_millis(500));
        scheduler.acquire(RequestPriority::Normal).await;

        let low_done = Arc::new(AtomicBool::new(false));

        let flood = {
            let scheduler = scheduler.clone();
            let low_done = low_done.clone();
            tokio::spawn(async move {
                while !low_done.load(Ordering::SeqCst) {
                    scheduler.acquire(RequestPriority::Normal).await;
                }
            })
        };

        let start = Instant::now();
        scheduler.acquire(RequestPriority::Low).await;
        let waited = start.elapsed();

        low_done.store(true, Ordering::SeqCst);
        flood.abort();

        assert!(
            waited <= Duration::from_secs(3),
            "low priority was starved: waited {waited:?}"
        );
    }
}