o402 0.1.5

//! In-process at-most-once settle: payload-hash set, spawned counter, `SettleOnce`.

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

use hashbrown::HashSet;
use tokio::sync::Notify;

/// Shared in-flight settle state for this process.
#[derive(Clone, Debug)]
pub(crate) struct InFlightSettles {
    inner: Arc<Inner>,
}

#[derive(Debug)]
struct Inner {
    /// Raw `Payment-Signature` bytes currently reserved for settle.
    hashes: std::sync::Mutex<HashSet<Vec<u8>>>,
    /// Spawned settle tasks that have not finished.
    spawned: AtomicUsize,
    /// Wakes [`InFlightSettles::wait_for_drain`] when `spawned` changes.
    notify: Notify,
}

/// RAII reservation of a signature. [`Drop`] removes it unless [`Self::disarm`] ran.
#[derive(Debug)]
pub(crate) struct HashReservation {
    inflight: InFlightSettles,
    signature: Vec<u8>,
    armed: bool,
}

impl InFlightSettles {
    /// Empty set, zero spawned settles.
    #[must_use]
    pub(crate) fn new() -> Self {
        Self {
            inner: Arc::new(Inner {
                hashes: std::sync::Mutex::new(HashSet::new()),
                spawned: AtomicUsize::new(0),
                notify: Notify::new(),
            }),
        }
    }

    /// Reserves `signature` before verify. `None` if it is already reserved.
    ///
    /// Drop of the returned guard releases the reservation unless [`HashReservation::disarm`]
    /// transfers removal to the spawned settle task.
    #[must_use]
    pub(crate) fn reserve(&self, signature: Vec<u8>) -> Option<HashReservation> {
        if !self.try_insert(&signature) {
            return None;
        }
        Some(HashReservation {
            inflight: self.clone(),
            signature,
            armed: true,
        })
    }

    /// Inserts `signature`. `false` if it is already reserved.
    #[must_use]
    pub(crate) fn try_insert(&self, signature: &[u8]) -> bool {
        let mut hashes = self
            .inner
            .hashes
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        hashes.insert(signature.to_vec())
    }

    /// Drops `signature` when this request will never settle, or after settle returns.
    pub(crate) fn remove(&self, signature: &[u8]) {
        let mut hashes = self
            .inner
            .hashes
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        hashes.remove(signature);
    }

    /// Counts a spawned settle before `tokio::spawn` so drain cannot miss it.
    pub(crate) fn increment_spawned(&self) {
        self.inner.spawned.fetch_add(1, Ordering::SeqCst);
    }

    /// After hash removal: drop the spawned count and wake drain waiters.
    pub(crate) fn decrement_spawned(&self) {
        self.inner.spawned.fetch_sub(1, Ordering::SeqCst);
        self.inner.notify.notify_waiters();
    }

    /// Waits until every spawned settle has finished, or `timeout` elapses.
    ///
    /// Subscribes to [`Notify`] before reading the counter so a decrement cannot
    /// be lost between the load and the wait.
    pub(crate) async fn wait_for_drain(&self, timeout: Duration) {
        let deadline = tokio::time::Instant::now() + timeout;
        loop {
            let notified = self.inner.notify.notified();
            if self.inner.spawned.load(Ordering::SeqCst) == 0 {
                return;
            }
            if tokio::time::timeout_at(deadline, notified).await.is_err() {
                tracing::warn!(
                    remaining = self.inner.spawned.load(Ordering::SeqCst),
                    "shutdown timeout elapsed with spawned settles still running"
                );
                return;
            }
        }
    }
}

impl HashReservation {
    /// Stops [`Drop`] from removing the hash. Caller must remove it after settle.
    #[must_use]
    pub(crate) fn disarm(mut self) -> (InFlightSettles, Vec<u8>) {
        self.armed = false;
        let signature = std::mem::take(&mut self.signature);
        (self.inflight.clone(), signature)
    }
}

impl Drop for HashReservation {
    fn drop(&mut self) {
        if self.armed {
            self.inflight.remove(&self.signature);
        }
    }
}

/// Per-response right to call settle. `take` is true for exactly one caller.
#[derive(Debug)]
pub(crate) struct SettleOnce {
    taken: AtomicBool,
}

impl SettleOnce {
    /// Unused settle right.
    #[must_use]
    pub(crate) const fn new() -> Self {
        Self {
            taken: AtomicBool::new(false),
        }
    }

    /// Consumes the settle right. True for the first caller only.
    #[must_use]
    pub(crate) fn take(&self) -> bool {
        !self.taken.swap(true, Ordering::SeqCst)
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::{InFlightSettles, SettleOnce};

    #[test]
    fn insert_rejects_duplicate_until_removed() {
        let inflight = InFlightSettles::new();
        let sig = b"sig";
        assert!(inflight.try_insert(sig), "first");
        assert!(!inflight.try_insert(sig), "duplicate");
        inflight.remove(sig);
        assert!(inflight.try_insert(sig), "after remove");
    }

    #[test]
    fn reservation_drop_releases_hash() {
        let inflight = InFlightSettles::new();
        {
            let reserved = inflight.reserve(b"sig".to_vec());
            assert!(reserved.is_some(), "reserve");
            assert!(!inflight.try_insert(b"sig"), "held");
        }
        assert!(inflight.try_insert(b"sig"), "after drop");
    }

    #[test]
    fn disarm_keeps_hash_until_remove() {
        let inflight = InFlightSettles::new();
        let reserved = inflight.reserve(b"sig".to_vec()).expect("reserve");
        let (handle, signature) = reserved.disarm();
        assert!(!inflight.try_insert(b"sig"), "held after disarm");
        handle.remove(&signature);
        assert!(inflight.try_insert(b"sig"), "after spawned remove");
    }

    #[test]
    fn settle_once_is_single_winner() {
        let once = SettleOnce::new();
        assert!(once.take(), "first");
        assert!(!once.take(), "second");
    }

    #[tokio::test]
    async fn wait_for_drain_returns_immediately_when_empty() {
        let inflight = InFlightSettles::new();
        let started = tokio::time::Instant::now();
        inflight.wait_for_drain(Duration::from_secs(2)).await;
        assert!(
            started.elapsed() < Duration::from_millis(200),
            "empty drain waited"
        );
    }

    #[tokio::test]
    async fn wait_for_drain_waits_until_decrement() {
        let inflight = InFlightSettles::new();
        inflight.increment_spawned();
        let worker = inflight.clone();
        drop(tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(50)).await;
            worker.decrement_spawned();
        }));
        let started = tokio::time::Instant::now();
        inflight.wait_for_drain(Duration::from_secs(2)).await;
        assert!(
            started.elapsed() >= Duration::from_millis(40),
            "drain returned before decrement"
        );
    }

    #[tokio::test]
    async fn wait_for_drain_times_out_while_spawned() {
        let inflight = InFlightSettles::new();
        inflight.increment_spawned();
        let started = tokio::time::Instant::now();
        inflight.wait_for_drain(Duration::from_millis(40)).await;
        assert!(
            started.elapsed() >= Duration::from_millis(30),
            "timeout returned early"
        );
        inflight.decrement_spawned();
    }
}