malice

//! Idle-exit tracking for the socket-activated static-only server: counts

//! in-flight requests and the time of the last one, so the process can shut

//! itself down once nothing has happened for `MALICE_IDLE_EXIT_SECS`.

use std::{

    sync::{

        Arc,

        atomic::{AtomicUsize, Ordering},

    },

    time::{Duration, Instant},

};

use axum::{extract::State, middleware::Next};

#[derive(Clone)]

pub struct IdleTracker {

    in_flight: Arc<AtomicUsize>,

    last_activity: Arc<std::sync::Mutex<Instant>>,

}

impl IdleTracker {

    pub fn new() -> Self {

        Self {

            in_flight: Arc::new(AtomicUsize::new(0)),

            last_activity: Arc::new(std::sync::Mutex::new(Instant::now())),

        }

    }

    fn touch(&self) {

        *self

            .last_activity

            .lock()

            .expect("idle tracker mutex poisoned") = Instant::now();

    }

    /// Resolves once no request has started or finished, and none is still

    /// in flight, for at least `idle`.

    pub async fn wait_idle(&self, idle: Duration) {

        loop {

            tokio::time::sleep(idle.min(Duration::from_millis(200))).await;

            if self.in_flight.load(Ordering::SeqCst) != 0 {

                continue;

            }

            let elapsed = self

                .last_activity

                .lock()

                .expect("idle tracker mutex poisoned")

                .elapsed();

            if elapsed >= idle {

                return;

            }

        }

    }

}

impl Default for IdleTracker {

    fn default() -> Self {

        Self::new()

    }

}

pub async fn track(

    State(idle): State<IdleTracker>,

    request: axum::extract::Request,

    next: Next,

) -> axum::response::Response {

    idle.in_flight.fetch_add(1, Ordering::SeqCst);

    idle.touch();

    let response = next.run(request).await;

    idle.touch();

    idle.in_flight.fetch_sub(1, Ordering::SeqCst);

    response

}