sim/src/cluster/transport.rs

//! Delivers entity messages between shard nodes over persistent WebSockets.
//!
//! Every peer gets one connection and one worker thread. Messages are queued in order and sent
//! in batches; the receiving shard answers every batch with one response per message, so
//! [`RemoteSender::flush`] can wait until everything sent so far has reached its mailbox.

use super::protocol::Peer;
use crate::world::{RemoteSender, Router};
use picard::remoting::{ActorAddress, RemoteRequest, RemoteResponse};
use std::collections::VecDeque;
use std::net::TcpStream;
use std::sync::{Arc, Condvar, Mutex};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
use tungstenite::stream::MaybeTlsStream;
use tungstenite::{Message, WebSocket};

const MAX_BATCH: usize = 512;
const FLUSH_TIMEOUT: Duration = Duration::from_secs(10);
const CONNECT_ATTEMPTS: u32 = 50;

#[derive(Default)]
struct Queue {
    waiting: VecDeque<RemoteRequest>,
    /// Messages queued or sent but not yet answered.
    unanswered: usize,
    closed: bool,
}

struct Link {
    queue: Mutex<Queue>,
    changed: Condvar,
}

impl Link {
    fn lock(&self) -> std::sync::MutexGuard<'_, Queue> {
        self.queue.lock().expect("transport queue lock")
    }
}

struct Connection {
    peer: Peer,
    link: Arc<Link>,
    worker: Option<JoinHandle<()>>,
}

pub struct ClusterTransport {
    connections: Vec<Option<Connection>>,
}

impl ClusterTransport {
    /// Connects lazily to every peer except `this` shard.
    pub fn new(this: usize, peers: Vec<Peer>) -> Arc<Self> {
        let connections = peers
            .into_iter()
            .enumerate()
            .map(|(index, peer)| {
                (index != this).then(|| {
                    let link = Arc::new(Link {
                        queue: Mutex::default(),
                        changed: Condvar::new(),
                    });
                    let worker = {
                        let (link, url) =
                            (Arc::clone(&link), format!("ws://{}/remote", peer.address));
                        std::thread::Builder::new()
                            .name(format!("shard-link-{index}"))
                            .spawn(move || run_link(&url, &link))
                            .expect("failed to start a shard link")
                    };
                    Connection {
                        peer,
                        link,
                        worker: Some(worker),
                    }
                })
            })
            .collect();
        Arc::new(Self { connections })
    }
}

impl RemoteSender for ClusterTransport {
    fn send(&self, shard: usize, path: &str, payload: serde_json::Value) {
        let Some(Some(connection)) = self.connections.get(shard) else {
            tracing::warn!(shard, path, "dropped a message for an unknown shard");
            return;
        };
        let destination = ActorAddress::local(connection.peer.system_id, path);
        let mut queue = connection.link.lock();
        queue
            .waiting
            .push_back(RemoteRequest::new(destination, None, payload));
        queue.unanswered += 1;
        connection.link.changed.notify_all();
    }

    fn flush(&self) {
        let deadline = Instant::now() + FLUSH_TIMEOUT;
        for (shard, connection) in self.connections.iter().enumerate() {
            let Some(connection) = connection else {
                continue;
            };
            let mut queue = connection.link.lock();
            while queue.unanswered > 0 && !queue.closed {
                let left = deadline.saturating_duration_since(Instant::now());
                if left.is_zero() {
                    tracing::warn!(
                        shard,
                        unanswered = queue.unanswered,
                        "shard flush timed out"
                    );
                    break;
                }
                queue = connection
                    .link
                    .changed
                    .wait_timeout(queue, left)
                    .expect("transport queue lock")
                    .0;
            }
        }
    }
}

impl Drop for ClusterTransport {
    fn drop(&mut self) {
        for connection in self.connections.iter_mut().flatten() {
            connection.link.lock().closed = true;
            connection.link.changed.notify_all();
            if let Some(worker) = connection.worker.take() {
                let _ = worker.join();
            }
        }
    }
}

type Socket = WebSocket<MaybeTlsStream<TcpStream>>;

fn connect(url: &str) -> Option<Socket> {
    for attempt in 0..CONNECT_ATTEMPTS {
        match tungstenite::connect(url) {
            Ok((socket, _)) => {
                if let MaybeTlsStream::Plain(stream) = socket.get_ref() {
                    let _ = stream.set_nodelay(true);
                }
                return Some(socket);
            }
            Err(error) if attempt + 1 == CONNECT_ATTEMPTS => {
                tracing::error!(url, %error, "failed to connect to a shard");
            }
            Err(_) => std::thread::sleep(Duration::from_millis(100)),
        }
    }
    None
}

/// Sends queued messages to one peer until the transport is dropped.
fn run_link(url: &str, link: &Link) {
    let mut socket: Option<Socket> = None;
    loop {
        let batch: Vec<RemoteRequest> = {
            let mut queue = link.lock();
            while queue.waiting.is_empty() && !queue.closed {
                queue = link.changed.wait(queue).expect("transport queue lock");
            }
            if queue.closed {
                if let Some(mut socket) = socket {
                    let _ = socket.close(None);
                }
                return;
            }
            let count = queue.waiting.len().min(MAX_BATCH);
            queue.waiting.drain(..count).collect()
        };
        let sent = batch.len();
        if socket.is_none() {
            socket = connect(url);
        }
        let delivered = socket
            .as_mut()
            .is_some_and(|socket| exchange(socket, &batch));
        if !delivered {
            tracing::warn!(url, dropped = sent, "lost messages to a shard");
            socket = None;
        }
        let mut queue = link.lock();
        queue.unanswered = queue.unanswered.saturating_sub(sent);
        link.changed.notify_all();
    }
}

fn exchange(socket: &mut Socket, batch: &[RemoteRequest]) -> bool {
    let Ok(frame) = serde_json::to_string(batch) else {
        return false;
    };
    if socket.send(Message::Text(frame.into())).is_err() {
        return false;
    }
    loop {
        match socket.read() {
            Ok(Message::Text(text)) => {
                let Ok(responses) = serde_json::from_str::<Vec<RemoteResponse>>(&text) else {
                    return false;
                };
                for response in responses.iter().filter(|response| !response.accepted) {
                    tracing::debug!(error = ?response.error, "a shard rejected a message");
                }
                return true;
            }
            Ok(Message::Ping(_) | Message::Pong(_)) => {}
            _ => return false,
        }
    }
}

/// Answers one batch of messages from another shard.
pub fn receive(router: &Router, peer: &str, frame: &str) -> String {
    let responses: Vec<RemoteResponse> = match serde_json::from_str::<Vec<RemoteRequest>>(frame) {
        Ok(requests) => requests
            .into_iter()
            .map(|request| router.handle(peer, request))
            .collect(),
        Err(error) => {
            tracing::warn!(%error, "received an unreadable batch from a shard");
            Vec::new()
        }
    };
    serde_json::to_string(&responses).unwrap_or_else(|_| "[]".into())
}

/// Runs one Picard console request on the shard at `address` and returns its JSON response.
#[cfg(feature = "console")]
pub fn console_query(address: &str, request: &str) -> Result<String, String> {
    let url = format!("ws://{address}/console");
    let (mut socket, _) =
        tungstenite::connect(&url).map_err(|e| format!("cannot reach {url}: {e}"))?;
    if let MaybeTlsStream::Plain(stream) = socket.get_ref() {
        let _ = stream.set_read_timeout(Some(FLUSH_TIMEOUT));
    }
    socket
        .send(Message::Text(request.to_owned().into()))
        .map_err(|e| e.to_string())?;
    loop {
        match socket.read() {
            Ok(Message::Text(text)) => {
                let _ = socket.close(None);
                return Ok(text.as_str().to_owned());
            }
            Ok(Message::Ping(_) | Message::Pong(_)) => {}
            Ok(_) => return Err("the shard closed the console".into()),
            Err(error) => return Err(error.to_string()),
        }
    }
}