salvo_core 0.94.0

use std::fmt::{self, Debug, Formatter};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};

use tokio::sync::Notify;

const RUNNING: u8 = 0;
const GRACEFUL_SHUTDOWN: u8 = 1;
const ABORT: u8 = 2;

/// Controls the lifetime of the transport connection serving a request.
///
/// A control is shared by every request multiplexed over the same connection.
/// Shutting down an HTTP/2 or HTTP/3 connection therefore also affects its
/// other in-flight request streams.
#[derive(Clone)]
pub struct ConnCtrl {
    inner: Arc<Inner>,
}

struct Inner {
    state: AtomicU8,
    notify: Notify,
    relax: AtomicBool,
}

impl Default for ConnCtrl {
    fn default() -> Self {
        Self::new()
    }
}

impl Debug for ConnCtrl {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("ConnCtrl")
            .field("state", &self.state())
            .finish()
    }
}

impl ConnCtrl {
    /// Creates a connection control in the running state.
    #[must_use]
    pub fn new() -> Self {
        Self {
            inner: Arc::new(Inner {
                state: AtomicU8::new(RUNNING),
                notify: Notify::new(),
                relax: AtomicBool::new(false),
            }),
        }
    }

    /// Stops accepting new requests and lets accepted requests finish.
    ///
    /// A later call to [`abort`](Self::abort) escalates graceful shutdown to an
    /// immediate abort.
    pub fn graceful_shutdown(&self) {
        if self
            .inner
            .state
            .compare_exchange(
                RUNNING,
                GRACEFUL_SHUTDOWN,
                Ordering::AcqRel,
                Ordering::Acquire,
            )
            .is_ok()
        {
            self.inner.notify.notify_waiters();
        }
    }

    /// Immediately aborts the underlying transport connection.
    ///
    /// The server will abruptly close the connection without sending any response.
    /// This causes clients to encounter errors such as `curl: (52) Empty reply from server`.
    ///
    /// Aborting the connection immediately frees up system resources allocated to the
    /// request flow. This forceful teardown should **only** be used in critical scenarios,
    /// such as mitigating active attacks, where maintaining the connection poses a security risk.
    pub fn abort(&self) {
        if self.inner.state.swap(ABORT, Ordering::AcqRel) != ABORT {
            self.inner.notify.notify_waiters();
        }
    }

    /// Returns `true` after graceful shutdown has been requested.
    #[must_use]
    pub fn is_graceful_shutdown(&self) -> bool {
        self.inner.state.load(Ordering::Acquire) == GRACEFUL_SHUTDOWN
    }

    /// Returns `true` after immediate abort has been requested.
    #[must_use]
    pub fn is_aborted(&self) -> bool {
        self.inner.state.load(Ordering::Acquire) == ABORT
    }

    /// Suspends the transport idle and write-stall fuse timeouts for this
    /// connection.
    ///
    /// The fuse [`connection_idle_timeout`](crate::fuse::FuseConfig::connection_idle_timeout)
    /// and [`write_stall_timeout`](crate::fuse::FuseConfig::write_stall_timeout) exist to
    /// close connections that stall mid-request. Long-lived protocols built on top of an HTTP
    /// upgrade — WebSocket, or a hand-rolled tunnel — legitimately spend long stretches with no
    /// transport activity and would otherwise trip those timers. A handler that hands the
    /// connection off to such a protocol should call this to keep the transport open.
    ///
    /// This does not affect [`graceful_shutdown`](Self::graceful_shutdown) or
    /// [`abort`](Self::abort); an aborted connection is still torn down.
    pub fn relax_timeouts(&self) {
        // A standalone flag with no other memory to order against; relaxed is enough,
        // and it is read on the transport poll path where we avoid needless fences.
        self.inner.relax.store(true, Ordering::Relaxed);
    }

    /// Returns `true` once [`relax_timeouts`](Self::relax_timeouts) has been called.
    #[must_use]
    pub fn is_relaxed(&self) -> bool {
        self.inner.relax.load(Ordering::Relaxed)
    }

    pub(crate) fn state(&self) -> ConnState {
        match self.inner.state.load(Ordering::Acquire) {
            GRACEFUL_SHUTDOWN => ConnState::GracefulShutdown,
            ABORT => ConnState::Abort,
            _ => ConnState::Running,
        }
    }

    #[cfg(any(feature = "http1", feature = "http2", feature = "quinn", test))]
    pub(crate) async fn notified(&self) -> ConnState {
        self.wait_for_state(false).await
    }

    #[cfg(any(feature = "http1", feature = "http2", feature = "quinn", test))]
    pub(crate) async fn aborted(&self) -> ConnState {
        self.wait_for_state(true).await
    }

    #[cfg(any(feature = "http1", feature = "http2", feature = "quinn", test))]
    async fn wait_for_state(&self, abort_only: bool) -> ConnState {
        loop {
            // Register this waiter before reading the state. This closes the
            // race where a transition could otherwise occur between the state
            // check and waiter registration. `Notify` keeps every registered
            // waiter and `notify_waiters` broadcasts to all of them.
            let notified = self.inner.notify.notified();
            tokio::pin!(notified);
            notified.as_mut().enable();

            let state = self.state();
            if state == ConnState::Abort || (!abort_only && state == ConnState::GracefulShutdown) {
                return state;
            }
            notified.await;
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum ConnState {
    Running,
    GracefulShutdown,
    Abort,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn abort_wakes_waiter() {
        let ctrl = ConnCtrl::new();
        let signal = ctrl.clone();
        let waiter = tokio::spawn(async move { signal.notified().await });
        ctrl.abort();
        assert_eq!(waiter.await.unwrap(), ConnState::Abort);
    }

    #[test]
    fn abort_escalates_graceful_shutdown() {
        let ctrl = ConnCtrl::new();
        ctrl.graceful_shutdown();
        assert!(ctrl.is_graceful_shutdown());
        ctrl.abort();
        assert!(ctrl.is_aborted());
    }

    #[tokio::test]
    async fn abort_waiter_ignores_graceful_shutdown() {
        let ctrl = ConnCtrl::new();
        ctrl.graceful_shutdown();
        assert!(
            tokio::time::timeout(std::time::Duration::from_millis(10), ctrl.aborted())
                .await
                .is_err()
        );
        ctrl.abort();
        assert_eq!(ctrl.aborted().await, ConnState::Abort);
    }

    #[tokio::test]
    async fn abort_wakes_every_waiter() {
        let ctrl = ConnCtrl::new();
        let connection_ctrl = ctrl.clone();
        let connection_waiter = tokio::spawn(async move { connection_ctrl.notified().await });
        let request_ctrl_a = ctrl.clone();
        let request_waiter_a = tokio::spawn(async move { request_ctrl_a.aborted().await });
        let request_ctrl_b = ctrl.clone();
        let request_waiter_b = tokio::spawn(async move { request_ctrl_b.aborted().await });

        // Give every future a chance to subscribe before broadcasting abort.
        tokio::task::yield_now().await;
        ctrl.abort();

        let all_waiters = async {
            assert_eq!(connection_waiter.await.unwrap(), ConnState::Abort);
            assert_eq!(request_waiter_a.await.unwrap(), ConnState::Abort);
            assert_eq!(request_waiter_b.await.unwrap(), ConnState::Abort);
        };
        tokio::time::timeout(std::time::Duration::from_secs(1), all_waiters)
            .await
            .expect("abort must wake every registered waiter");
    }
}