sim/src/cluster/gateway.rs

//! The gateway: starts the shard nodes, steps them in lockstep and presents their parts of the
//! world to the UI as one world.

use super::protocol::{Control, ControlReply, Peer};
use super::shard::LISTENING;
use crate::world::{EntityKey, Quote, Shard, regenerated_seed};
use serde_json::{Value, json};
use std::io::{BufRead, BufReader};
use std::net::TcpStream;
use std::process::{Child, Command, Stdio};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::sync::{mpsc, oneshot, watch};
use tungstenite::stream::MaybeTlsStream;
use tungstenite::{Message, WebSocket};
use uuid::Uuid;

const FRAME: Duration = Duration::from_millis(100);
const SAVE_INTERVAL: Duration = Duration::from_secs(1);
const SHARD_TIMEOUT: Duration = Duration::from_secs(120);

pub enum Msg {
    Edit(String),
    Script(String, oneshot::Sender<Result<(), String>>),
    #[cfg(feature = "console")]
    PicardConsole {
        request: String,
        shard: usize,
        reply: oneshot::Sender<String>,
    },
    Restart {
        regenerate: bool,
        reply: oneshot::Sender<u64>,
    },
}

pub struct GatewayOptions {
    pub shards: usize,
    pub seed: u64,
    /// Database URL of each shard; `{shard}` is replaced by the shard's index.
    pub database_template: String,
}

/// Where the gateway publishes the unified world for the HTTP handlers.
pub struct Outputs {
    pub terrain: Arc<Mutex<String>>,
    pub script: Arc<Mutex<String>>,
    pub state: watch::Sender<Arc<str>>,
}

/// Runs the cluster until the process exits; a failing shard takes the whole cluster down.
pub fn conduct(
    options: &GatewayOptions,
    outputs: &Outputs,
    mut messages: mpsc::UnboundedReceiver<Msg>,
) {
    let result = Cluster::start(options).and_then(|mut cluster| {
        let result = cluster.run(outputs, &mut messages);
        drop(cluster);
        result
    });
    if let Err(error) = result {
        tracing::error!(%error, "the cluster stopped");
        std::process::exit(1);
    }
}

struct ShardProcess {
    child: Child,
    address: String,
}

impl ShardProcess {
    fn start(index: usize, options: &GatewayOptions, cluster_token: &str) -> Result<Self, String> {
        let executable = std::env::current_exe().map_err(|e| e.to_string())?;
        let database = options
            .database_template
            .replace("{shard}", &index.to_string());
        let mut child = Command::new(executable)
            .args(["shard", "--attached"])
            .args(["--index", &index.to_string()])
            .args(["--count", &options.shards.to_string()])
            .args(["--seed", &options.seed.to_string()])
            .args(["--cluster-token", cluster_token])
            .args(["--database", &database])
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit())
            .spawn()
            .map_err(|e| format!("failed to start shard {index}: {e}"))?;
        let stdout = child.stdout.take().expect("piped shard stdout");
        let mut lines = BufReader::new(stdout).lines();
        let address = loop {
            match lines.next() {
                Some(Ok(line)) => match line.strip_prefix(LISTENING) {
                    Some(address) => break address.trim().to_owned(),
                    None => println!("[shard {index}] {line}"),
                },
                _ => {
                    let _ = child.kill();
                    return Err(format!("shard {index} exited before it was ready"));
                }
            }
        };
        std::thread::spawn(move || {
            for line in lines.map_while(Result::ok) {
                println!("[shard {index}] {line}");
            }
        });
        Ok(Self { child, address })
    }
}

impl Drop for ShardProcess {
    fn drop(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

type Socket = WebSocket<MaybeTlsStream<TcpStream>>;

/// The gateway's `/control` connection to one shard; commands are answered in order.
struct ShardLink {
    index: usize,
    socket: Socket,
}

impl ShardLink {
    fn connect(index: usize, address: &str) -> Result<Self, String> {
        let url = format!("ws://{address}/control");
        let (socket, _) = tungstenite::connect(&url)
            .map_err(|e| format!("failed to connect to shard {index} at {url}: {e}"))?;
        if let MaybeTlsStream::Plain(stream) = socket.get_ref() {
            let _ = stream.set_nodelay(true);
            let _ = stream.set_read_timeout(Some(SHARD_TIMEOUT));
        }
        Ok(Self { index, socket })
    }

    fn send(&mut self, control: &Control) -> Result<(), String> {
        let frame = serde_json::to_string(control).map_err(|e| e.to_string())?;
        self.socket
            .send(Message::Text(frame.into()))
            .map_err(|e| format!("lost shard {}: {e}", self.index))
    }

    fn receive(&mut self) -> Result<ControlReply, String> {
        loop {
            match self.socket.read() {
                Ok(Message::Text(text)) => {
                    return match serde_json::from_str(text.as_str()) {
                        Ok(ControlReply::Failed(error)) => {
                            Err(format!("shard {} failed: {error}", self.index))
                        }
                        Ok(reply) => Ok(reply),
                        Err(error) => Err(format!(
                            "unreadable reply from shard {}: {error}",
                            self.index
                        )),
                    };
                }
                Ok(Message::Ping(_) | Message::Pong(_)) => {}
                Ok(_) => return Err(format!("shard {} closed the connection", self.index)),
                Err(error) => return Err(format!("lost shard {}: {error}", self.index)),
            }
        }
    }

    fn call(&mut self, control: &Control) -> Result<ControlReply, String> {
        self.send(control)?;
        self.receive()
    }
}

struct Ready {
    system_id: Uuid,
    remote_address: String,
    seed: u64,
    tick: u64,
    restored: bool,
}

fn ready(reply: ControlReply) -> Result<Ready, String> {
    match reply {
        ControlReply::Ready {
            system_id,
            remote_address,
            seed,
            tick,
            restored,
        } => Ok(Ready {
            system_id,
            remote_address,
            seed,
            tick,
            restored,
        }),
        other => Err(unexpected(&other)),
    }
}

fn unexpected(reply: &ControlReply) -> String {
    format!("unexpected reply from a shard: {reply:?}")
}

struct Cluster {
    links: Vec<ShardLink>,
    _processes: Vec<ShardProcess>,
    seed: u64,
    cluster_token: String,
    /// Directory changes to publish with the next step.
    directory: Vec<(EntityKey, Option<(usize, String)>)>,
    board: Vec<Quote>,
    paused: bool,
    speed: u64,
    width: usize,
    terrain_rev: Option<u64>,
}

impl Cluster {
    fn start(options: &GatewayOptions) -> Result<Self, String> {
        let cluster_token = Uuid::new_v4().to_string();
        let processes = (0..options.shards)
            .map(|index| ShardProcess::start(index, options, &cluster_token))
            .collect::<Result<Vec<_>, _>>()?;
        let links = processes
            .iter()
            .enumerate()
            .map(|(index, process)| ShardLink::connect(index, &process.address))
            .collect::<Result<Vec<_>, _>>()?;
        let mut cluster = Self {
            links,
            _processes: processes,
            seed: options.seed,
            cluster_token,
            directory: Vec::new(),
            board: Vec::new(),
            paused: false,
            speed: 1,
            width: 0,
            terrain_rev: None,
        };
        let readies = cluster
            .broadcast(|_| Control::Hello)?
            .into_iter()
            .map(ready)
            .collect::<Result<Vec<_>, _>>()?;
        let consistent = readies
            .windows(2)
            .all(|pair| pair[0].tick == pair[1].tick && pair[0].seed == pair[1].seed);
        let restored = readies.iter().filter(|ready| ready.restored).count();
        if consistent && (restored == 0 || restored == readies.len()) {
            cluster.seed = readies[0].seed;
            cluster.connect_peers(&readies)?;
        } else {
            tracing::warn!("the shards' saved states don't match; starting a new world");
            cluster.restart(options.seed)?;
        }
        Ok(cluster)
    }

    /// Sends one command to every shard before waiting for any answer, so they work in parallel.
    fn broadcast(
        &mut self,
        control: impl Fn(usize) -> Control,
    ) -> Result<Vec<ControlReply>, String> {
        for link in &mut self.links {
            link.send(&control(link.index))?;
        }
        self.links.iter_mut().map(ShardLink::receive).collect()
    }

    fn broadcast_done(&mut self, control: impl Fn(usize) -> Control) -> Result<(), String> {
        for reply in self.broadcast(control)? {
            if !matches!(reply, ControlReply::Done) {
                return Err(unexpected(&reply));
            }
        }
        Ok(())
    }

    fn restart(&mut self, seed: u64) -> Result<(), String> {
        let readies = self
            .broadcast(|_| Control::Restart { seed })?
            .into_iter()
            .map(ready)
            .collect::<Result<Vec<_>, _>>()?;
        self.seed = seed;
        self.terrain_rev = None;
        self.connect_peers(&readies)
    }

    /// Tells every shard where the others are and publishes where every entity lives.
    fn connect_peers(&mut self, readies: &[Ready]) -> Result<(), String> {
        let peers: Vec<Peer> = readies
            .iter()
            .map(|ready| Peer {
                address: ready.remote_address.clone(),
                token: self.cluster_token.clone(),
                system_id: ready.system_id,
            })
            .collect();
        self.broadcast_done(|_| Control::Peers(peers.clone()))?;
        self.directory.clear();
        for (index, reply) in self.broadcast(|_| Control::Roster)?.into_iter().enumerate() {
            let ControlReply::Roster(keys) = reply else {
                return Err(unexpected(&reply));
            };
            self.directory.extend(
                keys.into_iter()
                    .map(|(key, path)| (key, Some((index, path)))),
            );
        }
        self.board.clear();
        Ok(())
    }

    fn step(&mut self) -> Result<(), String> {
        let directory = std::mem::take(&mut self.directory);
        let board = std::mem::take(&mut self.board);
        let replies = self.broadcast(|_| Control::Step {
            directory: directory.clone(),
            board: board.clone(),
        })?;
        let mut gone = Vec::new();
        for (index, reply) in replies.into_iter().enumerate() {
            let ControlReply::Stepped { report, quotes } = reply else {
                return Err(unexpected(&reply));
            };
            self.board.extend(quotes);
            self.directory.extend(
                report
                    .moved
                    .into_iter()
                    .map(|(key, to, path)| (key, Some((to, path)))),
            );
            self.directory.extend(
                report
                    .arrived
                    .into_iter()
                    .map(|(key, path)| (key, Some((index, path)))),
            );
            gone.extend(report.left.into_iter().map(|key| (key, None)));
        }
        // Departures last, so that an entity that moved and died in the same step stays gone.
        self.directory.extend(gone);
        Ok(())
    }

    fn run(
        &mut self,
        outputs: &Outputs,
        messages: &mut mpsc::UnboundedReceiver<Msg>,
    ) -> Result<(), String> {
        let mut last_saved = Instant::now();
        self.publish(outputs)?;
        loop {
            let frame_started = Instant::now();
            while let Ok(message) = messages.try_recv() {
                self.handle(message, outputs)?;
            }
            if !self.paused {
                for _ in 0..self.speed {
                    self.step()?;
                }
            }
            self.publish(outputs)?;
            if last_saved.elapsed() >= SAVE_INTERVAL {
                self.broadcast_done(|_| Control::Save)?;
                last_saved = Instant::now();
            }
            std::thread::sleep(FRAME.saturating_sub(frame_started.elapsed()));
        }
    }

    fn handle(&mut self, message: Msg, outputs: &Outputs) -> Result<(), String> {
        match message {
            Msg::Edit(command) => {
                for (index, command) in route_edit(&command, self.links.len()) {
                    let reply = self.links[index].call(&Control::Edit(command))?;
                    if !matches!(reply, ControlReply::Done) {
                        return Err(unexpected(&reply));
                    }
                }
            }
            Msg::Script(source, reply) => {
                let mut result = Ok(());
                for answer in self.broadcast(|_| Control::Script(source.clone()))? {
                    let ControlReply::Script(answer) = answer else {
                        return Err(unexpected(&answer));
                    };
                    if result.is_ok() {
                        result = answer;
                    }
                }
                if result.is_ok() {
                    *outputs.script.lock().expect("script lock") = source;
                }
                let _ = reply.send(result);
            }
            #[cfg(feature = "console")]
            Msg::PicardConsole {
                request,
                shard,
                reply,
            } => {
                let response = match self.links.get_mut(shard) {
                    Some(link) => match link.call(&Control::Console(request))? {
                        ControlReply::Console(response) => response,
                        other => return Err(unexpected(&other)),
                    },
                    None => json!({
                        "id": null,
                        "error": {"code": "invalid_request", "message": "no such shard"}
                    })
                    .to_string(),
                };
                let _ = reply.send(response);
            }
            Msg::Restart { regenerate, reply } => {
                let seed = if regenerate {
                    regenerated_seed(self.seed)
                } else {
                    self.seed
                };
                self.restart(seed)?;
                let _ = reply.send(seed);
            }
        }
        Ok(())
    }

    /// Merges the shards' views into one world for the UI.
    fn publish(&mut self, outputs: &Outputs) -> Result<(), String> {
        let mut states = Vec::with_capacity(self.links.len());
        for reply in self.broadcast(|_| Control::View { terrain: false })? {
            let ControlReply::View { state, .. } = reply else {
                return Err(unexpected(&reply));
            };
            states.push(serde_json::from_str::<Value>(&state).map_err(|e| e.to_string())?);
        }
        let first = &states[0];
        self.paused = first["paused"].as_bool().unwrap_or(false);
        self.speed = first["speed"].as_u64().unwrap_or(1);
        let rev = first["rev"].as_u64();
        if rev != self.terrain_rev {
            let reply = self.links[0].call(&Control::View { terrain: true })?;
            let ControlReply::View {
                terrain: Some(terrain),
                ..
            } = reply
            else {
                return Err(unexpected(&reply));
            };
            let parsed: Value = serde_json::from_str(&terrain).map_err(|e| e.to_string())?;
            self.width = parsed["w"]
                .as_u64()
                .and_then(|w| usize::try_from(w).ok())
                .unwrap_or(0);
            *outputs.terrain.lock().expect("terrain lock") = terrain;
            self.terrain_rev = rev;
        }
        let merged = merge_views(&states, self.width);
        outputs.state.send_replace(Arc::from(merged.to_string()));
        Ok(())
    }
}

/// Which shards an edit command goes to: people are found through shard 0's router, stock is
/// shared out, and everything else (including terrain, pause and speed) applies everywhere.
fn route_edit(command: &str, shards: usize) -> Vec<(usize, String)> {
    let words: Vec<&str> = command.split_whitespace().collect();
    match words.as_slice() {
        ["person", ..] => vec![(0, command.to_owned())],
        ["stock", name, value] => {
            let total = value
                .parse::<i64>()
                .map_or(0, |v| usize::try_from(v.max(0)).unwrap_or(usize::MAX));
            (0..shards)
                .filter_map(|index| Shard::new(index, shards))
                .map(|shard| (shard.index, format!("stock {name} {}", shard.share(total))))
                .collect()
        }
        _ => (0..shards)
            .map(|index| (index, command.to_owned()))
            .collect(),
    }
}

/// Combines the shards' state JSON: settings come from shard 0, entities are concatenated,
/// totals are summed, and `homes` and `shards` tell the UI where each shard lives.
fn merge_views(states: &[Value], width: usize) -> Value {
    let Some(Value::Object(first)) = states.first() else {
        return json!({});
    };
    let mut merged = first.clone();
    for (key, value) in &mut merged {
        match key.as_str() {
            "home" | "log" => {}
            "stock" | "actor_stats" => {
                if let Value::Object(totals) = value {
                    for (name, total) in totals.iter_mut() {
                        let sum: u64 = states
                            .iter()
                            .filter_map(|state| state[key.as_str()][name.as_str()].as_u64())
                            .sum();
                        *total = sum.into();
                    }
                }
            }
            _ => {
                if let Value::Array(items) = value {
                    for state in &states[1..] {
                        if let Some(more) = state[key.as_str()].as_array() {
                            items.extend(more.iter().cloned());
                        }
                    }
                }
            }
        }
    }
    let limit = states
        .iter()
        .filter_map(|state| state["log"].as_array().map(Vec::len))
        .max()
        .unwrap_or(0);
    let mut log: Vec<Value> = states
        .iter()
        .filter_map(|state| state["log"].as_array())
        .flatten()
        .cloned()
        .collect();
    log.sort_by_key(|entry| entry.as_str().and_then(log_day).unwrap_or(0));
    let log = log.split_off(log.len().saturating_sub(limit));
    merged.insert("log".into(), Value::Array(log));
    merged.insert(
        "homes".into(),
        states.iter().map(|state| state["home"].clone()).collect(),
    );
    let layout = Shard::new(0, states.len()).unwrap_or(Shard::SINGLE);
    merged.insert(
        "shards".into(),
        (0..states.len())
            .map(|index| {
                let columns = layout.columns_of(index, width);
                json!([columns.start, columns.end])
            })
            .collect(),
    );
    Value::Object(merged)
}

fn log_day(entry: &str) -> Option<u64> {
    entry.strip_prefix("[day ")?.split(']').next()?.parse().ok()
}

#[cfg(test)]
mod tests {
    use super::{merge_views, route_edit};
    use serde_json::json;

    #[test]
    fn views_merge_into_one_world() {
        let left = json!({
            "paused": false, "speed": 2, "tick": 10, "home": [5, 5],
            "stock": {"food": 3, "wood": 1},
            "actor_stats": {"live_actors": 4},
            "people": [[1, 5, 5]], "markets": [[1, 6, 6]],
            "log": ["[day 1] a", "[day 3] c"]
        });
        let right = json!({
            "paused": true, "speed": 9, "tick": 10, "home": [50, 5],
            "stock": {"food": 4, "wood": 2},
            "actor_stats": {"live_actors": 6},
            "people": [[2, 50, 5], [4, 51, 5]], "markets": [],
            "log": ["[day 2] b"]
        });
        let merged = merge_views(&[left, right], 96);
        assert_eq!(merged["speed"], 2);
        assert_eq!(merged["paused"], false);
        assert_eq!(merged["stock"], json!({"food": 7, "wood": 3}));
        assert_eq!(merged["actor_stats"]["live_actors"], 10);
        assert_eq!(merged["people"].as_array().unwrap().len(), 3);
        assert_eq!(merged["markets"].as_array().unwrap().len(), 1);
        assert_eq!(merged["log"], json!(["[day 2] b", "[day 3] c"]));
        assert_eq!(merged["homes"], json!([[5, 5], [50, 5]]));
        assert_eq!(merged["shards"], json!([[0, 48], [48, 96]]));
    }

    #[test]
    fn edits_go_where_they_apply() {
        assert_eq!(
            route_edit("person 4 coins 9", 3),
            vec![(0, "person 4 coins 9".into())]
        );
        assert_eq!(
            route_edit("stock food 10", 3),
            vec![
                (0, "stock food 4".into()),
                (1, "stock food 3".into()),
                (2, "stock food 3".into())
            ]
        );
        assert_eq!(route_edit("pause 1", 2).len(), 2);
    }
}