archify-vscode-ext-sync/crates/server/src/hub.rs

//! The live hub: one room per open space, holding its Yjs document and
//! awareness in memory. Clients speak the y-websocket protocol (sync step 1
//! and 2, updates, awareness). Every update is appended to PostgreSQL before
//! it is applied and passed on, under the room's lock, so a version's
//! `up_to` never drops an update the version does not hold.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use anyhow::{Context, Result};
use archify_sync_protocol::VersionKind;
use archify_sync_store::{Store, VersionRow};
use axum::extract::ws::{Message as WsMessage, WebSocket};
use bytes::Bytes;
use futures_util::{SinkExt, StreamExt};
use tokio::sync::{Mutex, broadcast};
use uuid::Uuid;
use yrs::sync::{Awareness, Message, SyncMessage};
use yrs::updates::decoder::Decode;
use yrs::updates::encoder::Encode;
use yrs::{ClientID, ReadTxn, Transact, Update};

use crate::doc;
use crate::metrics::{METRICS, Metrics};

/// An encoded update with no changes in it.
const EMPTY_UPDATE: &[u8] = &[0, 0];

pub struct Hub {
    store: Store,
    rooms: Mutex<HashMap<Uuid, Arc<Room>>>,
    snapshots: SnapshotPolicy,
    next_conn: AtomicU64,
}

/// When a room with edits gets a snapshot, which also empties its update log.
#[derive(Debug, Clone, Copy)]
pub struct SnapshotPolicy {
    /// After this long without edits.
    pub idle: Duration,
    /// After this long of editing without a pause, so a busy space's log stays short.
    pub max_age: Duration,
    /// After this many edits.
    pub max_edits: u64,
}

pub struct Room {
    pub space: Uuid,
    state: Mutex<RoomState>,
    /// Frames for every connection, tagged with the sender so it skips its own.
    tx: broadcast::Sender<(u64, Bytes)>,
}

struct RoomState {
    awareness: Awareness,
    /// Users holding the room: connections and short API calls.
    holders: usize,
    last_update: i64,
    latest: Option<VersionRow>,
    /// Edits since the latest version.
    dirty: Option<Dirty>,
}

#[derive(Debug, Clone, Copy)]
struct Dirty {
    since: Instant,
    last: Instant,
    edits: u64,
}

impl RoomState {
    fn edited(&mut self) {
        let now = Instant::now();
        let dirty = self.dirty.get_or_insert(Dirty {
            since: now,
            last: now,
            edits: 0,
        });
        dirty.last = now;
        dirty.edits += 1;
    }
}

impl SnapshotPolicy {
    fn due(&self, dirty: &Dirty) -> bool {
        dirty.last.elapsed() >= self.idle
            || dirty.since.elapsed() >= self.max_age
            || dirty.edits >= self.max_edits
    }
}

impl Hub {
    pub fn new(store: Store, snapshots: SnapshotPolicy) -> Self {
        Self {
            store,
            rooms: Mutex::new(HashMap::new()),
            snapshots,
            next_conn: AtomicU64::new(1),
        }
    }

    /// The space's room, loaded from the store on first use. Pair with [`Hub::release`].
    pub async fn acquire(&self, space: Uuid) -> Result<Arc<Room>> {
        let mut rooms = self.rooms.lock().await;
        if let Some(room) = rooms.get(&space) {
            room.state.lock().await.holders += 1;
            return Ok(room.clone());
        }
        let stored = self.store.load_doc(space).await?;
        let doc = doc::load(&stored).with_context(|| format!("space {space} does not load"))?;
        let (tx, _) = broadcast::channel(1024);
        let room = Arc::new(Room {
            space,
            tx,
            state: Mutex::new(RoomState {
                awareness: Awareness::new(doc),
                holders: 1,
                last_update: stored.last_update,
                latest: stored.latest,
                dirty: (!stored.updates.is_empty()).then(|| Dirty {
                    since: Instant::now(),
                    last: Instant::now(),
                    edits: stored.updates.len() as u64,
                }),
            }),
        });
        rooms.insert(space, room.clone());
        tracing::debug!(%space, "room opened");
        Ok(room)
    }

    /// Lets go of a room; the last one out takes a snapshot and closes it.
    pub async fn release(&self, room: &Arc<Room>) {
        let mut rooms = self.rooms.lock().await;
        let mut state = room.state.lock().await;
        state.holders = state.holders.saturating_sub(1);
        if state.holders > 0 {
            return;
        }
        if let Err(error) = self.snapshot_locked(room.space, &mut state).await {
            // The update log still has every edit; the room stays open for the next try.
            tracing::error!(space = %room.space, %error, "snapshot on close failed");
            return;
        }
        rooms.remove(&room.space);
        tracing::debug!(space = %room.space, "room closed");
    }

    /// Snapshots every room whose edits are due under the [`SnapshotPolicy`].
    pub async fn snapshot_idle_rooms(&self) {
        let rooms: Vec<Arc<Room>> = self.rooms.lock().await.values().cloned().collect();
        for room in rooms {
            let mut state = room.state.lock().await;
            if state.dirty.is_some_and(|dirty| self.snapshots.due(&dirty))
                && let Err(error) = self.snapshot_locked(room.space, &mut state).await
            {
                tracing::error!(space = %room.space, %error, "idle snapshot failed");
            }
        }
    }

    /// Snapshots every room with edits; on shutdown.
    pub async fn snapshot_all(&self) {
        let rooms: Vec<Arc<Room>> = self.rooms.lock().await.values().cloned().collect();
        for room in rooms {
            let mut state = room.state.lock().await;
            if let Err(error) = self.snapshot_locked(room.space, &mut state).await {
                tracing::error!(space = %room.space, %error, "snapshot on shutdown failed");
            }
        }
    }

    async fn snapshot_locked(
        &self,
        space: Uuid,
        state: &mut RoomState,
    ) -> Result<Option<VersionRow>> {
        if state.dirty.is_none() {
            return Ok(None);
        }
        self.version_locked(space, state, VersionKind::Snapshot, None, None, None)
            .await
            .map(Some)
    }

    async fn version_locked(
        &self,
        space: Uuid,
        state: &mut RoomState,
        kind: VersionKind,
        name: Option<&str>,
        restored_from: Option<i64>,
        by: Option<Uuid>,
    ) -> Result<VersionRow> {
        let (major, minor) = match (&state.latest, kind) {
            (None, VersionKind::Publish) => (1, 0),
            (None, _) => (0, 1),
            (Some(v), VersionKind::Publish) => (v.major + 1, 0),
            (Some(v), _) => (v.major, v.minor + 1),
        };
        let (kind_name, counter) = match kind {
            VersionKind::Snapshot => ("snapshot", &METRICS.snapshots_total),
            VersionKind::Publish => ("publish", &METRICS.publishes_total),
            VersionKind::Restore => ("restore", &METRICS.restores_total),
        };
        let encoded = doc::encode_state(state.awareness.doc());
        let row = self
            .store
            .insert_version(
                space,
                major,
                minor,
                kind_name,
                name,
                &encoded,
                state.last_update,
                restored_from,
                by,
            )
            .await?;
        tracing::info!(%space, version = %format!("{major}.{minor}"), kind = kind_name, "version stored");
        Metrics::inc(counter);
        state.latest = Some(row.clone());
        state.dirty = None;
        Ok(row)
    }

    /// Marks the current state as a named, major version.
    pub async fn publish(&self, space: Uuid, name: &str, by: Uuid) -> Result<VersionRow> {
        let room = self.acquire(space).await?;
        let result = {
            let mut state = room.state.lock().await;
            self.version_locked(
                space,
                &mut state,
                VersionKind::Publish,
                Some(name),
                None,
                Some(by),
            )
            .await
        };
        self.release(&room).await;
        result
    }

    /// Brings back the files of an earlier version as a new edit, then stores a restore version.
    pub async fn restore(&self, space: Uuid, version: i64, by: Uuid) -> Result<Option<VersionRow>> {
        let Some((_, old_state)) = self.store.version(space, version).await? else {
            return Ok(None);
        };
        let target = doc::files(&doc::from_state(&old_state)?);
        let room = self.acquire(space).await?;
        let result = async {
            let mut state = room.state.lock().await;
            let update = doc::replace_files(state.awareness.doc(), &target);
            if update != EMPTY_UPDATE {
                state.last_update = self.store.append_update(space, &update, Some(by)).await?;
                state.edited();
                let frame = Message::Sync(SyncMessage::Update(update)).encode_v1();
                let _ = room.tx.send((0, Bytes::from(frame)));
            }
            self.version_locked(
                space,
                &mut state,
                VersionKind::Restore,
                None,
                Some(version),
                Some(by),
            )
            .await
        }
        .await;
        self.release(&room).await;
        result.map(Some)
    }

    /// Serves one WebSocket until it closes.
    pub async fn serve(self: Arc<Self>, socket: WebSocket, space: Uuid, account: Uuid) {
        let room = match self.acquire(space).await {
            Ok(room) => room,
            Err(error) => {
                tracing::error!(%space, %error, "room does not open");
                return;
            }
        };
        let conn = self.next_conn.fetch_add(1, Ordering::Relaxed);
        Metrics::inc(&METRICS.connections_total);
        METRICS.connections_open.fetch_add(1, Ordering::Relaxed);
        let mut clients = HashSet::new();
        if let Err(error) = self.run(&room, socket, conn, account, &mut clients).await {
            tracing::debug!(%space, conn, %error, "connection ended");
        }
        METRICS.connections_open.fetch_sub(1, Ordering::Relaxed);
        self.leave_awareness(&room, conn, clients).await;
        self.release(&room).await;
    }

    async fn run(
        &self,
        room: &Arc<Room>,
        socket: WebSocket,
        conn: u64,
        account: Uuid,
        clients: &mut HashSet<ClientID>,
    ) -> Result<()> {
        let (mut sink, mut stream) = socket.split();
        let mut rx = room.tx.subscribe();

        // Ask for what the client has that the server lacks, and show who is here.
        let hello = {
            let state = room.state.lock().await;
            let sv = state.awareness.doc().transact().state_vector();
            let mut frames = vec![Message::Sync(SyncMessage::SyncStep1(sv)).encode_v1()];
            let awareness = state.awareness.update()?;
            if !awareness.clients.is_empty() {
                frames.push(Message::Awareness(awareness).encode_v1());
            }
            frames
        };
        for frame in hello {
            sink.send(WsMessage::Binary(frame.into())).await?;
        }

        loop {
            tokio::select! {
                incoming = stream.next() => {
                    let Some(incoming) = incoming else { return Ok(()) };
                    match incoming? {
                        WsMessage::Binary(data) => {
                            for reply in self.handle(room, conn, account, &data, clients).await? {
                                sink.send(WsMessage::Binary(reply.into())).await?;
                            }
                        }
                        WsMessage::Close(_) => return Ok(()),
                        _ => {}
                    }
                }
                outgoing = rx.recv() => match outgoing {
                    Ok((from, frame)) if from != conn => sink.send(WsMessage::Binary(frame)).await?,
                    Ok(_) => {}
                    // Too slow to keep up: close, and the client catches up on reconnect.
                    Err(broadcast::error::RecvError::Lagged(_)) => {
                        Metrics::inc(&METRICS.lagged_connections_total);
                        anyhow::bail!("connection lagged behind")
                    }
                    Err(broadcast::error::RecvError::Closed) => return Ok(()),
                },
            }
        }
    }

    /// One frame from a client; returns the frames to answer with.
    async fn handle(
        &self,
        room: &Arc<Room>,
        conn: u64,
        account: Uuid,
        data: &[u8],
        clients: &mut HashSet<ClientID>,
    ) -> Result<Vec<Vec<u8>>> {
        match Message::decode_v1(data)? {
            Message::Sync(SyncMessage::SyncStep1(sv)) => {
                let state = room.state.lock().await;
                let update = state
                    .awareness
                    .doc()
                    .transact()
                    .encode_state_as_update_v1(&sv);
                Ok(vec![
                    Message::Sync(SyncMessage::SyncStep2(update)).encode_v1(),
                ])
            }
            Message::Sync(SyncMessage::SyncStep2(update))
            | Message::Sync(SyncMessage::Update(update)) => {
                if update.as_slice() == EMPTY_UPDATE {
                    return Ok(vec![]);
                }
                let decoded = Update::decode_v1(&update)?;
                let mut state = room.state.lock().await;
                // Durable first: the edit is in PostgreSQL before anyone else sees it.
                let started = Instant::now();
                state.last_update = self
                    .store
                    .append_update(room.space, &update, Some(account))
                    .await?;
                METRICS.update_stored(update.len(), started.elapsed());
                state.awareness.doc().transact_mut().apply_update(decoded)?;
                state.edited();
                let frame = Message::Sync(SyncMessage::Update(update)).encode_v1();
                let _ = room.tx.send((conn, Bytes::from(frame)));
                Ok(vec![])
            }
            Message::Awareness(update) => {
                let mut state = room.state.lock().await;
                if let Some(summary) = state.awareness.apply_update_summary(update)? {
                    clients.extend(summary.added.iter().chain(&summary.updated));
                }
                let _ = room.tx.send((conn, Bytes::copy_from_slice(data)));
                Ok(vec![])
            }
            Message::AwarenessQuery => {
                let state = room.state.lock().await;
                Ok(vec![
                    Message::Awareness(state.awareness.update()?).encode_v1(),
                ])
            }
            Message::Auth(_) | Message::Custom(..) => Ok(vec![]),
        }
    }

    /// Clears the awareness states a connection brought, and tells the others.
    async fn leave_awareness(&self, room: &Arc<Room>, conn: u64, clients: HashSet<ClientID>) {
        if clients.is_empty() {
            return;
        }
        let mut state = room.state.lock().await;
        for client in &clients {
            state.awareness.remove_state(*client);
        }
        if let Ok(update) = state.awareness.update_with_clients(clients) {
            let _ = room
                .tx
                .send((conn, Bytes::from(Message::Awareness(update).encode_v1())));
        }
    }

    pub async fn open_rooms(&self) -> usize {
        self.rooms.lock().await.len()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn a_busy_space_still_gets_snapshots() {
        let policy = SnapshotPolicy {
            idle: Duration::from_secs(30),
            max_age: Duration::from_secs(600),
            max_edits: 2000,
        };
        let now = Instant::now();
        let typing = Dirty {
            since: now,
            last: now,
            edits: 10,
        };
        assert!(!policy.due(&typing));
        assert!(policy.due(&Dirty {
            last: now - Duration::from_secs(31),
            ..typing
        }));
        assert!(policy.due(&Dirty {
            since: now - Duration::from_secs(601),
            ..typing
        }));
        assert!(policy.due(&Dirty {
            edits: 2000,
            ..typing
        }));
    }
}