moonpool_transport/rpc/

failure_monitor.rs

//! `FailureMonitor`: Reactive failure tracking for addresses and endpoints.
//!
//! Tracks two levels of failure:
//! - **Address-level**: Is a remote machine reachable? (Missing = Failed)
//! - **Endpoint-level**: Is a specific endpoint permanently dead?
//!
//! Producers (`connection_task`) call [`FailureMonitor::set_status`] and
//! [`FailureMonitor::notify_disconnect`]. Consumers (delivery mode functions)
//! poll [`FailureMonitor::on_disconnect_or_failure`] to race replies against
//! disconnect signals.
//!
//! # FDB Reference
//! `SimpleFailureMonitor` from `FailureMonitor.h:146`, `FailureMonitor.actor.cpp`

use std::collections::{BTreeMap, BTreeSet};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::RwLock;
use std::task::{Poll, Waker};
use std::time::Duration;

use moonpool_core::TimeProvider;

use crate::Endpoint;

/// Type-erased sleep closure capturing a concrete [`TimeProvider`].
///
/// Keeps [`FailureMonitor`] non-generic while still able to call into the
/// active time provider (real or simulated).
type SleepFn = Box<dyn Fn(Duration) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;

/// Maximum number of permanently failed endpoints before clearing the map.
///
/// Matches FDB: `failedEndpoints.size() > 100000` triggers clear.
const MAX_FAILED_ENDPOINTS: usize = 100_000;

/// Status of a network address or endpoint.
///
/// # FDB Reference
/// `FailureStatus` from `FailureMonitor.h:34-60`
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FailureStatus {
    /// Address is reachable.
    Available,
    /// Address is unreachable (default for unknown addresses).
    Failed,
}

/// Reactive failure monitor for address and endpoint tracking.
///
/// Thread-safe (`Send + Sync`) via `Arc<RwLock<…>>` for interior mutability.
/// Consumers register wakers via `on_disconnect_or_failure` and similar methods;
/// producers wake them via `set_status`, `notify_disconnect`, `endpoint_not_found`.
///
/// Use [`on_failed_for`](Self::on_failed_for) to wait for a sustained failure
/// duration before giving up — the canonical building block behind
/// [`get_reply_unless_failed_for`](crate::rpc::delivery::get_reply_unless_failed_for).
///
/// # FDB Reference
/// `SimpleFailureMonitor` from `FailureMonitor.h:146`, `FailureMonitor.actor.cpp`
pub struct FailureMonitor {
    inner: RwLock<FailureMonitorInner>,
    sleep_fn: SleepFn,
}

struct FailureMonitorInner {
    /// Address-level status. Missing entry = Failed (FDB default).
    address_status: BTreeMap<String, FailureStatus>,
    /// Permanently failed endpoints (e.g., endpoint not found on remote).
    failed_endpoints: BTreeSet<Endpoint>,
    /// Wakers waiting for endpoint state changes, keyed by address.
    /// Woken on: `set_status` change, `notify_disconnect`, `endpoint_not_found`.
    endpoint_watchers: BTreeMap<String, Vec<Waker>>,
    /// Wakers waiting for disconnect events, keyed by address.
    /// Woken only on: `notify_disconnect`.
    disconnect_watchers: BTreeMap<String, Vec<Waker>>,
}

impl FailureMonitor {
    /// Create a new failure monitor with empty state.
    ///
    /// All unknown addresses default to [`FailureStatus::Failed`] until
    /// a connection succeeds and calls [`set_status`](Self::set_status).
    ///
    /// The `time` provider is captured behind a type-erased sleep closure
    /// so [`FailureMonitor`] itself stays non-generic.
    pub fn new<T: TimeProvider + 'static>(time: T) -> Self {
        let sleep_fn: SleepFn = Box::new(move |duration| {
            let time = time.clone();
            Box::pin(async move {
                let _ = time.sleep(duration).await;
            })
        });
        Self {
            inner: RwLock::new(FailureMonitorInner {
                address_status: BTreeMap::new(),
                failed_endpoints: BTreeSet::new(),
                endpoint_watchers: BTreeMap::new(),
                disconnect_watchers: BTreeMap::new(),
            }),
            sleep_fn,
        }
    }

    // =========================================================================
    // Producer methods (called by connection_task)
    // =========================================================================

    /// Update the status of an address.
    ///
    /// Called by `connection_task` on successful connect (`Available`) or
    /// connection failure (`Failed`).
    ///
    /// Wakes all endpoint watchers for this address on status change.
    ///
    /// # FDB Reference
    /// `SimpleFailureMonitor::setStatus` (FailureMonitor.actor.cpp:83-115)
    ///
    /// # Panics
    ///
    /// Panics if the internal `RwLock` is poisoned (only possible if a prior
    /// task panicked while holding the lock).
    pub fn set_status(&self, address: &str, status: FailureStatus) {
        let mut inner = self
            .inner
            .write()
            .expect("RwLock poisoned: prior task panicked");

        let changed = match status {
            FailureStatus::Available => {
                let prev = inner.address_status.insert(address.to_string(), status);
                prev != Some(FailureStatus::Available)
            }
            FailureStatus::Failed => {
                // Missing = Failed, so remove the entry
                inner.address_status.remove(address).is_some()
            }
        };

        if changed {
            wake_all(&mut inner.endpoint_watchers, address);
        }
    }

    /// Signal that a connection to the given address has been lost.
    ///
    /// Wakes both endpoint watchers and disconnect watchers for this address.
    /// Called by `connection_task` after `disconnect_notify.notify_waiters()`.
    ///
    /// # FDB Reference
    /// `SimpleFailureMonitor::notifyDisconnect` (FailureMonitor.actor.cpp:150-154)
    ///
    /// # Panics
    ///
    /// Panics if the internal `RwLock` is poisoned (only possible if a prior
    /// task panicked while holding the lock).
    pub fn notify_disconnect(&self, address: &str) {
        let mut inner = self
            .inner
            .write()
            .expect("RwLock poisoned: prior task panicked");
        wake_all(&mut inner.endpoint_watchers, address);
        wake_all(&mut inner.disconnect_watchers, address);
    }

    /// Mark an endpoint as permanently failed (e.g., not found on remote).
    ///
    /// Permanently failed endpoints are never automatically recovered.
    /// Wakes endpoint watchers for the endpoint's address.
    ///
    /// Skips well-known tokens (system endpoints that always exist).
    ///
    /// # FDB Reference
    /// `SimpleFailureMonitor::endpointNotFound` (FailureMonitor.actor.cpp:117-139)
    ///
    /// # Panics
    ///
    /// Panics if the internal `RwLock` is poisoned (only possible if a prior
    /// task panicked while holding the lock).
    pub fn endpoint_not_found(&self, endpoint: &Endpoint) {
        // Skip well-known tokens (FDB: `if token.first() == -1 return`)
        if endpoint.token.is_well_known() {
            return;
        }

        let mut inner = self
            .inner
            .write()
            .expect("RwLock poisoned: prior task panicked");

        // Cap is informational: every entry here is a permanent NotFound, so
        // there is nothing to evict. Drop the new addition with a warning when
        // we hit the limit rather than thrashing the set.
        if inner.failed_endpoints.len() >= MAX_FAILED_ENDPOINTS
            && !inner.failed_endpoints.contains(endpoint)
        {
            tracing::warn!(
                cap = MAX_FAILED_ENDPOINTS,
                ?endpoint,
                "FailureMonitor: failed-endpoint set at cap, dropping new NotFound marker"
            );
            return;
        }

        inner.failed_endpoints.insert(endpoint.clone());

        let address = endpoint.address.to_string();
        wake_all(&mut inner.endpoint_watchers, &address);
    }

    // =========================================================================
    // Consumer methods (used by delivery mode functions)
    // =========================================================================

    /// Get the current failure status of an endpoint.
    ///
    /// Returns [`FailureStatus::Failed`] if:
    /// - The endpoint is permanently failed, OR
    /// - The endpoint's address is unknown or failed
    ///
    /// # FDB Reference
    /// `SimpleFailureMonitor::getState(Endpoint)` (FailureMonitor.actor.cpp:196-206)
    ///
    /// # Panics
    ///
    /// Panics if the internal `RwLock` is poisoned (only possible if a prior
    /// task panicked while holding the lock).
    pub fn state(&self, endpoint: &Endpoint) -> FailureStatus {
        let inner = self
            .inner
            .read()
            .expect("RwLock poisoned: prior task panicked");

        if inner.failed_endpoints.contains(endpoint) {
            return FailureStatus::Failed;
        }

        let address = endpoint.address.to_string();
        inner
            .address_status
            .get(&address)
            .copied()
            .unwrap_or(FailureStatus::Failed) // Missing = Failed
    }

    /// Check if an endpoint is permanently failed.
    ///
    /// # FDB Reference
    /// `SimpleFailureMonitor::permanentlyFailed` (FailureMonitor.h:226-228)
    ///
    /// # Panics
    ///
    /// Panics if the internal `RwLock` is poisoned (only possible if a prior
    /// task panicked while holding the lock).
    pub fn permanently_failed(&self, endpoint: &Endpoint) -> bool {
        self.inner
            .read()
            .expect("RwLock poisoned: prior task panicked")
            .failed_endpoints
            .contains(endpoint)
    }

    /// Returns a future that resolves when the endpoint's address disconnects
    /// or the endpoint becomes permanently failed.
    ///
    /// **Fast path**: Returns `Ready` immediately if already failed.
    ///
    /// Used by `try_get_reply()` to race reply vs disconnect.
    ///
    /// # FDB Reference
    /// `SimpleFailureMonitor::onDisconnectOrFailure` (FailureMonitor.actor.cpp:156-178)
    ///
    /// # Panics
    ///
    /// The returned future panics if the internal `RwLock` is poisoned (only
    /// possible if a prior task panicked while holding the lock).
    pub fn on_disconnect_or_failure(
        self: &Arc<Self>,
        endpoint: &Endpoint,
    ) -> impl Future<Output = ()> + Send {
        let fm = Arc::clone(self);
        let address = endpoint.address.to_string();
        let endpoint = endpoint.clone();

        std::future::poll_fn(move |cx| {
            let inner = fm
                .inner
                .read()
                .expect("RwLock poisoned: prior task panicked");

            // Fast path: already failed
            if inner.failed_endpoints.contains(&endpoint) {
                return Poll::Ready(());
            }
            if !inner
                .address_status
                .get(&address)
                .is_some_and(|s| *s == FailureStatus::Available)
            {
                // Missing or Failed → already failed
                return Poll::Ready(());
            }

            // Slow path: register waker
            drop(inner);
            let mut inner = fm
                .inner
                .write()
                .expect("RwLock poisoned: prior task panicked");
            inner
                .endpoint_watchers
                .entry(address.clone())
                .or_default()
                .push(cx.waker().clone());
            Poll::Pending
        })
    }

    /// Returns a future that resolves when the given address disconnects.
    ///
    /// Only triggered by explicit disconnect events, not status changes.
    ///
    /// # FDB Reference
    /// `SimpleFailureMonitor::onDisconnect` (FailureMonitor.actor.cpp:180-182)
    ///
    /// # Panics
    ///
    /// The returned future panics if the internal `RwLock` is poisoned (only
    /// possible if a prior task panicked while holding the lock).
    pub fn on_disconnect(self: &Arc<Self>, address: &str) -> impl Future<Output = ()> + Send {
        let fm = Arc::clone(self);
        let address = address.to_string();
        let registered = Arc::new(std::sync::atomic::AtomicBool::new(false));

        std::future::poll_fn(move |cx| {
            // If we already registered and got woken, the disconnect happened
            if registered.load(std::sync::atomic::Ordering::Acquire) {
                return Poll::Ready(());
            }

            // First poll: register waker and mark as registered
            registered.store(true, std::sync::atomic::Ordering::Release);
            let mut inner = fm
                .inner
                .write()
                .expect("RwLock poisoned: prior task panicked");
            inner
                .disconnect_watchers
                .entry(address.clone())
                .or_default()
                .push(cx.waker().clone());
            Poll::Pending
        })
    }

    /// Waits for failure to be detected, then sleeps `duration` before returning.
    ///
    /// Two-phase behaviour, **not** sustained-failure detection:
    ///
    /// 1. Block via [`on_disconnect_or_failure`](Self::on_disconnect_or_failure)
    ///    until a failure is observed at least once.
    /// 2. Sleep `duration` using the time provider captured at construction,
    ///    regardless of whether the endpoint recovers in the meantime.
    ///
    /// The recovery case is intentionally ignored — callers like
    /// [`get_reply_unless_failed_for`](crate::rpc::delivery::get_reply_unless_failed_for)
    /// race this future against the reply future, so a recovered endpoint that
    /// answers in time wins the race anyway. Treat this as a deadline timer
    /// armed by the first observed failure, not a continuous health check.
    pub async fn on_failed_for(self: &Arc<Self>, endpoint: &Endpoint, duration: Duration) {
        self.on_disconnect_or_failure(endpoint).await;
        (self.sleep_fn)(duration).await;
    }
}

impl std::fmt::Debug for FailureMonitor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let inner = self
            .inner
            .read()
            .expect("RwLock poisoned: prior task panicked");
        f.debug_struct("FailureMonitor")
            .field("addresses_available", &inner.address_status.len())
            .field("endpoints_failed", &inner.failed_endpoints.len())
            .field("endpoint_watchers", &inner.endpoint_watchers.len())
            .field("disconnect_watchers", &inner.disconnect_watchers.len())
            // `sleep_fn` is a captured closure with no useful Debug representation.
            .finish_non_exhaustive()
    }
}

/// Drain and wake all wakers registered for the given address.
fn wake_all(watchers: &mut BTreeMap<String, Vec<Waker>>, address: &str) {
    if let Some(wakers) = watchers.remove(address) {
        for waker in wakers {
            waker.wake();
        }
    }
}

#[cfg(test)]
mod tests {
    use std::net::{IpAddr, Ipv4Addr};

    use moonpool_core::TokioTimeProvider;

    use super::*;
    use crate::{NetworkAddress, UID};

    fn test_addr() -> NetworkAddress {
        NetworkAddress::new(IpAddr::V4(Ipv4Addr::new(10, 0, 1, 1)), 4500)
    }

    fn test_endpoint() -> Endpoint {
        Endpoint::new(test_addr(), UID::new(42, 1))
    }

    fn make_fm() -> FailureMonitor {
        FailureMonitor::new(TokioTimeProvider::new())
    }

    #[test]
    fn test_unknown_address_is_failed() {
        let fm = make_fm();
        let ep = test_endpoint();
        assert_eq!(fm.state(&ep), FailureStatus::Failed);
    }

    #[test]
    fn test_set_status_available() {
        let fm = make_fm();
        let ep = test_endpoint();
        fm.set_status("10.0.1.1:4500", FailureStatus::Available);
        assert_eq!(fm.state(&ep), FailureStatus::Available);
    }

    #[test]
    fn test_endpoint_not_found_marks_permanent() {
        let fm = make_fm();
        let ep = test_endpoint();
        fm.set_status("10.0.1.1:4500", FailureStatus::Available);
        fm.endpoint_not_found(&ep);
        assert!(fm.permanently_failed(&ep));
        assert_eq!(fm.state(&ep), FailureStatus::Failed);
    }

    #[test]
    fn test_endpoint_not_found_skips_well_known() {
        let fm = make_fm();
        let ep = Endpoint::well_known(test_addr(), crate::WellKnownToken::Ping);
        fm.endpoint_not_found(&ep);
        assert!(!fm.permanently_failed(&ep));
    }

    #[tokio::test]
    async fn test_on_disconnect_or_failure_fast_path_unknown() {
        let fm = Arc::new(make_fm());
        let ep = test_endpoint();
        // Unknown address → already failed → should resolve immediately
        fm.on_disconnect_or_failure(&ep).await;
    }

    #[tokio::test]
    async fn test_on_disconnect_or_failure_fast_path_permanent() {
        let fm = Arc::new(make_fm());
        let ep = test_endpoint();
        fm.set_status("10.0.1.1:4500", FailureStatus::Available);
        fm.endpoint_not_found(&ep);
        // Permanently failed → should resolve immediately
        fm.on_disconnect_or_failure(&ep).await;
    }

    #[test]
    fn test_on_disconnect_or_failure_wakes_on_status_change() {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("build runtime");
        rt.block_on(async {
            let fm = Arc::new(make_fm());
            let ep = test_endpoint();
            fm.set_status("10.0.1.1:4500", FailureStatus::Available);

            let fm2 = Arc::clone(&fm);
            let handle = tokio::spawn(async move {
                fm2.on_disconnect_or_failure(&ep).await;
            });

            // Yield to let the future register its waker
            tokio::task::yield_now().await;

            // Trigger disconnect → should wake the future
            fm.set_status("10.0.1.1:4500", FailureStatus::Failed);

            handle.await.expect("task should complete");
        });
    }

    #[test]
    fn test_on_disconnect_wakes_on_notify() {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("build runtime");
        rt.block_on(async {
            let fm = Arc::new(make_fm());
            fm.set_status("10.0.1.1:4500", FailureStatus::Available);

            let fm2 = Arc::clone(&fm);
            let handle = tokio::spawn(async move {
                fm2.on_disconnect("10.0.1.1:4500").await;
            });

            tokio::task::yield_now().await;

            fm.notify_disconnect("10.0.1.1:4500");

            handle.await.expect("task should complete");
        });
    }

    #[test]
    fn test_debug_impl() {
        let fm = make_fm();
        fm.set_status("10.0.1.1:4500", FailureStatus::Available);
        let debug = format!("{fm:?}");
        assert!(debug.contains("FailureMonitor"));
        assert!(debug.contains("addresses_available: 1"));
    }

    #[tokio::test]
    async fn test_on_failed_for_resolves_after_sustained_failure() {
        let fm = Arc::new(make_fm());
        let addr = NetworkAddress::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), 4500);
        let token = UID::new(0xAAAA, 0xBBBB);
        let endpoint = Endpoint::new(addr.clone(), token);

        // Mark address failed BEFORE awaiting on_failed_for so the disconnect
        // detection step resolves immediately and we measure only the sleep.
        fm.set_status(&addr.to_string(), FailureStatus::Failed);

        let start = tokio::time::Instant::now();
        fm.on_failed_for(&endpoint, Duration::from_millis(50)).await;
        let elapsed = start.elapsed();

        assert!(
            elapsed >= Duration::from_millis(50),
            "should sleep at least 50ms, got {elapsed:?}"
        );
        assert!(
            elapsed < Duration::from_millis(500),
            "should not exceed 500ms wall clock, got {elapsed:?}"
        );
    }
}