rstango 0.4.0

//! Stopping the servers of this process when it is asked to terminate, see [`terminate`]
//! and [`super::ServerBuilder::stop_on_terminate`].
//!
//! The process terminating is a flag that is never cleared: every server that asked is
//! stopped, the same way the `Kill` command of an admin device does, a boot in progress is
//! given up, and any later boot fails right away. Terminating again stops the servers
//! again, which does nothing more.

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::task::{Context, Poll, Waker};
use std::future::Future;
use std::pin::Pin;

use futures::executor::block_on;
use futures::future::join_all;
use tracing::{info, warn};

use super::server::Server;


/// Whether the process is terminating, see [`terminate`].
static TERMINATING: AtomicBool = AtomicBool::new(false);

/// The boots waiting for the process to terminate, see [`Terminating`], each woken once it
/// does, under the id of its future.
static WAKERS: Mutex<Vec<(u64, Waker)>> = Mutex::new(Vec::new());

/// The id of the next [`Terminating`] future.
static NEXT_ID: AtomicU64 = AtomicU64::new(0);

/// The servers to stop on termination, held weakly: a server that is gone has nothing
/// left to stop, and is pruned on the next registration.
static SERVERS: Mutex<Vec<Weak<Server>>> = Mutex::new(Vec::new());

/// Terminate the servers of this process that are stopped on termination, see
/// [`super::ServerBuilder::stop_on_terminate`]: those served are stopped, as the `Kill`
/// command of their admin device does, which deletes their devices, while those booting
/// give up their boot, as does any served afterward.
///
/// This is what the signal handler of the `cli` feature calls, see
/// [`super::ServerBuilder::terminate_on_signals`], and what an application handling signals
/// itself calls from its own handler. This blocks until every server is stopped.
pub fn terminate() {

    // Raised before the wakers and the servers are taken, which a boot polling or a server
    // registering past them then sees.
    TERMINATING.store(true, Ordering::SeqCst);
    for (_, waker) in WAKERS.lock().unwrap().drain(..) {
        waker.wake();
    }

    let servers = SERVERS.lock().unwrap().iter()
        .filter_map(Weak::upgrade)
        .collect::<Vec<_>>();

    info!(servers = servers.len(), "Terminating the servers of the process");

    let stopped = block_on(join_all(servers.iter().map(|server| server.stop_async())));
    for (server, stopped) in servers.iter().zip(stopped) {
        if let Err(err) = stopped {
            warn!(server = %server.name(), %err, "Failed to stop the server on termination");
        }
    }

}

/// Resolve once the process is terminating, see [`terminate`], which a boot races.
pub(super) fn terminating() -> Terminating {
    Terminating { id: NEXT_ID.fetch_add(1, Ordering::Relaxed) }
}

/// See [`terminating`].
pub(super) struct Terminating {
    id: u64,
}

impl Future for Terminating {

    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {

        // Checked under the lock, which the wakers are taken under once the flag is raised.
        let mut wakers = WAKERS.lock().unwrap();
        if TERMINATING.load(Ordering::SeqCst) {
            return Poll::Ready(());
        }

        match wakers.iter_mut().find(|(id, _)| *id == self.id) {
            Some((_, waker)) => waker.clone_from(cx.waker()),
            None => wakers.push((self.id, cx.waker().clone())),
        }

        Poll::Pending

    }

}

impl Drop for Terminating {
    fn drop(&mut self) {
        WAKERS.lock().unwrap().retain(|(id, _)| *id != self.id);
    }
}

/// Stop that server on termination, see [`terminate`], returning whether the process is
/// terminating already: the server may then have been missed, and is its caller's to stop.
pub(super) fn register(server: &Arc<Server>) -> bool {
    let mut servers = SERVERS.lock().unwrap();
    servers.retain(|server| server.strong_count() > 0);
    servers.push(Arc::downgrade(server));
    drop(servers);
    TERMINATING.load(Ordering::SeqCst)
}