Restarts drain: SIGTERM stops taking connections and lets requests finish c92a99490cba

//! Stopping without cutting anyone off (oxygenhub#233). On SIGTERM (a

//! `systemctl restart`) or Ctrl-C the hub stops taking connections, ends

//! its event streams (browsers reconnect by themselves), and gives the

//! requests in flight up to `DRAIN` to finish. Given a listening socket by

//! systemd (socket activation, docs/deploy.md), it serves that rather than

//! binding: the socket outlives the restart, so a request that arrives

//! between the old process and the new one waits in its queue instead of

//! being refused.

use std::sync::OnceLock;

use tokio::sync::watch;

/// How long the requests in flight get to finish.

pub const DRAIN: std::time::Duration = std::time::Duration::from_secs(30);

fn sender() -> &'static watch::Sender<bool> {

    static STOP: OnceLock<watch::Sender<bool>> = OnceLock::new();

    STOP.get_or_init(|| watch::channel(false).0)

}

/// Begin stopping.

pub fn begin() {

    sender().send_replace(true);

}

pub fn is_stopping() -> bool {

    *sender().borrow()

}

/// Done when the hub begins stopping: for event streams to end on, and for

/// the servers' graceful shutdown.

pub async fn stopping() {

    let mut rx = sender().subscribe();

    let _ = rx.wait_for(|s| *s).await;

}

/// Wait for SIGTERM or Ctrl-C, then begin stopping. Installed before the

/// servers start, so a signal never kills the process outright.

pub fn on_signal() -> anyhow::Result<impl std::future::Future<Output = ()>> {

    #[cfg(unix)]

    let mut term = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?;

    Ok(async move {

        #[cfg(unix)]

        tokio::select! {

            _ = term.recv() => tracing::info!("SIGTERM: stopping (requests in flight may take {} s)", DRAIN.as_secs()),

            _ = tokio::signal::ctrl_c() => tracing::info!("stopping"),

        }

        #[cfg(not(unix))]

        {

            let _ = tokio::signal::ctrl_c().await;

            tracing::info!("stopping");

        }

        begin();

    })

}

/// The listening socket systemd passed (LISTEN_FDS, first descriptor 3),

/// when it passed one to this process.

pub fn systemd_listener() -> Option<std::net::TcpListener> {

    #[cfg(unix)]

    {

        use std::os::fd::FromRawFd;

        let pid: u32 = std::env::var("LISTEN_PID").ok()?.parse().ok()?;

        let fds: u32 = std::env::var("LISTEN_FDS").ok()?.parse().ok()?;

        if pid != std::process::id() || fds < 1 {

            return None;

        }

        // SAFETY: systemd hands this process descriptor 3 as an open,

        // listening socket (sd_listen_fds(3)); nothing else owns it.

        let listener = unsafe { std::net::TcpListener::from_raw_fd(3) };

        listener.set_nonblocking(true).ok()?;

        Some(listener)

    }

    #[cfg(not(unix))]

    None

}

//! Stopping drains (oxygenhub#233): its own test binary, as stopping is

//! for the whole process.

use oxygenhub::config::Config;

use oxygenhub::db;

use tokio::io::{AsyncReadExt, AsyncWriteExt};

/// A live event stream ends when the hub begins stopping, so it cannot

/// hold the drain open, and `serve` returns well within the drain's bound.

#[tokio::test(flavor = "multi_thread")]

async fn stopping_ends_event_streams_and_serve_returns() {

    let home = tempfile::tempdir().unwrap();

    std::env::set_var("OXY_HOME", home.path());

    let dir = tempfile::tempdir().unwrap();

    let mut config = Config::load(dir.path()).unwrap();

    let port = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port();

    config.server.bind = format!("127.0.0.1:{port}");

    config.server.base_url = format!("http://127.0.0.1:{port}");

    let app = oxygenhub::build(config).await.unwrap();

    let repo = db::repos::create(&app.db, "andy", "demo", "", hub_api::Visibility::Public).await.unwrap();

    let server = tokio::spawn(oxygenhub::serve(app.clone()));

    // An open event stream, as a browser on a live page holds.

    let mut stream = None;

    for _ in 0..50 {

        if let Ok(s) = tokio::net::TcpStream::connect(("127.0.0.1", port)).await {

            stream = Some(s);

            break;

        }

        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

    }

    let mut stream = stream.expect("the hub listens");

    let request = format!("GET /live?topics=repo:{} HTTP/1.1\r\nHost: 127.0.0.1\r\nAccept: text/event-stream\r\n\r\n", repo.id);

    stream.write_all(request.as_bytes()).await.unwrap();

    let mut head = vec![0u8; 512];

    let n = stream.read(&mut head).await.unwrap();

    let head = String::from_utf8_lossy(&head[..n]).to_string();

    assert!(head.starts_with("HTTP/1.1 200") && head.contains("text/event-stream"), "{head}");

    let started = std::time::Instant::now();

    oxygenhub::shutdown::begin();

    // The stream ends (the server closes it) rather than staying open.

    let mut rest = Vec::new();

    tokio::time::timeout(std::time::Duration::from_secs(10), stream.read_to_end(&mut rest))

        .await

        .expect("the event stream ends when stopping begins")

        .unwrap();

    tokio::time::timeout(std::time::Duration::from_secs(10), server).await.expect("serve returns once drained").unwrap().unwrap();

    assert!(started.elapsed() < oxygenhub::shutdown::DRAIN, "drained in {:?}", started.elapsed());

    assert!(oxygenhub::shutdown::is_stopping());

}

pub mod shutdown;

    // Installed first: from here a SIGTERM drains rather than kills.

    let signal = shutdown::on_signal()?;

        let listener = match shutdown::systemd_listener() {

            Some(l) => {

                tracing::info!("http on the socket systemd holds ({})", l.local_addr().map(|a| a.to_string()).unwrap_or_default());

                tokio::net::TcpListener::from_std(l)?

            }

            None => {

                tracing::info!("http on {addr}");

                tokio::net::TcpListener::bind(addr).await?

            }

        };

            axum::serve(listener, r.into_make_service_with_connect_info::<std::net::SocketAddr>())

                .with_graceful_shutdown(shutdown::stopping())

                .await

                .map_err(anyhow::Error::from)

        let handle = axum_server::Handle::new();

        let draining = handle.clone();

        tokio::spawn(async move {

            shutdown::stopping().await;

            draining.graceful_shutdown(Some(shutdown::DRAIN));

        });

                .handle(handle)

        _ = signal => {

            // The servers stop taking connections and finish what they

            // have; whatever is still running after DRAIN is cut off.

            let drained = tokio::time::timeout(shutdown::DRAIN, async { while tasks.join_next().await.is_some() {} }).await;

            if drained.is_err() {

                tracing::warn!("requests still running after {} s: stopping anyway", shutdown::DRAIN.as_secs());

            }

            tracing::info!("stopped");

            Ok(())

        }

    Sse::new(futures::StreamExt::take_until(futures::stream::iter(hello.chain(replay)).chain(live), crate::shutdown::stopping()))

        .keep_alive(KeepAlive::default())

        .into_response()

    Ok(axum::response::Sse::new(futures::StreamExt::take_until(stream, crate::shutdown::stopping()))

        .keep_alive(axum::response::sse::KeepAlive::default()))

    Ok(axum::response::Sse::new(futures::StreamExt::take_until(first.chain(rest), crate::shutdown::stopping()))

        .keep_alive(axum::response::sse::KeepAlive::default()))

## Restarts without errors (oxygenhub#233)

A deploy restarts the hub. Since #233, on SIGTERM (what `systemctl restart`

sends) the hub stops taking connections, ends its live event streams

(browsers reconnect by themselves) and gives the requests in flight up to

30 seconds to finish. Requests that arrive *between* the old process and the

new one are refused unless systemd holds the listening socket across the

restart: socket activation. The hub serves a socket systemd passes

(`LISTEN_FDS`) and binds `server.bind` itself otherwise.

To turn it on, add `/etc/systemd/system/oxygenhub.socket`:

```ini

[Unit]

Description=Oxygenhub's listening socket, held across restarts

[Socket]

ListenStream=0.0.0.0:7300

# A restart's queue: connections wait here, not refused.

Backlog=512

NoDelay=yes

[Install]

WantedBy=sockets.target

```

Then give `oxygenhub.service` `Requires=oxygenhub.socket`,

`After=oxygenhub.socket` and `TimeoutStopSec=40`, and run:

```sh

sudo systemctl daemon-reload

sudo systemctl stop oxygenhub

sudo systemctl enable --now oxygenhub.socket

sudo systemctl start oxygenhub

```

The service's log then says "http on the socket systemd holds". From then on

`systemctl restart oxygenhub` keeps the socket, so nobody sees a refused

connection or a Cloudflare 502. HTTPS (7343, the developer edition's

self-signed listener) still binds itself.