From 232cc1f2e4b457d252223b9cd1c635f60a956f9d Mon Sep 17 00:00:00 2001 From: Wolf Mermelstein Date: Tue, 27 Oct 2026 19:13:15 +0000 Subject: [PATCH 1/3] feat(sync-service): host sync routing in the connection gateway Embed a router per authenticated gateway socket so multiplexed document sync ships through the existing gateway deployment. Remove the standalone Redis relay and bind downstream cancellation to socket ownership. --- .github/workspace-dep-closures.json | 18 +- Cargo.lock | 16 +- Cargo.toml | 2 +- apps/web/AGENTS.md | 6 +- crates/connection_gateway_models/Cargo.toml | 4 - .../connection_gateway_models/src/fanout.rs | 134 --------- .../src/fanout/test.rs | 46 --- crates/connection_gateway_models/src/lib.rs | 2 - crates/sync_router/AGENTS.md | 38 +++ .../sync_router}/CLAUDE.md | 0 .../sync_router}/Cargo.toml | 11 +- .../sync_router}/src/domain/envelope.rs | 0 .../sync_router}/src/domain/envelope/test.rs | 0 .../sync_router}/src/domain/mod.rs | 0 .../sync_router}/src/domain/models.rs | 17 +- .../sync_router}/src/domain/ports.rs | 32 +- .../sync_router}/src/domain/router.rs | 54 ++-- .../sync_router}/src/domain/router/test.rs | 60 ++-- .../sync_router}/src/lib.rs | 3 +- .../src/outbound/do_downstream.rs | 87 +++--- .../sync_router}/src/outbound/mod.rs | 1 - docker/docker-compose.yml | 17 +- docs/README.md | 2 + docs/agent-guide/documents.md | 6 +- docs/architecture/sync-gateway.md | 45 +++ .../collaboration/src/sync-service/gateway.ts | 4 +- services/connection_gateway/AGENTS.md | 17 +- services/connection_gateway/Cargo.toml | 7 +- services/connection_gateway/README.md | 14 + .../connection_gateway/src/api/connection.rs | 34 ++- .../src/api/connection/messages.rs | 26 +- services/connection_gateway/src/context.rs | 5 +- services/connection_gateway/src/main.rs | 16 +- .../connection_gateway/src/model/message.rs | 2 +- .../connection_gateway/src/service/fanout.rs | 166 ----------- .../connection_gateway/src/service/mod.rs | 2 +- .../connection_gateway/src/service/sync.rs | 124 ++++++++ .../src/service/sync/test.rs | 280 ++++++++++++++++++ services/sync-router/AGENTS.md | 46 --- services/sync-router/src/inbound/mod.rs | 3 - .../sync-router/src/inbound/redis_edge.rs | 129 -------- .../src/inbound/redis_edge/test.rs | 75 ----- services/sync-router/src/main.rs | 57 ---- .../sync-router/src/outbound/redis_sink.rs | 39 --- .../crates/xtask_local/src/local/inventory.rs | 11 - .../xtask_local/src/local/inventory/test.rs | 3 +- 46 files changed, 728 insertions(+), 933 deletions(-) delete mode 100644 crates/connection_gateway_models/src/fanout.rs delete mode 100644 crates/connection_gateway_models/src/fanout/test.rs create mode 100644 crates/sync_router/AGENTS.md rename {services/sync-router => crates/sync_router}/CLAUDE.md (100%) rename {services/sync-router => crates/sync_router}/Cargo.toml (51%) rename {services/sync-router => crates/sync_router}/src/domain/envelope.rs (100%) rename {services/sync-router => crates/sync_router}/src/domain/envelope/test.rs (100%) rename {services/sync-router => crates/sync_router}/src/domain/mod.rs (100%) rename {services/sync-router => crates/sync_router}/src/domain/models.rs (72%) rename {services/sync-router => crates/sync_router}/src/domain/ports.rs (63%) rename {services/sync-router => crates/sync_router}/src/domain/router.rs (82%) rename {services/sync-router => crates/sync_router}/src/domain/router/test.rs (82%) rename {services/sync-router => crates/sync_router}/src/lib.rs (89%) rename {services/sync-router => crates/sync_router}/src/outbound/do_downstream.rs (72%) rename {services/sync-router => crates/sync_router}/src/outbound/mod.rs (80%) create mode 100644 docs/architecture/sync-gateway.md delete mode 100644 services/connection_gateway/src/service/fanout.rs create mode 100644 services/connection_gateway/src/service/sync.rs create mode 100644 services/connection_gateway/src/service/sync/test.rs delete mode 100644 services/sync-router/AGENTS.md delete mode 100644 services/sync-router/src/inbound/mod.rs delete mode 100644 services/sync-router/src/inbound/redis_edge.rs delete mode 100644 services/sync-router/src/inbound/redis_edge/test.rs delete mode 100644 services/sync-router/src/main.rs delete mode 100644 services/sync-router/src/outbound/redis_sink.rs diff --git a/.github/workspace-dep-closures.json b/.github/workspace-dep-closures.json index 9367cc3d2bc..d5e863e6e72 100644 --- a/.github/workspace-dep-closures.json +++ b/.github/workspace-dep-closures.json @@ -884,6 +884,7 @@ "crates/static_file", "crates/static_file_service_client", "crates/stream", + "crates/sync_router", "crates/sync_service", "crates/sync_service_client", "crates/system_properties", @@ -897,6 +898,7 @@ "services/agent_harness_service", "services/connection_gateway", "services/mcp_auth_proxy", + "services/sync-service/bebop-schema", "services/unfurl_service" ], "agent_inmem": [ @@ -4248,6 +4250,7 @@ "crates/macro_entrypoint", "crates/macro_env", "crates/macro_env_var", + "crates/macro_service_urls", "crates/macro_tower_layers", "crates/macro_user_id", "crates/macro_uuid", @@ -4266,22 +4269,22 @@ "crates/remote_env_var", "crates/secretsmanager_client", "crates/stream", + "crates/sync_router", "crates/workspace-hack", - "services/connection_gateway" + "services/connection_gateway", + "services/sync-service/bebop-schema" ], "connection_gateway_client": [ "crates/connection_gateway_client", "crates/connection_gateway_models", "crates/cowlike", "crates/macro_tower_layers", - "crates/macro_user_id", "crates/model-entity", "crates/workspace-hack" ], "connection_gateway_models": [ "crates/connection_gateway_models", "crates/cowlike", - "crates/macro_user_id", "crates/model-entity", "crates/workspace-hack" ], @@ -13286,15 +13289,8 @@ "crates/workspace-hack" ], "sync_router": [ - "crates/connection_gateway_models", - "crates/cowlike", - "crates/macro_entrypoint", - "crates/macro_env", - "crates/macro_env_var", - "crates/macro_user_id", - "crates/model-entity", + "crates/sync_router", "crates/workspace-hack", - "services/sync-router", "services/sync-service/bebop-schema" ], "sync_service": [ diff --git a/Cargo.lock b/Cargo.lock index cc30803203c..e8aadffcc9b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4872,6 +4872,7 @@ dependencies = [ "aws-sdk-dynamodb", "aws-sdk-secretsmanager", "axum", + "bebop", "chrono", "connection_gateway_client", "connection_gateway_models", @@ -4892,6 +4893,7 @@ dependencies = [ "macro_entrypoint", "macro_env", "macro_env_var", + "macro_service_urls", "macro_tower_layers", "macro_user_id", "mockall", @@ -4899,7 +4901,6 @@ dependencies = [ "model-entity", "models_bulk_upload", "opentelemetry", - "postcard", "redis", "secretsmanager_client", "serde", @@ -4909,7 +4910,11 @@ dependencies = [ "stream", "strum 0.27.2", "strum_macros 0.27.2", + "sync_router", + "sync_service_bebop_schema", "tokio", + "tokio-tungstenite 0.28.0", + "tokio-util", "tower 0.5.3", "tower-http", "tracing", @@ -4940,9 +4945,7 @@ dependencies = [ name = "connection_gateway_models" version = "0.1.0" dependencies = [ - "macro_user_id", "model-entity", - "postcard", "serde", "serde_json", "utoipa", @@ -19463,20 +19466,17 @@ version = "0.1.0" dependencies = [ "anyhow", "bebop", - "connection_gateway_models", "futures", - "macro_entrypoint", - "macro_env_var", "mockall", "opentelemetry", - "postcard", - "redis", "sync_service_bebop_schema", "thiserror 2.0.18", "tokio", "tokio-tungstenite 0.28.0", + "tokio-util", "tracing", "tracing-opentelemetry", + "url", "workspace-hack", ] diff --git a/Cargo.toml b/Cargo.toml index 563c6eca686..672304d52b3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -175,7 +175,7 @@ members = [ "services/sha_cleanup_worker", "services/slack_import_worker", "services/static_file_service", - "services/sync-router", + "crates/sync_router", "services/sync-service", "services/sync-service/bebop-schema", "services/unfurl_service", diff --git a/apps/web/AGENTS.md b/apps/web/AGENTS.md index fdb72f1d5cc..25bd67374b5 100644 --- a/apps/web/AGENTS.md +++ b/apps/web/AGENTS.md @@ -116,12 +116,12 @@ Then trigger the interaction and read `window.__inst.log`. `'1,2,3' → '' → ' Document sync uses the `sync-via-gateway` PostHog flag in both deployed dev and production. Off or unknown flag values use direct sync sockets. -Local stacks with `sync_router` can opt in with `VITE_SYNC_VIA_GATEWAY=true`; +Local stacks with the updated connection gateway can opt in with `VITE_SYNC_VIA_GATEWAY=true`; `false` forces direct sync. The choice is made when a sync source is created and stays fixed for its lifetime. PostHog can serve a cached value on reload; recreate document sessions after the client receives the updated flag to change -transport. This is not an immediate kill switch for open sessions. Provision -the hosted router before enabling. +transport. This is not an immediate kill switch for open sessions. Finish deploying the updated +connection gateway before enabling; routing is embedded in that service. Saved canvases use the shared Loro runtime; the stored JSON is only the original migration seed. Keep whole-file simple-save out of scene editing. Canvas diff --git a/crates/connection_gateway_models/Cargo.toml b/crates/connection_gateway_models/Cargo.toml index 3ca6ada5dfb..da13111e3ce 100644 --- a/crates/connection_gateway_models/Cargo.toml +++ b/crates/connection_gateway_models/Cargo.toml @@ -5,12 +5,8 @@ publish = false version = "0.1.0" [dependencies] -macro_user_id = { path = "../macro_user_id" } model-entity = { path = "../model-entity" } serde = { workspace = true } serde_json = { workspace = true } utoipa = { workspace = true } workspace-hack = { version = "0.1", path = "../workspace-hack" } - -[dev-dependencies] -postcard = { workspace = true } diff --git a/crates/connection_gateway_models/src/fanout.rs b/crates/connection_gateway_models/src/fanout.rs deleted file mode 100644 index bbe18b80751..00000000000 --- a/crates/connection_gateway_models/src/fanout.rs +++ /dev/null @@ -1,134 +0,0 @@ -//! Wire types for the gateway's inbound fan-out. -//! -//! Every gateway instance republishes its entire client websocket traffic — -//! connection lifecycle and every frame, text or binary, unparsed — to Redis -//! on [`INBOUND_CHANNEL`]. Consumer services subscribe and pick out what they -//! care about (the sync tier consumes binary frames; others may follow). A -//! consumer replies to a specific connection on that gateway instance's -//! [`outbound_channel`]. -//! -//! Encoding on both channels is `postcard`; Redis pub/sub is binary-safe. -//! -//! Delivery is fire-and-forget. Consumers detect a dead gateway instance by -//! its [`FromGateway::Heartbeat`] going quiet, not by any delivery guarantee. -//! In particular, a crashed gateway never publishes its connections' -//! [`FromGateway::Disconnected`] — consumers must treat heartbeat silence as -//! the disconnect of everything that instance held, and must tolerate -//! duplicate or missing lifecycle messages generally. - -#[cfg(test)] -mod test; - -use macro_user_id::user_id::MacroUserIdStr; -use serde::{Deserialize, Serialize}; -use std::fmt; - -/// Boot-unique id of a gateway instance. Minted fresh at gateway startup, so -/// a restarted instance is a new id and the old one simply goes quiet. -#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] -#[serde(transparent)] -pub struct GatewayId(pub String); - -impl fmt::Display for GatewayId { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(&self.0) - } -} - -// For string literals in tests and constructors at the wire boundary. -impl From<&str> for GatewayId { - fn from(id: &str) -> Self { - Self(id.to_string()) - } -} - -/// Gateway-local id of one websocket connection. Only meaningful together -/// with the [`GatewayId`] that minted it. -#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] -#[serde(transparent)] -pub struct ConnId(pub String); - -impl fmt::Display for ConnId { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(&self.0) - } -} - -// For string literals in tests and constructors at the wire boundary. -impl From<&str> for ConnId { - fn from(id: &str) -> Self { - Self(id.to_string()) - } -} - -/// Channel every gateway instance publishes its inbound traffic to. -pub const INBOUND_CHANNEL: &str = "connection_gateway.fanout.inbound"; - -/// Interval between [`FromGateway::Heartbeat`] messages. -pub const HEARTBEAT_INTERVAL_SECS: u64 = 5; - -/// The per-instance channel consumers reply on. `gateway` is minted fresh at -/// gateway boot, so a restarted instance is a new peer and the old id simply -/// goes quiet. -pub fn outbound_channel(gateway: &GatewayId) -> String { - format!("connection_gateway.fanout.outbound.{gateway}") -} - -/// Everything a gateway instance publishes on [`INBOUND_CHANNEL`]. -#[derive(Debug, Serialize, Deserialize)] -pub enum FromGateway { - /// A websocket connection was accepted and its user authenticated. - Connected { - /// Boot-unique id of the publishing gateway instance. - gateway: GatewayId, - /// Gateway-local id of the websocket connection. - conn: ConnId, - /// User the gateway authenticated at the socket edge. - user_id: MacroUserIdStr<'static>, - }, - /// One client frame, in per-connection order, unparsed. - Frame { - /// Boot-unique id of the publishing gateway instance. - gateway: GatewayId, - /// Gateway-local id of the websocket connection. - conn: ConnId, - /// Whether the frame was a text websocket message (else binary). - text: bool, - /// The client's bytes, unparsed. - payload: Vec, - }, - /// The websocket connection closed. - Disconnected { - /// Boot-unique id of the publishing gateway instance. - gateway: GatewayId, - /// Gateway-local id of the websocket connection. - conn: ConnId, - }, - /// Liveness beacon, published every [`HEARTBEAT_INTERVAL_SECS`]. - /// Consumers should drop all state for a gateway id that goes quiet. - Heartbeat { - /// Boot-unique id of the publishing gateway instance. - gateway: GatewayId, - }, -} - -/// Messages a consumer publishes on a gateway's [`outbound_channel`]. -#[derive(Debug, Serialize, Deserialize)] -pub enum ToGateway { - /// Deliver a frame to a connection's websocket. - Frame { - /// Gateway-local id of the target websocket connection. - conn: ConnId, - /// Whether to send as a text websocket message (else binary). - text: bool, - /// The bytes to send, unparsed. - payload: Vec, - }, - /// Close a connection's websocket. - Close { - /// Gateway-local id of the target websocket connection. - conn: ConnId, - /// Websocket close code to send. - code: u16, - }, -} diff --git a/crates/connection_gateway_models/src/fanout/test.rs b/crates/connection_gateway_models/src/fanout/test.rs deleted file mode 100644 index 664d46c75be..00000000000 --- a/crates/connection_gateway_models/src/fanout/test.rs +++ /dev/null @@ -1,46 +0,0 @@ -use super::*; - -/// The gateway encodes and the sync tier decodes (and vice versa); this pins -/// the postcard round trip both directions. -#[test] -fn from_gateway_round_trips() { - let message = FromGateway::Frame { - gateway: GatewayId("g-1".into()), - conn: ConnId("c-1".into()), - text: false, - payload: vec![0xff, 0x00, 0x7f], - }; - let bytes = postcard::to_stdvec(&message).unwrap(); - let decoded: FromGateway = postcard::from_bytes(&bytes).unwrap(); - match decoded { - FromGateway::Frame { - gateway, - conn, - text, - payload, - } => { - assert_eq!(gateway, GatewayId("g-1".into())); - assert_eq!(conn, ConnId("c-1".into())); - assert!(!text); - assert_eq!(payload, vec![0xff, 0x00, 0x7f]); - } - other => panic!("wrong variant: {other:?}"), - } -} - -#[test] -fn to_gateway_round_trips() { - let message = ToGateway::Close { - conn: ConnId("c-9".into()), - code: 4000, - }; - let bytes = postcard::to_stdvec(&message).unwrap(); - let decoded: ToGateway = postcard::from_bytes(&bytes).unwrap(); - match decoded { - ToGateway::Close { conn, code } => { - assert_eq!(conn, ConnId("c-9".into())); - assert_eq!(code, 4000); - } - other => panic!("wrong variant: {other:?}"), - } -} diff --git a/crates/connection_gateway_models/src/lib.rs b/crates/connection_gateway_models/src/lib.rs index 49d907ae0da..13b39e88983 100644 --- a/crates/connection_gateway_models/src/lib.rs +++ b/crates/connection_gateway_models/src/lib.rs @@ -1,8 +1,6 @@ #![deny(missing_docs)] //! Shared types for the connection gateway service and client. -pub mod fanout; - use model_entity::Entity; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; diff --git a/crates/sync_router/AGENTS.md b/crates/sync_router/AGENTS.md new file mode 100644 index 00000000000..2e8e99a5ec8 --- /dev/null +++ b/crates/sync_router/AGENTS.md @@ -0,0 +1,38 @@ +# sync_router + +## What it is + +Document sync routing library embedded in `connection_gateway`. Each authenticated +WebSocket owns a router and its downstream connections; there is no separate +service, Redis subscriber, or Pulumi stack. + +## Start here + +- `src/domain/router.rs`: document subscriptions, bounded queues, route epochs. +- `src/domain/ports.rs`: socket delivery and downstream lifecycle contracts. +- `src/outbound/do_downstream.rs`: scoped-token WebSocket dialing and frame pumps. +- `../../services/connection_gateway/src/service/sync.rs`: socket-owned task + tracking, cancellation, and local delivery adapter. + +## Scope + +- Owns routing and cancellation, not document state or authorization. +- `services/sync-service` validates each document-scoped token and owns CRDT state. +- Browser envelopes are generated from `services/sync-service/bebop/schema.bop`. + +## Checks + +- `cargo test -p sync_router` (no live external services required). +- `cargo test -p connection_gateway --lib service::sync` exercises real local + downstream sockets, isolation, backpressure, and cancellation. +- After dependency changes, regenerate and check `cargo x deps`. + +## Gotchas + +- Preserve inner chunk framing byte-for-byte. A full queue must close its route + or socket and trigger resync; never drop a chunk and continue the stream. +- Dropping a downstream route cancels its task, including pending dials/writes. +- Every gateway replica handles only its own sockets. No distributed route + ownership or singleton deployment is required. +- Deploy the updated gateway before enabling PostHog `sync-via-gateway`. + `VITE_SYNC_VIA_GATEWAY=true` opts local clients in. diff --git a/services/sync-router/CLAUDE.md b/crates/sync_router/CLAUDE.md similarity index 100% rename from services/sync-router/CLAUDE.md rename to crates/sync_router/CLAUDE.md diff --git a/services/sync-router/Cargo.toml b/crates/sync_router/Cargo.toml similarity index 51% rename from services/sync-router/Cargo.toml rename to crates/sync_router/Cargo.toml index 84ed82c1956..497e369dd22 100644 --- a/services/sync-router/Cargo.toml +++ b/crates/sync_router/Cargo.toml @@ -7,20 +7,17 @@ version = "0.1.0" [dependencies] anyhow = { workspace = true } bebop = { workspace = true } -connection_gateway_models = { path = "../../crates/connection_gateway_models" } futures = { workspace = true } -macro_entrypoint = { path = "../../crates/macro_entrypoint" } -macro_env_var = { path = "../../crates/macro_env_var" } opentelemetry = { workspace = true } -postcard = { workspace = true } -redis = { workspace = true, features = ["aio", "tokio-comp", "tokio-rustls-comp"] } -sync_service_bebop_schema = { path = "../sync-service/bebop-schema", features = ["bebop-owned-all"] } +sync_service_bebop_schema = { path = "../../services/sync-service/bebop-schema", features = ["bebop-owned-all"] } thiserror = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread", "sync", "time"] } +tokio-util = { workspace = true, features = ["rt"] } +url = { workspace = true } tokio-tungstenite = { workspace = true } tracing = { workspace = true } tracing-opentelemetry = { workspace = true } -workspace-hack = { version = "0.1", path = "../../crates/workspace-hack" } +workspace-hack = { version = "0.1", path = "../workspace-hack" } [dev-dependencies] mockall = { workspace = true } diff --git a/services/sync-router/src/domain/envelope.rs b/crates/sync_router/src/domain/envelope.rs similarity index 100% rename from services/sync-router/src/domain/envelope.rs rename to crates/sync_router/src/domain/envelope.rs diff --git a/services/sync-router/src/domain/envelope/test.rs b/crates/sync_router/src/domain/envelope/test.rs similarity index 100% rename from services/sync-router/src/domain/envelope/test.rs rename to crates/sync_router/src/domain/envelope/test.rs diff --git a/services/sync-router/src/domain/mod.rs b/crates/sync_router/src/domain/mod.rs similarity index 100% rename from services/sync-router/src/domain/mod.rs rename to crates/sync_router/src/domain/mod.rs diff --git a/services/sync-router/src/domain/models.rs b/crates/sync_router/src/domain/models.rs similarity index 72% rename from services/sync-router/src/domain/models.rs rename to crates/sync_router/src/domain/models.rs index 698af897ffc..8211268a621 100644 --- a/services/sync-router/src/domain/models.rs +++ b/crates/sync_router/src/domain/models.rs @@ -1,16 +1,8 @@ //! Core types shared across the router. -pub use connection_gateway_models::fanout::{ConnId, GatewayId}; - -/// A client connection, globally identified by the gateway instance that owns -/// its socket plus that gateway's local connection id. +/// A gateway-generated client connection id. Each socket has its own router. #[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct ConnectionId { - /// Boot-unique id of the gateway instance holding the socket. - pub gateway: GatewayId, - /// Gateway-local websocket connection id. - pub conn: ConnId, -} +pub struct ConnectionId(pub String); /// A document id, as it appears in envelope frames and downstream URLs. #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -39,11 +31,6 @@ pub enum EdgeEvent { /// The closed connection. conn: ConnectionId, }, - /// A gateway instance stopped heartbeating: every connection it held is gone. - GatewayLost { - /// The dead gateway instance's boot-unique id. - gateway: GatewayId, - }, } /// Everything the router loop reacts to. diff --git a/services/sync-router/src/domain/ports.rs b/crates/sync_router/src/domain/ports.rs similarity index 63% rename from services/sync-router/src/domain/ports.rs rename to crates/sync_router/src/domain/ports.rs index 0a77bb7ec08..d6707fb385b 100644 --- a/services/sync-router/src/domain/ports.rs +++ b/crates/sync_router/src/domain/ports.rs @@ -3,9 +3,10 @@ use crate::domain::models::{ConnectionId, DocId}; use std::fmt::Debug; use tokio::sync::mpsc; +use tokio_util::sync::{CancellationToken, DropGuard}; /// Delivers an already-serialized `FromRouter` envelope to a client -/// connection. Implemented over the gateway's Redis outbound channel. +/// connection. The gateway implements this with the owning socket's queue. #[cfg_attr(test, mockall::automock(type Err = anyhow::Error;))] pub trait EdgeSink: Send + Sync + 'static { /// The error type returned by delivery. @@ -35,11 +36,26 @@ pub trait DownstreamFactory: Send + Sync + 'static { /// (already-serialized `FromPeer`) frames. `epoch` identifies this route /// incarnation and must be echoed in the resulting /// [`crate::domain::models::Event::DownstreamClosed`]. - fn open( - &self, - conn: ConnectionId, - doc: DocId, - token: String, - epoch: u64, - ) -> mpsc::Sender>; + fn open(&self, conn: ConnectionId, doc: DocId, token: String, epoch: u64) -> Downstream; +} + +/// A downstream route whose work is cancelled as soon as the route is removed. +pub struct Downstream { + sender: mpsc::Sender>, + _cancel_on_drop: DropGuard, +} + +impl Downstream { + /// Own the frame queue and cancellation of its downstream pump. + pub fn new(sender: mpsc::Sender>, cancel: CancellationToken) -> Self { + Self { + sender, + _cancel_on_drop: cancel.drop_guard(), + } + } + + /// Offer a frame without blocking the router on a slow downstream. + pub fn try_send(&self, frame: Vec) -> Result<(), mpsc::error::TrySendError>> { + self.sender.try_send(frame) + } } diff --git a/services/sync-router/src/domain/router.rs b/crates/sync_router/src/domain/router.rs similarity index 82% rename from services/sync-router/src/domain/router.rs rename to crates/sync_router/src/domain/router.rs index cc44f65ba1c..362194cccfe 100644 --- a/services/sync-router/src/domain/router.rs +++ b/crates/sync_router/src/domain/router.rs @@ -7,11 +7,12 @@ mod test; use crate::domain::{ envelope::{self, ClientEnvelope}, models::{ConnectionId, DocId, EdgeEvent, Event}, - ports::{DownstreamFactory, EdgeSink}, + ports::{Downstream, DownstreamFactory, EdgeSink}, }; use std::collections::{HashMap, HashSet}; use std::sync::Arc; use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; use tracing::{debug, warn}; /// Routes envelope frames between gateway connections and per-document @@ -29,7 +30,7 @@ pub struct Router { /// One live (or dialing) downstream route. struct Route { epoch: u64, - sender: mpsc::Sender>, + sender: Downstream, } /// Documents one connection may hold open at once. The downstream token is @@ -50,9 +51,16 @@ impl Router { } /// Drive the router from an event stream until it closes. - pub async fn run(mut self, mut events: mpsc::Receiver) { - while let Some(event) = events.recv().await { - self.handle(event).await; + pub async fn run(mut self, mut events: mpsc::Receiver, cancel: CancellationToken) { + loop { + tokio::select! { + biased; + _ = cancel.cancelled() => break, + event = events.recv() => match event { + Some(event) => self.handle(event).await, + None => break, + }, + } } } @@ -69,22 +77,6 @@ impl Router { } } Event::Edge(EdgeEvent::Disconnected { conn }) => self.drop_conn(&conn), - Event::Edge(EdgeEvent::GatewayLost { gateway }) => { - let conns: Vec = self - .by_conn - .keys() - .filter(|conn| conn.gateway == gateway) - .cloned() - .collect(); - debug!( - gateway = %gateway, - count = conns.len(), - "gateway lost; dropping its connections" - ); - for conn in conns { - self.drop_conn(&conn); - } - } Event::DownstreamClosed { conn, doc, epoch } => { // The pump already told the client; forget the route — but // only if it is still the incarnation that died. A stale @@ -100,7 +92,7 @@ impl Router { } } - #[tracing::instrument(skip_all, fields(conn = %conn.conn, gateway = %conn.gateway))] + #[tracing::instrument(skip_all, fields(conn = %conn.0))] async fn handle_client(&mut self, conn: ConnectionId, frame: ClientEnvelope) { match frame { ClientEnvelope::RouterSubscribe { doc_id, token } => { @@ -151,12 +143,16 @@ impl Router { ); return; }; - // try_send: a full buffer means a downstream that can't keep - // up; drop the frame — the sync protocol self-heals via - // catch-up requests. A closed channel means the downstream - // died; DownstreamClosed will clean the route up. + // Losing a chunk can corrupt the inner framing stream. Tear + // down this route so the client reconnects and resynchronizes. if let Err(error) = sender.sender.try_send(payload) { - warn!(doc = doc.as_str(), error = %error, "dropping frame for downstream"); + warn!(doc = doc.as_str(), error = %error, "downstream queue unavailable; closing route"); + self.forget(&conn, &doc); + self.deliver( + &conn, + envelope::doc_closed(doc.as_str(), "downstream queue unavailable"), + ) + .await; } } } @@ -166,7 +162,7 @@ impl Router { let Some(docs) = self.by_conn.remove(conn) else { return; }; - debug!(conn = %conn.conn, count = docs.len(), "dropping connection's downstreams"); + debug!(conn = %conn.0, count = docs.len(), "dropping connection's downstreams"); for doc in docs { self.routes.remove(&(conn.clone(), doc)); } @@ -189,7 +185,7 @@ impl Router { .await .map_err(Into::into) .inspect_err(|error: &anyhow::Error| { - warn!(error = ?error, conn = %conn.conn, "failed to deliver frame to client"); + warn!(error = ?error, conn = %conn.0, "failed to deliver frame to client"); }) .ok(); } diff --git a/services/sync-router/src/domain/router/test.rs b/crates/sync_router/src/domain/router/test.rs similarity index 82% rename from services/sync-router/src/domain/router/test.rs rename to crates/sync_router/src/domain/router/test.rs index 903ed89c2e5..30dc31a7ffc 100644 --- a/services/sync-router/src/domain/router/test.rs +++ b/crates/sync_router/src/domain/router/test.rs @@ -4,10 +4,7 @@ use bebop::{Record, SliceWrapper}; use sync_service_bebop_schema::ToRouter; fn conn(gateway: &str, id: &str) -> ConnectionId { - ConnectionId { - gateway: gateway.into(), - conn: id.into(), - } + ConnectionId(format!("{gateway}:{id}")) } fn frame_event(c: &ConnectionId, inner: ToRouter<'_>) -> Event { @@ -54,7 +51,7 @@ fn recording_factory( .returning(move |_conn, doc, _token, _epoch| { let (tx, rx) = mpsc::channel(capacity); record.lock().unwrap().push((doc, rx)); - tx + Downstream::new(tx, CancellationToken::new()) }); (factory, opened) } @@ -117,34 +114,6 @@ async fn disconnect_tears_down_every_route_for_the_connection() { router.handle(doc_frame(&c, "doc-a", b"late")).await; } -#[tokio::test] -async fn gateway_lost_drops_only_that_gateways_connections() { - let (factory, opened) = recording_factory(8); - let mut router = Router::new(std::sync::Arc::new(MockEdgeSink::new()), factory); - - let dead = conn("g1", "c1"); - let alive = conn("g2", "c2"); - router.handle(subscribe(&dead, "doc-a")).await; - router.handle(subscribe(&alive, "doc-b")).await; - router - .handle(Event::Edge(EdgeEvent::GatewayLost { - gateway: "g1".into(), - })) - .await; - - let mut opened = opened.lock().unwrap(); - let (_, dead_rx) = &mut opened[0]; - assert!(matches!( - dead_rx.try_recv(), - Err(mpsc::error::TryRecvError::Disconnected) - )); - let (_, alive_rx) = &mut opened[1]; - assert!(matches!( - alive_rx.try_recv(), - Err(mpsc::error::TryRecvError::Empty) - )); -} - #[tokio::test] async fn unknown_doc_frames_and_garbage_are_dropped_without_panic() { let (factory, _opened) = recording_factory(8); @@ -181,3 +150,28 @@ async fn downstream_closed_forgets_the_route_so_resubscribe_reopens() { assert_eq!(opened.lock().unwrap().len(), 2); } + +#[tokio::test] +async fn full_downstream_queue_closes_the_route_instead_of_corrupting_chunk_framing() { + let (factory, opened) = recording_factory(1); + let mut sink = MockEdgeSink::new(); + sink.expect_deliver().times(1).returning(|_, bytes| { + assert_eq!( + bytes, + envelope::doc_closed("doc-a", "downstream queue unavailable") + ); + Box::pin(async { Ok(()) }) + }); + let mut router = Router::new(std::sync::Arc::new(sink), factory); + let c = conn("g1", "c1"); + router.handle(subscribe(&c, "doc-a")).await; + router.handle(doc_frame(&c, "doc-a", &[0, 1])).await; + router.handle(doc_frame(&c, "doc-a", &[1, 2])).await; + assert!( + !router + .routes + .contains_key(&(c.clone(), DocId("doc-a".into()))) + ); + router.handle(subscribe(&c, "doc-a")).await; + assert_eq!(opened.lock().unwrap().len(), 2); +} diff --git a/services/sync-router/src/lib.rs b/crates/sync_router/src/lib.rs similarity index 89% rename from services/sync-router/src/lib.rs rename to crates/sync_router/src/lib.rs index 3007b904e33..bd85d8caebf 100644 --- a/services/sync-router/src/lib.rs +++ b/crates/sync_router/src/lib.rs @@ -1,6 +1,6 @@ #![deny(missing_docs)] //! The sync-router: forwards sync-protocol frames between the connection -//! gateway's Redis fanout and a per-document downstream. +//! gateway's local sockets and a per-document downstream. //! //! Chapter 1: the downstream is the existing Cloudflare Durable Object //! sync-service, dialed per `(connection, document)` so the DO sees each user @@ -13,5 +13,4 @@ //! router core does not change. pub mod domain; -pub mod inbound; pub mod outbound; diff --git a/services/sync-router/src/outbound/do_downstream.rs b/crates/sync_router/src/outbound/do_downstream.rs similarity index 72% rename from services/sync-router/src/outbound/do_downstream.rs rename to crates/sync_router/src/outbound/do_downstream.rs index 3dae0ce9514..5b086217d89 100644 --- a/services/sync-router/src/outbound/do_downstream.rs +++ b/crates/sync_router/src/outbound/do_downstream.rs @@ -7,7 +7,7 @@ use crate::domain::envelope; use crate::domain::models::{ConnectionId, DocId, Event}; -use crate::domain::ports::{DownstreamFactory, EdgeSink}; +use crate::domain::ports::{Downstream, DownstreamFactory, EdgeSink}; use futures::{SinkExt, StreamExt}; use opentelemetry::trace::TraceContextExt; use std::sync::Arc; @@ -15,11 +15,13 @@ use std::time::Duration; use tokio::sync::mpsc; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message; -use tracing::{debug, warn}; +use tokio_util::{sync::CancellationToken, task::TaskTracker}; +use tracing::{Instrument, debug, warn}; use tracing_opentelemetry::OpenTelemetrySpanExt; +use url::Url; /// How many client frames may buffer while the dial is in flight (or the -/// upstream is slow). Overflow drops frames; the sync protocol self-heals. +/// upstream is slow). Overflow closes the route so the client resynchronizes. const DOWNSTREAM_BUFFER: usize = 256; /// Keep the upstream (and its Durable Object) alive, mirroring the ping the @@ -28,61 +30,74 @@ const UPSTREAM_PING_INTERVAL: Duration = Duration::from_secs(30); /// Dials `{base_url}/document/{doc}/connect?token=…&traceparent=…`. pub struct DoDownstreamFactory { - base_url: String, + base_url: Url, sink: Arc, events: mpsc::Sender, + cancel: CancellationToken, + tasks: TaskTracker, } impl DoDownstreamFactory { /// `base_url` is the sync-service origin, e.g. `http://localhost:8787`. - pub fn new(base_url: String, sink: Arc, events: mpsc::Sender) -> Self { + pub fn new( + base_url: Url, + sink: Arc, + events: mpsc::Sender, + cancel: CancellationToken, + tasks: TaskTracker, + ) -> Self { Self { base_url, sink, events, + cancel, + tasks, } } } impl DownstreamFactory for DoDownstreamFactory { #[tracing::instrument(skip(self, token), fields(document.id = doc.as_str()))] - fn open( - &self, - conn: ConnectionId, - doc: DocId, - token: String, - epoch: u64, - ) -> mpsc::Sender> { + fn open(&self, conn: ConnectionId, doc: DocId, token: String, epoch: u64) -> Downstream { let (sender, receiver) = mpsc::channel(DOWNSTREAM_BUFFER); - let ws_base = self - .base_url - .replacen("http://", "ws://", 1) - .replacen("https://", "wss://", 1); - let mut url = format!("{ws_base}/document/{}/connect?token={token}", doc.as_str()); - // The DO reads `traceparent` from the query string (browser websocket - // upgrades can't carry headers, so the DO accepts it there; we use the - // same mechanism). + let mut url = self.base_url.clone(); + // The composition root validates the base; URL APIs encode document ids + // and scoped tokens so neither can change the downstream destination. + let scheme = if url.scheme() == "http" { "ws" } else { "wss" }; + url.set_scheme(scheme) + .expect("HTTP sync origin accepts a WebSocket scheme"); + url.path_segments_mut() + .expect("sync origin is a base URL") + .pop_if_empty() + .extend(["document", doc.as_str(), "connect"]); + url.query_pairs_mut().clear().append_pair("token", &token); if let Some(traceparent) = current_traceparent() { - url.push_str(&format!("&traceparent={traceparent}")); + url.query_pairs_mut() + .append_pair("traceparent", &traceparent); } - tokio::spawn(pump( - url, - conn, - doc, - epoch, - receiver, - Arc::clone(&self.sink), - self.events.clone(), - )); - sender + let cancel = self.cancel.child_token(); + let pump_cancel = cancel.clone(); + let sink = Arc::clone(&self.sink); + let events = self.events.clone(); + self.tasks.spawn( + async move { + tokio::select! { + biased; + _ = pump_cancel.cancelled() => {}, + _ = pump(url.to_string(), conn, doc, epoch, receiver, sink, events) => {}, + } + } + .in_current_span(), + ); + Downstream::new(sender, cancel) } } /// The whole life of one downstream connection: dial, then pump frames both /// ways until either side hangs up. -#[tracing::instrument(skip_all, fields(document.id = doc.as_str(), conn = %conn.conn))] +#[tracing::instrument(skip_all, fields(document.id = doc.as_str(), conn = %conn.0))] async fn pump( url: String, conn: ConnectionId, @@ -106,10 +121,10 @@ async fn pump( } }; - let upstream = match connect_async(&url).await { - Ok((upstream, _response)) => upstream, - Err(error) => { - warn!(error = ?error, "downstream dial failed"); + let upstream = match tokio::time::timeout(Duration::from_secs(10), connect_async(&url)).await { + Ok(Ok((upstream, _response))) => upstream, + failure => { + warn!(timed_out = failure.is_err(), "downstream dial failed"); deliver(envelope::subscribe_failed(doc.as_str(), "dial failed")).await; events .send(Event::DownstreamClosed { conn, doc, epoch }) diff --git a/services/sync-router/src/outbound/mod.rs b/crates/sync_router/src/outbound/mod.rs similarity index 80% rename from services/sync-router/src/outbound/mod.rs rename to crates/sync_router/src/outbound/mod.rs index f64e965481b..0deeaa23500 100644 --- a/services/sync-router/src/outbound/mod.rs +++ b/crates/sync_router/src/outbound/mod.rs @@ -1,4 +1,3 @@ //! Driven adapters: implementations of the domain ports. pub mod do_downstream; -pub mod redis_sink; diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index 8c9e6580c73..3287dc3af9c 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -65,27 +65,18 @@ services: # ============================================================================ # connection_gateway (Port 8082) - # WebSocket gateway - Redis only + # WebSocket gateway and embedded document sync router # ============================================================================ - sync_router: - <<: [*common-env, *rust-services-image] - command: ["/app/out/sync_router"] - environment: - - SYNC_SERVICE_URL=http://sync-service:8787 - depends_on: - - redis - - sync_service - networks: - databases: - services: - connection_gateway: <<: [*common-env, *rust-services-image] command: ["/app/out/connection_gateway_service"] ports: - "8082:8080" + environment: + - OVERRIDE_SYNC_SERVICE_URL=http://sync-service:8787 depends_on: - redis + - sync_service networks: databases: services: diff --git a/docs/README.md b/docs/README.md index 0824da274f0..9c52aa355c3 100644 --- a/docs/README.md +++ b/docs/README.md @@ -13,3 +13,5 @@ Docs that belong to one app, service, or crate live next to its code. WhatsApp linked devices: [architecture](architecture/whatsapp-linked-devices.md) and [External view walkthrough](agent-guide/external.md). + +Document sync: [embedded gateway router](architecture/sync-gateway.md). diff --git a/docs/agent-guide/documents.md b/docs/agent-guide/documents.md index 6dd0d130304..55dec37625f 100644 --- a/docs/agent-guide/documents.md +++ b/docs/agent-guide/documents.md @@ -3,10 +3,10 @@ For gateway sync rollout checks, use PostHog flag `sync-via-gateway` (dev and production) and recreate document sessions after the client receives the updated flag. A reload can initially read PostHog's cached value; open sessions do not -change transport when the remote value changes. Off or unavailable uses a +change transport when the remote value changes. Off or unknown flag values use a direct sync-service WebSocket per document; on multiplexes document frames over -the connection gateway. Confirm the hosted `sync_router` is provisioned before -enabling. Local stacks can override with `VITE_SYNC_VIA_GATEWAY=true` or `false`. +the connection gateway. Confirm all gateway replicas have the embedded sync +router before enabling; no separate router service is needed. Local stacks can override with `VITE_SYNC_VIA_GATEWAY=true` or `false`. The document header has separate Share, Copy Share Link, and Side Panel buttons. They are borderless with a soft rounded background on hover. Share opens the diff --git a/docs/architecture/sync-gateway.md b/docs/architecture/sync-gateway.md new file mode 100644 index 00000000000..02c235b5ee4 --- /dev/null +++ b/docs/architecture/sync-gateway.md @@ -0,0 +1,45 @@ +# Document sync through the connection gateway + +The browser multiplexes document sync over one WebSocket when PostHog flag +`sync-via-gateway` is enabled. The existing connection gateway hosts the router +as the `sync_router` Rust library. There is no additional deployable service. + +```text +Browser: virtual document sockets + | + | WebSocket / Bebop ToRouter and FromRouter envelopes + v +connection_gateway replica + authenticated socket + -> socket-owned router + -> one downstream WebSocket per subscribed document + | + | existing chunked sync protocol and scoped document token + v +sync-service / document Durable Object + token authorization, CRDT state, persistence +``` + +Each authenticated client socket owns an independent router, outbound queue, and +set of tracked downstream tasks. Replicas do not subscribe to a shared broadcast +channel. Reconnecting to another replica creates a new router and resubscribes; +there is no shared routing state to transfer. Existing JSON gateway messages +retain their cross-instance Redis relay. + +Routing reads the document id from the outer Bebop envelope. The inner payload, +including the large-message chunk prefix, passes through unchanged. Tokens are +forwarded only to the configured sync origin; the sync service validates them. + +Unsubscribing cancels that document's downstream, including a pending handshake. +Client disconnect cancels and joins all routing work for that socket. A full +outbound socket queue closes the client connection; a full downstream queue +closes that document route. Clients reconnect and resynchronize, rather than +continuing a stream with missing chunks. Other sockets are unaffected. + +The gateway resolves `macro_service_urls::SyncServiceUrl` at startup. Hosted dev +and production use the existing sync-service origins; local Docker sets +`OVERRIDE_SYNC_SERVICE_URL=http://sync-service:8787`. No new secret or Pulumi +resource is needed. Deploy the updated gateway to every replica before enabling +`sync-via-gateway`. The frontend selects its transport when a sync source is +created; rollout updates do not move existing sessions. The env override +`VITE_SYNC_VIA_GATEWAY=true` enables local testing. diff --git a/packages/collaboration/src/sync-service/gateway.ts b/packages/collaboration/src/sync-service/gateway.ts index b19965d3f4c..8027380329d 100644 --- a/packages/collaboration/src/sync-service/gateway.ts +++ b/packages/collaboration/src/sync-service/gateway.ts @@ -28,7 +28,7 @@ import type { SyncSocket } from './socket'; * Wire shape: every frame is a Bebop `ToRouter`/`FromRouter` envelope naming * the document; the inner payloads are the same `FromPeer`/`FromRemote` * messages the direct socket speaks, unchanged. Server side, the - * connection-gateway fans inbound frames out to the sync-router, which holds + * connection-gateway routes inbound frames through its embedded router, which holds * one upstream connection per (socket, document). * * Each document gets a {@link SyncSocket}-compatible virtual socket, so @@ -98,7 +98,7 @@ export class GatewaySyncTransport { this.route(event.data); }); ws.addEventListener(WebsocketEvent.Reconnect, () => { - // A new gateway socket is a new fanout connection server-side: every + // A new gateway socket owns a new router session server-side: every // route was torn down, so every document must re-subscribe. for (const virtual of this.docs.values()) virtual.resubscribe(); }); diff --git a/services/connection_gateway/AGENTS.md b/services/connection_gateway/AGENTS.md index 2975e6ae053..854f5f4e57b 100644 --- a/services/connection_gateway/AGENTS.md +++ b/services/connection_gateway/AGENTS.md @@ -4,11 +4,13 @@ Shared infrastructure: manages per-entity WebSocket connections so other services can push messages to a specific connected entity without holding -the socket themselves. +the socket themselves. It also embeds document sync routing for each socket. ## Start here - `src/api/connection.rs`: WebSocket upgrade endpoint and connection lifecycle. +- `src/service/sync.rs`: per-socket sync router, bounded outbound delivery, + cancellation and task joining. Routing core lives in `crates/sync_router`. - `src/api/message.rs`: HTTP endpoint other services call to send a message to a connected entity. - `src/service/connection.rs` + `src/service/dynamodb.rs`: connection registry, @@ -21,9 +23,9 @@ the socket themselves. ## Scope - Owns: WebSocket connection tracking and message delivery to a connected - entity, across instances. -- Not here: the message/event payloads themselves, or who is authorized to - send to whom — callers decide that before calling in. + entity, across instances; multiplexed document sync for local sockets. +- Document state and document-token authorization remain in `sync-service`. + Other message/event payloads and send authorization remain with callers. ## Boundaries @@ -49,7 +51,12 @@ the socket themselves. - A send can target a connection on a different instance — it always goes through the Redis relay in `src/service/redis.rs`, never a direct in-process - lookup; don't bypass it when adding a new send path. + lookup. Sync replies are the exception: their router is owned by one local + socket, so they go directly to that socket's queue. +- Resolve `SyncServiceUrl` at startup; Docker uses + `OVERRIDE_SYNC_SERVICE_URL=http://sync-service:8787`. +- Complete the gateway rollout before enabling `sync-via-gateway` in PostHog; + there is no separate router deployment. See `docs/architecture/sync-gateway.md`. ## Related diff --git a/services/connection_gateway/Cargo.toml b/services/connection_gateway/Cargo.toml index f4e6d6486e7..de7771c8de9 100644 --- a/services/connection_gateway/Cargo.toml +++ b/services/connection_gateway/Cargo.toml @@ -59,9 +59,9 @@ macro_config = { path = "../../crates/macro_config" } macro_entrypoint = { path = "../../crates/macro_entrypoint" } macro_env = { path = "../../crates/macro_env" } macro_env_var = { path = "../../crates/macro_env_var" } +macro_service_urls = { path = "../../crates/macro_service_urls" } macro_tower_layers = { path = "../../crates/macro_tower_layers" } mockall.workspace = true -postcard = { workspace = true } model = { path = "../../crates/model" } model-entity = { path = "../../crates/model-entity" } models_bulk_upload = { path = "../../crates/models_bulk_upload" } @@ -79,7 +79,9 @@ sqlx = { workspace = true } strum = { workspace = true, features = ["derive"] } strum_macros = { workspace = true } stream = { path = "../../crates/stream" } +sync_router = { path = "../../crates/sync_router" } tokio = { workspace = true } +tokio-util = { workspace = true, features = ["rt"] } tower = { workspace = true } tower-http = { workspace = true } tracing = { workspace = true, features = ["attributes"] } @@ -92,5 +94,8 @@ workspace-hack = { version = "0.1", path = "../../crates/workspace-hack" } doppler-config = { workspace = true } [dev-dependencies] +bebop = { workspace = true } +sync_service_bebop_schema = { path = "../sync-service/bebop-schema" } +tokio-tungstenite = { workspace = true } cool_asserts = "2" macro_auth = { path = "../../crates/macro_auth", features = ["testing"] } diff --git a/services/connection_gateway/README.md b/services/connection_gateway/README.md index 17743775f7f..2cfc9599d8c 100644 --- a/services/connection_gateway/README.md +++ b/services/connection_gateway/README.md @@ -2,6 +2,20 @@ A service responsible for managing connections, and allowing services to send messages to a specific entity over a websocket connection. +## Document sync + +Binary WebSocket messages use the Bebop `ToRouter` / `FromRouter` envelope. +Each socket embeds `crates/sync_router`, which opens one downstream sync-service +connection per document and passes the inner chunked protocol unchanged. +Document tokens are validated by the sync service. Closing the client socket +cancels all of its downstream work, including in-flight handshakes. + +This ships in the existing gateway deployment. No additional service, Redis +fanout, or Pulumi stack is needed. `SyncServiceUrl` selects the existing dev/prod +sync origin; local Docker overrides it with `OVERRIDE_SYNC_SERVICE_URL`. +Deploy all gateway replicas before enabling PostHog `sync-via-gateway`. +See [sync gateway architecture](../../docs/architecture/sync-gateway.md). + ## Local Development *runs on port 8085* diff --git a/services/connection_gateway/src/api/connection.rs b/services/connection_gateway/src/api/connection.rs index 8b180879a41..44237a5e66e 100644 --- a/services/connection_gateway/src/api/connection.rs +++ b/services/connection_gateway/src/api/connection.rs @@ -53,8 +53,10 @@ pub async fn ws_handler( let macro_user_id = authorization.authorization.user.macro_user_id.clone(); let user_context = authorization.authorization.user.user_context.clone(); + let span = tracing::Span::current(); ws.on_upgrade(move |socket| { handle_websocket_connection(socket, ctx, config, macro_user_id, user_context) + .instrument(span) }) } @@ -72,13 +74,14 @@ async fn handle_websocket_connection( ) { let (sink, stream) = socket.split(); let (sender, receiver) = tokio::sync::mpsc::channel::(100); - let connection_id = uuid::Uuid::new_v4().to_string(); + let connection_id = uuid::Uuid::now_v7().to_string(); // Create guard that records last online time when websocket connection closes let last_online_guard = ctx.last_online_worker.new_guard(macro_user_id.clone()); let sender_connection_id = connection_id.clone(); - let sender_task = tokio::spawn(forwarder(sink, receiver, sender_connection_id)); + let mut sender_task = + tokio::spawn(forwarder(sink, receiver, sender_connection_id).in_current_span()); if let Err(err) = ctx .connection_manager @@ -103,12 +106,19 @@ async fn handle_websocket_connection( connection_id: &connection_id, }; - crate::service::fanout::connected(&connection_context, ¯o_user_id).await; - - let receiver_task = handle_websocket_stream(connection_context, stream, sender.clone()).fuse(); + let sync_session = crate::service::sync::SyncSession::new( + connection_id.clone(), + ctx.sync_service_url.clone(), + sender.clone(), + ); + let receiver_task = + handle_websocket_stream(connection_context, stream, sender.clone(), &sync_session).fuse(); tokio::select! { - res = sender_task => { + _ = sync_session.closed() => { + tracing::warn!(connection.id = %connection_id, "sync routing stopped; closing socket"); + } + res = &mut sender_task => { res.inspect(|_| tracing::debug!("sender task finished")) .inspect_err(|err| { tracing::error!(error=?err, "sender task failed"); @@ -122,8 +132,15 @@ async fn handle_websocket_connection( } } + // A blocked write must not keep the socket alive after routing stops. + sender_task.abort(); + if !sender_task.is_finished() { + let _ = sender_task.await; + } + sync_session.shutdown().await; + ctx.connection_manager - .remove_connection(&connection_id.clone()) + .remove_connection(&connection_id) .await .inspect(|_| tracing::debug!("connections removed")) .inspect_err(|e| { @@ -136,9 +153,6 @@ async fn handle_websocket_connection( }) .ok(); - // After remove_connection, so consumers never see "gone" while the - // gateway can still route to the socket. - crate::service::fanout::disconnected(&connection_context).await; drop(last_online_guard); } diff --git a/services/connection_gateway/src/api/connection/messages.rs b/services/connection_gateway/src/api/connection/messages.rs index 506ab63d0e9..ee1849863ea 100644 --- a/services/connection_gateway/src/api/connection/messages.rs +++ b/services/connection_gateway/src/api/connection/messages.rs @@ -5,7 +5,7 @@ use crate::{ tracking::{EntityConnectionExt, TrackingData}, websocket::ToWebsocketMessage, }, - service::{fanout, stream_event::handle_stream_events, tracker}, + service::{stream_event::handle_stream_events, sync::SyncSession, tracker}, }; use anyhow::{Context, Result}; use axum::extract::ws::{Message, WebSocket}; @@ -18,11 +18,13 @@ pub async fn handle_websocket_stream( connection_context: ConnectionContext<'_>, mut stream: SplitStream, sender: Sender, + sync_session: &SyncSession, ) -> Result<()> { while let Some(msg) = stream.next().await { match msg { Ok(msg) => { - if let Err(e) = handle_message(connection_context, msg, &sender).await { + if let Err(e) = handle_message(connection_context, msg, &sender, sync_session).await + { tracing::warn!(error=?e, "error handling message"); } } @@ -49,24 +51,10 @@ pub async fn handle_message( connection_context: ConnectionContext<'_>, message: Message, sender: &Sender, + sync_session: &SyncSession, ) -> Result<()> { - // Every client frame is republished to the fanout channel, unparsed, - // before any in-gateway handling. Backend services consume what they care - // about (the sync tier reads binary frames); the gateway stays a dumb tap. - // - // TODO: the in-gateway handling below (entity tracking / frecency, - // stream-event subscriptions) should probably migrate to services - // consuming this same fanout, keeping the gateway a pure edge. - match &message { - Message::Text(text) => { - fanout::frame(&connection_context, true, text.as_bytes().to_vec()).await; - } - Message::Binary(payload) => { - fanout::frame(&connection_context, false, payload.to_vec()).await; - // Binary frames have no in-gateway handling; fanout is it. - return Ok(()); - } - _ => {} + if let Message::Binary(payload) = message { + return sync_session.frame(payload.to_vec()).await; } let text_message = match message { diff --git a/services/connection_gateway/src/context.rs b/services/connection_gateway/src/context.rs index 2b48abdd894..eaf518aefda 100644 --- a/services/connection_gateway/src/context.rs +++ b/services/connection_gateway/src/context.rs @@ -18,9 +18,8 @@ pub type AuthorizationService = MacroAuthorizationServiceImpl, + /// Downstream sync origin, selected once at startup for this environment. + pub sync_service_url: url::Url, pub frecency_ingestor_service: EventIngestorImpl, pub redis_client: Arc, pub redis_connection: MultiplexedConnection, diff --git a/services/connection_gateway/src/main.rs b/services/connection_gateway/src/main.rs index 2ef4f22f891..8145ebc9b50 100644 --- a/services/connection_gateway/src/main.rs +++ b/services/connection_gateway/src/main.rs @@ -114,10 +114,17 @@ async fn main() -> Result<()> { let stream_service = RedisPostgresStreamRepo::new((*redis_client).clone(), pgpool.clone()); let stream_manager = RedisPostgresStreamManager::new(stream_service.obj()); - let fanout_gateway_id: Arc = Arc::from(uuid::Uuid::new_v4().to_string()); + let sync_service_url = url::Url::parse( + macro_service_urls::SyncServiceUrl::new_for_environment(config.environment)?.as_ref(), + ) + .context("invalid sync service URL")?; + anyhow::ensure!( + matches!(sync_service_url.scheme(), "http" | "https") && sync_service_url.has_host(), + "sync service URL must be an HTTP(S) origin", + ); let context = context::ApiContext { connection_manager, - fanout_gateway_id: Arc::clone(&fanout_gateway_id), + sync_service_url, redis_client: Arc::clone(&redis_client), redis_connection, frecency_ingestor_service: EventIngestorImpl::new(FrecencyPgStorage::new(pgpool.clone())), @@ -126,11 +133,6 @@ async fn main() -> Result<()> { }; tokio::spawn(poll_messages(context.clone())); - tokio::spawn(service::fanout::poll_outbound(context.clone())); - tokio::spawn(service::fanout::heartbeat_loop( - context.redis_connection.clone(), - fanout_gateway_id, - )); let config = Arc::new(config); let authorization_state = MacroAuthorizationState::new(Arc::new(AuthorizationService::new( diff --git a/services/connection_gateway/src/model/message.rs b/services/connection_gateway/src/model/message.rs index e3267430568..a99f692bfdc 100644 --- a/services/connection_gateway/src/model/message.rs +++ b/services/connection_gateway/src/model/message.rs @@ -60,7 +60,7 @@ impl TryFrom for Message { pub enum OutgoingMessage { Pong, Message(Message), - /// A raw binary frame (sync-protocol traffic from the fanout's consumers), + /// A raw binary frame from this socket's embedded sync router, /// sent to the client as a binary websocket message, unparsed. Binary(Vec), } diff --git a/services/connection_gateway/src/service/fanout.rs b/services/connection_gateway/src/service/fanout.rs deleted file mode 100644 index 51a0f385b58..00000000000 --- a/services/connection_gateway/src/service/fanout.rs +++ /dev/null @@ -1,166 +0,0 @@ -//! Republish this instance's entire inbound websocket traffic to Redis. -//! -//! Every accepted connection, every client frame (text or binary, unparsed), -//! and every disconnect is published on -//! [`connection_gateway_models::fanout::INBOUND_CHANNEL`], so backend services -//! can consume client traffic without the gateway knowing they exist. The -//! sync tier is the first consumer. -//! -//! Publishing is fire-and-forget by design: a publish failure is logged and -//! never fails the connection's own message handling. - -use crate::context::ApiContext; -use crate::model::connection::ConnectionContext; -use anyhow::{Context, Result}; -use connection_gateway_models::fanout::{ - ConnId, FromGateway, GatewayId, HEARTBEAT_INTERVAL_SECS, INBOUND_CHANNEL, ToGateway, - outbound_channel, -}; -use futures::StreamExt; -use macro_user_id::user_id::MacroUserIdStr; -use redis::AsyncCommands; - -async fn publish(ctx: &ConnectionContext<'_>, message: &FromGateway) -> Result<()> { - let payload = postcard::to_stdvec(message).context("failed to encode fanout message")?; - let mut connection = ctx.api_context.get_multiplexed_async_connection()?; - connection - .publish::<&str, &[u8], ()>(INBOUND_CHANNEL, &payload) - .await - .context("failed to publish fanout message")?; - Ok(()) -} - -/// Log-and-continue wrapper: fanout must never break connection handling. -async fn publish_best_effort(ctx: &ConnectionContext<'_>, message: &FromGateway) { - publish(ctx, message) - .await - .inspect_err(|error| tracing::warn!(error=?error, "failed to publish fanout message")) - .ok(); -} - -/// Announce an accepted, authenticated connection. -pub async fn connected(ctx: &ConnectionContext<'_>, user_id: &MacroUserIdStr<'static>) { - publish_best_effort( - ctx, - &FromGateway::Connected { - gateway: GatewayId(ctx.api_context.fanout_gateway_id.to_string()), - conn: ConnId(ctx.connection_id.to_string()), - user_id: user_id.clone(), - }, - ) - .await; -} - -/// Forward one client frame, unparsed. -pub async fn frame(ctx: &ConnectionContext<'_>, text: bool, payload: Vec) { - publish_best_effort( - ctx, - &FromGateway::Frame { - gateway: GatewayId(ctx.api_context.fanout_gateway_id.to_string()), - conn: ConnId(ctx.connection_id.to_string()), - text, - payload, - }, - ) - .await; -} - -/// Announce a closed connection. -pub async fn disconnected(ctx: &ConnectionContext<'_>) { - publish_best_effort( - ctx, - &FromGateway::Disconnected { - gateway: GatewayId(ctx.api_context.fanout_gateway_id.to_string()), - conn: ConnId(ctx.connection_id.to_string()), - }, - ) - .await; -} - -/// Liveness beacon so consumers can drop state for dead gateway instances. -/// Spawned once at boot. -pub async fn heartbeat_loop( - redis_connection: redis::aio::MultiplexedConnection, - gateway_id: std::sync::Arc, -) { - let mut tick = tokio::time::interval(std::time::Duration::from_secs(HEARTBEAT_INTERVAL_SECS)); - loop { - tick.tick().await; - let message = FromGateway::Heartbeat { - gateway: GatewayId(gateway_id.to_string()), - }; - let Ok(payload) = postcard::to_stdvec(&message) else { - continue; - }; - let mut connection = redis_connection.clone(); - connection - .publish::<&str, &[u8], ()>(INBOUND_CHANNEL, &payload) - .await - .inspect_err(|error| tracing::warn!(error=?error, "failed to publish fanout heartbeat")) - .ok(); - } -} - -/// Consume this instance's outbound fanout channel: frames published by -/// backend consumers (the sync tier) addressed to connections whose sockets -/// this instance holds. Mirrors [`crate::service::redis::poll_messages`]. -pub async fn poll_outbound(ctx: ApiContext) -> Result<()> { - let gateway_id = GatewayId(ctx.fanout_gateway_id.to_string()); - let channel = outbound_channel(&gateway_id); - - let (mut sink, mut stream) = ctx.redis_client.get_async_pubsub().await?.split(); - sink.subscribe(&channel) - .await - .context("failed to subscribe to fanout outbound channel")?; - tracing::debug!(channel, "subscribed to fanout outbound channel"); - - while let Some(message) = stream.next().await { - let payload: Vec = match message.get_payload() { - Ok(payload) => payload, - Err(error) => { - tracing::warn!(error = ?error, "unreadable outbound fanout payload; skipping"); - continue; - } - }; - match postcard::from_bytes::(&payload) { - Ok(ToGateway::Frame { - conn, - text: false, - payload, - }) => { - ctx.connection_manager - .send_binary(&conn.0, payload) - .await - .inspect_err(|error| { - tracing::warn!(error = ?error, conn = %conn, "failed to deliver outbound frame"); - }) - .ok(); - } - Ok(ToGateway::Frame { - conn, text: true, .. - }) => { - // No consumer sends text frames yet; refuse rather than guess - // at framing for the existing JSON protocol. - tracing::warn!(conn = %conn, "text outbound frames are not supported; dropping"); - } - Ok(ToGateway::Close { conn, code }) => { - // TODO: plumb a real close (with code) through the connection - // manager; for now drop the connection, which closes the socket. - tracing::debug!(conn = %conn, code, "closing connection at consumer request"); - ctx.connection_manager - .remove_connection(&conn.0) - .await - .inspect_err(|error| { - tracing::warn!(error = ?error, conn = %conn, "failed to close connection"); - }) - .ok(); - } - Err(error) => { - tracing::warn!(error = ?error, "undecodable outbound fanout message; skipping"); - } - } - } - - tracing::warn!("fanout outbound poller exited"); - Ok(()) -} diff --git a/services/connection_gateway/src/service/mod.rs b/services/connection_gateway/src/service/mod.rs index 02eb553f0c2..6395ab8dfd4 100644 --- a/services/connection_gateway/src/service/mod.rs +++ b/services/connection_gateway/src/service/mod.rs @@ -1,7 +1,7 @@ pub mod connection; pub mod dynamodb; -pub mod fanout; pub mod redis; pub mod sender; pub mod stream_event; +pub mod sync; pub mod tracker; diff --git a/services/connection_gateway/src/service/sync.rs b/services/connection_gateway/src/service/sync.rs new file mode 100644 index 00000000000..7d52ced54dd --- /dev/null +++ b/services/connection_gateway/src/service/sync.rs @@ -0,0 +1,124 @@ +//! Per-WebSocket sync routing. No shared subscription or cross-replica relay. + +use crate::model::message::OutgoingMessage; +use anyhow::{Context, Result}; +use std::sync::Arc; +use sync_router::{ + domain::{ + models::{ConnectionId, EdgeEvent, Event}, + ports::EdgeSink, + router::Router, + }, + outbound::do_downstream::DoDownstreamFactory, +}; +use tokio::sync::mpsc; +use tokio_util::{ + sync::{CancellationToken, DropGuard}, + task::TaskTracker, +}; +use tracing::Instrument; +use url::Url; + +#[cfg(test)] +mod test; + +const ROUTER_QUEUE_CAPACITY: usize = 64; + +/// Routing resources owned by one authenticated client WebSocket. +pub struct SyncSession { + connection: ConnectionId, + events: mpsc::Sender, + cancel: CancellationToken, + tasks: TaskTracker, + _cancel_on_drop: DropGuard, +} + +impl SyncSession { + /// Start a router that can deliver only to this socket's bounded queue. + pub fn new( + connection_id: String, + base_url: Url, + sender: mpsc::Sender, + ) -> Self { + let connection = ConnectionId(connection_id); + let cancel = CancellationToken::new(); + let tasks = TaskTracker::new(); + let (events, receiver) = mpsc::channel(ROUTER_QUEUE_CAPACITY); + let sink = Arc::new(SocketSink { + connection: connection.clone(), + sender, + cancel: cancel.clone(), + }); + let downstreams = DoDownstreamFactory::new( + base_url, + Arc::clone(&sink), + events.clone(), + cancel.clone(), + tasks.clone(), + ); + let router = Router::new(sink, downstreams); + tasks.spawn(router.run(receiver, cancel.clone()).in_current_span()); + Self { + connection, + events, + _cancel_on_drop: cancel.clone().drop_guard(), + cancel, + tasks, + } + } + + /// Apply backpressure only to this connection, preserving binary frame order. + pub async fn frame(&self, payload: Vec) -> Result<()> { + self.events + .send(Event::Edge(EdgeEvent::Frame { + conn: self.connection.clone(), + payload, + })) + .await + .context("sync router stopped") + } + + /// Stop the socket if routing fails or its outbound queue stops draining. + pub async fn closed(&self) { + tokio::select! { + _ = self.cancel.cancelled() => {}, + _ = self.events.closed() => {}, + } + } + + /// Cancel dialing and established downstreams, then join every routing task. + pub async fn shutdown(self) { + self.cancel.cancel(); + self.tasks.close(); + self.tasks.wait().await; + } +} + +struct SocketSink { + connection: ConnectionId, + sender: mpsc::Sender, + cancel: CancellationToken, +} + +impl EdgeSink for SocketSink { + type Err = anyhow::Error; + + #[tracing::instrument(skip_all, fields(connection.id = %conn.0), err)] + async fn deliver(&self, conn: &ConnectionId, frame: Vec) -> Result<()> { + anyhow::ensure!( + conn == &self.connection, + "sync frame targets another connection" + ); + if self + .sender + .try_send(OutgoingMessage::Binary(frame)) + .is_err() + { + // Dropping a chunk corrupts framing. Close this socket so its virtual + // sockets reconnect and resync instead of losing a partial message. + self.cancel.cancel(); + anyhow::bail!("sync client outbound queue is full or closed"); + } + Ok(()) + } +} diff --git a/services/connection_gateway/src/service/sync/test.rs b/services/connection_gateway/src/service/sync/test.rs new file mode 100644 index 00000000000..8d50929c131 --- /dev/null +++ b/services/connection_gateway/src/service/sync/test.rs @@ -0,0 +1,280 @@ +use super::*; +use bebop::{Record, SliceWrapper}; +use futures::{SinkExt, StreamExt}; +use sync_service_bebop_schema::{FromRouter, ToRouter}; +use tokio::{io::AsyncReadExt, net::TcpListener, sync::oneshot}; +use tokio_tungstenite::{accept_async, tungstenite::Message}; + +fn encode(message: ToRouter<'_>) -> Vec { + let mut bytes = Vec::new(); + message.serialize(&mut bytes).unwrap(); + bytes +} + +async fn next_frame(receiver: &mut mpsc::Receiver) -> Vec { + let OutgoingMessage::Binary(bytes) = receiver.recv().await.unwrap() else { + panic!("expected a binary sync envelope"); + }; + bytes +} + +#[tokio::test] +async fn separate_gateway_sessions_route_the_same_document_without_cross_delivery() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = Url::parse(&format!("http://{}", listener.local_addr().unwrap())).unwrap(); + let server = tokio::spawn(async move { + let (first, _) = listener.accept().await.unwrap(); + let mut first = accept_async(first).await.unwrap(); + let (second, _) = listener.accept().await.unwrap(); + let mut second = accept_async(second).await.unwrap(); + let first_bytes = loop { + if let Message::Binary(bytes) = first.next().await.unwrap().unwrap() { + break bytes; + } + }; + let second_bytes = loop { + if let Message::Binary(bytes) = second.next().await.unwrap().unwrap() { + break bytes; + } + }; + assert_eq!(first_bytes.as_ref(), &[0, 11, 12]); + assert_eq!(second_bytes.as_ref(), &[1, 21, 22]); + first.send(Message::Binary(first_bytes)).await.unwrap(); + second.send(Message::Binary(second_bytes)).await.unwrap(); + // Session shutdown must drop both downstream sockets. + let _ = first.next().await; + let _ = second.next().await; + }); + let (first_sender, mut first_receiver) = mpsc::channel(10); + let (second_sender, mut second_receiver) = mpsc::channel(10); + let first = SyncSession::new("gateway-a/socket-a".into(), base_url.clone(), first_sender); + let second = SyncSession::new("gateway-b/socket-b".into(), base_url, second_sender); + first + .frame(encode(ToRouter::RouterSubscribe { + doc_id: "same-document", + token: "first-token", + })) + .await + .unwrap(); + assert!(matches!( + FromRouter::deserialize(&next_frame(&mut first_receiver).await).unwrap(), + FromRouter::RouterSubscribed { + doc_id: "same-document" + } + )); + second + .frame(encode(ToRouter::RouterSubscribe { + doc_id: "same-document", + token: "second-token", + })) + .await + .unwrap(); + assert!(matches!( + FromRouter::deserialize(&next_frame(&mut second_receiver).await).unwrap(), + FromRouter::RouterSubscribed { + doc_id: "same-document" + } + )); + first + .frame(encode(ToRouter::RouterFrame { + doc_id: "same-document", + payload: SliceWrapper::Raw(&[0, 11, 12]), + })) + .await + .unwrap(); + second + .frame(encode(ToRouter::RouterFrame { + doc_id: "same-document", + payload: SliceWrapper::Raw(&[1, 21, 22]), + })) + .await + .unwrap(); + assert_eq!( + next_frame(&mut first_receiver).await, + sync_router::domain::envelope::doc_frame("same-document", &[0, 11, 12]) + ); + assert_eq!( + next_frame(&mut second_receiver).await, + sync_router::domain::envelope::doc_frame("same-document", &[1, 21, 22]) + ); + first.shutdown().await; + second.shutdown().await; + server.await.unwrap(); +} + +#[tokio::test] +async fn closing_a_session_cancels_a_downstream_that_never_completes_its_handshake() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = Url::parse(&format!("http://{}", listener.local_addr().unwrap())).unwrap(); + let (accepted, accepted_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + accepted.send(()).unwrap(); + let mut bytes = Vec::new(); + socket.read_to_end(&mut bytes).await.unwrap(); + }); + let (sender, _receiver) = mpsc::channel(10); + let session = SyncSession::new("socket".into(), base_url, sender); + session + .frame(encode(ToRouter::RouterSubscribe { + doc_id: "document", + token: "token", + })) + .await + .unwrap(); + accepted_rx.await.unwrap(); + session.shutdown().await; + tokio::time::timeout(std::time::Duration::from_secs(2), server) + .await + .unwrap() + .unwrap(); +} + +#[tokio::test] +async fn unsubscribe_cancels_a_pending_dial_without_stopping_the_session() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = Url::parse(&format!("http://{}", listener.local_addr().unwrap())).unwrap(); + let (accepted, accepted_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + accepted.send(()).unwrap(); + socket.read_to_end(&mut Vec::new()).await.unwrap(); + }); + let (sender, _receiver) = mpsc::channel(10); + let session = SyncSession::new("socket".into(), base_url, sender); + session + .frame(encode(ToRouter::RouterSubscribe { + doc_id: "document", + token: "token", + })) + .await + .unwrap(); + accepted_rx.await.unwrap(); + session + .frame(encode(ToRouter::RouterUnsubscribe { doc_id: "document" })) + .await + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), server) + .await + .unwrap() + .unwrap(); + assert!(!session.cancel.is_cancelled()); + session.shutdown().await; +} + +#[tokio::test] +async fn full_outbound_queue_cancels_only_its_owner_instead_of_dropping_a_chunk() { + let (sender, _receiver) = mpsc::channel(1); + sender.send(OutgoingMessage::Pong).await.unwrap(); + let cancel = CancellationToken::new(); + let other_cancel = CancellationToken::new(); + let conn = ConnectionId("socket".into()); + let sink = SocketSink { + connection: conn.clone(), + sender, + cancel: cancel.clone(), + }; + assert!(sink.deliver(&conn, vec![1, 2, 3]).await.is_err()); + assert!(cancel.is_cancelled()); + assert!(!other_cancel.is_cancelled()); +} + +#[tokio::test] +async fn socket_sink_rejects_frames_addressed_to_another_connection() { + let (sender, mut receiver) = mpsc::channel(1); + let cancel = CancellationToken::new(); + let sink = SocketSink { + connection: ConnectionId("owner".into()), + sender, + cancel: cancel.clone(), + }; + assert!( + sink.deliver(&ConnectionId("other".into()), vec![1]) + .await + .is_err() + ); + assert!(receiver.try_recv().is_err()); + assert!(!cancel.is_cancelled()); +} + +#[tokio::test] +async fn rejected_document_token_never_opens_a_virtual_socket() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = Url::parse(&format!("http://{}", listener.local_addr().unwrap())).unwrap(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let result = tokio_tungstenite::accept_hdr_async( + socket, + |_request: &tokio_tungstenite::tungstenite::handshake::server::Request, + _response: tokio_tungstenite::tungstenite::handshake::server::Response| { + Err(tokio_tungstenite::tungstenite::http::Response::builder() + .status(401) + .body(Some("rejected".into())) + .unwrap()) + }, + ) + .await; + assert!(result.is_err()); + }); + let (sender, mut receiver) = mpsc::channel(10); + let session = SyncSession::new("socket".into(), base_url, sender); + session + .frame(encode(ToRouter::RouterSubscribe { + doc_id: "document", + token: "invalid-token", + })) + .await + .unwrap(); + assert!(matches!( + FromRouter::deserialize(&next_frame(&mut receiver).await).unwrap(), + FromRouter::RouterSubscribeFailed { + doc_id: "document", + .. + } + )); + session.shutdown().await; + server.await.unwrap(); +} + +#[tokio::test] +async fn scoped_tokens_and_document_ids_cannot_change_the_downstream_url() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = Url::parse(&format!("http://{}", listener.local_addr().unwrap())).unwrap(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut socket = tokio_tungstenite::accept_hdr_async( + socket, + |request: &tokio_tungstenite::tungstenite::handshake::server::Request, + response: tokio_tungstenite::tungstenite::handshake::server::Response| { + assert_eq!( + request.uri().path(), + "/document/doc%2Fother%3Foverride/connect" + ); + let url = Url::parse(&format!("http://localhost{}", request.uri())).unwrap(); + assert_eq!( + url.query_pairs().collect::>(), + vec![("token".into(), "scoped&token=secret+value".into())] + ); + Ok(response) + }, + ) + .await + .unwrap(); + let _ = socket.next().await; + }); + let (sender, mut receiver) = mpsc::channel(10); + let session = SyncSession::new("socket".into(), base_url, sender); + session + .frame(encode(ToRouter::RouterSubscribe { + doc_id: "doc/other?override", + token: "scoped&token=secret+value", + })) + .await + .unwrap(); + assert!(matches!( + FromRouter::deserialize(&next_frame(&mut receiver).await).unwrap(), + FromRouter::RouterSubscribed { .. } + )); + session.shutdown().await; + server.await.unwrap(); +} diff --git a/services/sync-router/AGENTS.md b/services/sync-router/AGENTS.md deleted file mode 100644 index dbff05cd6ff..00000000000 --- a/services/sync-router/AGENTS.md +++ /dev/null @@ -1,46 +0,0 @@ -# sync_router - -## What it is - -Routes document collaboration from the connection gateway's Redis fanout to -Cloudflare Durable Object sync sockets. It owns routing, not document state. - -## Start here - -- `src/main.rs`: Redis and downstream adapter wiring, subscriber restart loop. -- `src/domain/router.rs`: one owner for routes keyed by gateway, connection, - and document; route epochs protect replacements from stale close events. -- `src/inbound/redis_edge.rs`: binary-frame filtering and gateway liveness. -- `src/outbound/do_downstream.rs`: downstream dialing, buffering, and replies. - -## Scope - -- Owns subscription routing and teardown when clients or gateways disappear. -- Document authorization and state belong to `services/sync-service`; forward - the scoped token to that service rather than interpreting it here. -- Redis wire types belong to `crates/connection_gateway_models`; browser - envelopes are generated from `services/sync-service/bebop/schema.bop`. - -## Checks - -- `cargo test -p sync_router` (domain and adapter tests need no live services). -- `cargo fmt --all -- --check` and `just clippy`. -- After changing dependencies, `cargo x deps --check`. - -## Gotchas - -- `REDIS_HOST` and `SYNC_SERVICE_URL` are required at startup. -- Document payloads retain the sync protocol's chunk framing. The router does - not decode or reassemble CRDT updates. -- Redis pub/sub is best effort. Do not assume replay or durable delivery. -- Every router subscribes to the same inbound channel; multiple instances do - not partition work. Add ownership coordination before horizontal scaling. -- Local Docker and service inventory wiring are not a hosted deployment. - Provision the router before enabling the `sync-via-gateway` PostHog flag. - Deployed dev and production defer to PostHog; local stacks can opt in with - `VITE_SYNC_VIA_GATEWAY=true`. - -## Related - -- Architecture: `../../.agents/skills/cloud-storage-hexagonal-architecture/SKILL.md` -- Rules: `../../.agents/rules/observability.md` diff --git a/services/sync-router/src/inbound/mod.rs b/services/sync-router/src/inbound/mod.rs deleted file mode 100644 index 247ae8a6866..00000000000 --- a/services/sync-router/src/inbound/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -//! Driving adapters: sources of [`crate::domain::models::Event`]. - -pub mod redis_edge; diff --git a/services/sync-router/src/inbound/redis_edge.rs b/services/sync-router/src/inbound/redis_edge.rs deleted file mode 100644 index 8bcd0896646..00000000000 --- a/services/sync-router/src/inbound/redis_edge.rs +++ /dev/null @@ -1,129 +0,0 @@ -//! Consume the connection gateway's Redis fanout and translate it into -//! [`EdgeEvent`]s. All Redis-isms (heartbeats, gateway liveness, text-frame -//! filtering) are absorbed here; the router core never sees them. - -#[cfg(test)] -mod test; - -use crate::domain::models::{ConnectionId, EdgeEvent, Event}; -use anyhow::{Context, Result}; -use connection_gateway_models::fanout::{ - FromGateway, GatewayId, HEARTBEAT_INTERVAL_SECS, INBOUND_CHANNEL, -}; -use futures::StreamExt; -use std::collections::HashMap; -use std::time::Duration; -use tokio::sync::mpsc; -use tokio::time::Instant; -use tracing::{debug, warn}; - -/// A gateway instance is declared dead after this long without a heartbeat -/// (or any other message from it). -const GATEWAY_STALE_AFTER: Duration = Duration::from_secs(HEARTBEAT_INTERVAL_SECS * 3); - -/// Subscribe to the fanout channel and pump events until Redis drops us. -/// The caller is expected to reconnect-and-rerun on exit. -pub async fn run(redis_client: &redis::Client, events: mpsc::Sender) -> Result<()> { - let mut pubsub = redis_client - .get_async_pubsub() - .await - .context("failed to open redis pubsub connection")?; - pubsub - .subscribe(INBOUND_CHANNEL) - .await - .context("failed to subscribe to fanout channel")?; - debug!(channel = INBOUND_CHANNEL, "subscribed to gateway fanout"); - - let mut liveness: HashMap = HashMap::new(); - let mut sweep = tokio::time::interval(Duration::from_secs(HEARTBEAT_INTERVAL_SECS)); - let mut stream = pubsub.on_message(); - - loop { - tokio::select! { - message = stream.next() => { - let Some(message) = message else { - // Redis connection died; surface it so main reconnects. - anyhow::bail!("fanout pubsub stream ended"); - }; - let payload: Vec = match message.get_payload() { - Ok(payload) => payload, - Err(error) => { - warn!(error = ?error, "unreadable fanout payload; skipping"); - continue; - } - }; - let decoded: FromGateway = match postcard::from_bytes(&payload) { - Ok(decoded) => decoded, - Err(error) => { - warn!(error = ?error, "undecodable fanout message; skipping"); - continue; - } - }; - if let Some(event) = translate(decoded, &mut liveness) - && events.send(event).await.is_err() - { - return Ok(()); // router shut down - } - } - _ = sweep.tick() => { - let now = Instant::now(); - let stale: Vec = liveness - .iter() - .filter(|(_, seen)| now.duration_since(**seen) > GATEWAY_STALE_AFTER) - .map(|(gateway, _)| gateway.clone()) - .collect(); - for gateway in stale { - warn!(gateway = %gateway, "gateway heartbeat went quiet; dropping its connections"); - liveness.remove(&gateway); - if events - .send(Event::Edge(EdgeEvent::GatewayLost { gateway })) - .await - .is_err() - { - return Ok(()); - } - } - } - } - } -} - -/// Turn one fanout message into an edge event, updating gateway liveness. -/// Returns `None` for messages the router doesn't act on (text frames, -/// heartbeats, `Connected` — the router keys everything off `Subscribe`). -fn translate(message: FromGateway, liveness: &mut HashMap) -> Option { - let mark = |liveness: &mut HashMap, gateway: &GatewayId| { - liveness.insert(gateway.clone(), Instant::now()); - }; - match message { - FromGateway::Heartbeat { gateway } => { - mark(liveness, &gateway); - None - } - FromGateway::Connected { gateway, .. } => { - mark(liveness, &gateway); - None - } - FromGateway::Frame { - gateway, - conn, - text, - payload, - } => { - mark(liveness, &gateway); - if text { - return None; // existing gateway JSON traffic, not ours - } - Some(Event::Edge(EdgeEvent::Frame { - conn: ConnectionId { gateway, conn }, - payload, - })) - } - FromGateway::Disconnected { gateway, conn } => { - mark(liveness, &gateway); - Some(Event::Edge(EdgeEvent::Disconnected { - conn: ConnectionId { gateway, conn }, - })) - } - } -} diff --git a/services/sync-router/src/inbound/redis_edge/test.rs b/services/sync-router/src/inbound/redis_edge/test.rs deleted file mode 100644 index 6b581657e86..00000000000 --- a/services/sync-router/src/inbound/redis_edge/test.rs +++ /dev/null @@ -1,75 +0,0 @@ -use super::*; -use connection_gateway_models::fanout::ConnId; - -fn gateway(id: &str) -> GatewayId { - GatewayId(id.into()) -} - -#[test] -fn text_frames_are_filtered_but_still_mark_liveness() { - let mut liveness = HashMap::new(); - let event = translate( - FromGateway::Frame { - gateway: gateway("g1"), - conn: ConnId("c1".into()), - text: true, - payload: b"{}".to_vec(), - }, - &mut liveness, - ); - assert!(event.is_none()); - assert!(liveness.contains_key(&gateway("g1"))); -} - -#[test] -fn binary_frames_become_edge_frames() { - let mut liveness = HashMap::new(); - let event = translate( - FromGateway::Frame { - gateway: gateway("g1"), - conn: ConnId("c1".into()), - text: false, - payload: vec![1, 2], - }, - &mut liveness, - ); - match event { - Some(Event::Edge(EdgeEvent::Frame { conn, payload })) => { - assert_eq!(conn.gateway, gateway("g1")); - assert_eq!(conn.conn, ConnId("c1".into())); - assert_eq!(payload, vec![1, 2]); - } - other => panic!("wrong event: {other:?}"), - } -} - -#[test] -fn heartbeats_and_connected_mark_liveness_without_events() { - let mut liveness = HashMap::new(); - assert!( - translate( - FromGateway::Heartbeat { - gateway: gateway("g1") - }, - &mut liveness - ) - .is_none() - ); - assert!(liveness.contains_key(&gateway("g1"))); -} - -#[test] -fn disconnected_becomes_an_edge_event() { - let mut liveness = HashMap::new(); - let event = translate( - FromGateway::Disconnected { - gateway: gateway("g1"), - conn: ConnId("c1".into()), - }, - &mut liveness, - ); - assert!(matches!( - event, - Some(Event::Edge(EdgeEvent::Disconnected { .. })) - )); -} diff --git a/services/sync-router/src/main.rs b/services/sync-router/src/main.rs deleted file mode 100644 index 4ec89b66e28..00000000000 --- a/services/sync-router/src/main.rs +++ /dev/null @@ -1,57 +0,0 @@ -//! sync-router: forwards sync frames between the connection gateway's Redis -//! fanout and per-document downstreams (Durable Objects, chapter 1). - -use anyhow::{Context, Result}; -use macro_entrypoint::MacroEntrypoint; -use macro_env_var::env_vars; -use std::sync::Arc; -use sync_router::domain::router::Router; -use sync_router::inbound::redis_edge; -use sync_router::outbound::{do_downstream::DoDownstreamFactory, redis_sink::RedisSink}; -use tokio::sync::mpsc; - -env_vars! { - pub struct RedisHost; - pub struct SyncServiceUrl; -} - -#[tokio::main] -async fn main() -> Result<()> { - MacroEntrypoint::default().init(); - - let redis_host = RedisHost::new().context("REDIS_HOST must be provided")?; - let sync_service_url = SyncServiceUrl::new().context("SYNC_SERVICE_URL must be provided")?; - - let redis_client = - redis::Client::open(redis_host.to_string()).context("failed to create redis client")?; - let redis_connection = redis_client - .get_multiplexed_async_connection() - .await - .context("failed to connect to redis")?; - - let (events_tx, events_rx) = mpsc::channel(4096); - let sink = Arc::new(RedisSink::new(redis_connection)); - let downstreams = DoDownstreamFactory::new( - sync_service_url.to_string(), - Arc::clone(&sink), - events_tx.clone(), - ); - - let router = Router::new(sink, downstreams); - tokio::spawn(router.run(events_rx)); - - tracing::info!("sync-router started"); - - // The subscriber is the process's spine: resubscribe forever on Redis - // failure. An Ok return means the router task itself is gone, which is - // unrecoverable — exit non-zero and let the supervisor restart us. - loop { - match redis_edge::run(&redis_client, events_tx.clone()).await { - Ok(()) => anyhow::bail!("router task shut down"), - Err(error) => { - tracing::error!(error = ?error, "fanout subscriber failed; reconnecting"); - } - } - tokio::time::sleep(std::time::Duration::from_secs(1)).await; - } -} diff --git a/services/sync-router/src/outbound/redis_sink.rs b/services/sync-router/src/outbound/redis_sink.rs deleted file mode 100644 index 07066c75091..00000000000 --- a/services/sync-router/src/outbound/redis_sink.rs +++ /dev/null @@ -1,39 +0,0 @@ -//! [`EdgeSink`] over the gateway's per-instance Redis outbound channel. - -use crate::domain::models::ConnectionId; -use crate::domain::ports::EdgeSink; -use anyhow::Context; -use connection_gateway_models::fanout::{ToGateway, outbound_channel}; -use redis::AsyncCommands; - -/// Publishes `FromRouter` envelope bytes to the gateway instance that holds -/// the target connection's socket. -pub struct RedisSink { - connection: redis::aio::MultiplexedConnection, -} - -impl RedisSink { - /// Wrap an established Redis connection. - pub fn new(connection: redis::aio::MultiplexedConnection) -> Self { - Self { connection } - } -} - -impl EdgeSink for RedisSink { - type Err = anyhow::Error; - - async fn deliver(&self, conn: &ConnectionId, frame: Vec) -> Result<(), Self::Err> { - let message = ToGateway::Frame { - conn: conn.conn.clone(), - text: false, - payload: frame, - }; - let payload = postcard::to_stdvec(&message).context("failed to encode ToGateway frame")?; - let mut connection = self.connection.clone(); - connection - .publish::(outbound_channel(&conn.gateway), &payload) - .await - .context("failed to publish frame to gateway outbound channel")?; - Ok(()) - } -} diff --git a/tooling/xtask/crates/xtask_local/src/local/inventory.rs b/tooling/xtask/crates/xtask_local/src/local/inventory.rs index 3fd60c7a4d0..28ab6ca0bc7 100644 --- a/tooling/xtask/crates/xtask_local/src/local/inventory.rs +++ b/tooling/xtask/crates/xtask_local/src/local/inventory.rs @@ -116,17 +116,6 @@ pub const RUST_SERVICES: &[RustService] = &[ // (`no_rate_limit`, `return_passwordless_code`) is additive. no_default_features: false, }, - RustService { - compose_name: "sync_router", - cargo_bin: "sync_router", - package: "sync_router", - host_port: None, - path_prefix: None, - is_websocket: false, - modes: &[Mode::Local, Mode::Dev], - opt_in: false, - no_default_features: false, - }, RustService { compose_name: "connection_gateway", cargo_bin: "connection_gateway_service", diff --git a/tooling/xtask/crates/xtask_local/src/local/inventory/test.rs b/tooling/xtask/crates/xtask_local/src/local/inventory/test.rs index 9fc98cc52fd..e37d56208c4 100644 --- a/tooling/xtask/crates/xtask_local/src/local/inventory/test.rs +++ b/tooling/xtask/crates/xtask_local/src/local/inventory/test.rs @@ -3,9 +3,8 @@ use super::*; #[test] fn local_binaries_are_unique_and_complete() { let bins = local_binaries(); - assert_eq!(bins.len(), 17, "{bins:?}"); + assert_eq!(bins.len(), 16, "{bins:?}"); assert!(bins.contains(&"whatsapp_session_service")); - assert!(bins.contains(&"sync_router")); assert!(bins.contains(&"pubsub_workers")); assert!(bins.contains(&"seed_cli")); assert!(bins.contains(&"document_upload_finalizer_local_worker")); From 34ad198b4949240c504dceb077f8a43c48c1e297 Mon Sep 17 00:00:00 2001 From: Wolf Mermelstein Date: Tue, 27 Oct 2026 19:29:28 +0000 Subject: [PATCH 2/3] fix(sync-service): require an origin for gateway sync destinations --- docs/architecture/sync-gateway.md | 9 +++++---- services/connection_gateway/src/main.rs | 8 +++++++- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/docs/architecture/sync-gateway.md b/docs/architecture/sync-gateway.md index 02c235b5ee4..514c41342fa 100644 --- a/docs/architecture/sync-gateway.md +++ b/docs/architecture/sync-gateway.md @@ -39,7 +39,8 @@ continuing a stream with missing chunks. Other sockets are unaffected. The gateway resolves `macro_service_urls::SyncServiceUrl` at startup. Hosted dev and production use the existing sync-service origins; local Docker sets `OVERRIDE_SYNC_SERVICE_URL=http://sync-service:8787`. No new secret or Pulumi -resource is needed. Deploy the updated gateway to every replica before enabling -`sync-via-gateway`. The frontend selects its transport when a sync source is -created; rollout updates do not move existing sessions. The env override -`VITE_SYNC_VIA_GATEWAY=true` enables local testing. +resource is needed. Overrides must be HTTP(S) origins without credentials, +non-root paths, queries, or fragments. Deploy the updated gateway to every +replica before enabling `sync-via-gateway`. The frontend selects its transport +when a sync source is created; rollout updates do not move existing sessions. +The env override `VITE_SYNC_VIA_GATEWAY=true` enables local testing. diff --git a/services/connection_gateway/src/main.rs b/services/connection_gateway/src/main.rs index 8145ebc9b50..a835a12d9d8 100644 --- a/services/connection_gateway/src/main.rs +++ b/services/connection_gateway/src/main.rs @@ -119,7 +119,13 @@ async fn main() -> Result<()> { ) .context("invalid sync service URL")?; anyhow::ensure!( - matches!(sync_service_url.scheme(), "http" | "https") && sync_service_url.has_host(), + matches!(sync_service_url.scheme(), "http" | "https") + && sync_service_url.has_host() + && sync_service_url.path() == "/" + && sync_service_url.query().is_none() + && sync_service_url.fragment().is_none() + && sync_service_url.username().is_empty() + && sync_service_url.password().is_none(), "sync service URL must be an HTTP(S) origin", ); let context = context::ApiContext { From ec5fe2fe5609dc8cac87b765cae39ffa67274ed2 Mon Sep 17 00:00:00 2001 From: Wolf Mermelstein Date: Tue, 27 Oct 2026 23:07:23 +0000 Subject: [PATCH 3/3] feat(sync-service): route sync through dedicated DSS sockets Host socket-owned routing in DSS so the multiplexed client can use the existing deployment independently of notification traffic. Bound transport queues, drain downstream work during shutdown, reject stale route failures, and release pending acknowledgements for reconnect replay. --- .github/workspace-dep-closures.json | 13 +- Cargo.lock | 15 +- apps/web/AGENTS.md | 5 +- .../web/src/lib/core/constant/featureFlags.ts | 4 +- .../service-sync/source/helpers.test.ts | 4 +- .../service-sync/source/helpers.ts | 14 +- .../service-sync/websocket-url.test.ts | 52 ++++++ .../service-sync/websocket-url.ts | 26 +++ crates/connection_gateway_models/Cargo.toml | 4 + .../connection_gateway_models/src/fanout.rs | 134 ++++++++++++++ .../src/fanout/test.rs | 46 +++++ crates/connection_gateway_models/src/lib.rs | 2 + crates/sync_router/AGENTS.md | 15 +- crates/sync_router/src/domain/models.rs | 21 ++- crates/sync_router/src/domain/ports.rs | 49 +++++- crates/sync_router/src/domain/router.rs | 42 +++-- crates/sync_router/src/domain/router/test.rs | 49 +++++- crates/sync_router/src/lib.rs | 6 +- .../sync_router/src/outbound/do_downstream.rs | 53 ++++-- .../sync_router/src/session.rs | 84 +++++++-- .../sync_router/src/session}/test.rs | 97 ++++++++-- docker/docker-compose.yml | 5 +- docs/README.md | 2 +- docs/agent-guide/documents.md | 8 +- docs/architecture/sync-gateway.md | 46 ----- docs/architecture/sync-router.md | 56 ++++++ .../src/sync-service/generated/schema.bop | 2 +- .../src/sync-service/generated/schema.ts | 2 +- .../{gateway.test.ts => router.test.ts} | 85 ++++++++- .../sync-service/{gateway.ts => router.ts} | 40 ++--- .../src/sync-service/source.test.ts | 26 +++ .../collaboration/src/sync-service/source.ts | 17 +- services/connection_gateway/AGENTS.md | 17 +- services/connection_gateway/Cargo.toml | 7 +- services/connection_gateway/README.md | 14 -- .../connection_gateway/src/api/connection.rs | 34 ++-- .../src/api/connection/messages.rs | 26 ++- services/connection_gateway/src/context.rs | 5 +- services/connection_gateway/src/main.rs | 22 +-- .../connection_gateway/src/model/message.rs | 2 +- .../connection_gateway/src/service/fanout.rs | 166 ++++++++++++++++++ .../connection_gateway/src/service/mod.rs | 2 +- services/document_storage_service/AGENTS.md | 4 + services/document_storage_service/Cargo.toml | 6 + services/document_storage_service/src/api.rs | 15 +- .../src/api/context.rs | 1 + .../document_storage_service/src/api/sync.rs | 114 ++++++++++++ .../src/api/sync/test.rs | 123 +++++++++++++ .../document_storage_service/src/server.rs | 3 + .../bebop-schema/src/generated.rs | 2 +- .../sync-service/bebop/generated/schema.ts | 2 +- services/sync-service/bebop/schema.bop | 2 +- tooling/just/check.just | 4 +- 53 files changed, 1312 insertions(+), 283 deletions(-) create mode 100644 apps/web/src/lib/service-clients/service-sync/websocket-url.test.ts create mode 100644 apps/web/src/lib/service-clients/service-sync/websocket-url.ts create mode 100644 crates/connection_gateway_models/src/fanout.rs create mode 100644 crates/connection_gateway_models/src/fanout/test.rs rename services/connection_gateway/src/service/sync.rs => crates/sync_router/src/session.rs (57%) rename {services/connection_gateway/src/service/sync => crates/sync_router/src/session}/test.rs (73%) delete mode 100644 docs/architecture/sync-gateway.md create mode 100644 docs/architecture/sync-router.md rename packages/collaboration/src/sync-service/{gateway.test.ts => router.test.ts} (63%) rename packages/collaboration/src/sync-service/{gateway.ts => router.ts} (91%) create mode 100644 services/connection_gateway/src/service/fanout.rs create mode 100644 services/document_storage_service/src/api/sync.rs create mode 100644 services/document_storage_service/src/api/sync/test.rs diff --git a/.github/workspace-dep-closures.json b/.github/workspace-dep-closures.json index 6c31fbe4abb..ec2f9302e93 100644 --- a/.github/workspace-dep-closures.json +++ b/.github/workspace-dep-closures.json @@ -886,7 +886,6 @@ "crates/static_file", "crates/static_file_service_client", "crates/stream", - "crates/sync_router", "crates/sync_service", "crates/sync_service_client", "crates/system_properties", @@ -901,7 +900,6 @@ "services/agent_harness_service", "services/connection_gateway", "services/mcp_auth_proxy", - "services/sync-service/bebop-schema", "services/unfurl_service" ], "agent_inmem": [ @@ -4253,7 +4251,6 @@ "crates/macro_entrypoint", "crates/macro_env", "crates/macro_env_var", - "crates/macro_service_urls", "crates/macro_tower_layers", "crates/macro_user_id", "crates/macro_uuid", @@ -4272,22 +4269,22 @@ "crates/remote_env_var", "crates/secretsmanager_client", "crates/stream", - "crates/sync_router", "crates/workspace-hack", - "services/connection_gateway", - "services/sync-service/bebop-schema" + "services/connection_gateway" ], "connection_gateway_client": [ "crates/connection_gateway_client", "crates/connection_gateway_models", "crates/cowlike", "crates/macro_tower_layers", + "crates/macro_user_id", "crates/model-entity", "crates/workspace-hack" ], "connection_gateway_models": [ "crates/connection_gateway_models", "crates/cowlike", + "crates/macro_user_id", "crates/model-entity", "crates/workspace-hack" ], @@ -5337,6 +5334,7 @@ "crates/sqs_worker", "crates/static_file", "crates/static_file_service_client", + "crates/sync_router", "crates/sync_service", "crates/sync_service_client", "crates/system_properties", @@ -5352,7 +5350,8 @@ "crates/whatsapp_linked", "crates/work_feed", "crates/workspace-hack", - "services/document_storage_service" + "services/document_storage_service", + "services/sync-service/bebop-schema" ], "document_storage_service_client": [ "crates/agent", diff --git a/Cargo.lock b/Cargo.lock index b29a10ad6fc..cde7b98fbab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4876,7 +4876,6 @@ dependencies = [ "aws-sdk-dynamodb", "aws-sdk-secretsmanager", "axum", - "bebop", "chrono", "connection_gateway_client", "connection_gateway_models", @@ -4897,7 +4896,6 @@ dependencies = [ "macro_entrypoint", "macro_env", "macro_env_var", - "macro_service_urls", "macro_tower_layers", "macro_user_id", "mockall", @@ -4905,6 +4903,7 @@ dependencies = [ "model-entity", "models_bulk_upload", "opentelemetry", + "postcard", "redis", "secretsmanager_client", "serde", @@ -4914,11 +4913,7 @@ dependencies = [ "stream", "strum 0.27.2", "strum_macros 0.27.2", - "sync_router", - "sync_service_bebop_schema", "tokio", - "tokio-tungstenite 0.28.0", - "tokio-util", "tower 0.5.3", "tower-http", "tracing", @@ -4949,7 +4944,9 @@ dependencies = [ name = "connection_gateway_models" version = "0.1.0" dependencies = [ + "macro_user_id", "model-entity", + "postcard", "serde", "serde_json", "utoipa", @@ -6469,6 +6466,7 @@ dependencies = [ "axum", "axum-extra", "base64 0.22.1", + "bebop", "bot_id", "bots", "bytes", @@ -6571,6 +6569,7 @@ dependencies = [ "models_soup", "notification", "opensearch_client", + "opentelemetry", "ordered-float", "projects", "properties", @@ -6600,6 +6599,8 @@ dependencies = [ "sqs_client", "sqs_worker", "static_file_service_client", + "sync_router", + "sync_service_bebop_schema", "sync_service_client", "sync_service_hex", "system_properties", @@ -6608,10 +6609,12 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-tungstenite 0.27.0", + "tokio-tungstenite 0.28.0", "tokio-util", "tower 0.5.3", "tower-http", "tracing", + "tracing-opentelemetry", "tungstenite 0.27.0", "unicode-segmentation", "urlencoding", diff --git a/apps/web/AGENTS.md b/apps/web/AGENTS.md index 8d70eff2898..dd4aa21dc7d 100644 --- a/apps/web/AGENTS.md +++ b/apps/web/AGENTS.md @@ -127,12 +127,13 @@ Then trigger the interaction and read `window.__inst.log`. `'1,2,3' → '' → ' Document sync uses the `sync-via-gateway` PostHog flag in both deployed dev and production. Off or unknown flag values use direct sync sockets. -Local stacks with the updated connection gateway can opt in with `VITE_SYNC_VIA_GATEWAY=true`; +Local stacks with the updated DSS can opt in with `VITE_SYNC_VIA_GATEWAY=true`; `false` forces direct sync. The choice is made when a sync source is created and stays fixed for its lifetime. PostHog can serve a cached value on reload; recreate document sessions after the client receives the updated flag to change transport. This is not an immediate kill switch for open sessions. Finish deploying the updated -connection gateway before enabling; routing is embedded in that service. +DSS before enabling; the dedicated `/sync/ws` socket hosts routing independently +of the connection-gateway notification socket. The PostHog key keeps its historical name. Saved canvases use the shared Loro runtime; the stored JSON is only the original migration seed. Keep whole-file simple-save out of scene editing. Canvas diff --git a/apps/web/src/lib/core/constant/featureFlags.ts b/apps/web/src/lib/core/constant/featureFlags.ts index 2584058473d..1c4cd683aad 100644 --- a/apps/web/src/lib/core/constant/featureFlags.ts +++ b/apps/web/src/lib/core/constant/featureFlags.ts @@ -216,8 +216,8 @@ export const ENABLE_SCRIPTING = defineFlag({ default: false, }).enabled; -// Deployed dev and production defer to PostHog: the router must be provisioned -// before enabling gateway sync. Local stacks can use the env override. +// Historical key retained at the deployed PostHog boundary. This selects the +// dedicated DSS sync socket; deploy DSS before enabling. Local env override supported. export const syncViaGateway = defineFlag({ key: 'sync-via-gateway', env: 'SYNC_VIA_GATEWAY', diff --git a/apps/web/src/lib/service-clients/service-sync/source/helpers.test.ts b/apps/web/src/lib/service-clients/service-sync/source/helpers.test.ts index 02e78027f18..c57e2ba2937 100644 --- a/apps/web/src/lib/service-clients/service-sync/source/helpers.test.ts +++ b/apps/web/src/lib/service-clients/service-sync/source/helpers.test.ts @@ -23,8 +23,8 @@ vi.mock('@core/constant/featureFlags', () => ({ syncViaGateway: { key: 'sync-via-gateway', override: undefined }, isFeatureEnabled: () => false, })); -vi.mock('@service-connection/websocket-url', () => ({ - resolveWsUrl: () => 'ws://gateway.test', +vi.mock('../websocket-url', () => ({ + resolveSyncRouterWsUrl: () => 'ws://dss.test/sync/ws', })); vi.mock('@service-storage/client', () => ({ storageServiceClient: {} })); diff --git a/apps/web/src/lib/service-clients/service-sync/source/helpers.ts b/apps/web/src/lib/service-clients/service-sync/source/helpers.ts index 88b4e020afe..e5d77b4916f 100644 --- a/apps/web/src/lib/service-clients/service-sync/source/helpers.ts +++ b/apps/web/src/lib/service-clients/service-sync/source/helpers.ts @@ -7,7 +7,7 @@ import type { SyncError, TimeoutError, } from '@macro-inc/collaboration/collab/source'; -import { GatewaySyncTransport } from '@macro-inc/collaboration/sync-service/gateway'; +import { SyncRouterTransport } from '@macro-inc/collaboration/sync-service/router'; import { createSyncSocket, type SyncSocket, @@ -19,9 +19,9 @@ import { } from '@macro-inc/collaboration/sync-service/source'; import type { UrlResolver } from '@macro-inc/collaboration/websocket'; import { createWebsocketStateSignal } from '@macro-inc/collaboration/websocket/solid/state-signal'; -import { resolveWsUrl as resolveGatewayWsUrl } from '@service-connection/websocket-url'; import { storageServiceClient } from '@service-storage/client'; import type { ResultAsync } from 'neverthrow'; +import { resolveSyncRouterWsUrl } from '../websocket-url'; import type { DocumentSyncAuthorization } from './authorization'; import { guardDocumentSyncSource } from './guarded-source'; @@ -107,10 +107,10 @@ export function createSyncServiceSocket( ); } -let gatewayTransport: GatewaySyncTransport | undefined; -function syncGatewayTransport(): GatewaySyncTransport { - gatewayTransport ??= new GatewaySyncTransport(resolveGatewayWsUrl); - return gatewayTransport; +let routerTransport: SyncRouterTransport | undefined; +function syncRouterTransport(): SyncRouterTransport { + routerTransport ??= new SyncRouterTransport(resolveSyncRouterWsUrl); + return routerTransport; } async function refreshPermissionToken(documentId: string): Promise { @@ -133,7 +133,7 @@ export const createSyncServiceSource = ( // Pick a transport for this source's lifetime. Later flag changes apply to // newly created sources without interrupting an editor's active sync session. const ws: SyncSocket = isFeatureEnabled(syncViaGateway) - ? syncGatewayTransport().attach( + ? syncRouterTransport().attach( documentId, authorization ? undefined : token, authorization?.getToken ?? (() => refreshPermissionToken(documentId)) diff --git a/apps/web/src/lib/service-clients/service-sync/websocket-url.test.ts b/apps/web/src/lib/service-clients/service-sync/websocket-url.test.ts new file mode 100644 index 00000000000..9631897291f --- /dev/null +++ b/apps/web/src/lib/service-clients/service-sync/websocket-url.test.ts @@ -0,0 +1,52 @@ +import { err, ok } from 'neverthrow'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const auth = vi.hoisted(() => ({ + bearer: false, + token: vi.fn(async () => 'token+/='), + cookie: vi.fn(), +})); +vi.mock('@core/constant/featureFlags', () => ({ + get ENABLE_BEARER_TOKEN_AUTH() { + return auth.bearer; + }, +})); +vi.mock('@core/constant/servers', () => ({ + SERVER_HOSTS: { 'document-storage-service': 'https://gateway.test/dss' }, +})); +vi.mock('@core/util/fetchWithToken', () => ({ fetchToken: auth.cookie })); +vi.mock('@service-auth/fetch', () => ({ getMacroApiToken: auth.token })); + +import { resolveSyncRouterWsUrl } from './websocket-url'; + +beforeEach(() => { + vi.clearAllMocks(); + auth.bearer = false; + auth.cookie.mockResolvedValue(ok(undefined)); +}); + +describe('dedicated DSS sync socket credentials', () => { + it('preserves the DSS gateway prefix and refreshes cookies on each connection', async () => { + expect(await resolveSyncRouterWsUrl()).toBe( + 'wss://gateway.test/dss/sync/ws' + ); + await resolveSyncRouterWsUrl(); + expect(auth.cookie).toHaveBeenCalledTimes(2); + expect(auth.token).not.toHaveBeenCalled(); + }); + + it('encodes a fresh bearer token in the upgrade query', async () => { + auth.bearer = true; + const url = new URL(await resolveSyncRouterWsUrl()); + expect(url.pathname).toBe('/dss/sync/ws'); + expect(url.searchParams.get('macro-api-token')).toBe('token+/='); + expect(auth.cookie).not.toHaveBeenCalled(); + }); + + it('does not connect after a failed cookie refresh', async () => { + auth.cookie.mockResolvedValue(err([{ code: 'UNAUTHORIZED' }])); + await expect(resolveSyncRouterWsUrl()).rejects.toThrow( + 'Unable to refresh sync session' + ); + }); +}); diff --git a/apps/web/src/lib/service-clients/service-sync/websocket-url.ts b/apps/web/src/lib/service-clients/service-sync/websocket-url.ts new file mode 100644 index 00000000000..41b95e81998 --- /dev/null +++ b/apps/web/src/lib/service-clients/service-sync/websocket-url.ts @@ -0,0 +1,26 @@ +import { ENABLE_BEARER_TOKEN_AUTH } from '@core/constant/featureFlags'; +import { SERVER_HOSTS } from '@core/constant/servers'; +import { fetchToken } from '@core/util/fetchWithToken'; +import { Telemetry } from '@macro-inc/observability'; +import { getMacroApiToken } from '@service-auth/fetch'; + +/** Resolve fresh account credentials for the dedicated DSS sync socket. */ +export async function resolveSyncRouterWsUrl(): Promise { + return Telemetry.span('sync.router.connect', async (span) => { + const url = new URL( + `${SERVER_HOSTS['document-storage-service'].replace(/\/$/, '')}/sync/ws` + ); + url.protocol = url.protocol === 'http:' ? 'ws:' : 'wss:'; + if (ENABLE_BEARER_TOKEN_AUTH) { + const token = await getMacroApiToken(); + if (!token) throw new Error('No Macro API token'); + url.searchParams.set('macro-api-token', token); + } else { + const refreshed = await fetchToken(); + if (refreshed.isErr()) throw new Error('Unable to refresh sync session'); + } + const traceparent = span.traceparent(); + if (traceparent) url.searchParams.set('traceparent', traceparent); + return url.toString(); + }); +} diff --git a/crates/connection_gateway_models/Cargo.toml b/crates/connection_gateway_models/Cargo.toml index da13111e3ce..3ca6ada5dfb 100644 --- a/crates/connection_gateway_models/Cargo.toml +++ b/crates/connection_gateway_models/Cargo.toml @@ -5,8 +5,12 @@ publish = false version = "0.1.0" [dependencies] +macro_user_id = { path = "../macro_user_id" } model-entity = { path = "../model-entity" } serde = { workspace = true } serde_json = { workspace = true } utoipa = { workspace = true } workspace-hack = { version = "0.1", path = "../workspace-hack" } + +[dev-dependencies] +postcard = { workspace = true } diff --git a/crates/connection_gateway_models/src/fanout.rs b/crates/connection_gateway_models/src/fanout.rs new file mode 100644 index 00000000000..bbe18b80751 --- /dev/null +++ b/crates/connection_gateway_models/src/fanout.rs @@ -0,0 +1,134 @@ +//! Wire types for the gateway's inbound fan-out. +//! +//! Every gateway instance republishes its entire client websocket traffic — +//! connection lifecycle and every frame, text or binary, unparsed — to Redis +//! on [`INBOUND_CHANNEL`]. Consumer services subscribe and pick out what they +//! care about (the sync tier consumes binary frames; others may follow). A +//! consumer replies to a specific connection on that gateway instance's +//! [`outbound_channel`]. +//! +//! Encoding on both channels is `postcard`; Redis pub/sub is binary-safe. +//! +//! Delivery is fire-and-forget. Consumers detect a dead gateway instance by +//! its [`FromGateway::Heartbeat`] going quiet, not by any delivery guarantee. +//! In particular, a crashed gateway never publishes its connections' +//! [`FromGateway::Disconnected`] — consumers must treat heartbeat silence as +//! the disconnect of everything that instance held, and must tolerate +//! duplicate or missing lifecycle messages generally. + +#[cfg(test)] +mod test; + +use macro_user_id::user_id::MacroUserIdStr; +use serde::{Deserialize, Serialize}; +use std::fmt; + +/// Boot-unique id of a gateway instance. Minted fresh at gateway startup, so +/// a restarted instance is a new id and the old one simply goes quiet. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct GatewayId(pub String); + +impl fmt::Display for GatewayId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +// For string literals in tests and constructors at the wire boundary. +impl From<&str> for GatewayId { + fn from(id: &str) -> Self { + Self(id.to_string()) + } +} + +/// Gateway-local id of one websocket connection. Only meaningful together +/// with the [`GatewayId`] that minted it. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct ConnId(pub String); + +impl fmt::Display for ConnId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +// For string literals in tests and constructors at the wire boundary. +impl From<&str> for ConnId { + fn from(id: &str) -> Self { + Self(id.to_string()) + } +} + +/// Channel every gateway instance publishes its inbound traffic to. +pub const INBOUND_CHANNEL: &str = "connection_gateway.fanout.inbound"; + +/// Interval between [`FromGateway::Heartbeat`] messages. +pub const HEARTBEAT_INTERVAL_SECS: u64 = 5; + +/// The per-instance channel consumers reply on. `gateway` is minted fresh at +/// gateway boot, so a restarted instance is a new peer and the old id simply +/// goes quiet. +pub fn outbound_channel(gateway: &GatewayId) -> String { + format!("connection_gateway.fanout.outbound.{gateway}") +} + +/// Everything a gateway instance publishes on [`INBOUND_CHANNEL`]. +#[derive(Debug, Serialize, Deserialize)] +pub enum FromGateway { + /// A websocket connection was accepted and its user authenticated. + Connected { + /// Boot-unique id of the publishing gateway instance. + gateway: GatewayId, + /// Gateway-local id of the websocket connection. + conn: ConnId, + /// User the gateway authenticated at the socket edge. + user_id: MacroUserIdStr<'static>, + }, + /// One client frame, in per-connection order, unparsed. + Frame { + /// Boot-unique id of the publishing gateway instance. + gateway: GatewayId, + /// Gateway-local id of the websocket connection. + conn: ConnId, + /// Whether the frame was a text websocket message (else binary). + text: bool, + /// The client's bytes, unparsed. + payload: Vec, + }, + /// The websocket connection closed. + Disconnected { + /// Boot-unique id of the publishing gateway instance. + gateway: GatewayId, + /// Gateway-local id of the websocket connection. + conn: ConnId, + }, + /// Liveness beacon, published every [`HEARTBEAT_INTERVAL_SECS`]. + /// Consumers should drop all state for a gateway id that goes quiet. + Heartbeat { + /// Boot-unique id of the publishing gateway instance. + gateway: GatewayId, + }, +} + +/// Messages a consumer publishes on a gateway's [`outbound_channel`]. +#[derive(Debug, Serialize, Deserialize)] +pub enum ToGateway { + /// Deliver a frame to a connection's websocket. + Frame { + /// Gateway-local id of the target websocket connection. + conn: ConnId, + /// Whether to send as a text websocket message (else binary). + text: bool, + /// The bytes to send, unparsed. + payload: Vec, + }, + /// Close a connection's websocket. + Close { + /// Gateway-local id of the target websocket connection. + conn: ConnId, + /// Websocket close code to send. + code: u16, + }, +} diff --git a/crates/connection_gateway_models/src/fanout/test.rs b/crates/connection_gateway_models/src/fanout/test.rs new file mode 100644 index 00000000000..664d46c75be --- /dev/null +++ b/crates/connection_gateway_models/src/fanout/test.rs @@ -0,0 +1,46 @@ +use super::*; + +/// The gateway encodes and the sync tier decodes (and vice versa); this pins +/// the postcard round trip both directions. +#[test] +fn from_gateway_round_trips() { + let message = FromGateway::Frame { + gateway: GatewayId("g-1".into()), + conn: ConnId("c-1".into()), + text: false, + payload: vec![0xff, 0x00, 0x7f], + }; + let bytes = postcard::to_stdvec(&message).unwrap(); + let decoded: FromGateway = postcard::from_bytes(&bytes).unwrap(); + match decoded { + FromGateway::Frame { + gateway, + conn, + text, + payload, + } => { + assert_eq!(gateway, GatewayId("g-1".into())); + assert_eq!(conn, ConnId("c-1".into())); + assert!(!text); + assert_eq!(payload, vec![0xff, 0x00, 0x7f]); + } + other => panic!("wrong variant: {other:?}"), + } +} + +#[test] +fn to_gateway_round_trips() { + let message = ToGateway::Close { + conn: ConnId("c-9".into()), + code: 4000, + }; + let bytes = postcard::to_stdvec(&message).unwrap(); + let decoded: ToGateway = postcard::from_bytes(&bytes).unwrap(); + match decoded { + ToGateway::Close { conn, code } => { + assert_eq!(conn, ConnId("c-9".into())); + assert_eq!(code, 4000); + } + other => panic!("wrong variant: {other:?}"), + } +} diff --git a/crates/connection_gateway_models/src/lib.rs b/crates/connection_gateway_models/src/lib.rs index 13b39e88983..49d907ae0da 100644 --- a/crates/connection_gateway_models/src/lib.rs +++ b/crates/connection_gateway_models/src/lib.rs @@ -1,6 +1,8 @@ #![deny(missing_docs)] //! Shared types for the connection gateway service and client. +pub mod fanout; + use model_entity::Entity; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; diff --git a/crates/sync_router/AGENTS.md b/crates/sync_router/AGENTS.md index 2e8e99a5ec8..46b95d02c1c 100644 --- a/crates/sync_router/AGENTS.md +++ b/crates/sync_router/AGENTS.md @@ -2,7 +2,7 @@ ## What it is -Document sync routing library embedded in `connection_gateway`. Each authenticated +Document sync routing library hosted by DSS at `/sync/ws`. Each authenticated WebSocket owns a router and its downstream connections; there is no separate service, Redis subscriber, or Pulumi stack. @@ -11,7 +11,7 @@ service, Redis subscriber, or Pulumi stack. - `src/domain/router.rs`: document subscriptions, bounded queues, route epochs. - `src/domain/ports.rs`: socket delivery and downstream lifecycle contracts. - `src/outbound/do_downstream.rs`: scoped-token WebSocket dialing and frame pumps. -- `../../services/connection_gateway/src/service/sync.rs`: socket-owned task +- `src/session.rs`: socket-owned task tracking, cancellation, and local delivery adapter. ## Scope @@ -23,16 +23,19 @@ service, Redis subscriber, or Pulumi stack. ## Checks - `cargo test -p sync_router` (no live external services required). -- `cargo test -p connection_gateway --lib service::sync` exercises real local - downstream sockets, isolation, backpressure, and cancellation. +- Session tests exercise real local downstream sockets, isolation, backpressure, + cancellation, and byte budgets. - After dependency changes, regenerate and check `cargo x deps`. ## Gotchas - Preserve inner chunk framing byte-for-byte. A full queue must close its route or socket and trigger resync; never drop a chunk and continue the stream. +- Inbound messages are capped at one MiB; downstream queues share an eight MiB + budget per session, retained through the write. Never increase count limits + without checking the byte bound. - Dropping a downstream route cancels its task, including pending dials/writes. -- Every gateway replica handles only its own sockets. No distributed route +- Every DSS replica handles only its own sockets. No distributed route ownership or singleton deployment is required. -- Deploy the updated gateway before enabling PostHog `sync-via-gateway`. +- Deploy the updated DSS before enabling PostHog `sync-via-gateway`. `VITE_SYNC_VIA_GATEWAY=true` opts local clients in. diff --git a/crates/sync_router/src/domain/models.rs b/crates/sync_router/src/domain/models.rs index 8211268a621..c2a5025ebcb 100644 --- a/crates/sync_router/src/domain/models.rs +++ b/crates/sync_router/src/domain/models.rs @@ -1,6 +1,6 @@ //! Core types shared across the router. -/// A gateway-generated client connection id. Each socket has its own router. +/// A host-generated client connection id. Each socket has its own router. #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ConnectionId(pub String); @@ -15,7 +15,7 @@ impl DocId { } } -/// Events from the gateway edge, already filtered to what the router handles +/// Events from the socket edge, already filtered to what the router handles /// (text frames never reach it). #[derive(Debug)] pub enum EdgeEvent { @@ -36,11 +36,11 @@ pub enum EdgeEvent { /// Everything the router loop reacts to. #[derive(Debug)] pub enum Event { - /// An event from the gateway edge. + /// An event from the socket edge. Edge(EdgeEvent), /// A downstream connection died (upstream close or dial failure). The - /// client was already sent `RouterDocClosed` / `RouterSubscribeFailed` by - /// the downstream pump; this only cleans up the route. + /// router checks the incarnation before notifying the client, so an old + /// pump cannot close a replacement virtual socket. DownstreamClosed { /// The connection whose downstream died. conn: ConnectionId, @@ -50,5 +50,16 @@ pub enum Event { /// dial failing after the client already re-subscribed) must not tear /// down the replacement route. epoch: u64, + /// Whether dialing failed or an established connection ended. + failure: DownstreamFailure, }, } + +/// Coarse downstream failure safe to expose to clients. +#[derive(Debug)] +pub enum DownstreamFailure { + /// Authentication, connection establishment, or its timeout failed. + DialFailed, + /// An established upstream connection ended. + Closed(&'static str), +} diff --git a/crates/sync_router/src/domain/ports.rs b/crates/sync_router/src/domain/ports.rs index d6707fb385b..21448c2d364 100644 --- a/crates/sync_router/src/domain/ports.rs +++ b/crates/sync_router/src/domain/ports.rs @@ -1,12 +1,12 @@ //! Ports the router core depends on. use crate::domain::models::{ConnectionId, DocId}; -use std::fmt::Debug; -use tokio::sync::mpsc; +use std::{fmt::Debug, sync::Arc}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; use tokio_util::sync::{CancellationToken, DropGuard}; /// Delivers an already-serialized `FromRouter` envelope to a client -/// connection. The gateway implements this with the owning socket's queue. +/// connection. The socket host implements this with the owning socket's queue. #[cfg_attr(test, mockall::automock(type Err = anyhow::Error;))] pub trait EdgeSink: Send + Sync + 'static { /// The error type returned by delivery. @@ -24,9 +24,8 @@ pub trait EdgeSink: Send + Sync + 'static { /// /// Returns a sender immediately; dialing happens in a background task, so /// frames sent before the dial completes buffer in the channel. The -/// implementation delivers `RouterSubscribed` / `RouterSubscribeFailed` / -/// `RouterDocFrame` / `RouterDocClosed` to the client itself and reports -/// route teardown via [`crate::domain::models::Event::DownstreamClosed`]. +/// implementation delivers `RouterSubscribed` / `RouterDocFrame` to the client +/// and reports failures for epoch-checked teardown and notification via [`crate::domain::models::Event::DownstreamClosed`]. /// /// Chapter 1 dials the Cloudflare Durable Object; chapter 2 consistent-hashes /// `doc` across native sync services. @@ -41,21 +40,53 @@ pub trait DownstreamFactory: Send + Sync + 'static { /// A downstream route whose work is cancelled as soon as the route is removed. pub struct Downstream { - sender: mpsc::Sender>, + sender: mpsc::Sender, _cancel_on_drop: DropGuard, + budget: Arc, } impl Downstream { /// Own the frame queue and cancellation of its downstream pump. - pub fn new(sender: mpsc::Sender>, cancel: CancellationToken) -> Self { + pub fn new( + sender: mpsc::Sender, + cancel: CancellationToken, + budget: Arc, + ) -> Self { Self { sender, + budget, _cancel_on_drop: cancel.drop_guard(), } } /// Offer a frame without blocking the router on a slow downstream. pub fn try_send(&self, frame: Vec) -> Result<(), mpsc::error::TrySendError>> { - self.sender.try_send(frame) + let Ok(size) = u32::try_from(frame.len()) else { + return Err(mpsc::error::TrySendError::Full(frame)); + }; + let Ok(permit) = self.budget.clone().try_acquire_many_owned(size) else { + return Err(mpsc::error::TrySendError::Full(frame)); + }; + self.sender + .try_send(QueuedFrame { + bytes: frame, + _permit: permit, + }) + .map_err(|error| match error { + mpsc::error::TrySendError::Full(frame) => { + mpsc::error::TrySendError::Full(frame.bytes) + } + mpsc::error::TrySendError::Closed(frame) => { + mpsc::error::TrySendError::Closed(frame.bytes) + } + }) } } + +/// Bytes charged to a session-wide budget until the downstream write completes. +#[derive(Debug)] +pub struct QueuedFrame { + /// The untouched inner sync frame. + pub bytes: Vec, + _permit: OwnedSemaphorePermit, +} diff --git a/crates/sync_router/src/domain/router.rs b/crates/sync_router/src/domain/router.rs index 362194cccfe..1f38eb18ed4 100644 --- a/crates/sync_router/src/domain/router.rs +++ b/crates/sync_router/src/domain/router.rs @@ -6,7 +6,7 @@ mod test; use crate::domain::{ envelope::{self, ClientEnvelope}, - models::{ConnectionId, DocId, EdgeEvent, Event}, + models::{ConnectionId, DocId, DownstreamFailure, EdgeEvent, Event}, ports::{Downstream, DownstreamFactory, EdgeSink}, }; use std::collections::{HashMap, HashSet}; @@ -15,7 +15,7 @@ use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use tracing::{debug, warn}; -/// Routes envelope frames between gateway connections and per-document +/// Routes envelope frames between client connections and per-document /// downstreams. Owned and driven by a single task; see [`Router::handle`]. pub struct Router { sink: Arc, @@ -77,16 +77,29 @@ impl Router { } } Event::Edge(EdgeEvent::Disconnected { conn }) => self.drop_conn(&conn), - Event::DownstreamClosed { conn, doc, epoch } => { - // The pump already told the client; forget the route — but - // only if it is still the incarnation that died. A stale - // close must not tear down a replacement route. + Event::DownstreamClosed { + conn, + doc, + epoch, + failure, + } => { + // Teardown and notification are atomic with respect to retries. + // A stale pump cannot affect a replacement route or virtual socket. let is_current = self .routes .get(&(conn.clone(), doc.clone())) .is_some_and(|route| route.epoch == epoch); if is_current { self.forget(&conn, &doc); + let frame = match failure { + DownstreamFailure::DialFailed => { + envelope::subscribe_failed(doc.as_str(), "dial failed") + } + DownstreamFailure::Closed(reason) => { + envelope::doc_closed(doc.as_str(), reason) + } + }; + self.deliver(&conn, frame).await; } } } @@ -96,14 +109,21 @@ impl Router { async fn handle_client(&mut self, conn: ConnectionId, frame: ClientEnvelope) { match frame { ClientEnvelope::RouterSubscribe { doc_id, token } => { + if doc_id.len() > 512 || token.len() > 16 * 1024 { + self.deliver( + &conn, + envelope::subscribe_failed(&doc_id, "invalid subscription"), + ) + .await; + return; + } let doc = DocId(doc_id); let key = (conn.clone(), doc.clone()); if self.routes.contains_key(&key) { - // Idempotent: the downstream is already up (or dialing). - // Re-ack so a retrying client settles. - self.deliver(&conn, envelope::subscribed(doc.as_str())) - .await; - return; + // Retry with fresh authority, cancelling the previous incarnation. + // The old pump may have sent its failure before its cleanup event + // reached us. Never acknowledge a still-pending handshake. + self.forget(&conn, &doc); } if self .by_conn diff --git a/crates/sync_router/src/domain/router/test.rs b/crates/sync_router/src/domain/router/test.rs index 30dc31a7ffc..3e3fef750e2 100644 --- a/crates/sync_router/src/domain/router/test.rs +++ b/crates/sync_router/src/domain/router/test.rs @@ -41,7 +41,9 @@ fn recording_factory( capacity: usize, ) -> ( MockDownstreamFactory, - std::sync::Arc>)>>>, + std::sync::Arc< + std::sync::Mutex)>>, + >, ) { let opened = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); let record = std::sync::Arc::clone(&opened); @@ -51,7 +53,11 @@ fn recording_factory( .returning(move |_conn, doc, _token, _epoch| { let (tx, rx) = mpsc::channel(capacity); record.lock().unwrap().push((doc, rx)); - Downstream::new(tx, CancellationToken::new()) + Downstream::new( + tx, + CancellationToken::new(), + Arc::new(tokio::sync::Semaphore::new(1024)), + ) }); (factory, opened) } @@ -69,17 +75,17 @@ async fn frames_buffer_into_the_downstream_channel_before_dial_completes() { let mut opened = opened.lock().unwrap(); let (doc, rx) = &mut opened[0]; assert_eq!(doc.as_str(), "doc-a"); - assert_eq!(rx.try_recv().unwrap(), b"first"); - assert_eq!(rx.try_recv().unwrap(), b"second"); + assert_eq!(rx.try_recv().unwrap().bytes, b"first"); + assert_eq!(rx.try_recv().unwrap().bytes, b"second"); } #[tokio::test] -async fn resubscribe_is_idempotent_and_reacks() { +async fn duplicate_subscribe_does_not_ack_before_authorization() { let (factory, opened) = recording_factory(8); let mut sink = MockEdgeSink::new(); - // Exactly one re-ack for the duplicate subscribe. + // Only the downstream can acknowledge a validated subscription. sink.expect_deliver() - .times(1) + .times(0) .returning(|_, _| Box::pin(async { Ok(()) })); let mut router = Router::new(std::sync::Arc::new(sink), factory); @@ -87,7 +93,7 @@ async fn resubscribe_is_idempotent_and_reacks() { router.handle(subscribe(&c, "doc-a")).await; router.handle(subscribe(&c, "doc-a")).await; - assert_eq!(opened.lock().unwrap().len(), 1); + assert_eq!(opened.lock().unwrap().len(), 2); } #[tokio::test] @@ -144,6 +150,7 @@ async fn downstream_closed_forgets_the_route_so_resubscribe_reopens() { conn: c.clone(), doc: DocId("doc-a".into()), epoch: 1, + failure: DownstreamFailure::DialFailed, }) .await; router.handle(subscribe(&c, "doc-a")).await; @@ -175,3 +182,29 @@ async fn full_downstream_queue_closes_the_route_instead_of_corrupting_chunk_fram router.handle(subscribe(&c, "doc-a")).await; assert_eq!(opened.lock().unwrap().len(), 2); } + +#[tokio::test] +async fn stale_failure_cannot_notify_or_remove_a_replacement_subscription() { + let (factory, opened) = recording_factory(8); + let mut sink = MockEdgeSink::new(); + sink.expect_deliver().times(0); + let mut router = Router::new(Arc::new(sink), factory); + let c = conn("host", "socket"); + router.handle(subscribe(&c, "doc-a")).await; + router.handle(subscribe(&c, "doc-a")).await; + router + .handle(Event::DownstreamClosed { + conn: c.clone(), + doc: DocId("doc-a".into()), + epoch: 1, + failure: DownstreamFailure::DialFailed, + }) + .await; + router.handle(doc_frame(&c, "doc-a", b"new-route")).await; + let mut opened = opened.lock().unwrap(); + assert!(matches!( + opened[0].1.try_recv(), + Err(mpsc::error::TryRecvError::Disconnected) + )); + assert_eq!(opened[1].1.try_recv().unwrap().bytes, b"new-route"); +} diff --git a/crates/sync_router/src/lib.rs b/crates/sync_router/src/lib.rs index bd85d8caebf..f56894e0ced 100644 --- a/crates/sync_router/src/lib.rs +++ b/crates/sync_router/src/lib.rs @@ -1,6 +1,5 @@ #![deny(missing_docs)] -//! The sync-router: forwards sync-protocol frames between the connection -//! gateway's local sockets and a per-document downstream. +//! The sync-router: forwards sync-protocol frames between the host's local sockets and a per-document downstream. //! //! Chapter 1: the downstream is the existing Cloudflare Durable Object //! sync-service, dialed per `(connection, document)` so the DO sees each user @@ -14,3 +13,6 @@ pub mod domain; pub mod outbound; + +/// Socket lifecycle and upstream adapter composition. +pub mod session; diff --git a/crates/sync_router/src/outbound/do_downstream.rs b/crates/sync_router/src/outbound/do_downstream.rs index 5b086217d89..c1787473163 100644 --- a/crates/sync_router/src/outbound/do_downstream.rs +++ b/crates/sync_router/src/outbound/do_downstream.rs @@ -6,15 +6,15 @@ //! id across native sync services; the router core is unaware. use crate::domain::envelope; -use crate::domain::models::{ConnectionId, DocId, Event}; -use crate::domain::ports::{Downstream, DownstreamFactory, EdgeSink}; +use crate::domain::models::{ConnectionId, DocId, DownstreamFailure, Event}; +use crate::domain::ports::{Downstream, DownstreamFactory, EdgeSink, QueuedFrame}; use futures::{SinkExt, StreamExt}; use opentelemetry::trace::TraceContextExt; use std::sync::Arc; use std::time::Duration; -use tokio::sync::mpsc; -use tokio_tungstenite::connect_async; +use tokio::sync::{Semaphore, mpsc}; use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::{connect_async_with_config, tungstenite::protocol::WebSocketConfig}; use tokio_util::{sync::CancellationToken, task::TaskTracker}; use tracing::{Instrument, debug, warn}; use tracing_opentelemetry::OpenTelemetrySpanExt; @@ -22,7 +22,9 @@ use url::Url; /// How many client frames may buffer while the dial is in flight (or the /// upstream is slow). Overflow closes the route so the client resynchronizes. -const DOWNSTREAM_BUFFER: usize = 256; +const DOWNSTREAM_BUFFER: usize = 8; +/// Shared by every document queue and in-flight upstream write in a session. +const DOWNSTREAM_BYTE_BUDGET: usize = 8 * 1024 * 1024; /// Keep the upstream (and its Durable Object) alive, mirroring the ping the /// browser used to send on its direct per-document socket. @@ -35,6 +37,7 @@ pub struct DoDownstreamFactory { events: mpsc::Sender, cancel: CancellationToken, tasks: TaskTracker, + budget: Arc, } impl DoDownstreamFactory { @@ -47,6 +50,7 @@ impl DoDownstreamFactory { tasks: TaskTracker, ) -> Self { Self { + budget: Arc::new(Semaphore::new(DOWNSTREAM_BYTE_BUDGET)), base_url, sink, events, @@ -91,7 +95,7 @@ impl DownstreamFactory for DoDownstreamFactory { } .in_current_span(), ); - Downstream::new(sender, cancel) + Downstream::new(sender, cancel, self.budget.clone()) } } @@ -103,7 +107,7 @@ async fn pump( conn: ConnectionId, doc: DocId, epoch: u64, - mut from_client: mpsc::Receiver>, + mut from_client: mpsc::Receiver, sink: Arc, events: mpsc::Sender, ) { @@ -121,13 +125,30 @@ async fn pump( } }; - let upstream = match tokio::time::timeout(Duration::from_secs(10), connect_async(&url)).await { + let upstream = match tokio::time::timeout( + Duration::from_secs(10), + connect_async_with_config( + &url, + Some( + WebSocketConfig::default() + .max_message_size(Some(crate::session::MAX_FRAME_BYTES - 1024)) + .max_frame_size(Some(crate::session::MAX_FRAME_BYTES - 1024)), + ), + false, + ), + ) + .await + { Ok(Ok((upstream, _response))) => upstream, failure => { warn!(timed_out = failure.is_err(), "downstream dial failed"); - deliver(envelope::subscribe_failed(doc.as_str(), "dial failed")).await; events - .send(Event::DownstreamClosed { conn, doc, epoch }) + .send(Event::DownstreamClosed { + conn, + doc, + epoch, + failure: DownstreamFailure::DialFailed, + }) .await .ok(); return; @@ -147,8 +168,8 @@ async fn pump( let closed_reason: Option<&str> = loop { tokio::select! { frame = from_client.recv() => match frame { - Some(bytes) => { - if let Err(error) = write.send(Message::Binary(bytes.into())).await { + Some(frame) => { + if let Err(error) = write.send(Message::Binary(frame.bytes.into())).await { warn!(error = ?error, "upstream write failed"); break Some("upstream write failed"); } @@ -180,9 +201,13 @@ async fn pump( match closed_reason { Some(reason) => { debug!(reason, "downstream closed"); - deliver(envelope::doc_closed(doc.as_str(), reason)).await; events - .send(Event::DownstreamClosed { conn, doc, epoch }) + .send(Event::DownstreamClosed { + conn, + doc, + epoch, + failure: DownstreamFailure::Closed(reason), + }) .await .ok(); } diff --git a/services/connection_gateway/src/service/sync.rs b/crates/sync_router/src/session.rs similarity index 57% rename from services/connection_gateway/src/service/sync.rs rename to crates/sync_router/src/session.rs index 7d52ced54dd..6e0f2d74e34 100644 --- a/services/connection_gateway/src/service/sync.rs +++ b/crates/sync_router/src/session.rs @@ -1,9 +1,6 @@ //! Per-WebSocket sync routing. No shared subscription or cross-replica relay. -use crate::model::message::OutgoingMessage; -use anyhow::{Context, Result}; -use std::sync::Arc; -use sync_router::{ +use crate::{ domain::{ models::{ConnectionId, EdgeEvent, Event}, ports::EdgeSink, @@ -11,10 +8,12 @@ use sync_router::{ }, outbound::do_downstream::DoDownstreamFactory, }; +use anyhow::{Context, Result}; +use std::sync::Arc; use tokio::sync::mpsc; use tokio_util::{ sync::{CancellationToken, DropGuard}, - task::TaskTracker, + task::{TaskTracker, task_tracker::TaskTrackerToken}, }; use tracing::Instrument; use url::Url; @@ -22,10 +21,11 @@ use url::Url; #[cfg(test)] mod test; -const ROUTER_QUEUE_CAPACITY: usize = 64; +const ROUTER_QUEUE_CAPACITY: usize = 8; /// Routing resources owned by one authenticated client WebSocket. pub struct SyncSession { + _active: TaskTrackerToken, connection: ConnectionId, events: mpsc::Sender, cancel: CancellationToken, @@ -35,13 +35,14 @@ pub struct SyncSession { impl SyncSession { /// Start a router that can deliver only to this socket's bounded queue. - pub fn new( + fn new( connection_id: String, base_url: Url, - sender: mpsc::Sender, + sender: mpsc::Sender>, + parent: &SyncSessionFactory, ) -> Self { let connection = ConnectionId(connection_id); - let cancel = CancellationToken::new(); + let cancel = parent.cancel.child_token(); let tasks = TaskTracker::new(); let (events, receiver) = mpsc::channel(ROUTER_QUEUE_CAPACITY); let sink = Arc::new(SocketSink { @@ -59,6 +60,7 @@ impl SyncSession { let router = Router::new(sink, downstreams); tasks.spawn(router.run(receiver, cancel.clone()).in_current_span()); Self { + _active: parent.active.token(), connection, events, _cancel_on_drop: cancel.clone().drop_guard(), @@ -69,6 +71,7 @@ impl SyncSession { /// Apply backpressure only to this connection, preserving binary frame order. pub async fn frame(&self, payload: Vec) -> Result<()> { + anyhow::ensure!(payload.len() <= MAX_FRAME_BYTES, "sync frame too large"); self.events .send(Event::Edge(EdgeEvent::Frame { conn: self.connection.clone(), @@ -96,7 +99,7 @@ impl SyncSession { struct SocketSink { connection: ConnectionId, - sender: mpsc::Sender, + sender: mpsc::Sender>, cancel: CancellationToken, } @@ -109,11 +112,11 @@ impl EdgeSink for SocketSink { conn == &self.connection, "sync frame targets another connection" ); - if self - .sender - .try_send(OutgoingMessage::Binary(frame)) - .is_err() - { + if frame.len() > MAX_FRAME_BYTES { + self.cancel.cancel(); + anyhow::bail!("sync response too large"); + } + if self.sender.try_send(frame).is_err() { // Dropping a chunk corrupts framing. Close this socket so its virtual // sockets reconnect and resync instead of losing a partial message. self.cancel.cancel(); @@ -122,3 +125,54 @@ impl EdgeSink for SocketSink { Ok(()) } } + +/// Maximum outer message size; inner sync chunks are at most 900,001 bytes. +pub const MAX_FRAME_BYTES: usize = 1024 * 1024; +/// The bounded client queue also caps queued response bytes at eight MiB. +pub const CLIENT_QUEUE_CAPACITY: usize = 8; + +/// Composition root shared by the host's sockets, including process shutdown. +#[derive(Clone)] +pub struct SyncSessionFactory { + base_url: Url, + cancel: CancellationToken, + active: TaskTracker, +} + +impl SyncSessionFactory { + /// Validate the upstream origin once at service startup. + pub fn new(base_url: Url) -> Result { + anyhow::ensure!( + matches!(base_url.scheme(), "http" | "https") + && base_url.has_host() + && base_url.path() == "/" + && base_url.query().is_none() + && base_url.fragment().is_none() + && base_url.username().is_empty() + && base_url.password().is_none(), + "sync service URL must be an HTTP(S) origin", + ); + Ok(Self { + base_url, + cancel: CancellationToken::new(), + active: TaskTracker::new(), + }) + } + + /// Each socket gets independent routes, cancellation, and bounded queues. + pub fn open(&self, connection_id: String, sender: mpsc::Sender>) -> SyncSession { + SyncSession::new(connection_id, self.base_url.clone(), sender, self) + } + + /// Begin draining while the HTTP host stops accepting new connections. + pub fn cancel(&self) { + self.cancel.cancel(); + self.active.close(); + } + + /// Drain every upgraded socket and downstream before the host exits. + pub async fn shutdown(&self) { + self.cancel(); + self.active.wait().await; + } +} diff --git a/services/connection_gateway/src/service/sync/test.rs b/crates/sync_router/src/session/test.rs similarity index 73% rename from services/connection_gateway/src/service/sync/test.rs rename to crates/sync_router/src/session/test.rs index 8d50929c131..503aa56c718 100644 --- a/services/connection_gateway/src/service/sync/test.rs +++ b/crates/sync_router/src/session/test.rs @@ -11,15 +11,12 @@ fn encode(message: ToRouter<'_>) -> Vec { bytes } -async fn next_frame(receiver: &mut mpsc::Receiver) -> Vec { - let OutgoingMessage::Binary(bytes) = receiver.recv().await.unwrap() else { - panic!("expected a binary sync envelope"); - }; - bytes +async fn next_frame(receiver: &mut mpsc::Receiver>) -> Vec { + receiver.recv().await.unwrap() } #[tokio::test] -async fn separate_gateway_sessions_route_the_same_document_without_cross_delivery() { +async fn separate_host_sessions_route_the_same_document_without_cross_delivery() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let base_url = Url::parse(&format!("http://{}", listener.local_addr().unwrap())).unwrap(); let server = tokio::spawn(async move { @@ -47,8 +44,12 @@ async fn separate_gateway_sessions_route_the_same_document_without_cross_deliver }); let (first_sender, mut first_receiver) = mpsc::channel(10); let (second_sender, mut second_receiver) = mpsc::channel(10); - let first = SyncSession::new("gateway-a/socket-a".into(), base_url.clone(), first_sender); - let second = SyncSession::new("gateway-b/socket-b".into(), base_url, second_sender); + let first = SyncSessionFactory::new(base_url.clone()) + .unwrap() + .open("gateway-a/socket-a".into(), first_sender); + let second = SyncSessionFactory::new(base_url) + .unwrap() + .open("gateway-b/socket-b".into(), second_sender); first .frame(encode(ToRouter::RouterSubscribe { doc_id: "same-document", @@ -91,11 +92,11 @@ async fn separate_gateway_sessions_route_the_same_document_without_cross_deliver .unwrap(); assert_eq!( next_frame(&mut first_receiver).await, - sync_router::domain::envelope::doc_frame("same-document", &[0, 11, 12]) + crate::domain::envelope::doc_frame("same-document", &[0, 11, 12]) ); assert_eq!( next_frame(&mut second_receiver).await, - sync_router::domain::envelope::doc_frame("same-document", &[1, 21, 22]) + crate::domain::envelope::doc_frame("same-document", &[1, 21, 22]) ); first.shutdown().await; second.shutdown().await; @@ -114,7 +115,9 @@ async fn closing_a_session_cancels_a_downstream_that_never_completes_its_handsha socket.read_to_end(&mut bytes).await.unwrap(); }); let (sender, _receiver) = mpsc::channel(10); - let session = SyncSession::new("socket".into(), base_url, sender); + let session = SyncSessionFactory::new(base_url) + .unwrap() + .open("socket".into(), sender); session .frame(encode(ToRouter::RouterSubscribe { doc_id: "document", @@ -141,7 +144,9 @@ async fn unsubscribe_cancels_a_pending_dial_without_stopping_the_session() { socket.read_to_end(&mut Vec::new()).await.unwrap(); }); let (sender, _receiver) = mpsc::channel(10); - let session = SyncSession::new("socket".into(), base_url, sender); + let session = SyncSessionFactory::new(base_url) + .unwrap() + .open("socket".into(), sender); session .frame(encode(ToRouter::RouterSubscribe { doc_id: "document", @@ -165,7 +170,7 @@ async fn unsubscribe_cancels_a_pending_dial_without_stopping_the_session() { #[tokio::test] async fn full_outbound_queue_cancels_only_its_owner_instead_of_dropping_a_chunk() { let (sender, _receiver) = mpsc::channel(1); - sender.send(OutgoingMessage::Pong).await.unwrap(); + sender.send(vec![]).await.unwrap(); let cancel = CancellationToken::new(); let other_cancel = CancellationToken::new(); let conn = ConnectionId("socket".into()); @@ -217,7 +222,9 @@ async fn rejected_document_token_never_opens_a_virtual_socket() { assert!(result.is_err()); }); let (sender, mut receiver) = mpsc::channel(10); - let session = SyncSession::new("socket".into(), base_url, sender); + let session = SyncSessionFactory::new(base_url) + .unwrap() + .open("socket".into(), sender); session .frame(encode(ToRouter::RouterSubscribe { doc_id: "document", @@ -263,7 +270,9 @@ async fn scoped_tokens_and_document_ids_cannot_change_the_downstream_url() { let _ = socket.next().await; }); let (sender, mut receiver) = mpsc::channel(10); - let session = SyncSession::new("socket".into(), base_url, sender); + let session = SyncSessionFactory::new(base_url) + .unwrap() + .open("socket".into(), sender); session .frame(encode(ToRouter::RouterSubscribe { doc_id: "doc/other?override", @@ -278,3 +287,61 @@ async fn scoped_tokens_and_document_ids_cannot_change_the_downstream_url() { session.shutdown().await; server.await.unwrap(); } + +#[tokio::test] +async fn host_shutdown_cancels_and_waits_for_all_sessions() { + let factory = SyncSessionFactory::new(Url::parse("http://127.0.0.1:1").unwrap()).unwrap(); + let (sender, _receiver) = mpsc::channel(1); + let session = factory.open("socket".into(), sender); + let shutdown = factory.shutdown(); + tokio::pin!(shutdown); + tokio::select! { + _ = &mut shutdown => panic!("must wait for socket cleanup"), + _ = session.closed() => {}, + } + assert!(!factory.active.is_empty()); + session.shutdown().await; + tokio::time::timeout(std::time::Duration::from_secs(1), shutdown) + .await + .unwrap(); + assert!(factory.active.is_empty()); +} + +#[tokio::test] +async fn oversized_frames_cannot_enter_client_or_router_queues() { + let factory = SyncSessionFactory::new(Url::parse("http://127.0.0.1:1").unwrap()).unwrap(); + let (sender, _receiver) = mpsc::channel(1); + let session = factory.open("socket".into(), sender.clone()); + assert!(session.frame(vec![0; MAX_FRAME_BYTES + 1]).await.is_err()); + let cancel = CancellationToken::new(); + let conn = ConnectionId("socket".into()); + let sink = SocketSink { + connection: conn.clone(), + sender, + cancel: cancel.clone(), + }; + assert!( + sink.deliver(&conn, vec![0; MAX_FRAME_BYTES + 1]) + .await + .is_err() + ); + assert!(cancel.is_cancelled()); + session.shutdown().await; +} + +#[tokio::test] +async fn all_document_queues_share_a_byte_budget_until_writes_finish() { + use crate::domain::ports::Downstream; + let budget = Arc::new(tokio::sync::Semaphore::new(10)); + let (first_tx, mut first_rx) = mpsc::channel(8); + let (second_tx, _second_rx) = mpsc::channel(8); + let first = Downstream::new(first_tx, CancellationToken::new(), budget.clone()); + let second = Downstream::new(second_tx, CancellationToken::new(), budget.clone()); + first.try_send(vec![0; 6]).unwrap(); + assert!(second.try_send(vec![0; 5]).is_err()); + let in_flight = first_rx.recv().await.unwrap(); + assert!(second.try_send(vec![0; 5]).is_err()); + drop(in_flight); + second.try_send(vec![0; 10]).unwrap(); + assert_eq!(budget.available_permits(), 0); +} diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index 3287dc3af9c..29c1774ff86 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -65,18 +65,15 @@ services: # ============================================================================ # connection_gateway (Port 8082) - # WebSocket gateway and embedded document sync router + # WebSocket gateway - Redis only # ============================================================================ connection_gateway: <<: [*common-env, *rust-services-image] command: ["/app/out/connection_gateway_service"] ports: - "8082:8080" - environment: - - OVERRIDE_SYNC_SERVICE_URL=http://sync-service:8787 depends_on: - redis - - sync_service networks: databases: services: diff --git a/docs/README.md b/docs/README.md index 9c52aa355c3..8b1ac46c66a 100644 --- a/docs/README.md +++ b/docs/README.md @@ -14,4 +14,4 @@ Docs that belong to one app, service, or crate live next to its code. WhatsApp linked devices: [architecture](architecture/whatsapp-linked-devices.md) and [External view walkthrough](agent-guide/external.md). -Document sync: [embedded gateway router](architecture/sync-gateway.md). +Document sync: [DSS sync router](architecture/sync-router.md). diff --git a/docs/agent-guide/documents.md b/docs/agent-guide/documents.md index 55dec37625f..7295e9d7f30 100644 --- a/docs/agent-guide/documents.md +++ b/docs/agent-guide/documents.md @@ -1,12 +1,14 @@ # Documents -For gateway sync rollout checks, use PostHog flag `sync-via-gateway` (dev and +For DSS sync rollout checks, use PostHog flag `sync-via-gateway` (dev and production) and recreate document sessions after the client receives the updated flag. A reload can initially read PostHog's cached value; open sessions do not change transport when the remote value changes. Off or unknown flag values use a direct sync-service WebSocket per document; on multiplexes document frames over -the connection gateway. Confirm all gateway replicas have the embedded sync -router before enabling; no separate router service is needed. Local stacks can override with `VITE_SYNC_VIA_GATEWAY=true` or `false`. +a dedicated DSS `/sync/ws` socket. Confirm all DSS replicas have the sync +endpoint before enabling; notification traffic still uses connection gateway. +Verify two clients converge after concurrent edits, a large update, and a forced +sync disconnect with pending edits. The historical flag name is unchanged. Local stacks can override with `VITE_SYNC_VIA_GATEWAY=true` or `false`. The document header has separate Share, Copy Share Link, and Side Panel buttons. They are borderless with a soft rounded background on hover. Share opens the diff --git a/docs/architecture/sync-gateway.md b/docs/architecture/sync-gateway.md deleted file mode 100644 index 514c41342fa..00000000000 --- a/docs/architecture/sync-gateway.md +++ /dev/null @@ -1,46 +0,0 @@ -# Document sync through the connection gateway - -The browser multiplexes document sync over one WebSocket when PostHog flag -`sync-via-gateway` is enabled. The existing connection gateway hosts the router -as the `sync_router` Rust library. There is no additional deployable service. - -```text -Browser: virtual document sockets - | - | WebSocket / Bebop ToRouter and FromRouter envelopes - v -connection_gateway replica - authenticated socket - -> socket-owned router - -> one downstream WebSocket per subscribed document - | - | existing chunked sync protocol and scoped document token - v -sync-service / document Durable Object - token authorization, CRDT state, persistence -``` - -Each authenticated client socket owns an independent router, outbound queue, and -set of tracked downstream tasks. Replicas do not subscribe to a shared broadcast -channel. Reconnecting to another replica creates a new router and resubscribes; -there is no shared routing state to transfer. Existing JSON gateway messages -retain their cross-instance Redis relay. - -Routing reads the document id from the outer Bebop envelope. The inner payload, -including the large-message chunk prefix, passes through unchanged. Tokens are -forwarded only to the configured sync origin; the sync service validates them. - -Unsubscribing cancels that document's downstream, including a pending handshake. -Client disconnect cancels and joins all routing work for that socket. A full -outbound socket queue closes the client connection; a full downstream queue -closes that document route. Clients reconnect and resynchronize, rather than -continuing a stream with missing chunks. Other sockets are unaffected. - -The gateway resolves `macro_service_urls::SyncServiceUrl` at startup. Hosted dev -and production use the existing sync-service origins; local Docker sets -`OVERRIDE_SYNC_SERVICE_URL=http://sync-service:8787`. No new secret or Pulumi -resource is needed. Overrides must be HTTP(S) origins without credentials, -non-root paths, queries, or fragments. Deploy the updated gateway to every -replica before enabling `sync-via-gateway`. The frontend selects its transport -when a sync source is created; rollout updates do not move existing sessions. -The env override `VITE_SYNC_VIA_GATEWAY=true` enables local testing. diff --git a/docs/architecture/sync-router.md b/docs/architecture/sync-router.md new file mode 100644 index 00000000000..ebf460ad9c9 --- /dev/null +++ b/docs/architecture/sync-router.md @@ -0,0 +1,56 @@ +# Document sync through DSS + +When PostHog `sync-via-gateway` is enabled, the browser opens a dedicated +WebSocket at DSS `/sync/ws` (public gateway prefix: `/dss/sync/ws`). The historical +flag key remains unchanged. Notifications retain their separate connection-gateway +socket. DSS hosts the reusable `crates/sync_router` library; no new service, +Pulumi stack, Redis relay, or distributed subscription registry is required. + +```text +Browser + +-- notification socket --> connection_gateway + | + +-- dedicated sync socket --> DSS /sync/ws + ToRouter / FromRouter socket-owned SyncSession + (docId, payload) +-- document A --> sync-service DO A + +-- document B --> sync-service DO B +``` + +The endpoint authenticates the account using the standard Macro cookie/API-token +extractor. Each subscription carries a separate document-scoped token. The +sync-service Durable Object validates that token and continues to own document +state, permission enforcement, and durable acknowledgements. Only a successful +upstream handshake produces `RouterSubscribed`; duplicate subscribes cannot +acknowledge a pending or rejected handshake. + +`SyncRouterTransport` multiplexes document virtual sockets using the generated +Bebop `ToRouter` and `FromRouter` envelopes. The router reads the outer document +id and forwards the inner bytes unchanged. Existing 900,000-byte chunking and +per-document reassembly remain in place under the envelope. Text `ping`/`pong` +heartbeats apply to the dedicated sync socket. + +Each client socket owns its routes and downstream tasks. Any DSS replica can +accept a connection; reconnecting to another replica creates new routes and +re-subscribes with refreshed document tokens. Replicas share no routing state. +The client sync runtime re-synchronizes unacknowledged edits from its document +state. A connection loss never counts as a durable acknowledgement. Pending ack waits +end immediately on disconnect, releasing the WAL flush before reconnect replay. +Trace context travels from the browser in the upgrade query and is propagated +through the DSS session into downstream connections. + +Both WebSocket receive directions have explicit size limits around one MiB. +The client and router queues hold at most eight messages each. All downstream +queues and in-flight writes share an eight MiB budget per session, independently +of the 64-document cap. Overflow closes the route or socket so clients resync; +no chunk is silently discarded from a continuing stream. Unsubscribe cancels +pending dials and writes. DSS shutdown cancels and drains upgraded sync sessions. + +Deploy DSS before enabling the flag. Off or unknown values retain direct +per-document sync-service sockets. Local stacks can set +`VITE_SYNC_VIA_GATEWAY=true` or `false`. A source selects its transport for its +lifetime; recreate document sessions after a flag update. Cached PostHog values +can survive reload, so the flag is not an immediate kill switch for open sessions. + +The router core depends on `EdgeSink` and `DownstreamFactory` ports. A future +standalone host or gRPC adapter can replace those adapters without changing the +browser envelope or document sync protocol. diff --git a/packages/collaboration/src/sync-service/generated/schema.bop b/packages/collaboration/src/sync-service/generated/schema.bop index 2cc08b4b069..15d60537f6e 100644 --- a/packages/collaboration/src/sync-service/generated/schema.bop +++ b/packages/collaboration/src/sync-service/generated/schema.bop @@ -58,7 +58,7 @@ struct InitializeFromSnapshotRequest { /* * Multiplex envelope between the browser and the sync-router (rides the - * connection-gateway socket as binary frames). Inner payloads are the + * dedicated DSS /sync/ws socket as binary frames). Inner payloads are the * existing FromPeer/FromRemote unions, already bebop-encoded — the router * forwards them without parsing. */ diff --git a/packages/collaboration/src/sync-service/generated/schema.ts b/packages/collaboration/src/sync-service/generated/schema.ts index eb33150f9e4..98a8c68f605 100644 --- a/packages/collaboration/src/sync-service/generated/schema.ts +++ b/packages/collaboration/src/sync-service/generated/schema.ts @@ -2257,7 +2257,7 @@ export class RouterFrame implements IRouterFrame { /** * Multiplex envelope between the browser and the sync-router (rides the - * connection-gateway socket as binary frames). Inner payloads are the + * dedicated DSS /sync/ws socket as binary frames). Inner payloads are the * existing FromPeer/FromRemote unions, already bebop-encoded — the router * forwards them without parsing. */ diff --git a/packages/collaboration/src/sync-service/gateway.test.ts b/packages/collaboration/src/sync-service/router.test.ts similarity index 63% rename from packages/collaboration/src/sync-service/gateway.test.ts rename to packages/collaboration/src/sync-service/router.test.ts index e06e5666495..348fe5f6007 100644 --- a/packages/collaboration/src/sync-service/gateway.test.ts +++ b/packages/collaboration/src/sync-service/router.test.ts @@ -2,15 +2,17 @@ import { once } from 'node:events'; import type { AddressInfo } from 'node:net'; import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest'; import { type WebSocket, WebSocketServer } from 'ws'; +import { InMemoryWALStore, WALSyncer } from '../collab/wal'; import { WebsocketEvent } from '../websocket'; import { intoFrames, Reassembler } from '../websocket/platform/framing'; -import { GatewaySyncTransport } from './gateway'; import { FromPeer, FromRemote, FromRouter, ToRouter } from './generated/schema'; +import { SyncRouterTransport } from './router'; import type { SyncSocket } from './socket'; +import { SyncServiceSource } from './source'; -describe('gateway document transport', () => { +describe('router document transport', () => { let server: WebSocketServer; - let transport: GatewaySyncTransport; + let transport: SyncRouterTransport; let peer: WebSocket; let received: ToRouter[]; let sockets: SyncSocket[]; @@ -27,7 +29,7 @@ describe('gateway document transport', () => { }); }); const port = (server.address() as AddressInfo).port; - transport = new GatewaySyncTransport(`ws://127.0.0.1:${port}`); + transport = new SyncRouterTransport(`ws://127.0.0.1:${port}`); }); afterEach(async () => { @@ -127,6 +129,81 @@ describe('gateway document transport', () => { } }); + test('replays an unacknowledged edit after the real shared socket disconnects', async () => { + let dropped = false; + const replayed = Promise.withResolvers(); + server.on('connection', (client) => { + const reassembler = new Reassembler(); + client.on('message', (bytes) => { + const message = ToRouter.decode(new Uint8Array(bytes as Buffer)); + if (message.isRouterSubscribe()) { + const docId = message.value.docId; + client.send( + FromRouter.encode(FromRouter.fromRouterSubscribed({ docId })) + ); + for (const payload of intoFrames( + FromRemote.encode( + FromRemote.fromRemoteInitialSync({ + snapshot: new Uint8Array(), + awareness: new Uint8Array(), + }) + ) + )) { + client.send( + FromRouter.encode( + FromRouter.fromRouterDocFrame({ docId, payload }) + ) + ); + } + } else if (message.isRouterFrame()) { + const complete = reassembler.push(message.value.payload); + if (!complete) return; + const update = FromPeer.decode(complete); + if (!update.isPeerUpdate()) return; + if (!dropped) { + dropped = true; + client.close(1012, 'restart before persistence'); + return; + } + expect(update.value.updates).toEqual([new Uint8Array([1, 2, 3])]); + for (const payload of intoFrames( + FromRemote.encode( + FromRemote.fromRemoteUpdateAck({ id: update.value.id }) + ) + )) { + client.send( + FromRouter.encode( + FromRouter.fromRouterDocFrame({ + docId: message.value.docId, + payload, + }) + ) + ); + } + replayed.resolve(); + } + }); + }); + const source = new SyncServiceSource(attach('doc', 'token'), 'doc'); + const wal = new WALSyncer(new InMemoryWALStore(), (updates) => + source.pushUpdate(updates) + ); + source.listen((event) => { + if (event.type === 'reconnect') void wal.flush(); + }); + try { + expect((await source.doInitialSync()).isOk()).toBe(true); + await wal.append(new Uint8Array([1, 2, 3])); + await replayed.promise; + await wal.pendingFlush; + expect((await wal.summary()).dirty).toBe(0); + expect(dropped).toBe(true); + } finally { + wal.destroy(); + source.cleanup(); + } + }); + test('retries a failed initial token refresh', async () => { const refresh = vi .fn() diff --git a/packages/collaboration/src/sync-service/gateway.ts b/packages/collaboration/src/sync-service/router.ts similarity index 91% rename from packages/collaboration/src/sync-service/gateway.ts rename to packages/collaboration/src/sync-service/router.ts index 8027380329d..1809942632a 100644 --- a/packages/collaboration/src/sync-service/gateway.ts +++ b/packages/collaboration/src/sync-service/router.ts @@ -23,18 +23,18 @@ import type { SyncSocket } from './socket'; /** * Multiplexes many documents' sync traffic over ONE websocket to the - * connection gateway, instead of one direct sync-service socket per document. + * DSS sync endpoint, instead of one direct sync-service socket per document. * * Wire shape: every frame is a Bebop `ToRouter`/`FromRouter` envelope naming * the document; the inner payloads are the same `FromPeer`/`FromRemote` * messages the direct socket speaks, unchanged. Server side, the - * connection-gateway routes inbound frames through its embedded router, which holds + * DSS host routes inbound frames through its embedded router, which holds * one upstream connection per (socket, document). * * Each document gets a {@link SyncSocket}-compatible virtual socket, so * `SyncServiceSource` works unmodified on top of this transport. */ -export class GatewaySyncTransport { +export class SyncRouterTransport { private readonly getUrl: UrlResolver; private ws: Websocket | undefined; private readonly docs = new Map(); @@ -47,7 +47,7 @@ export class GatewaySyncTransport { /** * Attach a document to the shared socket: subscribes immediately and * returns the virtual socket for it. `refreshToken` is consulted on every - * re-subscribe (upstream drop, gateway reconnect), mirroring how the direct + * re-subscribe (upstream drop, router reconnect), mirroring how the direct * socket refetches a permission token per reconnect. */ public attach( @@ -74,7 +74,7 @@ export class GatewaySyncTransport { this.docs.delete(documentId); } - /** The shared gateway websocket, built lazily on first use. */ + /** The shared router websocket, built lazily on first use. */ public socket(): Websocket { if (this.ws) return this.ws; const ws = new WebsocketBuilder(this.getUrl) @@ -98,16 +98,16 @@ export class GatewaySyncTransport { this.route(event.data); }); ws.addEventListener(WebsocketEvent.Reconnect, () => { - // A new gateway socket owns a new router session server-side: every + // A new router socket owns a new router session server-side: every // route was torn down, so every document must re-subscribe. for (const virtual of this.docs.values()) virtual.resubscribe(); }); ws.addEventListener(WebsocketEvent.Close, (_ws, event) => { - for (const virtual of this.docs.values()) virtual.onGatewayClose(event); + for (const virtual of this.docs.values()) virtual.onRouterClose(event); }); ws.addEventListener(WebsocketEvent.HeartbeatMissed, (_ws, event) => { for (const virtual of this.docs.values()) { - virtual.onGatewayHeartbeatMissed(event); + virtual.onRouterHeartbeatMissed(event); } }); @@ -147,16 +147,16 @@ const RESUBSCRIBE_BASE_MS = 500; const RESUBSCRIBE_MAX_MS = 8_000; /** - * A {@link SyncSocket} for one document, riding the shared gateway socket. + * A {@link SyncSocket} for one document, riding the shared router socket. * - * Lifecycle mapping (virtual event ← gateway/envelope cause): + * Lifecycle mapping (virtual event ← router/envelope cause): * - Open ← first `RouterSubscribed` * - Message ← `RouterDocFrame` (decoded to `FromRemote`) - * - Close ← `RouterSubscribeFailed` / `RouterDocClosed` / gateway close - * - Reconnect ← any later `RouterSubscribed` (upstream or gateway recovery) + * - Close ← `RouterSubscribeFailed` / `RouterDocClosed` / router close + * - Reconnect ← any later `RouterSubscribed` (upstream or router recovery) */ class VirtualSyncSocket implements SyncSocket { - private readonly transport: GatewaySyncTransport; + private readonly transport: SyncRouterTransport; private readonly documentId: string; private readonly refreshToken: () => Promise; private initialToken: string | undefined; @@ -172,7 +172,7 @@ class VirtualSyncSocket implements SyncSocket { private pending: IFromPeer[] = []; public constructor( - transport: GatewaySyncTransport, + transport: SyncRouterTransport, documentId: string, initialToken: string | undefined, refreshToken: () => Promise @@ -184,8 +184,8 @@ class VirtualSyncSocket implements SyncSocket { } public get connectionState(): WebsocketConnectionState { - const gateway = this.transport.socket().connectionState; - if (gateway !== WebsocketConnectionState.Open) return gateway; + const router = this.transport.socket().connectionState; + if (router !== WebsocketConnectionState.Open) return router; return match(this.state) .with('open', () => WebsocketConnectionState.Open) .with('subscribing', () => WebsocketConnectionState.Connecting) @@ -303,7 +303,7 @@ class VirtualSyncSocket implements SyncSocket { this.dispatch( WebsocketEvent.Reconnect, new CustomEvent(WebsocketEvent.Reconnect, { - detail: { retries: 0, lastConnection: undefined, url: 'gateway' }, + detail: { retries: 0, lastConnection: undefined, url: 'router' }, }) ); } @@ -332,9 +332,9 @@ class VirtualSyncSocket implements SyncSocket { this.scheduleResubscribe(); } - public onGatewayClose(event: CloseEvent): void { + public onRouterClose(event: CloseEvent): void { if (this.state === 'closed') return; - // The gateway socket owns reconnection; we just report the drop and wait + // The router socket owns reconnection; we just report the drop and wait // for the transport's Reconnect handler to re-subscribe us. this.clearResubscribeTimer(); this.state = 'dropped'; @@ -343,7 +343,7 @@ class VirtualSyncSocket implements SyncSocket { this.dispatch(WebsocketEvent.Close, event); } - public onGatewayHeartbeatMissed( + public onRouterHeartbeatMissed( event: CustomEvent<{ missedHeartbeats: number; willReconnect: boolean }> ): void { this.dispatch( diff --git a/packages/collaboration/src/sync-service/source.test.ts b/packages/collaboration/src/sync-service/source.test.ts index 69a9aaac1e4..704ecb77756 100644 --- a/packages/collaboration/src/sync-service/source.test.ts +++ b/packages/collaboration/src/sync-service/source.test.ts @@ -291,6 +291,32 @@ describe('SyncServiceSource', () => { expect(await acked).toBe(true); }); + it('releases unacknowledged edits on disconnect before a fast resubscribe', async () => { + vi.useFakeTimers(); + try { + const ws = new FakeSocket(); + const src = new SyncServiceSource(ws, 'doc1', { newId: () => 'retry' }); + ws.deliver(remote.initialSync(snap, aw)); + await src.doInitialSync(); + let delivered: boolean | undefined; + void src.pushUpdate([new Uint8Array([1])]).then((result) => { + delivered = result; + }); + ws.fireClose(1012, 'restart'); + await vi.advanceTimersByTimeAsync(0); + expect(delivered).toBe(false); + expect(await src.pushUpdate([new Uint8Array([1])])).toBe(false); + ws.fireReconnect(); + ws.deliver(remote.initialSync(snap, aw)); + const retried = src.pushUpdate([new Uint8Array([1])]); + ws.deliver(remote.ack('retry')); + expect(await retried).toBe(true); + src.cleanup(); + } finally { + vi.useRealTimers(); + } + }); + it('pushUpdate resolves false when the ack times out', async () => { vi.useFakeTimers(); try { diff --git a/packages/collaboration/src/sync-service/source.ts b/packages/collaboration/src/sync-service/source.ts index 036d908d0fe..5956e642b8e 100644 --- a/packages/collaboration/src/sync-service/source.ts +++ b/packages/collaboration/src/sync-service/source.ts @@ -66,6 +66,7 @@ function syncEventForMessage(message: FromRemote): SyncSourceEvent | null { interface MessageWaiter { match: (message: FromRemote) => boolean; resolve: (message: FromRemote) => void; + cancelOnDisconnect?: () => void; } export interface SyncServiceSourceOptions { @@ -179,7 +180,8 @@ export class SyncServiceSource implements LiveSyncSource { // can't be missed. const acked = this.awaitMessage( (message) => message.isRemoteUpdateAck() && message.value.id === id, - TIMEOUTS.ACK + TIMEOUTS.ACK, + true ).then( () => true, () => false @@ -285,11 +287,18 @@ export class SyncServiceSource implements LiveSyncSource { */ private awaitMessage( match: (message: FromRemote) => boolean, - timeoutMs: number + timeoutMs: number, + cancelOnDisconnect = false ): Promise { const { promise, resolve, reject } = Promise.withResolvers(); const waiter: MessageWaiter = { match, + cancelOnDisconnect: cancelOnDisconnect + ? () => { + remove(); + reject(SyncError.connectionFailed()); + } + : undefined, resolve: (message) => { remove(); resolve(message); @@ -337,6 +346,10 @@ export class SyncServiceSource implements LiveSyncSource { IFromPeer, FromRemote > = (_ws, event) => { + this.initialSyncReceived = false; + // Release pending WAL pushes before a fast reconnect can request replay. + // A disconnected socket cannot supply a durable acknowledgement. + for (const waiter of [...this.waiters]) waiter.cancelOnDisconnect?.(); logSyncService({ documentId: this.documentId, level: 'warn', diff --git a/services/connection_gateway/AGENTS.md b/services/connection_gateway/AGENTS.md index 854f5f4e57b..2975e6ae053 100644 --- a/services/connection_gateway/AGENTS.md +++ b/services/connection_gateway/AGENTS.md @@ -4,13 +4,11 @@ Shared infrastructure: manages per-entity WebSocket connections so other services can push messages to a specific connected entity without holding -the socket themselves. It also embeds document sync routing for each socket. +the socket themselves. ## Start here - `src/api/connection.rs`: WebSocket upgrade endpoint and connection lifecycle. -- `src/service/sync.rs`: per-socket sync router, bounded outbound delivery, - cancellation and task joining. Routing core lives in `crates/sync_router`. - `src/api/message.rs`: HTTP endpoint other services call to send a message to a connected entity. - `src/service/connection.rs` + `src/service/dynamodb.rs`: connection registry, @@ -23,9 +21,9 @@ the socket themselves. It also embeds document sync routing for each socket. ## Scope - Owns: WebSocket connection tracking and message delivery to a connected - entity, across instances; multiplexed document sync for local sockets. -- Document state and document-token authorization remain in `sync-service`. - Other message/event payloads and send authorization remain with callers. + entity, across instances. +- Not here: the message/event payloads themselves, or who is authorized to + send to whom — callers decide that before calling in. ## Boundaries @@ -51,12 +49,7 @@ the socket themselves. It also embeds document sync routing for each socket. - A send can target a connection on a different instance — it always goes through the Redis relay in `src/service/redis.rs`, never a direct in-process - lookup. Sync replies are the exception: their router is owned by one local - socket, so they go directly to that socket's queue. -- Resolve `SyncServiceUrl` at startup; Docker uses - `OVERRIDE_SYNC_SERVICE_URL=http://sync-service:8787`. -- Complete the gateway rollout before enabling `sync-via-gateway` in PostHog; - there is no separate router deployment. See `docs/architecture/sync-gateway.md`. + lookup; don't bypass it when adding a new send path. ## Related diff --git a/services/connection_gateway/Cargo.toml b/services/connection_gateway/Cargo.toml index de7771c8de9..f4e6d6486e7 100644 --- a/services/connection_gateway/Cargo.toml +++ b/services/connection_gateway/Cargo.toml @@ -59,9 +59,9 @@ macro_config = { path = "../../crates/macro_config" } macro_entrypoint = { path = "../../crates/macro_entrypoint" } macro_env = { path = "../../crates/macro_env" } macro_env_var = { path = "../../crates/macro_env_var" } -macro_service_urls = { path = "../../crates/macro_service_urls" } macro_tower_layers = { path = "../../crates/macro_tower_layers" } mockall.workspace = true +postcard = { workspace = true } model = { path = "../../crates/model" } model-entity = { path = "../../crates/model-entity" } models_bulk_upload = { path = "../../crates/models_bulk_upload" } @@ -79,9 +79,7 @@ sqlx = { workspace = true } strum = { workspace = true, features = ["derive"] } strum_macros = { workspace = true } stream = { path = "../../crates/stream" } -sync_router = { path = "../../crates/sync_router" } tokio = { workspace = true } -tokio-util = { workspace = true, features = ["rt"] } tower = { workspace = true } tower-http = { workspace = true } tracing = { workspace = true, features = ["attributes"] } @@ -94,8 +92,5 @@ workspace-hack = { version = "0.1", path = "../../crates/workspace-hack" } doppler-config = { workspace = true } [dev-dependencies] -bebop = { workspace = true } -sync_service_bebop_schema = { path = "../sync-service/bebop-schema" } -tokio-tungstenite = { workspace = true } cool_asserts = "2" macro_auth = { path = "../../crates/macro_auth", features = ["testing"] } diff --git a/services/connection_gateway/README.md b/services/connection_gateway/README.md index 2cfc9599d8c..17743775f7f 100644 --- a/services/connection_gateway/README.md +++ b/services/connection_gateway/README.md @@ -2,20 +2,6 @@ A service responsible for managing connections, and allowing services to send messages to a specific entity over a websocket connection. -## Document sync - -Binary WebSocket messages use the Bebop `ToRouter` / `FromRouter` envelope. -Each socket embeds `crates/sync_router`, which opens one downstream sync-service -connection per document and passes the inner chunked protocol unchanged. -Document tokens are validated by the sync service. Closing the client socket -cancels all of its downstream work, including in-flight handshakes. - -This ships in the existing gateway deployment. No additional service, Redis -fanout, or Pulumi stack is needed. `SyncServiceUrl` selects the existing dev/prod -sync origin; local Docker overrides it with `OVERRIDE_SYNC_SERVICE_URL`. -Deploy all gateway replicas before enabling PostHog `sync-via-gateway`. -See [sync gateway architecture](../../docs/architecture/sync-gateway.md). - ## Local Development *runs on port 8085* diff --git a/services/connection_gateway/src/api/connection.rs b/services/connection_gateway/src/api/connection.rs index 44237a5e66e..8b180879a41 100644 --- a/services/connection_gateway/src/api/connection.rs +++ b/services/connection_gateway/src/api/connection.rs @@ -53,10 +53,8 @@ pub async fn ws_handler( let macro_user_id = authorization.authorization.user.macro_user_id.clone(); let user_context = authorization.authorization.user.user_context.clone(); - let span = tracing::Span::current(); ws.on_upgrade(move |socket| { handle_websocket_connection(socket, ctx, config, macro_user_id, user_context) - .instrument(span) }) } @@ -74,14 +72,13 @@ async fn handle_websocket_connection( ) { let (sink, stream) = socket.split(); let (sender, receiver) = tokio::sync::mpsc::channel::(100); - let connection_id = uuid::Uuid::now_v7().to_string(); + let connection_id = uuid::Uuid::new_v4().to_string(); // Create guard that records last online time when websocket connection closes let last_online_guard = ctx.last_online_worker.new_guard(macro_user_id.clone()); let sender_connection_id = connection_id.clone(); - let mut sender_task = - tokio::spawn(forwarder(sink, receiver, sender_connection_id).in_current_span()); + let sender_task = tokio::spawn(forwarder(sink, receiver, sender_connection_id)); if let Err(err) = ctx .connection_manager @@ -106,19 +103,12 @@ async fn handle_websocket_connection( connection_id: &connection_id, }; - let sync_session = crate::service::sync::SyncSession::new( - connection_id.clone(), - ctx.sync_service_url.clone(), - sender.clone(), - ); - let receiver_task = - handle_websocket_stream(connection_context, stream, sender.clone(), &sync_session).fuse(); + crate::service::fanout::connected(&connection_context, ¯o_user_id).await; + + let receiver_task = handle_websocket_stream(connection_context, stream, sender.clone()).fuse(); tokio::select! { - _ = sync_session.closed() => { - tracing::warn!(connection.id = %connection_id, "sync routing stopped; closing socket"); - } - res = &mut sender_task => { + res = sender_task => { res.inspect(|_| tracing::debug!("sender task finished")) .inspect_err(|err| { tracing::error!(error=?err, "sender task failed"); @@ -132,15 +122,8 @@ async fn handle_websocket_connection( } } - // A blocked write must not keep the socket alive after routing stops. - sender_task.abort(); - if !sender_task.is_finished() { - let _ = sender_task.await; - } - sync_session.shutdown().await; - ctx.connection_manager - .remove_connection(&connection_id) + .remove_connection(&connection_id.clone()) .await .inspect(|_| tracing::debug!("connections removed")) .inspect_err(|e| { @@ -153,6 +136,9 @@ async fn handle_websocket_connection( }) .ok(); + // After remove_connection, so consumers never see "gone" while the + // gateway can still route to the socket. + crate::service::fanout::disconnected(&connection_context).await; drop(last_online_guard); } diff --git a/services/connection_gateway/src/api/connection/messages.rs b/services/connection_gateway/src/api/connection/messages.rs index ee1849863ea..506ab63d0e9 100644 --- a/services/connection_gateway/src/api/connection/messages.rs +++ b/services/connection_gateway/src/api/connection/messages.rs @@ -5,7 +5,7 @@ use crate::{ tracking::{EntityConnectionExt, TrackingData}, websocket::ToWebsocketMessage, }, - service::{stream_event::handle_stream_events, sync::SyncSession, tracker}, + service::{fanout, stream_event::handle_stream_events, tracker}, }; use anyhow::{Context, Result}; use axum::extract::ws::{Message, WebSocket}; @@ -18,13 +18,11 @@ pub async fn handle_websocket_stream( connection_context: ConnectionContext<'_>, mut stream: SplitStream, sender: Sender, - sync_session: &SyncSession, ) -> Result<()> { while let Some(msg) = stream.next().await { match msg { Ok(msg) => { - if let Err(e) = handle_message(connection_context, msg, &sender, sync_session).await - { + if let Err(e) = handle_message(connection_context, msg, &sender).await { tracing::warn!(error=?e, "error handling message"); } } @@ -51,10 +49,24 @@ pub async fn handle_message( connection_context: ConnectionContext<'_>, message: Message, sender: &Sender, - sync_session: &SyncSession, ) -> Result<()> { - if let Message::Binary(payload) = message { - return sync_session.frame(payload.to_vec()).await; + // Every client frame is republished to the fanout channel, unparsed, + // before any in-gateway handling. Backend services consume what they care + // about (the sync tier reads binary frames); the gateway stays a dumb tap. + // + // TODO: the in-gateway handling below (entity tracking / frecency, + // stream-event subscriptions) should probably migrate to services + // consuming this same fanout, keeping the gateway a pure edge. + match &message { + Message::Text(text) => { + fanout::frame(&connection_context, true, text.as_bytes().to_vec()).await; + } + Message::Binary(payload) => { + fanout::frame(&connection_context, false, payload.to_vec()).await; + // Binary frames have no in-gateway handling; fanout is it. + return Ok(()); + } + _ => {} } let text_message = match message { diff --git a/services/connection_gateway/src/context.rs b/services/connection_gateway/src/context.rs index eaf518aefda..2b48abdd894 100644 --- a/services/connection_gateway/src/context.rs +++ b/services/connection_gateway/src/context.rs @@ -18,8 +18,9 @@ pub type AuthorizationService = MacroAuthorizationServiceImpl, pub frecency_ingestor_service: EventIngestorImpl, pub redis_client: Arc, pub redis_connection: MultiplexedConnection, diff --git a/services/connection_gateway/src/main.rs b/services/connection_gateway/src/main.rs index a835a12d9d8..2ef4f22f891 100644 --- a/services/connection_gateway/src/main.rs +++ b/services/connection_gateway/src/main.rs @@ -114,23 +114,10 @@ async fn main() -> Result<()> { let stream_service = RedisPostgresStreamRepo::new((*redis_client).clone(), pgpool.clone()); let stream_manager = RedisPostgresStreamManager::new(stream_service.obj()); - let sync_service_url = url::Url::parse( - macro_service_urls::SyncServiceUrl::new_for_environment(config.environment)?.as_ref(), - ) - .context("invalid sync service URL")?; - anyhow::ensure!( - matches!(sync_service_url.scheme(), "http" | "https") - && sync_service_url.has_host() - && sync_service_url.path() == "/" - && sync_service_url.query().is_none() - && sync_service_url.fragment().is_none() - && sync_service_url.username().is_empty() - && sync_service_url.password().is_none(), - "sync service URL must be an HTTP(S) origin", - ); + let fanout_gateway_id: Arc = Arc::from(uuid::Uuid::new_v4().to_string()); let context = context::ApiContext { connection_manager, - sync_service_url, + fanout_gateway_id: Arc::clone(&fanout_gateway_id), redis_client: Arc::clone(&redis_client), redis_connection, frecency_ingestor_service: EventIngestorImpl::new(FrecencyPgStorage::new(pgpool.clone())), @@ -139,6 +126,11 @@ async fn main() -> Result<()> { }; tokio::spawn(poll_messages(context.clone())); + tokio::spawn(service::fanout::poll_outbound(context.clone())); + tokio::spawn(service::fanout::heartbeat_loop( + context.redis_connection.clone(), + fanout_gateway_id, + )); let config = Arc::new(config); let authorization_state = MacroAuthorizationState::new(Arc::new(AuthorizationService::new( diff --git a/services/connection_gateway/src/model/message.rs b/services/connection_gateway/src/model/message.rs index a99f692bfdc..e3267430568 100644 --- a/services/connection_gateway/src/model/message.rs +++ b/services/connection_gateway/src/model/message.rs @@ -60,7 +60,7 @@ impl TryFrom for Message { pub enum OutgoingMessage { Pong, Message(Message), - /// A raw binary frame from this socket's embedded sync router, + /// A raw binary frame (sync-protocol traffic from the fanout's consumers), /// sent to the client as a binary websocket message, unparsed. Binary(Vec), } diff --git a/services/connection_gateway/src/service/fanout.rs b/services/connection_gateway/src/service/fanout.rs new file mode 100644 index 00000000000..51a0f385b58 --- /dev/null +++ b/services/connection_gateway/src/service/fanout.rs @@ -0,0 +1,166 @@ +//! Republish this instance's entire inbound websocket traffic to Redis. +//! +//! Every accepted connection, every client frame (text or binary, unparsed), +//! and every disconnect is published on +//! [`connection_gateway_models::fanout::INBOUND_CHANNEL`], so backend services +//! can consume client traffic without the gateway knowing they exist. The +//! sync tier is the first consumer. +//! +//! Publishing is fire-and-forget by design: a publish failure is logged and +//! never fails the connection's own message handling. + +use crate::context::ApiContext; +use crate::model::connection::ConnectionContext; +use anyhow::{Context, Result}; +use connection_gateway_models::fanout::{ + ConnId, FromGateway, GatewayId, HEARTBEAT_INTERVAL_SECS, INBOUND_CHANNEL, ToGateway, + outbound_channel, +}; +use futures::StreamExt; +use macro_user_id::user_id::MacroUserIdStr; +use redis::AsyncCommands; + +async fn publish(ctx: &ConnectionContext<'_>, message: &FromGateway) -> Result<()> { + let payload = postcard::to_stdvec(message).context("failed to encode fanout message")?; + let mut connection = ctx.api_context.get_multiplexed_async_connection()?; + connection + .publish::<&str, &[u8], ()>(INBOUND_CHANNEL, &payload) + .await + .context("failed to publish fanout message")?; + Ok(()) +} + +/// Log-and-continue wrapper: fanout must never break connection handling. +async fn publish_best_effort(ctx: &ConnectionContext<'_>, message: &FromGateway) { + publish(ctx, message) + .await + .inspect_err(|error| tracing::warn!(error=?error, "failed to publish fanout message")) + .ok(); +} + +/// Announce an accepted, authenticated connection. +pub async fn connected(ctx: &ConnectionContext<'_>, user_id: &MacroUserIdStr<'static>) { + publish_best_effort( + ctx, + &FromGateway::Connected { + gateway: GatewayId(ctx.api_context.fanout_gateway_id.to_string()), + conn: ConnId(ctx.connection_id.to_string()), + user_id: user_id.clone(), + }, + ) + .await; +} + +/// Forward one client frame, unparsed. +pub async fn frame(ctx: &ConnectionContext<'_>, text: bool, payload: Vec) { + publish_best_effort( + ctx, + &FromGateway::Frame { + gateway: GatewayId(ctx.api_context.fanout_gateway_id.to_string()), + conn: ConnId(ctx.connection_id.to_string()), + text, + payload, + }, + ) + .await; +} + +/// Announce a closed connection. +pub async fn disconnected(ctx: &ConnectionContext<'_>) { + publish_best_effort( + ctx, + &FromGateway::Disconnected { + gateway: GatewayId(ctx.api_context.fanout_gateway_id.to_string()), + conn: ConnId(ctx.connection_id.to_string()), + }, + ) + .await; +} + +/// Liveness beacon so consumers can drop state for dead gateway instances. +/// Spawned once at boot. +pub async fn heartbeat_loop( + redis_connection: redis::aio::MultiplexedConnection, + gateway_id: std::sync::Arc, +) { + let mut tick = tokio::time::interval(std::time::Duration::from_secs(HEARTBEAT_INTERVAL_SECS)); + loop { + tick.tick().await; + let message = FromGateway::Heartbeat { + gateway: GatewayId(gateway_id.to_string()), + }; + let Ok(payload) = postcard::to_stdvec(&message) else { + continue; + }; + let mut connection = redis_connection.clone(); + connection + .publish::<&str, &[u8], ()>(INBOUND_CHANNEL, &payload) + .await + .inspect_err(|error| tracing::warn!(error=?error, "failed to publish fanout heartbeat")) + .ok(); + } +} + +/// Consume this instance's outbound fanout channel: frames published by +/// backend consumers (the sync tier) addressed to connections whose sockets +/// this instance holds. Mirrors [`crate::service::redis::poll_messages`]. +pub async fn poll_outbound(ctx: ApiContext) -> Result<()> { + let gateway_id = GatewayId(ctx.fanout_gateway_id.to_string()); + let channel = outbound_channel(&gateway_id); + + let (mut sink, mut stream) = ctx.redis_client.get_async_pubsub().await?.split(); + sink.subscribe(&channel) + .await + .context("failed to subscribe to fanout outbound channel")?; + tracing::debug!(channel, "subscribed to fanout outbound channel"); + + while let Some(message) = stream.next().await { + let payload: Vec = match message.get_payload() { + Ok(payload) => payload, + Err(error) => { + tracing::warn!(error = ?error, "unreadable outbound fanout payload; skipping"); + continue; + } + }; + match postcard::from_bytes::(&payload) { + Ok(ToGateway::Frame { + conn, + text: false, + payload, + }) => { + ctx.connection_manager + .send_binary(&conn.0, payload) + .await + .inspect_err(|error| { + tracing::warn!(error = ?error, conn = %conn, "failed to deliver outbound frame"); + }) + .ok(); + } + Ok(ToGateway::Frame { + conn, text: true, .. + }) => { + // No consumer sends text frames yet; refuse rather than guess + // at framing for the existing JSON protocol. + tracing::warn!(conn = %conn, "text outbound frames are not supported; dropping"); + } + Ok(ToGateway::Close { conn, code }) => { + // TODO: plumb a real close (with code) through the connection + // manager; for now drop the connection, which closes the socket. + tracing::debug!(conn = %conn, code, "closing connection at consumer request"); + ctx.connection_manager + .remove_connection(&conn.0) + .await + .inspect_err(|error| { + tracing::warn!(error = ?error, conn = %conn, "failed to close connection"); + }) + .ok(); + } + Err(error) => { + tracing::warn!(error = ?error, "undecodable outbound fanout message; skipping"); + } + } + } + + tracing::warn!("fanout outbound poller exited"); + Ok(()) +} diff --git a/services/connection_gateway/src/service/mod.rs b/services/connection_gateway/src/service/mod.rs index 6395ab8dfd4..02eb553f0c2 100644 --- a/services/connection_gateway/src/service/mod.rs +++ b/services/connection_gateway/src/service/mod.rs @@ -1,7 +1,7 @@ pub mod connection; pub mod dynamodb; +pub mod fanout; pub mod redis; pub mod sender; pub mod stream_event; -pub mod sync; pub mod tracker; diff --git a/services/document_storage_service/AGENTS.md b/services/document_storage_service/AGENTS.md index 6612d2707a7..18e1ba6c41d 100644 --- a/services/document_storage_service/AGENTS.md +++ b/services/document_storage_service/AGENTS.md @@ -8,6 +8,10 @@ threads, pins, recents, notifications, folders, activity. ## Start here +- `src/api/sync.rs`: dedicated authenticated `/sync/ws` endpoint; `crates/sync_router` + owns multiplex routing and cancellation. Sync-service still owns document state + and validates scoped tokens. DSS shutdown drains these upgraded sockets. + This binary WebSocket protocol is intentionally outside the REST OpenAPI/SDK surface. - `src/api.rs`: route module list — `documents`, `items`, `entity`, `folders_query`, `graphql_soup`, `threads`, `annotations`, `history`, `pins`, `recents`, `notification`, `instructions`, `internal`. diff --git a/services/document_storage_service/Cargo.toml b/services/document_storage_service/Cargo.toml index a639f8facf2..90d2f8188b8 100644 --- a/services/document_storage_service/Cargo.toml +++ b/services/document_storage_service/Cargo.toml @@ -190,6 +190,8 @@ macro_queues = { path = "../../crates/macro_queues" } macro_service_urls = { path = "../../crates/macro_service_urls" } macro_sha_count_client = { path = "../../crates/macro_sha_count_client" } macro_sync_service_jwt = { path = "../../crates/macro_sync_service_jwt" } +opentelemetry = { workspace = true } +tracing-opentelemetry = { workspace = true } macro_tower_layers = { path = "../../crates/macro_tower_layers" } macro_user_id = { path = "../../crates/macro_user_id" } macro_uuid = { path = "../../crates/macro_uuid" } @@ -266,6 +268,7 @@ sqs_client = { path = "../../crates/sqs_client", default-features = false, featu ] } sqs_worker = { path = "../../crates/sqs_worker" } static_file_service_client = { path = "../../crates/static_file_service_client" } +sync_router = { path = "../../crates/sync_router" } sync_service_client = { path = "../../crates/sync_service_client" } sync_service_hex = { path = "../../crates/sync_service", package = "sync_service_hex", features = [ "inbound", @@ -299,6 +302,9 @@ workspace-hack = { version = "0.1", path = "../../crates/workspace-hack" } zip = { workspace = true } [dev-dependencies] +bebop = { workspace = true } +sync_service_bebop_schema = { path = "../sync-service/bebop-schema" } +tokio-tungstenite = { workspace = true } cool_asserts = "2" flate2 = "1" frecency = { path = "../../crates/frecency", features = ["mock"] } diff --git a/services/document_storage_service/src/api.rs b/services/document_storage_service/src/api.rs index bcce70a4cf7..f1cc07aa111 100644 --- a/services/document_storage_service/src/api.rs +++ b/services/document_storage_service/src/api.rs @@ -19,6 +19,7 @@ use utoipa_swagger_ui::SwaggerUi; // Utilities pub(crate) mod context; mod saved_views; +mod sync; pub(crate) mod util; // Middleware @@ -49,6 +50,7 @@ mod threads; pub const MACRO_INTERNAL_USER_ID: &str = "macro|INTERNAL@macro.com"; pub async fn setup_and_serve(state: ApiContext) -> anyhow::Result<()> { + let sync_sessions = state.sync_sessions.clone(); let app = api_router(state.clone()) .layer( ServiceBuilder::new() @@ -77,10 +79,16 @@ pub async fn setup_and_serve(state: ApiContext) -> anyhow::Result<()> { &state.config.environment, &state.config.port ); - axum::serve(listener, app.into_make_service()) - .with_graceful_shutdown(macro_entrypoint::shutdown_signal()) + let shutdown_sessions = sync_sessions.clone(); + let result = axum::serve(listener, app.into_make_service()) + .with_graceful_shutdown(async move { + macro_entrypoint::shutdown_signal().await; + shutdown_sessions.cancel(); + }) .await - .context("error starting service") + .context("error starting service"); + sync_sessions.shutdown().await; + result } fn items_router(state: ApiContext) -> Router { @@ -124,6 +132,7 @@ fn api_router(state: ApiContext) -> Router { ); let internal_router = Router::new() + .merge(sync::router()) .nest( "/dictation", dictation::inbound::axum_router::dictation_router(state.dictation_state.clone()), diff --git a/services/document_storage_service/src/api/context.rs b/services/document_storage_service/src/api/context.rs index 3e6432ae739..e4af048cdf4 100644 --- a/services/document_storage_service/src/api/context.rs +++ b/services/document_storage_service/src/api/context.rs @@ -901,6 +901,7 @@ pub(crate) struct ApiContext { pub notification_ingress_service: Arc, pub conn_gateway_client: Arc, pub sync_service_client: Arc, + pub sync_sessions: sync_router::session::SyncSessionFactory, pub system_properties_service: Arc, pub properties_service: Arc, pub opensearch_client: Arc, diff --git a/services/document_storage_service/src/api/sync.rs b/services/document_storage_service/src/api/sync.rs new file mode 100644 index 00000000000..fad64a43ea7 --- /dev/null +++ b/services/document_storage_service/src/api/sync.rs @@ -0,0 +1,114 @@ +//! Dedicated multiplexed document-sync transport. Authorization of each document +//! remains in sync-service's scoped-token handshake. + +use super::context::{ApiContext, AuthorizationService}; +use axum::{ + Router, + extract::{ + Query, State, WebSocketUpgrade, + ws::{Message, WebSocket}, + }, + http::StatusCode, + response::{IntoResponse, Response}, + routing::get, +}; +use axum_extra::extract::Cached; +use macro_authorization::{OptionalMacroAuthorizationExtractor, UserOrInternalService}; +use opentelemetry::trace::TraceContextExt; +use std::time::Duration; +use sync_router::session::{ + CLIENT_QUEUE_CAPACITY, MAX_FRAME_BYTES, SyncSession, SyncSessionFactory, +}; +use tokio::sync::mpsc; +use tracing::Instrument; +use tracing_opentelemetry::OpenTelemetrySpanExt; + +#[cfg(test)] +mod test; + +pub(super) fn router() -> Router { + Router::new().route("/sync/ws", get(upgrade)) +} + +#[derive(Default, serde::Deserialize)] +struct TraceQuery { + traceparent: Option, +} + +async fn upgrade( + State(state): State, + Cached(auth): Cached< + OptionalMacroAuthorizationExtractor, + >, + Query(trace): Query, + upgrade: WebSocketUpgrade, +) -> Response { + if auth + .authorization + .as_ref() + .and_then(|auth| auth.acting_user()) + .is_none() + { + return (StatusCode::UNAUTHORIZED, "authentication required for sync").into_response(); + } + upgrade_session(upgrade, &state.sync_sessions, trace.traceparent.as_deref()) +} + +fn upgrade_session( + upgrade: WebSocketUpgrade, + factory: &SyncSessionFactory, + traceparent: Option<&str>, +) -> Response { + let connection_id = uuid::Uuid::now_v7().to_string(); + let span = tracing::info_span!("sync.websocket", connection.id = %connection_id); + if let Some(traceparent) = traceparent { + let carrier = + std::collections::HashMap::from([("traceparent".to_string(), traceparent.to_string())]); + let parent = opentelemetry::global::get_text_map_propagator(|propagator| { + propagator.extract(&carrier) + }); + if parent.span().span_context().is_valid() { + let _ = span.set_parent(parent); + } + } + + let (sender, receiver) = mpsc::channel(CLIENT_QUEUE_CAPACITY); + let session = span.in_scope(|| factory.open(connection_id, sender)); + upgrade + .max_message_size(MAX_FRAME_BYTES) + .max_frame_size(MAX_FRAME_BYTES) + .max_write_buffer_size(2 * MAX_FRAME_BYTES) + .on_upgrade(move |socket| serve(socket, session, receiver).instrument(span)) +} + +async fn serve(mut socket: WebSocket, session: SyncSession, mut frames: mpsc::Receiver>) { + // Cancellation surrounds all socket IO, including a blocked write or event + // queue. Upgraded sockets therefore drain when DSS receives SIGTERM. + tokio::select! { + biased; + _ = session.closed() => {}, + _ = async { + loop { + tokio::select! { + incoming = socket.recv() => match incoming { + Some(Ok(Message::Binary(bytes))) => { + if session.frame(bytes.to_vec()).await.is_err() { break; } + } + Some(Ok(Message::Text(text))) if text == "ping" => { + if socket.send(Message::Text("pong".into())).await.is_err() { break; } + } + Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => {}, + _ => break, + }, + frame = frames.recv() => { + let Some(frame) = frame else { break; }; + if socket.send(Message::Binary(frame.into())).await.is_err() { break; } + }, + _ = tokio::time::sleep(Duration::from_secs(60)) => break, + } + } + } => {}, + } + let _ = tokio::time::timeout(Duration::from_secs(1), socket.send(Message::Close(None))).await; + session.shutdown().await; +} diff --git a/services/document_storage_service/src/api/sync/test.rs b/services/document_storage_service/src/api/sync/test.rs new file mode 100644 index 00000000000..8dc059184a1 --- /dev/null +++ b/services/document_storage_service/src/api/sync/test.rs @@ -0,0 +1,123 @@ +use super::*; +use bebop::{Record, SliceWrapper}; +use futures::{SinkExt, StreamExt}; +use sync_service_bebop_schema::{FromRouter, ToRouter}; +use tokio::net::TcpListener; +use tokio_tungstenite::{accept_async, connect_async, tungstenite::Message as WsMessage}; + +#[tokio::test] +async fn dedicated_endpoint_preserves_large_chunks_heartbeats_and_drains_on_shutdown() { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let factory = SyncSessionFactory::new( + format!("http://{}", upstream.local_addr().unwrap()) + .parse() + .unwrap(), + ) + .unwrap(); + let downstream = tokio::spawn(async move { + let (socket, _) = upstream.accept().await.unwrap(); + let mut socket = accept_async(socket).await.unwrap(); + while let Some(Ok(message)) = socket.next().await { + if let WsMessage::Binary(bytes) = message { + socket.send(WsMessage::Binary(bytes)).await.unwrap(); + } + } + }); + // Authentication is exercised separately against the deployed route. This + // mounts the same upgrade/IO adapter with a session already authorized. + let host = factory.clone(); + let app = Router::new().route( + "/sync/ws", + get(move |upgrade: WebSocketUpgrade| async move { upgrade_session(upgrade, &host, None) }), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let (mut client, _) = connect_async(format!("ws://{address}/sync/ws")) + .await + .unwrap(); + client.send(WsMessage::Text("ping".into())).await.unwrap(); + assert_eq!( + client.next().await.unwrap().unwrap(), + WsMessage::Text("pong".into()) + ); + let mut subscribe = Vec::new(); + ToRouter::RouterSubscribe { + doc_id: "document", + token: "scoped-token", + } + .serialize(&mut subscribe) + .unwrap(); + client + .send(WsMessage::Binary(subscribe.into())) + .await + .unwrap(); + let ack = client.next().await.unwrap().unwrap().into_data(); + assert!(matches!( + FromRouter::deserialize(&ack).unwrap(), + FromRouter::RouterSubscribed { .. } + )); + let chunk = vec![1; 900_001]; + let mut frame = Vec::new(); + ToRouter::RouterFrame { + doc_id: "document", + payload: SliceWrapper::Raw(&chunk), + } + .serialize(&mut frame) + .unwrap(); + client.send(WsMessage::Binary(frame.into())).await.unwrap(); + let response = client.next().await.unwrap().unwrap().into_data(); + assert_eq!( + response.as_ref(), + sync_router::domain::envelope::doc_frame("document", &chunk) + ); + tokio::time::timeout(Duration::from_secs(2), factory.shutdown()) + .await + .unwrap(); + assert!(matches!( + client.next().await, + Some(Ok(WsMessage::Close(_))) | None + )); + downstream.await.unwrap(); + server.abort(); +} + +#[tokio::test] +async fn oversized_websocket_message_closes_only_that_session() { + let factory = SyncSessionFactory::new("http://127.0.0.1:1".parse().unwrap()).unwrap(); + let host = factory.clone(); + let app = Router::new().route( + "/sync/ws", + get(move |upgrade: WebSocketUpgrade| async move { upgrade_session(upgrade, &host, None) }), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let (mut client, _) = connect_async(format!("ws://{address}/sync/ws")) + .await + .unwrap(); + let (mut other, _) = connect_async(format!("ws://{address}/sync/ws")) + .await + .unwrap(); + client + .send(WsMessage::Binary(vec![0; MAX_FRAME_BYTES + 1].into())) + .await + .unwrap(); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(2), client.next()) + .await + .unwrap(), + Some(Ok(WsMessage::Close(_))) | Some(Err(_)) | None + )); + other.send(WsMessage::Text("ping".into())).await.unwrap(); + assert_eq!( + other.next().await.unwrap().unwrap(), + WsMessage::Text("pong".into()) + ); + factory.shutdown().await; + server.abort(); +} diff --git a/services/document_storage_service/src/server.rs b/services/document_storage_service/src/server.rs index d121a6ab30d..cb137b8dce9 100644 --- a/services/document_storage_service/src/server.rs +++ b/services/document_storage_service/src/server.rs @@ -2126,7 +2126,10 @@ pub async fn run() -> anyhow::Result<()> { complete_graph::EntityAccessWorkFeedViewers(entity_access_service.clone()), ); + let sync_sessions = + sync_router::session::SyncSessionFactory::new(SyncServiceUrl::new()?.as_ref().parse()?)?; let api_context = ApiContext { + sync_sessions, dictation_state, graphql_work_feed_context, contacts_ingress: contacts_ingress.clone(), diff --git a/services/sync-service/bebop-schema/src/generated.rs b/services/sync-service/bebop-schema/src/generated.rs index 9c8f984588b..7c33607d941 100644 --- a/services/sync-service/bebop-schema/src/generated.rs +++ b/services/sync-service/bebop-schema/src/generated.rs @@ -545,7 +545,7 @@ impl<'raw> ::bebop::SubRecord<'raw> for InitializeFromSnapshotRequest<'raw> { impl<'raw> ::bebop::Record<'raw> for InitializeFromSnapshotRequest<'raw> {} /// Multiplex envelope between the browser and the sync-router (rides the -/// connection-gateway socket as binary frames). Inner payloads are the +/// dedicated DSS /sync/ws socket as binary frames). Inner payloads are the /// existing FromPeer/FromRemote unions, already bebop-encoded — the router /// forwards them without parsing. #[derive(Clone, Debug, PartialEq)] diff --git a/services/sync-service/bebop/generated/schema.ts b/services/sync-service/bebop/generated/schema.ts index 18ffe7ff9d1..fd309141cf5 100644 --- a/services/sync-service/bebop/generated/schema.ts +++ b/services/sync-service/bebop/generated/schema.ts @@ -2257,7 +2257,7 @@ export class RouterFrame implements IRouterFrame { /** * Multiplex envelope between the browser and the sync-router (rides the - * connection-gateway socket as binary frames). Inner payloads are the + * dedicated DSS /sync/ws socket as binary frames). Inner payloads are the * existing FromPeer/FromRemote unions, already bebop-encoded — the router * forwards them without parsing. */ diff --git a/services/sync-service/bebop/schema.bop b/services/sync-service/bebop/schema.bop index 2cc08b4b069..15d60537f6e 100644 --- a/services/sync-service/bebop/schema.bop +++ b/services/sync-service/bebop/schema.bop @@ -58,7 +58,7 @@ struct InitializeFromSnapshotRequest { /* * Multiplex envelope between the browser and the sync-router (rides the - * connection-gateway socket as binary frames). Inner payloads are the + * dedicated DSS /sync/ws socket as binary frames). Inner payloads are the * existing FromPeer/FromRemote unions, already bebop-encoded — the router * forwards them without parsing. */ diff --git a/tooling/just/check.just b/tooling/just/check.just index fb1eba27aa0..076c1586450 100644 --- a/tooling/just/check.just +++ b/tooling/just/check.just @@ -29,8 +29,8 @@ check *ARGS: if [ -z "$changed" ]; then echo "no changed files — nothing to check"; exit 0; fi rs=$(grep -E '^(crates|services|tooling)/.*\.rs$' <<<"$changed" || true) - # Match packages/sdk/biome.jsonc: preserve generator output verbatim. - js=$(grep -E '^(apps/(web|marketing)|packages|services/(analytics-proxy|ai-editing-worker|lexical-service|websocket-service))/.*\.(ts|tsx)$' <<<"$changed" | grep -v '^packages/sdk/generated/' || true) + # Match the SDK and collaboration Biome configs: preserve generator output verbatim. + js=$(grep -E '^(apps/(web|marketing)|packages|services/(analytics-proxy|ai-editing-worker|lexical-service|websocket-service))/.*\.(ts|tsx)$' <<<"$changed" | grep -Ev '^packages/(sdk/generated|collaboration/src/sync-service/generated)/' || true) sg=$(grep -E '^(apps|crates|packages|services|tooling)/.*\.(rs|ts|tsx)$' <<<"$changed" || true) fail=0