MetaCrate/crates/metacrate-grid-agent/tests/session_supervisor.rs

use libremetaverse_types::compat::CancellationToken;
use metacrate_grid_agent::{
    GridSession, GridSessionBackend, ReconnectPolicy, SessionFailure, SessionFailureKind,
    SessionFuture, SessionSignal, SessionState, SessionSupervisor, SessionSupervisorHandle,
    SessionWork, WorkDisposition, WorkKind,
};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

#[derive(Default)]
struct ResourceStats {
    attempts: AtomicUsize,
    logouts: AtomicUsize,
    flushes: AtomicUsize,
    sessions: AtomicUsize,
    callbacks: AtomicUsize,
    workers: AtomicUsize,
    sockets: AtomicUsize,
    max_sessions: AtomicUsize,
    max_callbacks: AtomicUsize,
    max_workers: AtomicUsize,
    max_sockets: AtomicUsize,
}

impl ResourceStats {
    fn acquire(&self) {
        acquire(&self.sessions, &self.max_sessions, 1);
        acquire(&self.callbacks, &self.max_callbacks, 3);
        acquire(&self.workers, &self.max_workers, 2);
        acquire(&self.sockets, &self.max_sockets, 1);
    }

    fn release(&self) {
        self.sessions.fetch_sub(1, Ordering::AcqRel);
        self.callbacks.fetch_sub(3, Ordering::AcqRel);
        self.workers.fetch_sub(2, Ordering::AcqRel);
        self.sockets.fetch_sub(1, Ordering::AcqRel);
    }
}

fn acquire(active: &AtomicUsize, maximum: &AtomicUsize, amount: usize) {
    let next = active.fetch_add(amount, Ordering::AcqRel) + amount;
    maximum.fetch_max(next, Ordering::AcqRel);
}

struct ReconnectingFakeGrid {
    stats: Arc<ResourceStats>,
}

impl GridSessionBackend for ReconnectingFakeGrid {
    fn login(
        &self,
        generation: u64,
        _cancellation: CancellationToken,
    ) -> SessionFuture<'_, Result<Box<dyn GridSession>, SessionFailure>> {
        let attempt = self.stats.attempts.fetch_add(1, Ordering::AcqRel) + 1;
        self.stats.acquire();
        let stats = Arc::clone(&self.stats);
        Box::pin(async move {
            let session: Box<dyn GridSession> = Box::new(FakeConnectedSession {
                generation,
                attempt,
                step: 0,
                stats,
                released: false,
            });
            Ok(session)
        })
    }

    fn flush_audit(
        &self,
        _cancellation: CancellationToken,
    ) -> SessionFuture<'_, Result<(), SessionFailure>> {
        self.stats.flushes.fetch_add(1, Ordering::AcqRel);
        Box::pin(async { Ok(()) })
    }
}

struct FakeConnectedSession {
    generation: u64,
    attempt: usize,
    step: u8,
    stats: Arc<ResourceStats>,
    released: bool,
}

impl FakeConnectedSession {
    fn release(&mut self) {
        if !self.released {
            self.released = true;
            self.stats.release();
        }
    }
}

impl GridSession for FakeConnectedSession {
    fn generation(&self) -> u64 {
        self.generation
    }

    fn next_signal(&mut self, cancellation: CancellationToken) -> SessionFuture<'_, SessionSignal> {
        let attempt = self.attempt;
        let step = self.step;
        self.step = self.step.saturating_add(1);
        Box::pin(async move {
            if step == 0 {
                return SessionSignal::Ready;
            }
            if attempt <= 2 {
                tokio::select! {
                    () = tokio::time::sleep(Duration::from_secs(1)) => {
                        let kind = if attempt == 1 {
                            SessionFailureKind::Maintenance
                        } else {
                            SessionFailureKind::Kicked
                        };
                        SessionSignal::Disconnected(SessionFailure::new(kind))
                    }
                    () = cancellation.cancelled() => SessionSignal::Disconnected(
                        SessionFailure::new(SessionFailureKind::TransientTransport),
                    ),
                }
            } else {
                cancellation.cancelled().await;
                SessionSignal::Disconnected(SessionFailure::new(
                    SessionFailureKind::TransientTransport,
                ))
            }
        })
    }

    fn logout(
        mut self: Box<Self>,
        _cancellation: CancellationToken,
    ) -> SessionFuture<'static, Result<(), SessionFailure>> {
        self.stats.logouts.fetch_add(1, Ordering::AcqRel);
        self.release();
        Box::pin(async { Ok(()) })
    }
}

impl Drop for FakeConnectedSession {
    fn drop(&mut self) {
        self.release();
    }
}

async fn settle() {
    for _ in 0..30 {
        tokio::task::yield_now().await;
    }
}

async fn wait_online(handle: &SessionSupervisorHandle) {
    for _ in 0..100 {
        if handle.state() == SessionState::Online {
            return;
        }
        tokio::task::yield_now().await;
    }
    panic!("fake grid did not become ready: {:?}", handle.status());
}

#[tokio::test(start_paused = true)]
async fn repeated_reconnects_have_one_resource_set_no_replay_and_no_shutdown_leaks() {
    let stats = Arc::new(ResourceStats::default());
    let backend: Arc<dyn GridSessionBackend> = Arc::new(ReconnectingFakeGrid {
        stats: Arc::clone(&stats),
    });
    let policy = ReconnectPolicy {
        initial_delay: Duration::from_secs(1),
        maximum_delay: Duration::from_secs(4),
        readiness_timeout: Duration::from_secs(30),
        stable_reset_after: Duration::from_secs(30),
        shutdown_deadline: Duration::from_secs(3),
        jitter_basis_points: 0,
        instance_seed: 1,
        offline_work_capacity: 8,
    };
    let mut handle = SessionSupervisor::new(backend, policy, 32, 256)
        .expect("supervisor")
        .start();

    wait_online(&handle).await;
    let stale_generation = handle.status().generation;
    assert_eq!(
        handle
            .submit(SessionWork::new("response-1", WorkKind::ReadOnly).expect("work"))
            .await
            .expect("accepted"),
        WorkDisposition::Accepted {
            generation: stale_generation
        }
    );
    tokio::time::advance(Duration::from_secs(1)).await;
    settle().await;
    tokio::time::advance(Duration::from_secs(1)).await;
    settle().await;
    wait_online(&handle).await;
    tokio::time::advance(Duration::from_secs(1)).await;
    settle().await;
    tokio::time::advance(Duration::from_secs(2)).await;
    settle().await;
    wait_online(&handle).await;

    assert!(!handle.accepts_result(stale_generation));
    assert_eq!(
        handle
            .submit(SessionWork::new("response-1", WorkKind::ReadOnly).expect("work"))
            .await
            .expect("dedupe"),
        WorkDisposition::RejectedDuplicate
    );
    assert_eq!(stats.max_sessions.load(Ordering::Acquire), 1);
    assert_eq!(stats.max_callbacks.load(Ordering::Acquire), 3);
    assert_eq!(stats.max_workers.load(Ordering::Acquire), 2);
    assert_eq!(stats.max_sockets.load(Ordering::Acquire), 1);

    handle.shutdown().await.expect("shutdown");
    assert_eq!(stats.sessions.load(Ordering::Acquire), 0);
    assert_eq!(stats.callbacks.load(Ordering::Acquire), 0);
    assert_eq!(stats.workers.load(Ordering::Acquire), 0);
    assert_eq!(stats.sockets.load(Ordering::Acquire), 0);
    assert_eq!(stats.logouts.load(Ordering::Acquire), 3);
    assert_eq!(stats.flushes.load(Ordering::Acquire), 1);
}