saddle_runtime/

request.rs

use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};

use saddle_core::{ErrorKind, Result, SaddleError};
use tokio::sync::Notify;

/// The externally visible phase of the Saddle application lifecycle.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ApplicationPhase {
    Starting = 0,
    Ready = 1,
    Draining = 2,
    Stopped = 3,
}

const PHASE_SHIFT: u32 = usize::BITS - 2;
const COUNT_MASK: usize = (1 << PHASE_SHIFT) - 1;

#[derive(Debug)]
struct Shared {
    /// The top two bits hold `ApplicationPhase`; the remaining bits hold the
    /// number of admitted requests. Updating both in one CAS closes the race
    /// between request admission and the transition to draining.
    state: AtomicUsize,
    drained: Notify,
}

/// Coordinates request admission with graceful application shutdown.
///
/// Service adapters must acquire a [`RequestGuard`] before dispatching an
/// accepted request. Once shutdown begins, new acquisitions fail while guards
/// already issued remain valid until the corresponding request completes.
#[derive(Clone, Debug)]
pub struct RequestLifecycle {
    shared: Arc<Shared>,
}

impl RequestLifecycle {
    pub(crate) fn new() -> Self {
        Self {
            shared: Arc::new(Shared {
                state: AtomicUsize::new(encode(ApplicationPhase::Starting, 0)),
                drained: Notify::new(),
            }),
        }
    }

    /// Returns the application's current lifecycle phase.
    pub fn phase(&self) -> ApplicationPhase {
        phase(self.shared.state.load(Ordering::Acquire))
    }

    /// Admits one request if the application is ready and accepting work.
    ///
    /// The returned guard must be held for the complete request execution. Its
    /// `Drop` implementation records completion even when the request future is
    /// cancelled or unwinds.
    pub fn try_accept(&self) -> Result<RequestGuard> {
        let mut current = self.shared.state.load(Ordering::Acquire);
        loop {
            match phase(current) {
                ApplicationPhase::Ready => {
                    if count(current) == COUNT_MASK {
                        return Err(SaddleError::new(
                            ErrorKind::Internal,
                            "runtime.request_count_overflow",
                            "request accounting capacity exhausted",
                        ));
                    }
                    match self.shared.state.compare_exchange_weak(
                        current,
                        current + 1,
                        Ordering::AcqRel,
                        Ordering::Acquire,
                    ) {
                        Ok(_) => {
                            return Ok(RequestGuard {
                                shared: Arc::clone(&self.shared),
                            });
                        }
                        Err(observed) => current = observed,
                    }
                }
                ApplicationPhase::Starting => {
                    return Err(SaddleError::new(
                        ErrorKind::Unavailable,
                        "runtime.not_ready",
                        "application is not ready",
                    ));
                }
                ApplicationPhase::Draining => {
                    return Err(SaddleError::new(
                        ErrorKind::Unavailable,
                        "runtime.shutting_down",
                        "application is shutting down",
                    ));
                }
                ApplicationPhase::Stopped => {
                    return Err(SaddleError::new(
                        ErrorKind::Unavailable,
                        "runtime.stopped",
                        "application is stopped",
                    ));
                }
            }
        }
    }

    pub(crate) fn mark_ready(&self) {
        let result = self.shared.state.compare_exchange(
            encode(ApplicationPhase::Starting, 0),
            encode(ApplicationPhase::Ready, 0),
            Ordering::AcqRel,
            Ordering::Acquire,
        );
        debug_assert!(result.is_ok());
    }

    pub(crate) fn begin_draining(&self) {
        let mut current = self.shared.state.load(Ordering::Acquire);
        while matches!(
            phase(current),
            ApplicationPhase::Starting | ApplicationPhase::Ready
        ) {
            let draining = encode(ApplicationPhase::Draining, count(current));
            match self.shared.state.compare_exchange_weak(
                current,
                draining,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => {
                    current = draining;
                    break;
                }
                Err(observed) => current = observed,
            }
        }
        if count(current) == 0 {
            self.shared.drained.notify_one();
        }
    }

    pub(crate) async fn wait_until_drained(&self) {
        loop {
            if count(self.shared.state.load(Ordering::Acquire)) == 0 {
                return;
            }
            self.shared.drained.notified().await;
        }
    }

    pub(crate) fn mark_stopped(&self) {
        let result = self.shared.state.compare_exchange(
            encode(ApplicationPhase::Draining, 0),
            encode(ApplicationPhase::Stopped, 0),
            Ordering::AcqRel,
            Ordering::Acquire,
        );
        debug_assert!(result.is_ok());
    }
}

const fn encode(phase: ApplicationPhase, count: usize) -> usize {
    ((phase as usize) << PHASE_SHIFT) | count
}

const fn phase(state: usize) -> ApplicationPhase {
    match state >> PHASE_SHIFT {
        0 => ApplicationPhase::Starting,
        1 => ApplicationPhase::Ready,
        2 => ApplicationPhase::Draining,
        3 => ApplicationPhase::Stopped,
        _ => unreachable!(),
    }
}

const fn count(state: usize) -> usize {
    state & COUNT_MASK
}

/// Proof that one request was admitted by Saddle before shutdown began.
#[derive(Debug)]
pub struct RequestGuard {
    shared: Arc<Shared>,
}

impl Drop for RequestGuard {
    fn drop(&mut self) {
        let previous = self.shared.state.fetch_sub(1, Ordering::AcqRel);
        debug_assert!(count(previous) > 0);
        if count(previous) == 1 && phase(previous) == ApplicationPhase::Draining {
            self.shared.drained.notify_one();
        }
    }
}

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

    use super::*;

    fn test_runtime() -> tokio::runtime::Runtime {
        tokio::runtime::Builder::new_current_thread()
            .build()
            .expect("test runtime must build")
    }

    #[test]
    fn only_ready_applications_accept_requests() {
        let requests = RequestLifecycle::new();
        assert_eq!(
            requests.try_accept().unwrap_err().code(),
            "runtime.not_ready"
        );

        requests.mark_ready();
        let request = requests.try_accept().expect("ready request is admitted");
        requests.begin_draining();

        assert_eq!(requests.phase(), ApplicationPhase::Draining);
        assert_eq!(
            requests.try_accept().unwrap_err().code(),
            "runtime.shutting_down"
        );

        drop(request);
        test_runtime().block_on(requests.wait_until_drained());
        requests.mark_stopped();
        assert_eq!(requests.phase(), ApplicationPhase::Stopped);
    }

    #[test]
    fn draining_waits_for_every_admitted_request() {
        test_runtime().block_on(async {
            let requests = RequestLifecycle::new();
            requests.mark_ready();
            let first = requests.try_accept().unwrap();
            let second = requests.try_accept().unwrap();
            requests.begin_draining();

            let requests_for_waiter = requests.clone();
            let waiter = tokio::spawn(async move {
                requests_for_waiter.wait_until_drained().await;
            });
            tokio::task::yield_now().await;
            assert!(!waiter.is_finished());

            drop(first);
            tokio::task::yield_now().await;
            assert!(!waiter.is_finished());

            drop(second);
            waiter.await.unwrap();
        });
    }

    #[test]
    fn admission_racing_with_drain_never_leaks_a_request() {
        const WORKERS: usize = 8;
        let requests = RequestLifecycle::new();
        requests.mark_ready();
        let barrier = Arc::new(Barrier::new(WORKERS + 1));
        let workers: Vec<_> = (0..WORKERS)
            .map(|_| {
                let requests = requests.clone();
                let barrier = Arc::clone(&barrier);
                std::thread::spawn(move || {
                    let admitted_before_drain = requests.try_accept().unwrap();
                    barrier.wait();
                    loop {
                        match requests.try_accept() {
                            Ok(request) => drop(request),
                            Err(error) => {
                                assert_eq!(error.code(), "runtime.shutting_down");
                                drop(admitted_before_drain);
                                break;
                            }
                        }
                    }
                })
            })
            .collect();

        barrier.wait();
        requests.begin_draining();
        for worker in workers {
            worker.join().unwrap();
        }
        test_runtime().block_on(requests.wait_until_drained());
        assert_eq!(count(requests.shared.state.load(Ordering::Acquire)), 0);
    }

    #[test]
    fn guard_completion_is_safe_from_another_thread() {
        let requests = RequestLifecycle::new();
        requests.mark_ready();
        let guard = requests.try_accept().unwrap();
        requests.begin_draining();

        std::thread::spawn(move || drop(guard)).join().unwrap();
        test_runtime().block_on(requests.wait_until_drained());

        // Keep this assertion explicit: the thread above is runtime-internal
        // test coverage, not a business-facing execution API.
        assert_eq!(Arc::strong_count(&requests.shared), 1);
    }
}