liminal-server 0.24.2

use super::*;
use liminal::protocol::{ProtocolVersion, decode, encode, encoded_len};
use std::io::{Read, Write};
use std::net::TcpListener;

type TestResult<T = ()> = Result<T, Box<dyn std::error::Error>>;

#[derive(Clone, Debug, Default)]
struct Calls {
    detached: Vec<(u64, PassPrincipal)>,
    unregistered: Vec<u64>,
    order: Vec<&'static str>,
    shutdowns: usize,
}

#[derive(Debug, Default)]
struct Individual(Mutex<Calls>);

impl ConnectionNotifier for Individual {
    fn on_worker_registered(&self, _: u64, _: &WorkerRegistration) -> Result<(), ServerError> {
        Ok(())
    }
    fn on_worker_unregistered(&self, pid: u64) {
        let mut calls = recover_lock(&self.0);
        calls.unregistered.push(pid);
        calls.order.push("worker");
    }
    fn on_pass_detached(&self, pid: u64, principal: &PassPrincipal) {
        let mut calls = recover_lock(&self.0);
        calls.detached.push((pid, principal.clone()));
        calls.order.push("pass");
    }
}

#[derive(Debug, Default)]
struct Batched {
    individual: Individual,
    batch: Mutex<Calls>,
    runtime: Mutex<Option<Weak<ConnectionRuntime>>>,
    detached_signal: Mutex<Option<Sender<()>>>,
}

impl ConnectionNotifier for Batched {
    fn on_worker_registered(&self, _: u64, _: &WorkerRegistration) -> Result<(), ServerError> {
        Ok(())
    }
    fn on_worker_unregistered(&self, pid: u64) {
        self.individual.on_worker_unregistered(pid);
    }
    fn on_pass_detached(&self, pid: u64, principal: &PassPrincipal) {
        self.individual.on_pass_detached(pid, principal);
        let signal = recover_lock(&self.detached_signal).clone();
        if let Some(signal) = signal {
            assert!(signal.send(()).is_ok());
        }
    }
    fn on_shutdown(&self, detached: &[(u64, &PassPrincipal)], unregistered: &[u64]) {
        let runtime = recover_lock(&self.runtime).as_ref().and_then(Weak::upgrade);
        if let Some(runtime) = runtime {
            assert_eq!(
                runtime.active_count(),
                0,
                "all records leave before the hook"
            );
        }
        let mut batch = recover_lock(&self.batch);
        batch.shutdowns += 1;
        batch.detached.extend(
            detached
                .iter()
                .map(|(pid, principal)| (*pid, (*principal).clone())),
        );
        batch.unregistered.extend_from_slice(unregistered);
    }
}

fn supervisor(notifier: Arc<dyn ConnectionNotifier>) -> TestResult<ConnectionSupervisor> {
    Ok(ConnectionSupervisor::with_services_and_notifier(
        Arc::new(LiminalConnectionServices::empty()?),
        notifier,
    )?)
}

fn send(client: &mut TcpStream, frame: &Frame) -> TestResult {
    let mut bytes = vec![0; encoded_len(frame)?];
    let written = encode(frame, &mut bytes)?;
    client.write_all(bytes.get(..written).ok_or("invalid encoded length")?)?;
    Ok(())
}

fn open(supervisor: &ConnectionSupervisor) -> TestResult<(u64, TcpStream)> {
    let listener = TcpListener::bind("127.0.0.1:0")?;
    let mut client = TcpStream::connect(listener.local_addr()?)?;
    let (server, _) = listener.accept()?;
    let pid = supervisor.spawn_connection(server)?.pid();
    send(
        &mut client,
        &Frame::Connect {
            flags: 0,
            min_version: ProtocolVersion::new(1, 0),
            max_version: ProtocolVersion::new(1, 0),
            auth_token: Vec::new(),
        },
    )?;
    let mut bytes = Vec::new();
    loop {
        let mut chunk = [0; 1024];
        let read = client.read(&mut chunk)?;
        if read == 0 {
            return Err("connection closed before ConnectAck".into());
        }
        bytes.extend_from_slice(chunk.get(..read).ok_or("invalid read length")?);
        match decode(&bytes) {
            Ok((Frame::ConnectAck { .. }, _)) => return Ok((pid, client)),
            Ok((other, _)) => return Err(format!("expected ConnectAck, got {other:?}").into()),
            Err(
                liminal::protocol::ProtocolError::IncompleteHeader { .. }
                | liminal::protocol::ProtocolError::TruncatedPayload { .. },
            ) => {}
            Err(error) => return Err(error.into()),
        }
    }
}

fn principal(pid: u64) -> PassPrincipal {
    PassPrincipal {
        participant: pid.to_be_bytes().to_vec(),
        public_key: [7; 32],
        conversations: std::collections::BTreeSet::from([pid]),
        live: format!("seat.{pid}."),
        may_enroll: true,
    }
}

struct Population {
    clients: Vec<TcpStream>,
    detached: Vec<(u64, PassPrincipal)>,
    workers: Vec<u64>,
}

fn population(
    supervisor: &ConnectionSupervisor,
    passes: usize,
    workers: usize,
) -> TestResult<Population> {
    let mut clients = Vec::new();
    let mut detached = Vec::new();
    let mut unregistered = Vec::new();
    for _ in 0..passes {
        let (pid, client) = open(supervisor)?;
        let principal = principal(pid);
        supervisor
            .inner
            .runtime
            .attach_pass_principal(pid, &principal)?;
        detached.push((pid, principal));
        clients.push(client);
    }
    for _ in 0..workers {
        let (pid, client) = open(supervisor)?;
        supervisor.inner.runtime.set_registration(
            pid,
            WorkerRegistration {
                namespaces: vec!["default".into()],
                task_queue: format!("queue-{pid}"),
                node: None,
                activity_types: Vec::new(),
                identity: format!("worker-{pid}"),
                activities: Vec::new(),
            },
        )?;
        unregistered.push(pid);
        clients.push(client);
    }
    Ok(Population {
        clients,
        detached,
        workers: unregistered,
    })
}

#[test]
fn shutdown_batches_eleven_passes_and_two_workers_after_removal() -> TestResult {
    let notifier = Arc::new(Batched::default());
    let supervisor = supervisor(notifier.clone())?;
    *recover_lock(&notifier.runtime) = Some(Arc::downgrade(&supervisor.inner.runtime));
    let Population {
        clients,
        mut detached,
        mut workers,
    } = population(&supervisor, 11, 2)?;
    supervisor.shutdown();
    assert_eq!(supervisor.inner.runtime.active_count(), 0);
    let batch = recover_lock(&notifier.batch).clone();
    assert_eq!(batch.shutdowns, 1);
    let mut actual = batch.detached.clone();
    actual.sort_by_key(|(pid, _)| *pid);
    detached.sort_by_key(|(pid, _)| *pid);
    assert_eq!(actual, detached);
    let mut actual = batch.unregistered;
    actual.sort_unstable();
    workers.sort_unstable();
    assert_eq!(actual, workers);
    let individual = recover_lock(&notifier.individual.0).clone();
    assert!(individual.detached.is_empty());
    assert!(individual.unregistered.is_empty());
    drop(clients);
    Ok(())
}

#[test]
fn shutdown_default_keeps_eleven_detaches_and_two_unregisters_once() -> TestResult {
    let notifier = Arc::new(Individual::default());
    let supervisor = supervisor(notifier.clone())?;
    let Population {
        clients,
        mut detached,
        mut workers,
    } = population(&supervisor, 11, 2)?;
    supervisor.shutdown();
    for (pid, _) in &detached {
        supervisor.inner.runtime.finish(*pid);
    }
    let calls = recover_lock(&notifier.0).clone();
    let mut actual = calls.detached.clone();
    actual.sort_by_key(|(pid, _)| *pid);
    detached.sort_by_key(|(pid, _)| *pid);
    assert_eq!(actual, detached);
    let mut actual = calls.unregistered.clone();
    actual.sort_unstable();
    workers.sort_unstable();
    assert_eq!(actual, workers);
    assert_eq!(calls.order, [vec!["pass"; 11], vec!["worker"; 2]].concat());
    drop(clients);
    Ok(())
}

#[test]
fn disconnect_still_detaches_once_without_shutdown_notification() -> TestResult {
    let notifier = Arc::new(Batched::default());
    let (signal, detached) = channel();
    *recover_lock(&notifier.detached_signal) = Some(signal);
    let supervisor = supervisor(notifier.clone())?;
    let (pid, mut client) = open(&supervisor)?;
    let principal = principal(pid);
    supervisor
        .inner
        .runtime
        .attach_pass_principal(pid, &principal)?;
    send(&mut client, &Frame::Disconnect { flags: 0 })?;
    detached.recv()?;
    let mut bytes = [0; 1024];
    while client.read(&mut bytes)? != 0 {}
    assert_eq!(
        recover_lock(&notifier.individual.0).detached,
        [(pid, principal)]
    );
    assert_eq!(recover_lock(&notifier.batch).shutdowns, 0);
    supervisor.shutdown();
    Ok(())
}

#[derive(Debug)]
struct Noop;
impl ConnectionNotifier for Noop {
    fn on_worker_registered(&self, _: u64, _: &WorkerRegistration) -> Result<(), ServerError> {
        Ok(())
    }
    fn on_worker_unregistered(&self, _: u64) {}
    fn on_shutdown(&self, _: &[(u64, &PassPrincipal)], _: &[u64]) {}
}

#[test]
fn shutdown_noop_wall_time_is_reported() -> TestResult {
    for connections in [1, 11, 50] {
        let supervisor = supervisor(Arc::new(Noop))?;
        let Population { clients, .. } = population(&supervisor, connections, 0)?;
        let started = Instant::now();
        supervisor.shutdown();
        println!(
            "SHUTDOWN-NOOP n={connections} elapsed_ms={:.6}",
            started.elapsed().as_secs_f64() * 1000.0
        );
        drop(clients);
    }
    Ok(())
}