Archify Sync server, Dockerfile and Helm chart

Rust server (axum, yrs, sqlx/PostgreSQL): live shared models over the y-websocket protocol, versions, browser sign-in with PKCE, invites, single sign-on through OIDC, metrics. Docker image on distroless, compose file, and a Helm chart with one pod, optional bundled PostgreSQL and secrets from existing Secrets.

crates/server/src/hub.rs

@@ -0,0 +1,466 @@
+//! 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
+        }));
+    }
+}