diff --git a/src/connection.rs b/src/connection.rs index d74b333..4f512b9 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -6,23 +6,26 @@ use std::sync::Arc; use tokio::sync::mpsc; use uuid::Uuid; +// --------------------------------------------------------------------------- +// RenderTrigger +// --------------------------------------------------------------------------- + +/// Signal sent to a connection task to request a rerender. +pub enum RenderTrigger { + Invalidated, +} + // --------------------------------------------------------------------------- // Connection state // --------------------------------------------------------------------------- /// Per-connection state stored in the manager. pub struct ConnectionState { - /// Sender half of the per-connection SSE channel. - /// - /// The fanout task sends fully-constructed axum `Event`s here; the SSE - /// handler streams them directly to the client without re-wrapping. - pub sse_sender: mpsc::Sender, - - /// The full signal JSON from the most recent render for this connection. - /// - /// Stored so the fanout task can reconstruct a `SignalStore` and build a - /// `RenderContext` for server-initiated rerenders, without a client - /// request carrying the signals. + /// Sends render triggers to the connection task. + pub trigger_tx: mpsc::Sender, + /// Sends new SSE pipe senders to the connection task on reconnect. + pub pipe_tx: mpsc::Sender>, + /// Last known signals JSON, updated by action handlers. pub last_signals_json: Option, } @@ -35,8 +38,8 @@ pub struct ConnectionState { /// Accessed from: /// - `GET /sse` handler (insert) /// - `POST /action` handler (update signals) -/// - SSE drop cleanup (remove) -/// - Fanout task (read sender + signals for rerender) +/// - Fanout task (get trigger sender for rerender) +/// - Connection task (cleanup on grace period expiry) pub struct ConnectionManager { connections: DashMap, } @@ -48,26 +51,29 @@ impl ConnectionManager { } } - /// Register a new connection with its SSE sender. - pub fn insert(&self, id: Uuid, sender: mpsc::Sender) { + /// Register a new connection with its trigger and pipe senders. + pub fn insert(&self, id: Uuid, trigger_tx: mpsc::Sender, pipe_tx: mpsc::Sender>) { self.connections.insert( id, ConnectionState { - sse_sender: sender, + trigger_tx, + pipe_tx, last_signals_json: None, }, ); } - /// Get the SSE sender for a connection. - pub fn get_sender(&self, id: &Uuid) -> Option> { - self.connections.get(id).map(|c| c.sse_sender.clone()) + /// Get the render trigger sender for a connection. + pub fn get_trigger_sender(&self, id: &Uuid) -> Option> { + self.connections.get(id).map(|c| c.trigger_tx.clone()) + } + + /// Get the pipe sender for a connection (used to hand off a new SSE pipe on reconnect). + pub fn get_pipe_sender(&self, id: &Uuid) -> Option>> { + self.connections.get(id).map(|c| c.pipe_tx.clone()) } /// Get the last signals JSON for a connection. - /// - /// Called by the fanout task to reconstruct signal state for a - /// server-initiated rerender. pub fn get_signals(&self, id: &Uuid) -> Option { self.connections .get(id) @@ -75,15 +81,13 @@ impl ConnectionManager { } /// Update the stored signals for a connection. - /// - /// Called by `POST /action` after ingesting the client's signal blob. pub fn update_signals(&self, id: &Uuid, signals_json: String) { if let Some(mut conn) = self.connections.get_mut(id) { conn.last_signals_json = Some(signals_json); } } - /// Remove a connection (called on SSE pipe drop). + /// Remove a connection. pub fn remove(&self, id: &Uuid) -> Option<(Uuid, ConnectionState)> { self.connections.remove(id) } @@ -168,4 +172,4 @@ impl SubscriptionRegistry { .map(|c| c.clone()) .unwrap_or_default() } -} \ No newline at end of file +} diff --git a/src/lib.rs b/src/lib.rs index 642dea2..91af1ef 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -30,7 +30,7 @@ pub use store::{Store, KVStore}; #[cfg(feature = "pg_replication")] pub use pg_query::PgQuery; pub use context::{RenderContext, ComponentTree, ComponentNode}; -pub use connection::{ConnectionManager, ConnectionState, SubscriptionRegistry}; +pub use connection::{ConnectionManager, ConnectionState, SubscriptionRegistry, RenderTrigger}; pub use sse::{format_patch_elements, format_patch_signals}; pub use server::{AppState, SharedState, RenderFn, sse_handler, spawn_fanout, ingest_signals, ActionParams, SseParams}; pub use config::{LiveVueConfig, ServerConfig, BrotliConfig, SharedConfig, load_config, load_and_watch_config}; diff --git a/src/server.rs b/src/server.rs index 6f82e63..348c8ad 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,5 +1,5 @@ use crate::config::{LiveVueConfig, SharedConfig}; -use crate::connection::{ConnectionManager, SubscriptionRegistry}; +use crate::connection::{ConnectionManager, RenderTrigger, SubscriptionRegistry}; use crate::context::{ComponentTree, RenderContext}; use crate::query::CacheKey; use crate::signal::SignalStore; @@ -12,9 +12,13 @@ use std::convert::Infallible; use std::sync::Arc; use tokio::sync::mpsc; use tokio::sync::RwLock; +use tokio::time::{sleep, Duration}; use tokio_stream::wrappers::ReceiverStream; use uuid::Uuid; +const DEBOUNCE_MS: u64 = 2; +const RECONNECT_GRACE_SECS: u64 = 30; + // --------------------------------------------------------------------------- // RenderFn // --------------------------------------------------------------------------- @@ -99,6 +103,151 @@ impl AppState { } } +// --------------------------------------------------------------------------- +// render_and_send helper +// --------------------------------------------------------------------------- + +/// Parse signals, render, and send SSE events for one connection. +/// +/// On any error, logs and returns without panicking. +async fn render_and_send( + conn_id: Uuid, + signals_json: &str, + sse_tx: &mpsc::Sender, + state: &AppState, +) { + let signals = match serde_json::from_str::(signals_json) { + Ok(v) => SignalStore::from_json(v), + Err(e) => { + tracing::error!("render_and_send: failed to parse signals for {}: {}", conn_id, e); + return; + } + }; + + let tree: ComponentTree = signals + .get("tree") + .and_then(|v| serde_json::from_value(v.clone()).ok()) + .unwrap_or_default(); + + let cx = RenderContext::new(conn_id, signals, tree, state.store.clone()); + match (state.render_fn)(cx).await { + Ok((html, keys)) => { + state.subscriptions.update(conn_id, keys); + + let data_payload: String = html + .lines() + .map(|line| format!("elements {}", line)) + .collect::>() + .join("\n"); + + let html_event = Event::default() + .event("datastar-patch-elements") + .data(data_payload); + + if sse_tx.send(html_event).await.is_err() { + tracing::debug!("render_and_send: SSE channel closed for {}", conn_id); + return; + } + + let signals_event = Event::default() + .event("datastar-patch-signals") + .data("signals ".to_owned() + signals_json); + + if sse_tx.send(signals_event).await.is_err() { + tracing::debug!("render_and_send: SSE channel closed for {}", conn_id); + } + } + Err(e) => { + tracing::error!("render_and_send: render error for {}: {}", conn_id, e); + } + } +} + +// --------------------------------------------------------------------------- +// connection_task +// --------------------------------------------------------------------------- + +/// Long-lived task that owns the SSE pipe for one logical connection. +/// +/// Survives SSE disconnects for a grace period (RECONNECT_GRACE_SECS), allowing +/// the client to reconnect cleanly. Debounces render triggers by DEBOUNCE_MS. +async fn connection_task( + conn_id: Uuid, + mut inbox: mpsc::Receiver, + mut pipe_rx: mpsc::Receiver>, + initial_signals_json: String, + state: SharedState, +) { + let mut signals_json = initial_signals_json; + + loop { + // ---------------------------------------------------------------- + // Wait for first/next SSE pipe (with reconnect grace period). + // ---------------------------------------------------------------- + let sse_tx: mpsc::Sender = tokio::select! { + maybe_pipe = pipe_rx.recv() => { + match maybe_pipe { + Some(tx) => tx, + None => { + // pipe channel closed — shut down + return; + } + } + } + _ = sleep(Duration::from_secs(RECONNECT_GRACE_SECS)) => { + tracing::debug!( + "connection_task: grace period expired for {}, cleaning up", + conn_id + ); + state.connections.remove(&conn_id); + state.subscriptions.remove_connection(&conn_id); + return; + } + }; + + // Do an immediate render on pipe arrival. + render_and_send(conn_id, &signals_json, &sse_tx, &state).await; + + // ---------------------------------------------------------------- + // Main event loop for this pipe. + // ---------------------------------------------------------------- + loop { + tokio::select! { + // SSE pipe dropped by the HTTP layer — wait for reconnect. + _ = sse_tx.closed() => { + tracing::debug!("connection_task: SSE pipe closed for {}, waiting for reconnect", conn_id); + break; // back to outer loop to wait for new pipe + } + + // Trigger channel closed — task is being shut down. + maybe_trigger = inbox.recv() => { + match maybe_trigger { + None => { + // Trigger sender dropped — clean up and exit. + state.connections.remove(&conn_id); + state.subscriptions.remove_connection(&conn_id); + return; + } + Some(_) => { + // Debounce: drain any additional triggers within DEBOUNCE_MS. + sleep(Duration::from_millis(DEBOUNCE_MS)).await; + while inbox.try_recv().is_ok() {} + + // Read fresh signals and render. + signals_json = state + .connections + .get_signals(&conn_id) + .unwrap_or_else(|| signals_json.clone()); + + render_and_send(conn_id, &signals_json, &sse_tx, &state).await; + } + } + } + } + } + } +} + // --------------------------------------------------------------------------- // Fanout task // --------------------------------------------------------------------------- @@ -108,16 +257,11 @@ impl AppState { /// Must be called once at startup, after building `AppState`. The task owns /// a `broadcast::Receiver` and runs for the lifetime of the server. /// -/// On each invalidation event: -/// 1. Look up subscribed connections via `SubscriptionRegistry`. -/// 2. For each connection, fetch `last_signals_json` + SSE sender. -/// 3. Reconstruct `SignalStore` + `ComponentTree`, build `RenderContext`. -/// 4. Call `render_fn` to produce HTML. -/// 5. Frame as SSE and send down the connection's SSE channel. -/// 6. Update subscription keys and `last_signals_json` from the new render. +/// On each invalidation event, routes a `RenderTrigger::Invalidated` to each +/// subscribed connection's long-lived connection task. /// /// ```ignore -/// let state = AppState::new(store.clone(), render_fn); +/// let state = AppState::new(store.clone(), render_fn, config); /// spawn_fanout(state.clone()); /// ``` pub fn spawn_fanout(state: SharedState) { @@ -140,69 +284,9 @@ pub fn spawn_fanout(state: SharedState) { let conn_ids = state.subscriptions.get_connections_for_key(&key); for conn_id in conn_ids { - let signals_json = state - .connections - .get_signals(&conn_id) - .unwrap_or_else(|| "{}".to_string()); - let Some(sse_tx) = state.connections.get_sender(&conn_id) else { - continue; - }; - - let signals = match serde_json::from_str::(&signals_json) { - Ok(v) => SignalStore::from_json(v), - Err(e) => { - tracing::error!("fanout: failed to parse signals for {}: {}", conn_id, e); - continue; - } - }; - - // Parse the component tree from the `tree` signal. - let tree = signals - .get("tree") - .and_then(|v| serde_json::from_value(v.clone()).ok()) - .unwrap_or_default(); - - let mut cx = RenderContext::new(conn_id, signals, tree, state.store.clone()); - - let render_fn = state.render_fn.clone(); - let state_ref = state.clone(); - - tokio::spawn(async move { - match render_fn(cx).await { - Ok((html, keys)) => { - state_ref.subscriptions.update(conn_id, keys); - - // Build a datastar-patch-elements SSE event. - // Each HTML line gets an `elements ` prefix; axum's - // Event::data() splits on \n and writes a `data:` - // line for each, which is correct SSE multiline format. - let data_payload: String = html - .lines() - .map(|line| format!("elements {}", line)) - .collect::>() - .join("\n"); - - let html_event = Event::default() - .event("datastar-patch-elements") - .data(data_payload); - - if sse_tx.send(html_event).await.is_err() { - tracing::debug!("fanout: SSE channel closed for {}", conn_id); - } - - let signals_event = Event::default() - .event("datastar-patch-signals") - .data("signals ".to_owned() + &signals_json); - - if sse_tx.send(signals_event).await.is_err() { - tracing::debug!("fanout: SSE channel closed for {}", conn_id); - } - } - Err(e) => { - tracing::error!("fanout: render error for {}: {}", conn_id, e); - } - } - }); + if let Some(tx) = state.connections.get_trigger_sender(&conn_id) { + tx.send(RenderTrigger::Invalidated).await.ok(); + } } } }); @@ -220,73 +304,44 @@ pub struct SseParams { /// `GET /sse?conn=` — Long-lived SSE pipe for server-initiated push. /// -/// The conn_id is generated by `GET /` and pre-seeded with initial signals -/// in `ConnectionManager` before this handler is called. On connect, the -/// handler immediately triggers a render and pushes the result down the pipe, -/// so the client sees content without any further round-trip. +/// On first connect, spawns a `connection_task` that owns the lifecycle of the +/// connection. On reconnect with the same conn_id, hands a new SSE pipe to the +/// already-running task. The task survives disconnects for RECONNECT_GRACE_SECS. pub async fn sse_handler( State(state): State, AxumQuery(params): AxumQuery, ) -> Sse>> { let conn_id = params.conn; - let (tx, rx) = mpsc::channel::(32); - let tx_watch = tx.clone(); - state.connections.insert(conn_id, tx.clone()); + let pipe_tx: mpsc::Sender> = if !state.connections.contains(&conn_id) { + // Fresh connection — create channels and spawn the connection task. + let (trigger_tx, trigger_rx) = mpsc::channel::(32); + let (pipe_tx, pipe_rx) = mpsc::channel::>(4); - // Cleanup on disconnect - let state_cleanup = state.clone(); - tokio::spawn(async move { - tx_watch.closed().await; - state_cleanup.connections.remove(&conn_id); - state_cleanup.subscriptions.remove_connection(&conn_id); - tracing::debug!("cleaned up connection {}", conn_id); - }); - println!("SSE Handler called"); + let initial_signals = state + .connections + .get_signals(&conn_id) + .unwrap_or_else(|| "{}".to_string()); - // Trigger first render immediately using the pre-seeded signals. - let signals_json = state - .connections - .get_signals(&conn_id) - .unwrap_or_else(|| "{}".to_string()); + state.connections.insert(conn_id, trigger_tx, pipe_tx.clone()); - let render_fn = state.render_fn.clone(); - let store = state.store.clone(); - let subscriptions = state.subscriptions.clone(); - let tx_init = tx.clone(); + let state_task = state.clone(); + tokio::spawn(async move { + connection_task(conn_id, trigger_rx, pipe_rx, initial_signals, state_task).await; + }); - tokio::spawn(async move { - let signals = match serde_json::from_str::(&signals_json) { - Ok(v) => SignalStore::from_json(v), - Err(e) => { - tracing::error!("sse_handler: failed to parse seeded signals: {}", e); - return; - } - }; - let tree = signals - .get("tree") - .and_then(|v| serde_json::from_value(v.clone()).ok()) - .unwrap_or_default(); + pipe_tx + } else { + // Reconnect — retrieve existing pipe sender. + state.connections.get_pipe_sender(&conn_id) + .expect("connection exists but has no pipe_tx") + }; - let cx = RenderContext::new(conn_id, signals, tree, store); - match render_fn(cx).await { - Ok((html, keys)) => { - subscriptions.update(conn_id, keys); - let data_payload = html - .lines() - .map(|line| format!("elements {}", line)) - .collect::>() - .join("\n"); - let event = Event::default() - .event("datastar-patch-elements") - .data(data_payload); - tx_init.send(event).await.ok(); - } - Err(e) => tracing::error!("sse_handler: initial render failed: {}", e), - } - }); + // Create a fresh SSE channel for this HTTP connection's lifetime. + let (sse_tx, sse_rx) = mpsc::channel::(32); + pipe_tx.send(sse_tx).await.ok(); - let stream = ReceiverStream::new(rx); + let stream = ReceiverStream::new(sse_rx); let event_stream = tokio_stream::StreamExt::map(stream, |event| Ok::(event)); Sse::new(event_stream).keep_alive(KeepAlive::default()) diff --git a/src/store.rs b/src/store.rs index 96df5c4..d3ed2c4 100644 --- a/src/store.rs +++ b/src/store.rs @@ -108,7 +108,7 @@ impl Store { /// that can be buffered before slow receivers start lagging. The fanout /// task should be fast enough that this is never hit in practice. pub fn new(db: sqlx::SqlitePool) -> Arc { - let (events_tx, _) = broadcast::channel(256); + let (events_tx, _) = broadcast::channel(16384); let kv = Arc::new(KVStore::new(events_tx.clone())); Arc::new(Self { db,