Skip to main content

arete_server/websocket/
server.rs

1use crate::bus::{BusManager, BusMessage, StateUpdate};
2use crate::cache::{cmp_seq, EntityCache, SnapshotBatchConfig};
3use crate::compression::maybe_compress;
4use crate::view::{ViewIndex, ViewSpec};
5use crate::websocket::auth::{
6    AuthContext, AuthDecision, AuthDeny, ConnectionAuthRequest, WebSocketAuthPlugin,
7};
8use crate::websocket::client_manager::{ClientManager, RateLimitConfig};
9use crate::websocket::frame::{
10    apply_wire_format, Frame, Mode, SnapshotEntity, SnapshotFrame, SortConfig, SortOrder,
11    SubscribedFrame, UnsubscribedFrame,
12};
13use crate::websocket::subscription::{
14    ClientMessage, RefreshAuthRequest, RefreshAuthResponse, SocketIssueMessage, Subscription,
15    SubscriptionQuery, Unsubscription, PROTOCOL_VERSION,
16};
17use crate::websocket::usage::{WebSocketUsageEmitter, WebSocketUsageEvent};
18use crate::WebSocketDeliveryConfig;
19use anyhow::Result;
20use bytes::Bytes;
21use futures_util::StreamExt;
22use serde::Serialize;
23use serde_json::Value;
24use std::collections::{HashMap, HashSet, VecDeque};
25use std::future::Future;
26use std::net::SocketAddr;
27use std::sync::Arc;
28use std::time::{Duration, Instant};
29use tokio::net::{TcpListener, TcpStream};
30use tokio::sync::{broadcast, watch};
31use tokio_tungstenite::{
32    accept_hdr_async,
33    tungstenite::{
34        handshake::server::{ErrorResponse as HandshakeErrorResponse, Request, Response},
35        http::{header::CONTENT_TYPE, StatusCode},
36        Error as WsError,
37    },
38};
39use tokio_util::sync::CancellationToken;
40use tokio_util::task::TaskTracker;
41use tracing::{debug, error, info, info_span, warn, Instrument};
42use uuid::Uuid;
43
44#[cfg(feature = "otel")]
45use crate::metrics::Metrics;
46
47#[derive(Clone, Default)]
48struct WsMetrics {
49    #[cfg(feature = "otel")]
50    inner: Option<Arc<Metrics>>,
51    #[cfg(test)]
52    probe: Option<Arc<DeliveryProbe>>,
53}
54
55/// Test-only mirror of the delivery instruments recorded through
56/// [`WsMetrics`], so the real-socket load harness can report them without the
57/// `otel` feature. Each counter has the meaning of the metric named beside it.
58#[cfg(test)]
59#[derive(Debug, Default)]
60pub(crate) struct DeliveryProbe {
61    /// `arete.ws.messages.sent`
62    pub(crate) messages_sent: std::sync::atomic::AtomicU64,
63    /// `arete.ws.subscription.lagged`
64    pub(crate) lag_events: std::sync::atomic::AtomicU64,
65    /// `arete.ws.subscription.dropped_updates`
66    pub(crate) dropped_updates: std::sync::atomic::AtomicU64,
67    /// `arete.ws.subscription.resnapshots`
68    pub(crate) resnapshots: std::sync::atomic::AtomicU64,
69    /// `arete.ws.collection.coalesced_updates`
70    pub(crate) coalesced_updates: std::sync::atomic::AtomicU64,
71    /// `arete.ws.collection.coalesced_flushes`
72    pub(crate) coalesced_flushes: std::sync::atomic::AtomicU64,
73    /// `arete.ws.delivery.stopped`, by reason
74    pub(crate) delivery_stopped: std::sync::Mutex<std::collections::BTreeMap<&'static str, u64>>,
75}
76
77#[cfg(test)]
78impl DeliveryProbe {
79    fn add(counter: &std::sync::atomic::AtomicU64, value: u64) {
80        counter.fetch_add(value, std::sync::atomic::Ordering::Relaxed);
81    }
82}
83
84impl WsMetrics {
85    #[cfg(feature = "otel")]
86    fn new(inner: Option<Arc<Metrics>>) -> Self {
87        Self {
88            inner,
89            #[cfg(test)]
90            probe: None,
91        }
92    }
93
94    #[cfg(test)]
95    fn probe(&self, record: impl FnOnce(&DeliveryProbe)) {
96        if let Some(probe) = &self.probe {
97            record(probe);
98        }
99    }
100
101    fn connection_opened(&self, metering_key: Option<&str>) {
102        #[cfg(not(feature = "otel"))]
103        let _ = metering_key;
104        #[cfg(feature = "otel")]
105        if let Some(metrics) = &self.inner {
106            if let Some(metering_key) = metering_key {
107                metrics.record_ws_connection_with_metering(metering_key);
108            } else {
109                metrics.record_ws_connection();
110            }
111        }
112    }
113
114    fn connection_closed(&self, duration_secs: f64, metering_key: Option<&str>) {
115        #[cfg(not(feature = "otel"))]
116        let _ = (duration_secs, metering_key);
117        #[cfg(feature = "otel")]
118        if let Some(metrics) = &self.inner {
119            if let Some(metering_key) = metering_key {
120                metrics.record_ws_disconnection_with_metering(duration_secs, metering_key);
121            } else {
122                metrics.record_ws_disconnection(duration_secs);
123            }
124        }
125    }
126
127    fn message_received(&self, metering_key: Option<&str>) {
128        #[cfg(not(feature = "otel"))]
129        let _ = metering_key;
130        #[cfg(feature = "otel")]
131        if let Some(metrics) = &self.inner {
132            if let Some(metering_key) = metering_key {
133                metrics.record_ws_message_received_with_metering(metering_key);
134            } else {
135                metrics.record_ws_message_received();
136            }
137        }
138    }
139
140    fn message_sent(&self) {
141        #[cfg(test)]
142        self.probe(|probe| DeliveryProbe::add(&probe.messages_sent, 1));
143        #[cfg(feature = "otel")]
144        if let Some(metrics) = &self.inner {
145            metrics.record_ws_message_sent();
146        }
147    }
148
149    fn subscription_created(&self, view: &str, metering_key: Option<&str>) {
150        #[cfg(not(feature = "otel"))]
151        let _ = (view, metering_key);
152        #[cfg(feature = "otel")]
153        if let Some(metrics) = &self.inner {
154            if let Some(metering_key) = metering_key {
155                metrics.record_subscription_created_with_metering(view, metering_key);
156            } else {
157                metrics.record_subscription_created(view);
158            }
159        }
160    }
161
162    fn subscription_removed(&self, view: &str, metering_key: Option<&str>) {
163        #[cfg(not(feature = "otel"))]
164        let _ = (view, metering_key);
165        #[cfg(feature = "otel")]
166        if let Some(metrics) = &self.inner {
167            if let Some(metering_key) = metering_key {
168                metrics.record_subscription_removed_with_metering(view, metering_key);
169            } else {
170                metrics.record_subscription_removed(view);
171            }
172        }
173    }
174
175    fn protocol_error(&self, code: &str) {
176        #[cfg(not(feature = "otel"))]
177        let _ = code;
178        #[cfg(feature = "otel")]
179        if let Some(metrics) = &self.inner {
180            metrics.record_ws_protocol_error(code);
181        }
182    }
183
184    fn subscription_lagged(&self, view: &str, skipped: u64) {
185        #[cfg(test)]
186        self.probe(|probe| {
187            DeliveryProbe::add(&probe.lag_events, 1);
188            DeliveryProbe::add(&probe.dropped_updates, skipped);
189        });
190        #[cfg(not(feature = "otel"))]
191        let _ = (view, skipped);
192        #[cfg(feature = "otel")]
193        if let Some(metrics) = &self.inner {
194            metrics.record_ws_subscription_lagged(view, skipped);
195        }
196    }
197
198    fn subscription_resnapshot(&self, view: &str) {
199        #[cfg(test)]
200        self.probe(|probe| DeliveryProbe::add(&probe.resnapshots, 1));
201        #[cfg(not(feature = "otel"))]
202        let _ = view;
203        #[cfg(feature = "otel")]
204        if let Some(metrics) = &self.inner {
205            metrics.record_ws_subscription_resnapshot(view);
206        }
207    }
208
209    fn collection_coalesced(&self, view: &str, updates: u64) {
210        #[cfg(test)]
211        self.probe(|probe| {
212            DeliveryProbe::add(&probe.coalesced_updates, updates);
213            DeliveryProbe::add(&probe.coalesced_flushes, 1);
214        });
215        #[cfg(not(feature = "otel"))]
216        let _ = (view, updates);
217        #[cfg(feature = "otel")]
218        if let Some(metrics) = &self.inner {
219            metrics.record_ws_collection_coalesced(view, updates);
220        }
221    }
222
223    fn delivery_stopped(&self, view: &str, reason: &'static str) {
224        #[cfg(test)]
225        self.probe(|probe| {
226            *probe
227                .delivery_stopped
228                .lock()
229                .expect("delivery probe lock poisoned")
230                .entry(reason)
231                .or_default() += 1;
232        });
233        #[cfg(not(feature = "otel"))]
234        let _ = (view, reason);
235        #[cfg(feature = "otel")]
236        if let Some(metrics) = &self.inner {
237            metrics.record_ws_delivery_stopped(view, reason);
238        }
239    }
240}
241
242async fn handle_refresh_auth(
243    client_id: Uuid,
244    refresh_req: &RefreshAuthRequest,
245    client_manager: &ClientManager,
246    auth_plugin: &Arc<dyn WebSocketAuthPlugin>,
247) {
248    let refresh_result: Result<AuthContext, String> = if let Some(signed_plugin) = auth_plugin
249        .as_any()
250        .downcast_ref::<crate::websocket::auth::SignedSessionAuthPlugin>()
251    {
252        signed_plugin
253            .verify_refresh_token(&refresh_req.token)
254            .await
255            .map_err(|error| error.reason)
256    } else {
257        Err("In-band auth refresh not supported with current auth plugin".to_string())
258    };
259
260    let response = match refresh_result {
261        Ok(new_context) => {
262            let expires_at = new_context.expires_at;
263            if client_manager.update_client_auth(client_id, new_context) {
264                RefreshAuthResponse {
265                    success: true,
266                    error: None,
267                    expires_at: Some(expires_at),
268                }
269            } else {
270                RefreshAuthResponse {
271                    success: false,
272                    error: Some("client-not-found".to_string()),
273                    expires_at: None,
274                }
275            }
276        }
277        Err(error) => {
278            let code = if error.contains("expired") {
279                "token-expired"
280            } else if error.contains("signature") {
281                "token-invalid-signature"
282            } else if error.contains("issuer") {
283                "token-invalid-issuer"
284            } else if error.contains("audience") {
285                "token-invalid-audience"
286            } else {
287                "token-invalid"
288            };
289            RefreshAuthResponse {
290                success: false,
291                error: Some(code.to_string()),
292                expires_at: None,
293            }
294        }
295    };
296
297    if let Ok(json) = serde_json::to_string(&response) {
298        let _ = client_manager.send_text_to_client(client_id, json).await;
299    }
300}
301
302async fn send_socket_issue(
303    client_id: Uuid,
304    client_manager: &ClientManager,
305    deny: &AuthDeny,
306    fatal: bool,
307    subscription_id: Option<String>,
308) {
309    let message = SocketIssueMessage::from_auth_deny(deny, fatal, subscription_id);
310    if let Ok(json) = serde_json::to_string(&message) {
311        let _ = client_manager.send_text_to_client(client_id, json).await;
312    }
313}
314
315async fn send_protocol_issue(
316    client_id: Uuid,
317    client_manager: &ClientManager,
318    metrics: &WsMetrics,
319    subscription_id: Option<String>,
320    code: &str,
321    message: impl Into<String>,
322) {
323    metrics.protocol_error(code);
324    let issue = SocketIssueMessage::protocol(subscription_id, code, message);
325    if let Ok(json) = serde_json::to_string(&issue) {
326        let _ = client_manager.send_text_to_client(client_id, json).await;
327    }
328}
329
330/// Send an already-built issue, for refusals that carry structured detail.
331async fn send_prepared_issue(
332    client_id: Uuid,
333    client_manager: &ClientManager,
334    metrics: &WsMetrics,
335    issue: SocketIssueMessage,
336) {
337    metrics.protocol_error(&issue.code);
338    if let Ok(json) = serde_json::to_string(&issue) {
339        let _ = client_manager.send_text_to_client(client_id, json).await;
340    }
341}
342
343fn key_class_label(key_class: arete_auth::KeyClass) -> &'static str {
344    match key_class {
345        arete_auth::KeyClass::Secret => "secret",
346        arete_auth::KeyClass::Publishable => "publishable",
347    }
348}
349
350fn emit_usage_event(
351    usage_emitter: &Option<Arc<dyn WebSocketUsageEmitter>>,
352    event: WebSocketUsageEvent,
353) {
354    if let Some(emitter) = usage_emitter.clone() {
355        tokio::spawn(async move {
356            emitter.emit(event).await;
357        });
358    }
359}
360
361fn usage_identity(
362    auth_context: Option<&AuthContext>,
363) -> (
364    Option<String>,
365    Option<String>,
366    Option<String>,
367    Option<String>,
368) {
369    match auth_context {
370        Some(context) => (
371            Some(context.metering_key.clone()),
372            Some(context.subject.clone()),
373            Some(key_class_label(context.key_class).to_string()),
374            context.deployment_id.clone(),
375        ),
376        None => (None, None, None, None),
377    }
378}
379
380fn emit_update_sent_for_client(
381    usage_emitter: &Option<Arc<dyn WebSocketUsageEmitter>>,
382    client_manager: &ClientManager,
383    client_id: Uuid,
384    view_id: &str,
385    bytes: usize,
386) {
387    let auth_context = client_manager.get_auth_context(client_id);
388    let (metering_key, subject, _, deployment_id) = usage_identity(auth_context.as_ref());
389    emit_usage_event(
390        usage_emitter,
391        WebSocketUsageEvent::UpdateSent {
392            client_id: client_id.to_string(),
393            deployment_id,
394            metering_key,
395            subject,
396            view_id: view_id.to_string(),
397            messages: 1,
398            bytes: bytes as u64,
399        },
400    );
401}
402
403#[derive(Clone)]
404struct SubscriptionContext {
405    client_id: Uuid,
406    client_manager: ClientManager,
407    bus_manager: BusManager,
408    entity_cache: EntityCache,
409    view_index: Arc<ViewIndex>,
410    usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
411    journal: Option<Arc<crate::journal::EventJournal>>,
412    metrics: WsMetrics,
413    delivery: WebSocketDeliveryConfig,
414    /// Cancelled when the server stops; every session ends through its normal
415    /// cleanup path rather than being dropped mid-flight.
416    shutdown: CancellationToken,
417}
418
419pub struct WebSocketServer {
420    bind_addr: SocketAddr,
421    client_manager: ClientManager,
422    bus_manager: BusManager,
423    entity_cache: EntityCache,
424    view_index: Arc<ViewIndex>,
425    max_clients: usize,
426    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
427    usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
428    rate_limit_config: Option<RateLimitConfig>,
429    journal: Option<Arc<crate::journal::EventJournal>>,
430    delivery: WebSocketDeliveryConfig,
431    #[cfg(feature = "otel")]
432    metrics: Option<Arc<Metrics>>,
433}
434
435impl WebSocketServer {
436    #[cfg(feature = "otel")]
437    pub fn new(
438        bind_addr: SocketAddr,
439        bus_manager: BusManager,
440        entity_cache: EntityCache,
441        view_index: Arc<ViewIndex>,
442        metrics: Option<Arc<Metrics>>,
443    ) -> Self {
444        Self {
445            bind_addr,
446            client_manager: ClientManager::new(),
447            bus_manager,
448            entity_cache,
449            view_index,
450            max_clients: 10_000,
451            auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
452            usage_emitter: None,
453            rate_limit_config: None,
454            journal: None,
455            delivery: WebSocketDeliveryConfig::default(),
456            metrics,
457        }
458    }
459
460    #[cfg(not(feature = "otel"))]
461    pub fn new(
462        bind_addr: SocketAddr,
463        bus_manager: BusManager,
464        entity_cache: EntityCache,
465        view_index: Arc<ViewIndex>,
466    ) -> Self {
467        Self {
468            bind_addr,
469            client_manager: ClientManager::new(),
470            bus_manager,
471            entity_cache,
472            view_index,
473            max_clients: 10_000,
474            auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
475            usage_emitter: None,
476            rate_limit_config: None,
477            journal: None,
478            delivery: WebSocketDeliveryConfig::default(),
479        }
480    }
481
482    pub fn with_max_clients(mut self, max_clients: usize) -> Self {
483        self.max_clients = max_clients;
484        self
485    }
486
487    pub fn with_auth_plugin(mut self, auth_plugin: Arc<dyn WebSocketAuthPlugin>) -> Self {
488        self.auth_plugin = auth_plugin;
489        self
490    }
491
492    pub fn with_usage_emitter(mut self, usage_emitter: Arc<dyn WebSocketUsageEmitter>) -> Self {
493        self.usage_emitter = Some(usage_emitter);
494        self
495    }
496
497    /// Serve replayable append subscriptions from the retained event journal.
498    pub fn with_journal(mut self, journal: Arc<crate::journal::EventJournal>) -> Self {
499        self.journal = Some(journal);
500        self
501    }
502
503    pub fn with_rate_limit_config(mut self, config: RateLimitConfig) -> Self {
504        self.rate_limit_config = Some(config);
505        self
506    }
507
508    pub fn with_delivery_config(mut self, config: WebSocketDeliveryConfig) -> Self {
509        self.delivery = config;
510        self
511    }
512
513    /// Bind the configured address and serve connections until the task is
514    /// dropped. Equivalent to [`into_acceptor`](Self::into_acceptor) followed by
515    /// [`ConnectionAcceptor::serve_listener`].
516    pub async fn start(self) -> Result<()> {
517        info!(
518            "Starting WebSocket server on {} (max_clients: {})",
519            self.bind_addr, self.max_clients
520        );
521        let listener = TcpListener::bind(&self.bind_addr).await?;
522        let (acceptor, _cleanup) = self.into_acceptor();
523        acceptor.serve_listener(listener).await
524    }
525
526    /// Split this server into the part that serves connections and the
527    /// client-manager cleanup task, leaving the caller to own the listener.
528    ///
529    /// The cleanup handle is returned rather than detached so a caller that
530    /// stops serving can stop it too.
531    pub(crate) fn into_acceptor(self) -> (ConnectionAcceptor, tokio::task::JoinHandle<()>) {
532        let client_manager = self
533            .rate_limit_config
534            .map(ClientManager::with_config)
535            .unwrap_or(self.client_manager);
536        let cleanup = client_manager.start_cleanup_task();
537
538        #[cfg(feature = "otel")]
539        let metrics = WsMetrics::new(self.metrics.clone());
540        #[cfg(not(feature = "otel"))]
541        let metrics = WsMetrics::default();
542
543        let acceptor = ConnectionAcceptor {
544            client_manager,
545            bus_manager: self.bus_manager,
546            entity_cache: self.entity_cache,
547            view_index: self.view_index,
548            max_clients: self.max_clients,
549            auth_plugin: self.auth_plugin,
550            usage_emitter: self.usage_emitter,
551            journal: self.journal,
552            delivery: self.delivery,
553            metrics,
554            shutdown: CancellationToken::new(),
555            sessions: TaskTracker::new(),
556        };
557        (acceptor, cleanup)
558    }
559}
560
561/// Serves already-accepted TCP connections against one server's buses, cache
562/// and views.
563///
564/// This is what [`WebSocketServer::start`] runs behind its listener, separated
565/// so that a caller that owns the listener (an application that terminates
566/// TLS itself, a test with an ephemeral port) can hand streams in without the
567/// server binding anything.
568#[derive(Clone)]
569pub(crate) struct ConnectionAcceptor {
570    client_manager: ClientManager,
571    bus_manager: BusManager,
572    entity_cache: EntityCache,
573    view_index: Arc<ViewIndex>,
574    max_clients: usize,
575    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
576    usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
577    journal: Option<Arc<crate::journal::EventJournal>>,
578    delivery: WebSocketDeliveryConfig,
579    metrics: WsMetrics,
580    shutdown: CancellationToken,
581    /// Sessions spawned by [`serve_listener`](Self::serve_listener), so a
582    /// stop can wait for them. Sessions a caller serves on its own tasks are
583    /// the caller's to wait for.
584    sessions: TaskTracker,
585}
586
587impl ConnectionAcceptor {
588    /// Mirror this acceptor's delivery instruments into `probe`. Must be set
589    /// before serving: each session copies the metrics handle when it starts.
590    #[cfg(test)]
591    pub(crate) fn with_delivery_probe(mut self, probe: Arc<DeliveryProbe>) -> Self {
592        self.metrics.probe = Some(probe);
593        self
594    }
595
596    /// Number of clients currently connected to this server.
597    pub(crate) fn client_count(&self) -> usize {
598        self.client_manager.client_count()
599    }
600
601    /// End every session this acceptor is serving and stop accepting.
602    ///
603    /// Sessions notice on their next poll and leave through the same cleanup
604    /// as a client disconnect, so the client manager, buses and usage events
605    /// see an ordinary close.
606    pub(crate) fn shutdown(&self) {
607        self.shutdown.cancel();
608        self.sessions.close();
609    }
610
611    /// Resolves once every listener-spawned session has finished cleaning
612    /// up. Call after [`shutdown`](Self::shutdown).
613    pub(crate) async fn wait_for_sessions(&self) {
614        self.sessions.wait().await;
615    }
616
617    /// Serve one accepted connection: WebSocket handshake, authentication,
618    /// then the subscription session until the peer disconnects.
619    ///
620    /// Returns `Ok(())` without serving when the server is at its client
621    /// limit, exactly as the listener loop does.
622    pub(crate) async fn serve(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
623        if self.client_manager.client_count() >= self.max_clients {
624            warn!(
625                "Rejecting connection from {}: max clients reached",
626                remote_addr
627            );
628            return Ok(());
629        }
630
631        let context = SubscriptionContext {
632            client_id: Uuid::nil(),
633            client_manager: self.client_manager.clone(),
634            bus_manager: self.bus_manager.clone(),
635            entity_cache: self.entity_cache.clone(),
636            view_index: self.view_index.clone(),
637            usage_emitter: self.usage_emitter.clone(),
638            journal: self.journal.clone(),
639            delivery: self.delivery.clone(),
640            metrics: self.metrics.clone(),
641            shutdown: self.shutdown.clone(),
642        };
643        handle_connection(stream, context, remote_addr, self.auth_plugin.clone()).await
644    }
645
646    /// Accept from `listener` until [`shutdown`](Self::shutdown), serving each
647    /// connection on its own task.
648    pub(crate) async fn serve_listener(self, listener: TcpListener) -> Result<()> {
649        loop {
650            let accepted = tokio::select! {
651                _ = self.shutdown.cancelled() => return Ok(()),
652                accepted = listener.accept() => accepted,
653            };
654            match accepted {
655                Ok((stream, addr)) => {
656                    let acceptor = self.clone();
657                    self.sessions.spawn(
658                        async move {
659                            if let Err(error) = acceptor.serve(stream, addr).await {
660                                error!("WebSocket connection error: {}", error);
661                            }
662                        }
663                        .instrument(info_span!("ws.connection", %addr)),
664                    );
665                }
666                Err(error) => error!("Failed to accept connection: {}", error),
667            }
668        }
669    }
670}
671
672#[derive(Debug, Clone)]
673struct HandshakeReject {
674    status: StatusCode,
675    body: crate::websocket::auth::ErrorResponse,
676    error_code: String,
677    retry_after_secs: Option<u64>,
678}
679
680impl HandshakeReject {
681    fn from_deny(deny: &AuthDeny) -> Self {
682        let retry_after_secs = match deny.retry_policy {
683            crate::websocket::auth::RetryPolicy::RetryAfter(duration) => Some(duration.as_secs()),
684            _ => None,
685        };
686        Self {
687            status: StatusCode::from_u16(deny.http_status).unwrap_or(StatusCode::UNAUTHORIZED),
688            body: deny.to_error_response(),
689            error_code: deny.code.to_string(),
690            retry_after_secs,
691        }
692    }
693}
694
695fn build_handshake_error_response(
696    response: &Response,
697    reject: &HandshakeReject,
698) -> HandshakeErrorResponse {
699    let mut builder = Response::builder()
700        .status(reject.status)
701        .version(response.version())
702        .header(CONTENT_TYPE, "application/json; charset=utf-8")
703        .header("X-Error-Code", &reject.error_code)
704        .header("Cache-Control", "no-store");
705    if let Some(retry_after_secs) = reject.retry_after_secs {
706        builder = builder.header("Retry-After", retry_after_secs.to_string());
707    }
708    let body = serde_json::to_string(&reject.body).unwrap_or_else(|_| {
709        format!(
710            r#"{{"error":"{}","message":"{}","code":"{}","retryable":false}}"#,
711            reject.body.error, reject.body.message, reject.body.code
712        )
713    });
714    builder
715        .body(Some(body))
716        .expect("handshake rejection response should build")
717}
718
719#[allow(clippy::result_large_err)]
720async fn accept_authorized_connection(
721    stream: TcpStream,
722    remote_addr: SocketAddr,
723    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
724    client_manager: ClientManager,
725) -> Result<Option<(tokio_tungstenite::WebSocketStream<TcpStream>, AuthContext)>> {
726    use std::sync::Mutex;
727
728    let capture: Arc<Mutex<Option<Result<AuthContext, HandshakeReject>>>> =
729        Arc::new(Mutex::new(None));
730    let capture_ref = capture.clone();
731    let auth_plugin_ref = auth_plugin.clone();
732    let manager_ref = client_manager.clone();
733
734    let handshake_result = accept_hdr_async(stream, move |request: &Request, response| {
735        let request = ConnectionAuthRequest::from_http_request(remote_addr, request);
736        let result = tokio::task::block_in_place(|| {
737            tokio::runtime::Handle::current().block_on(async {
738                match auth_plugin_ref.authorize(&request).await {
739                    AuthDecision::Allow(context) => manager_ref
740                        .check_connection_allowed(remote_addr, &Some(context.clone()))
741                        .await
742                        .map(|()| context)
743                        .map_err(|deny| HandshakeReject::from_deny(&deny)),
744                    AuthDecision::Deny(deny) => Err(HandshakeReject::from_deny(&deny)),
745                }
746            })
747        });
748        *capture_ref.lock().expect("capture lock poisoned") = Some(result.clone());
749        match result {
750            Ok(_) => Ok(response),
751            Err(reject) => Err(build_handshake_error_response(&response, &reject)),
752        }
753    })
754    .await;
755
756    let auth_result = capture.lock().expect("capture lock poisoned").take();
757    match handshake_result {
758        Ok(stream) => match auth_result {
759            Some(Ok(context)) => Ok(Some((stream, context))),
760            Some(Err(reject)) => Err(anyhow::anyhow!(
761                "handshake unexpectedly succeeded after rejection: {}",
762                reject.body.message
763            )),
764            None => Err(anyhow::anyhow!("no auth result captured during handshake")),
765        },
766        Err(WsError::Http(_)) => Ok(None),
767        Err(error) => Err(error.into()),
768    }
769}
770
771async fn handle_connection(
772    stream: TcpStream,
773    mut context: SubscriptionContext,
774    remote_addr: SocketAddr,
775    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
776) -> Result<()> {
777    // The handshake is raced against shutdown too: a peer that stalls it must
778    // not keep a task alive after the server has stopped.
779    let accepted = tokio::select! {
780        _ = context.shutdown.cancelled() => return Ok(()),
781        accepted = accept_authorized_connection(
782            stream,
783            remote_addr,
784            auth_plugin.clone(),
785            context.client_manager.clone(),
786        ) => accepted?,
787    };
788    let Some((ws_stream, auth_context)) = accepted else {
789        return Ok(());
790    };
791
792    let client_id = Uuid::new_v4();
793    context.client_id = client_id;
794    let connection_start = Instant::now();
795    let (metering_key, subject, key_class, deployment_id) = usage_identity(Some(&auth_context));
796    context.metrics.connection_opened(metering_key.as_deref());
797
798    let (ws_sender, mut ws_receiver) = ws_stream.split();
799    context
800        .client_manager
801        .add_client(client_id, ws_sender, Some(auth_context), remote_addr);
802    emit_usage_event(
803        &context.usage_emitter,
804        WebSocketUsageEvent::ConnectionEstablished {
805            client_id: client_id.to_string(),
806            remote_addr: remote_addr.to_string(),
807            deployment_id: deployment_id.clone(),
808            metering_key: metering_key.clone(),
809            subject: subject.clone(),
810            key_class,
811        },
812    );
813
814    let mut active_subscriptions: HashMap<String, String> = HashMap::new();
815    loop {
816        let message = tokio::select! {
817            _ = context.shutdown.cancelled() => break,
818            next = ws_receiver.next() => match next {
819                Some(message) => message,
820                None => break,
821            },
822        };
823        let message = match message {
824            Ok(message) => message,
825            Err(error) => {
826                warn!("WebSocket error for client {}: {}", client_id, error);
827                break;
828            }
829        };
830        if message.is_close() {
831            break;
832        }
833        context.client_manager.update_client_last_seen(client_id);
834        if !message.is_text() {
835            continue;
836        }
837        if let Err(deny) = context
838            .client_manager
839            .check_inbound_message_allowed(client_id)
840        {
841            send_socket_issue(client_id, &context.client_manager, &deny, true, None).await;
842            break;
843        }
844        context.metrics.message_received(metering_key.as_deref());
845
846        let text = match message.to_text() {
847            Ok(text) => text,
848            Err(_) => continue,
849        };
850        let client_message = match serde_json::from_str::<ClientMessage>(text) {
851            Ok(message) => message,
852            Err(parse_error) => {
853                let subscription_id = extract_subscription_id(text);
854                send_protocol_issue(
855                    client_id,
856                    &context.client_manager,
857                    &context.metrics,
858                    subscription_id,
859                    "malformed-message",
860                    format!("invalid protocol v2 message: {parse_error}"),
861                )
862                .await;
863                continue;
864            }
865        };
866
867        match client_message {
868            ClientMessage::Subscribe(subscription) => {
869                let subscription_id = subscription.subscription_id.clone();
870                if let Err(message) = subscription.validate() {
871                    send_protocol_issue(
872                        client_id,
873                        &context.client_manager,
874                        &context.metrics,
875                        Some(subscription_id),
876                        "invalid-subscription",
877                        message,
878                    )
879                    .await;
880                    continue;
881                }
882                if let Err(deny) = context
883                    .client_manager
884                    .check_subscription_allowed(client_id)
885                    .await
886                {
887                    send_socket_issue(
888                        client_id,
889                        &context.client_manager,
890                        &deny,
891                        false,
892                        Some(subscription_id),
893                    )
894                    .await;
895                    continue;
896                }
897
898                let cancel_token = CancellationToken::new();
899                if !context
900                    .client_manager
901                    .add_client_subscription(
902                        client_id,
903                        subscription_id.clone(),
904                        cancel_token.clone(),
905                    )
906                    .await
907                {
908                    send_protocol_issue(
909                        client_id,
910                        &context.client_manager,
911                        &context.metrics,
912                        Some(subscription_id),
913                        "duplicate-subscription-id",
914                        "subscriptionId is already active on this connection",
915                    )
916                    .await;
917                    continue;
918                }
919
920                let view = subscription.query.view.clone();
921                if let Err(error) = attach_client_to_bus(&context, subscription, cancel_token).await
922                {
923                    context
924                        .client_manager
925                        .remove_client_subscription(client_id, &subscription_id)
926                        .await;
927                    // A refusal that already knows what to tell the client
928                    // (an expired cursor, a changed epoch) keeps its own
929                    // frame; anything else is a generic rejection.
930                    match error.downcast::<RejectedSubscription>() {
931                        Ok(rejected) => {
932                            send_prepared_issue(
933                                client_id,
934                                &context.client_manager,
935                                &context.metrics,
936                                rejected.0,
937                            )
938                            .await;
939                        }
940                        Err(error) => {
941                            send_protocol_issue(
942                                client_id,
943                                &context.client_manager,
944                                &context.metrics,
945                                Some(subscription_id),
946                                "subscription-rejected",
947                                error.to_string(),
948                            )
949                            .await;
950                        }
951                    }
952                    continue;
953                }
954
955                active_subscriptions.insert(subscription_id, view.clone());
956                context
957                    .metrics
958                    .subscription_created(&view, metering_key.as_deref());
959                emit_usage_event(
960                    &context.usage_emitter,
961                    WebSocketUsageEvent::SubscriptionCreated {
962                        client_id: client_id.to_string(),
963                        deployment_id: deployment_id.clone(),
964                        metering_key: metering_key.clone(),
965                        subject: subject.clone(),
966                        view_id: view,
967                    },
968                );
969            }
970            ClientMessage::Unsubscribe(unsubscription) => {
971                handle_unsubscribe(
972                    &context,
973                    unsubscription,
974                    &mut active_subscriptions,
975                    metering_key.as_deref(),
976                    &deployment_id,
977                    &metering_key,
978                    &subject,
979                )
980                .await;
981            }
982            ClientMessage::Ping => debug!("Received ping from client {}", client_id),
983            ClientMessage::RefreshAuth(request) => {
984                handle_refresh_auth(client_id, &request, &context.client_manager, &auth_plugin)
985                    .await;
986            }
987        }
988    }
989
990    context
991        .client_manager
992        .cancel_all_client_subscriptions(client_id)
993        .await;
994    context.client_manager.remove_client(client_id);
995    if let Some(rate_limiter) = context.client_manager.rate_limiter().cloned() {
996        rate_limiter.remove_client_buckets(client_id).await;
997    }
998    for view in active_subscriptions.values() {
999        context
1000            .metrics
1001            .subscription_removed(view, metering_key.as_deref());
1002        emit_usage_event(
1003            &context.usage_emitter,
1004            WebSocketUsageEvent::SubscriptionRemoved {
1005                client_id: client_id.to_string(),
1006                deployment_id: deployment_id.clone(),
1007                metering_key: metering_key.clone(),
1008                subject: subject.clone(),
1009                view_id: view.clone(),
1010            },
1011        );
1012    }
1013    let duration = connection_start.elapsed().as_secs_f64();
1014    context
1015        .metrics
1016        .connection_closed(duration, metering_key.as_deref());
1017    emit_usage_event(
1018        &context.usage_emitter,
1019        WebSocketUsageEvent::ConnectionClosed {
1020            client_id: client_id.to_string(),
1021            deployment_id,
1022            metering_key,
1023            subject,
1024            duration_secs: Some(duration),
1025            subscription_count: u32::try_from(active_subscriptions.len()).unwrap_or(u32::MAX),
1026        },
1027    );
1028    Ok(())
1029}
1030
1031#[allow(clippy::too_many_arguments)]
1032async fn handle_unsubscribe(
1033    context: &SubscriptionContext,
1034    unsubscription: Unsubscription,
1035    active_subscriptions: &mut HashMap<String, String>,
1036    metrics_metering_key: Option<&str>,
1037    deployment_id: &Option<String>,
1038    usage_metering_key: &Option<String>,
1039    subject: &Option<String>,
1040) {
1041    let subscription_id = unsubscription.subscription_id.clone();
1042    if let Err(message) = unsubscription.validate() {
1043        send_protocol_issue(
1044            context.client_id,
1045            &context.client_manager,
1046            &context.metrics,
1047            Some(subscription_id),
1048            "invalid-unsubscription",
1049            message,
1050        )
1051        .await;
1052        return;
1053    }
1054
1055    if !context
1056        .client_manager
1057        .remove_client_subscription(context.client_id, &subscription_id)
1058        .await
1059    {
1060        send_protocol_issue(
1061            context.client_id,
1062            &context.client_manager,
1063            &context.metrics,
1064            Some(subscription_id),
1065            "unknown-subscription-id",
1066            "subscriptionId is not active on this connection",
1067        )
1068        .await;
1069        return;
1070    }
1071
1072    let Some(view) = active_subscriptions.remove(&subscription_id) else {
1073        return;
1074    };
1075    let _ = send_control_frame(context, &UnsubscribedFrame::new(subscription_id), &view);
1076    context
1077        .metrics
1078        .subscription_removed(&view, metrics_metering_key);
1079    emit_usage_event(
1080        &context.usage_emitter,
1081        WebSocketUsageEvent::SubscriptionRemoved {
1082            client_id: context.client_id.to_string(),
1083            deployment_id: deployment_id.clone(),
1084            metering_key: usage_metering_key.clone(),
1085            subject: subject.clone(),
1086            view_id: view,
1087        },
1088    );
1089}
1090
1091fn extract_subscription_id(text: &str) -> Option<String> {
1092    serde_json::from_str::<Value>(text)
1093        .ok()?
1094        .get("subscriptionId")?
1095        .as_str()
1096        .map(str::to_string)
1097}
1098
1099struct SnapshotMetadata<'a> {
1100    subscription_id: &'a str,
1101    snapshot_id: &'a str,
1102    authoritative: bool,
1103    mode: Mode,
1104    view_id: &'a str,
1105    key: Option<&'a str>,
1106}
1107
1108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1109enum SnapshotPurpose {
1110    Initial,
1111    Recovery,
1112}
1113
1114impl SnapshotPurpose {
1115    fn authoritative(self, subscription: &Subscription) -> bool {
1116        match self {
1117            Self::Initial => subscription.query.after.is_none(),
1118            // Recovery replaces the exact query membership even when `after`
1119            // made the initial snapshot incremental. A merge-only snapshot
1120            // cannot remove members whose delete frames were skipped.
1121            Self::Recovery => true,
1122        }
1123    }
1124}
1125
1126fn create_snapshot_batches(
1127    entities: &[SnapshotEntity],
1128    metadata: SnapshotMetadata<'_>,
1129    batch_config: &SnapshotBatchConfig,
1130) -> Vec<SnapshotFrame> {
1131    if entities.is_empty() {
1132        return vec![SnapshotFrame {
1133            protocol_version: PROTOCOL_VERSION,
1134            subscription_id: metadata.subscription_id.to_string(),
1135            snapshot_id: metadata.snapshot_id.to_string(),
1136            authoritative: metadata.authoritative,
1137            mode: metadata.mode,
1138            export: metadata.view_id.to_string(),
1139            op: "snapshot",
1140            key: metadata.key.map(str::to_string),
1141            data: vec![],
1142            complete: true,
1143        }];
1144    }
1145
1146    let mut batches = Vec::new();
1147    let mut offset = 0;
1148    while offset < entities.len() {
1149        let configured_size = if offset == 0 {
1150            batch_config.initial_batch_size
1151        } else {
1152            batch_config.subsequent_batch_size
1153        };
1154        let end = (offset + configured_size.max(1)).min(entities.len());
1155        batches.push(SnapshotFrame {
1156            protocol_version: PROTOCOL_VERSION,
1157            subscription_id: metadata.subscription_id.to_string(),
1158            snapshot_id: metadata.snapshot_id.to_string(),
1159            authoritative: metadata.authoritative,
1160            mode: metadata.mode,
1161            export: metadata.view_id.to_string(),
1162            op: "snapshot",
1163            key: metadata.key.map(str::to_string),
1164            data: entities[offset..end].to_vec(),
1165            complete: end == entities.len(),
1166        });
1167        offset = end;
1168    }
1169    batches
1170}
1171
1172async fn send_snapshot_batches(
1173    context: &SubscriptionContext,
1174    subscription: &Subscription,
1175    entities: &[SnapshotEntity],
1176    mode: Mode,
1177    purpose: SnapshotPurpose,
1178    batch_config: &SnapshotBatchConfig,
1179) -> Result<()> {
1180    let snapshot_id = Uuid::new_v4().to_string();
1181    let authoritative = purpose.authoritative(subscription);
1182    let frames = create_snapshot_batches(
1183        entities,
1184        SnapshotMetadata {
1185            subscription_id: &subscription.subscription_id,
1186            snapshot_id: &snapshot_id,
1187            authoritative,
1188            mode,
1189            view_id: &subscription.query.view,
1190            key: subscription.query.key.as_deref(),
1191        },
1192        batch_config,
1193    );
1194
1195    for frame in frames {
1196        let rows = frame.data.len() as u32;
1197        let json = serde_json::to_vec(&frame)?;
1198        let payload = maybe_compress(&json);
1199        let bytes = payload.as_bytes().len() as u64;
1200        context
1201            .client_manager
1202            .send_compressed_async(context.client_id, payload)
1203            .await
1204            .map_err(|error| anyhow::anyhow!("failed to send snapshot: {error}"))?;
1205        context.metrics.message_sent();
1206
1207        let auth_context = context.client_manager.get_auth_context(context.client_id);
1208        let (metering_key, subject, _, deployment_id) = usage_identity(auth_context.as_ref());
1209        emit_usage_event(
1210            &context.usage_emitter,
1211            WebSocketUsageEvent::SnapshotSent {
1212                client_id: context.client_id.to_string(),
1213                deployment_id,
1214                metering_key,
1215                subject,
1216                view_id: subscription.query.view.clone(),
1217                rows,
1218                messages: 1,
1219                bytes,
1220            },
1221        );
1222    }
1223    Ok(())
1224}
1225
1226fn extract_sort_config(view_spec: &ViewSpec) -> Option<SortConfig> {
1227    if let Some(sort) = view_spec
1228        .pipeline
1229        .as_ref()
1230        .and_then(|pipeline| pipeline.sort.as_ref())
1231    {
1232        return Some(SortConfig {
1233            field: sort.field_path.clone(),
1234            order: match sort.order {
1235                crate::materialized_view::SortOrder::Asc => SortOrder::Asc,
1236                crate::materialized_view::SortOrder::Desc => SortOrder::Desc,
1237            },
1238        });
1239    }
1240    (view_spec.mode == Mode::List).then(|| SortConfig {
1241        field: vec!["_seq".to_string()],
1242        order: SortOrder::Desc,
1243    })
1244}
1245
1246fn send_control_frame<T: Serialize>(
1247    context: &SubscriptionContext,
1248    frame: &T,
1249    view_id: &str,
1250) -> Result<()> {
1251    let json = serde_json::to_vec(frame)?;
1252    let bytes = json.len();
1253    context
1254        .client_manager
1255        .send_to_client(context.client_id, Arc::new(Bytes::from(json)))
1256        .map_err(|error| anyhow::anyhow!("failed to send control frame: {error}"))?;
1257    context.metrics.message_sent();
1258    emit_update_sent_for_client(
1259        &context.usage_emitter,
1260        &context.client_manager,
1261        context.client_id,
1262        view_id,
1263        bytes,
1264    );
1265    Ok(())
1266}
1267
1268fn send_subscribed_frame(
1269    context: &SubscriptionContext,
1270    subscription: &Subscription,
1271    view_spec: &ViewSpec,
1272) -> Result<()> {
1273    let frame = SubscribedFrame::new(
1274        subscription.subscription_id.clone(),
1275        subscription.query.clone(),
1276        view_spec.mode,
1277        extract_sort_config(view_spec),
1278    );
1279    send_control_frame(context, &frame, &subscription.query.view)
1280}
1281
1282fn enforce_snapshot_limit(context: &SubscriptionContext, rows: usize) -> Result<()> {
1283    context
1284        .client_manager
1285        .check_snapshot_allowed(context.client_id, u32::try_from(rows).unwrap_or(u32::MAX))
1286        .map_err(|deny| anyhow::anyhow!(deny.reason))
1287}
1288
1289async fn subscribe_state_then_snapshot<F, Fut, T>(
1290    bus_manager: &BusManager,
1291    view_id: &str,
1292    key: &str,
1293    snapshot: F,
1294) -> (watch::Receiver<StateUpdate>, u64, T)
1295where
1296    F: FnOnce() -> Fut,
1297    Fut: Future<Output = T>,
1298{
1299    let mut receiver = bus_manager.get_or_create_state_bus(view_id, key).await;
1300    let published = receiver.borrow_and_update().published;
1301    let snapshot = snapshot().await;
1302    (receiver, published, snapshot)
1303}
1304
1305async fn subscribe_list_then_snapshot<F, Fut, T>(
1306    bus_manager: &BusManager,
1307    view_id: &str,
1308    snapshot: F,
1309) -> (broadcast::Receiver<Arc<BusMessage>>, T)
1310where
1311    F: FnOnce() -> Fut,
1312    Fut: Future<Output = T>,
1313{
1314    let receiver = bus_manager.get_or_create_list_bus(view_id).await;
1315    let snapshot = snapshot().await;
1316    (receiver, snapshot)
1317}
1318
1319async fn attach_client_to_bus(
1320    context: &SubscriptionContext,
1321    mut subscription: Subscription,
1322    cancel_token: CancellationToken,
1323) -> Result<()> {
1324    let view_spec = context
1325        .view_index
1326        .get_view(&subscription.query.view)
1327        .cloned()
1328        .ok_or_else(|| anyhow::anyhow!("unknown view: {}", subscription.query.view))?;
1329
1330    if view_spec.mode == Mode::State && !view_spec.is_derived() && subscription.query.key.is_none()
1331    {
1332        return Err(anyhow::anyhow!("state subscriptions require query.key"));
1333    }
1334    if view_spec.is_derived() && subscription.query.take.is_none() {
1335        subscription.query.take = view_spec
1336            .pipeline
1337            .as_ref()
1338            .and_then(|pipeline| pipeline.limit);
1339    }
1340
1341    // A retained tape takes precedence for append views: it is the only
1342    // delivery that can honour a cursor. Without one, fall through to the
1343    // previous latest-state behaviour.
1344    let journal = context
1345        .journal
1346        .clone()
1347        .filter(|journal| journal.is_enabled() && view_spec.mode == Mode::Append);
1348    if let Some(journal) = journal {
1349        return attach_journal_subscription(
1350            context,
1351            subscription,
1352            view_spec,
1353            journal,
1354            cancel_token,
1355        )
1356        .await;
1357    }
1358
1359    if view_spec.mode == Mode::State && !view_spec.is_derived() {
1360        attach_state_subscription(context, subscription, view_spec, cancel_token).await
1361    } else {
1362        attach_collection_subscription(context, subscription, view_spec, cancel_token).await
1363    }
1364}
1365
1366async fn attach_state_subscription(
1367    context: &SubscriptionContext,
1368    subscription: Subscription,
1369    view_spec: ViewSpec,
1370    cancel_token: CancellationToken,
1371) -> Result<()> {
1372    let view_id = subscription.query.view.clone();
1373    let key = subscription.query.key.clone().unwrap_or_default();
1374    let query = subscription.query.clone();
1375    let cache = context.entity_cache.clone();
1376    let view_spec_for_snapshot = view_spec.clone();
1377    let (mut receiver, mut seen, initial) =
1378        subscribe_state_then_snapshot(&context.bus_manager, &view_id, &key, move || async move {
1379            load_query_entities(&cache, None, &view_spec_for_snapshot, &query, false).await
1380        })
1381        .await;
1382
1383    let mut snapshot_entities = initial;
1384    if let Some(limit) = subscription.query.snapshot_limit {
1385        snapshot_entities.truncate(limit);
1386    }
1387    enforce_snapshot_limit(context, snapshot_entities.len())?;
1388    send_subscribed_frame(context, &subscription, &view_spec)?;
1389    // The client holds the entity only if the snapshot actually carried it.
1390    // A disabled snapshot, or one `snapshotLimit` cut to nothing, leaves it
1391    // without a copy, so its first change must be a full `upsert`: a patch
1392    // would have nothing to merge into.
1393    let delivered = subscription.snapshot.enabled && !snapshot_entities.is_empty();
1394    let snapshot_rows = to_wire_snapshot_entities(snapshot_entities, &view_spec);
1395    // The entity as the client last received it whole, which a catch-up is
1396    // measured against.
1397    let synced = if delivered {
1398        snapshot_rows.first().map(|row| row.data.clone())
1399    } else {
1400        None
1401    };
1402    if subscription.snapshot.enabled {
1403        send_snapshot_batches(
1404            context,
1405            &subscription,
1406            &snapshot_rows,
1407            view_spec.mode,
1408            SnapshotPurpose::Initial,
1409            &context.entity_cache.snapshot_config(),
1410        )
1411        .await?;
1412    }
1413
1414    let task_context = context.clone();
1415    let subscription_id = subscription.subscription_id.clone();
1416    let query = subscription.query.clone();
1417    let view_spec_task = view_spec.clone();
1418    let span_view = view_id.clone();
1419    let span_key = key.clone();
1420    tokio::spawn(
1421        async move {
1422            // Whether the client holds this key, not whether it exists.
1423            let mut member = delivered;
1424            // The entity as the client last received it whole (its snapshot
1425            // row, an upsert or the last catch-up), in wire format. A catch-up
1426            // sends what changed since, and every field a patch forwarded in
1427            // between set (`touched`, see `mark_fields`): the entity may have
1428            // set it back, and a client may have dropped the patch.
1429            let mut synced = synced;
1430            let mut touched = Value::Null;
1431            // Whether the client lacks fields of frames the bus overwrote
1432            // before this task read them. The cached entity is sent in their
1433            // place, but while the cache lacks the key there is none, so the
1434            // debt waits for the next frame that finds it cached.
1435            let mut behind = false;
1436            // Whether, while behind, a frame the client missed replaced the
1437            // entity (a delete, or a whole upsert): its copy may then hold
1438            // fields the entity no longer has, which merging keeps, so the
1439            // catch-up must replace it.
1440            let mut replaced = false;
1441            loop {
1442                tokio::select! {
1443                    _ = cancel_token.cancelled() => break,
1444                    changed = receiver.changed() => {
1445                        if changed.is_err() {
1446                            break;
1447                        }
1448                        let (payload, published, replaced_at) = {
1449                            let update = receiver.borrow_and_update();
1450                            (update.payload.clone(), update.published, update.replaced)
1451                        };
1452                        // The bus keeps only the latest frame. If more than one
1453                        // was published since the last read, the earlier ones
1454                        // were overwritten and the latest patch alone would
1455                        // drop their fields, so send the cached entity instead.
1456                        behind |= published > seen + 1;
1457                        replaced |= behind && replaced_at > seen;
1458                        seen = published;
1459                        let metadata = source_frame_metadata(&payload);
1460                        if metadata.op == "delete" {
1461                            task_context.entity_cache.remove(&query.view, &key).await;
1462                            if member && send_membership_frame(
1463                                &task_context,
1464                                &subscription_id,
1465                                &view_spec_task,
1466                                "delete",
1467                                &key,
1468                                Value::Null,
1469                                metadata.seq,
1470                            ).is_err() {
1471                                break;
1472                            }
1473                            member = false;
1474                            behind = false;
1475                            replaced = false;
1476                            synced = None;
1477                            touched = Value::Null;
1478                            continue;
1479                        }
1480
1481                        // A key missing from the cache was evicted there or
1482                        // never held whole, not deleted (deletes arrive
1483                        // above): the cache refused this patch and the whole
1484                        // entity is on its way (see `EntityResync`). The
1485                        // client's copy is still good, so a holder gets the
1486                        // patch and stays a holder; `remove` is only for an
1487                        // entity that stopped matching. If it is behind, it
1488                        // stays behind: the resend that brings the entity
1489                        // back carries the seq of the latest change, which it
1490                        // already has, so only the catch-up below reaches it.
1491                        let Some(cached) = task_context
1492                            .entity_cache
1493                            .get(&view_spec_task.id, &key)
1494                            .await
1495                        else {
1496                            // A copy of a replaced entity takes no patch: it
1497                            // waits for the whole entity to replace it.
1498                            if member && !replaced {
1499                                record_forwarded(&payload, &metadata.op, &mut synced, &mut touched);
1500                                if send_scoped_source_payload(
1501                                    &task_context,
1502                                    &subscription_id,
1503                                    &query.view,
1504                                    payload,
1505                                ).is_err() {
1506                                    break;
1507                                }
1508                            }
1509                            continue;
1510                        };
1511                        let selected =
1512                            select_query_entities(vec![(key.clone(), cached)], &query, true, false);
1513                        let is_member = !selected.is_empty();
1514                        let caught_up = !std::mem::take(&mut behind);
1515                        let replace = std::mem::take(&mut replaced);
1516                        let result = match (member, selected.into_iter().next()) {
1517                            (true, Some(_)) if caught_up => {
1518                                record_forwarded(&payload, &metadata.op, &mut synced, &mut touched);
1519                                send_scoped_source_payload(
1520                                    &task_context,
1521                                    &subscription_id,
1522                                    &query.view,
1523                                    payload,
1524                                )
1525                            }
1526                            // The catch-up: a patch of the cached entity's
1527                            // fields that changed since the client last received
1528                            // it whole or that a forwarded patch set, which
1529                            // merges into what the client holds.
1530                            // If a missed frame replaced the entity, it is the
1531                            // whole entity as an upsert instead, which replaces
1532                            // the client's copy: merging would keep the fields of
1533                            // the entity that was deleted or replaced. Either
1534                            // carries no seq, because it is newer than anything
1535                            // this subscriber was sent and seqs are not ordered
1536                            // within a slot: account updates and instructions
1537                            // number themselves differently, so the latest
1538                            // patch's seq can sort below one already delivered.
1539                            (true, Some((entity_key, data))) => {
1540                                let mut current = data;
1541                                apply_wire_format(&mut current, &view_spec_task.wire_format);
1542                                let body = match (&synced, replace) {
1543                                    (Some(base), false) => changed_fields(base, &current, &touched),
1544                                    _ => Some(current.clone()),
1545                                };
1546                                synced = Some(current);
1547                                touched = Value::Null;
1548                                match body {
1549                                    Some(body) => send_membership_frame(
1550                                        &task_context,
1551                                        &subscription_id,
1552                                        &view_spec_task,
1553                                        if replace { "upsert" } else { "patch" },
1554                                        &entity_key,
1555                                        body,
1556                                        None,
1557                                    ),
1558                                    None => Ok(()),
1559                                }
1560                            }
1561                            (false, Some((entity_key, data))) => {
1562                                let mut current = data;
1563                                apply_wire_format(&mut current, &view_spec_task.wire_format);
1564                                synced = Some(current.clone());
1565                                touched = Value::Null;
1566                                send_membership_frame(
1567                                    &task_context,
1568                                    &subscription_id,
1569                                    &view_spec_task,
1570                                    "upsert",
1571                                    &entity_key,
1572                                    current,
1573                                    metadata.seq,
1574                                )
1575                            }
1576                            (true, None) => {
1577                                synced = None;
1578                                touched = Value::Null;
1579                                send_membership_frame(
1580                                &task_context,
1581                                &subscription_id,
1582                                &view_spec_task,
1583                                "remove",
1584                                &key,
1585                                Value::Null,
1586                                metadata.seq,
1587                                )
1588                            }
1589                            (false, None) => Ok(()),
1590                        };
1591                        if result.is_err() {
1592                            break;
1593                        }
1594                        member = is_member;
1595                    }
1596                }
1597            }
1598        }
1599        .instrument(info_span!("ws.subscribe.state", client_id = %context.client_id, view = %span_view, key = %span_key)),
1600    );
1601    Ok(())
1602}
1603
1604async fn attach_collection_subscription(
1605    context: &SubscriptionContext,
1606    subscription: Subscription,
1607    view_spec: ViewSpec,
1608    cancel_token: CancellationToken,
1609) -> Result<()> {
1610    let view_id = subscription.query.view.clone();
1611    let source_view_id = view_spec
1612        .source_view
1613        .clone()
1614        .unwrap_or_else(|| view_id.clone());
1615    let (mut receiver, initial_membership) = subscribe_collection_then_snapshot(
1616        context,
1617        &source_view_id,
1618        &view_spec,
1619        &subscription.query,
1620    )
1621    .await;
1622
1623    let mut snapshot_entities = initial_membership.clone();
1624    if let Some(limit) = subscription.query.snapshot_limit {
1625        snapshot_entities.truncate(limit);
1626    }
1627    enforce_snapshot_limit(context, snapshot_entities.len())?;
1628    send_subscribed_frame(context, &subscription, &view_spec)?;
1629    // Window members the client was never sent: everything past a
1630    // `snapshotLimit` cut, or the whole window when the snapshot is disabled.
1631    // They stay in the window (take/skip positions still count them) but the
1632    // client holds no copy, so none of them may ever be sent a patch.
1633    let delivered_rows = if subscription.snapshot.enabled {
1634        snapshot_entities.len()
1635    } else {
1636        0
1637    };
1638    let mut undelivered: HashSet<String> = initial_membership
1639        .iter()
1640        .skip(delivered_rows)
1641        .map(|(key, _)| key.clone())
1642        .collect();
1643    if subscription.snapshot.enabled {
1644        send_snapshot_batches(
1645            context,
1646            &subscription,
1647            &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1648            view_spec.mode,
1649            SnapshotPurpose::Initial,
1650            &context.entity_cache.snapshot_config(),
1651        )
1652        .await?;
1653    }
1654
1655    let task_context = context.clone();
1656    let task_subscription = subscription.clone();
1657    let subscription_id = subscription.subscription_id.clone();
1658    let query = subscription.query.clone();
1659    let view_spec_task = view_spec.clone();
1660    let span_view = view_id.clone();
1661    tokio::spawn(
1662        async move {
1663            let mut current = initial_membership;
1664            // Append views are event tapes: collapsing two records would lose
1665            // observable history. Coalescing is only valid for latest-state
1666            // list membership, where a full final entity preserves meaning.
1667            let coalesce_ms = (view_spec_task.mode == Mode::List)
1668                .then(|| {
1669                    view_spec_task
1670                        .delivery
1671                        .coalesce_ms
1672                        .or(task_context.delivery.collection_coalesce_ms)
1673                        .filter(|milliseconds| *milliseconds > 0)
1674                })
1675                .flatten();
1676            let mut flush_interval = coalesce_ms.map(|milliseconds| {
1677                let period = Duration::from_millis(milliseconds);
1678                let mut interval = tokio::time::interval_at(
1679                    tokio::time::Instant::now() + period,
1680                    period,
1681                );
1682                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
1683                interval
1684            });
1685            let mut pending = HashMap::<String, Arc<BusMessage>>::new();
1686            let mut pending_updates = 0_u64;
1687
1688            loop {
1689                tokio::select! {
1690                    _ = cancel_token.cancelled() => break,
1691                    _ = async {
1692                        flush_interval
1693                            .as_mut()
1694                            .expect("coalescing interval is guarded")
1695                            .tick()
1696                            .await;
1697                    }, if flush_interval.is_some() => {
1698                        if pending.is_empty() {
1699                            continue;
1700                        }
1701                        let sorted_caches = view_spec_task
1702                            .is_derived()
1703                            .then(|| task_context.view_index.sorted_caches());
1704                        let next = load_query_entities(
1705                            &task_context.entity_cache,
1706                            sorted_caches,
1707                            &view_spec_task,
1708                            &query,
1709                            false,
1710                        ).await;
1711                        if emit_coalesced_collection_delta(
1712                            &task_context,
1713                            &subscription_id,
1714                            &view_spec_task,
1715                            &current,
1716                            &next,
1717                            &pending,
1718                            &mut undelivered,
1719                        ).is_err() {
1720                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
1721                            break;
1722                        }
1723                        task_context
1724                            .metrics
1725                            .collection_coalesced(&view_id, pending_updates);
1726                        current = next;
1727                        pending.clear();
1728                        pending_updates = 0;
1729                    }
1730                    received = receiver.recv() => {
1731                        let envelope = match received {
1732                            Ok(envelope) => envelope,
1733                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
1734                                task_context.metrics.subscription_lagged(&view_id, skipped);
1735                                warn!(
1736                                    "Subscription {} lagged by {} updates",
1737                                    subscription_id, skipped
1738                                );
1739                                if view_spec_task.mode == Mode::Append {
1740                                    let _ = send_control_frame(
1741                                        &task_context,
1742                                        &SocketIssueMessage::append_subscription_lagged(
1743                                            subscription_id.clone(),
1744                                            skipped,
1745                                        ),
1746                                        &view_id,
1747                                    );
1748                                    task_context.metrics.delivery_stopped(&view_id, "append-lagged-without-replay");
1749                                    break;
1750                                }
1751                                if !task_subscription.snapshot.enabled {
1752                                    let _ = send_control_frame(
1753                                        &task_context,
1754                                        &SocketIssueMessage::subscription_lagged(
1755                                            subscription_id.clone(),
1756                                            skipped,
1757                                        ),
1758                                        &view_id,
1759                                    );
1760                                    task_context.metrics.delivery_stopped(&view_id, "lagged-without-snapshot");
1761                                    break;
1762                                }
1763                                info!(
1764                                    "Subscription {} is recovering from an authoritative snapshot",
1765                                    subscription_id
1766                                );
1767                                match recover_collection_subscription(
1768                                    &task_context,
1769                                    &task_subscription,
1770                                    &view_spec_task,
1771                                    &source_view_id,
1772                                ).await {
1773                                    Ok((next_receiver, recovered)) => {
1774                                        receiver = next_receiver;
1775                                        current = recovered;
1776                                        // Recovery replaces the whole live
1777                                        // membership, untruncated.
1778                                        undelivered.clear();
1779                                        pending.clear();
1780                                        pending_updates = 0;
1781                                        task_context.metrics.subscription_resnapshot(&view_id);
1782                                        continue;
1783                                    }
1784                                    Err(error) => {
1785                                        warn!(
1786                                            "Subscription {} failed to recover from lag: {error:#}",
1787                                            subscription_id
1788                                        );
1789                                        task_context.metrics.delivery_stopped(&view_id, "resnapshot-failed");
1790                                        break;
1791                                    }
1792                                }
1793                            }
1794                            Err(broadcast::error::RecvError::Closed) => break,
1795                        };
1796
1797                        apply_collection_source_event(
1798                            &task_context,
1799                            &source_view_id,
1800                            &view_spec_task,
1801                            &query,
1802                            &envelope,
1803                        ).await;
1804
1805                        if flush_interval.is_some() {
1806                            pending_updates = pending_updates.saturating_add(1);
1807                            pending.insert(envelope.key.clone(), envelope);
1808                            continue;
1809                        }
1810
1811                        let metadata = source_frame_metadata(&envelope.payload);
1812                        let sorted_caches = view_spec_task
1813                            .is_derived()
1814                            .then(|| task_context.view_index.sorted_caches());
1815                        let next = load_query_entities(
1816                            &task_context.entity_cache,
1817                            sorted_caches,
1818                            &view_spec_task,
1819                            &query,
1820                            false,
1821                        ).await;
1822                        if emit_collection_delta(
1823                            &task_context,
1824                            &subscription_id,
1825                            &view_spec_task,
1826                            &current,
1827                            &next,
1828                            &envelope,
1829                            &metadata,
1830                            &mut undelivered,
1831                        ).is_err() {
1832                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
1833                            break;
1834                        }
1835                        current = next;
1836                    }
1837                }
1838            }
1839        }
1840        .instrument(info_span!("ws.subscribe.collection", client_id = %context.client_id, view = %span_view)),
1841    );
1842    Ok(())
1843}
1844
1845async fn subscribe_collection_then_snapshot(
1846    context: &SubscriptionContext,
1847    source_view_id: &str,
1848    view_spec: &ViewSpec,
1849    query: &SubscriptionQuery,
1850) -> (broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>) {
1851    let cache = context.entity_cache.clone();
1852    let sorted_caches = view_spec
1853        .is_derived()
1854        .then(|| context.view_index.sorted_caches());
1855    let view_spec = view_spec.clone();
1856    let query = query.clone();
1857    subscribe_list_then_snapshot(&context.bus_manager, source_view_id, move || async move {
1858        load_query_entities(&cache, sorted_caches, &view_spec, &query, false).await
1859    })
1860    .await
1861}
1862
1863async fn recover_collection_subscription(
1864    context: &SubscriptionContext,
1865    subscription: &Subscription,
1866    view_spec: &ViewSpec,
1867    source_view_id: &str,
1868) -> Result<(broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>)> {
1869    let (receiver, membership) =
1870        subscribe_collection_then_snapshot(context, source_view_id, view_spec, &subscription.query)
1871            .await;
1872    // `snapshotLimit` caps only the initial transfer. Recovery has to replace
1873    // the complete live membership; truncating it would make the replacement
1874    // authoritative while immediately omitting members the server still
1875    // considers current.
1876    enforce_snapshot_limit(context, membership.len())?;
1877    send_snapshot_batches(
1878        context,
1879        subscription,
1880        &to_wire_snapshot_entities(membership.clone(), view_spec),
1881        view_spec.mode,
1882        SnapshotPurpose::Recovery,
1883        &context.entity_cache.snapshot_config(),
1884    )
1885    .await?;
1886    Ok((receiver, membership))
1887}
1888
1889async fn apply_collection_source_event(
1890    context: &SubscriptionContext,
1891    source_view_id: &str,
1892    view_spec: &ViewSpec,
1893    query: &SubscriptionQuery,
1894    envelope: &BusMessage,
1895) {
1896    let metadata = source_frame_metadata(&envelope.payload);
1897    if metadata.op != "delete" {
1898        return;
1899    }
1900    // A slow subscription can observe an old delete after the projector has
1901    // already recreated the key. Never let subscriber-local lag erase newer
1902    // shared cache state.
1903    let current = context
1904        .entity_cache
1905        .get(source_view_id, &envelope.key)
1906        .await;
1907    if source_delete_is_stale(current.as_ref(), metadata.seq.as_deref()) {
1908        return;
1909    }
1910    context
1911        .entity_cache
1912        .remove(source_view_id, &envelope.key)
1913        .await;
1914    if view_spec.is_derived() {
1915        let caches = context.view_index.sorted_caches();
1916        let mut guard = caches.write().await;
1917        if let Some(cache) = guard.get_mut(&query.view) {
1918            cache.remove(&envelope.key);
1919        }
1920    }
1921}
1922
1923fn source_delete_is_stale(current: Option<&Value>, delete_seq: Option<&str>) -> bool {
1924    match (current, delete_seq) {
1925        (Some(current), Some(delete_seq)) => current
1926            .get("_seq")
1927            .and_then(Value::as_str)
1928            .is_some_and(|current_seq| {
1929                cmp_seq(current_seq, delete_seq) == std::cmp::Ordering::Greater
1930            }),
1931        _ => false,
1932    }
1933}
1934
1935/// A subscription refused for a reason the client needs spelled out.
1936///
1937/// Attach paths that return this get the registration released by the
1938/// connection loop, the same as any other failure, while the client still
1939/// receives the specific error rather than a generic `subscription-rejected`.
1940/// Sending the frame and returning `Ok` instead would leave a registered
1941/// subscription with nothing attached: it would hold a slot against the
1942/// client's limit and make the advertised "resubscribe" remediation fail with
1943/// `duplicate-subscription-id`.
1944#[derive(Debug)]
1945pub(crate) struct RejectedSubscription(pub SocketIssueMessage);
1946
1947impl RejectedSubscription {
1948    fn into_error(issue: SocketIssueMessage) -> anyhow::Error {
1949        anyhow::Error::new(Self(issue))
1950    }
1951}
1952
1953impl std::fmt::Display for RejectedSubscription {
1954    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1955        formatter.write_str(&self.0.message)
1956    }
1957}
1958
1959impl std::error::Error for RejectedSubscription {}
1960
1961/// Deliver an append view as an event tape: replay the retained records after
1962/// the cursor, then forward live frames.
1963///
1964/// This deliberately does not recompute membership from the entity cache the
1965/// way [`attach_collection_subscription`] does. The cache folds each patch
1966/// into the resident entity, so a membership diff cannot express "these three
1967/// events happened"; the retained records can.
1968///
1969/// For the same reason records are never promoted to a full `upsert` the
1970/// first time this subscription sees their key, the way a window entry is.
1971/// The cache only holds the latest state, not the state as of a record's
1972/// offset, so substituting it would hand a consumer state from after its
1973/// cursor. Records go out exactly as retained, `patch` included: a tape is
1974/// events, and the consumer that resumes from a cursor holds whatever came
1975/// before it. Clients therefore apply frames that carry an `offset` as events
1976/// and do not discard them as partial entities (see "Partial entities" in
1977/// `docs/websocket-v2-protocol.md`).
1978async fn attach_journal_subscription(
1979    context: &SubscriptionContext,
1980    subscription: Subscription,
1981    view_spec: ViewSpec,
1982    journal: Arc<crate::journal::EventJournal>,
1983    cancel_token: CancellationToken,
1984) -> Result<()> {
1985    let view_id = view_spec.id.clone();
1986    let subscription_id = subscription.subscription_id.clone();
1987
1988    // A tape has no membership window, so `take`/`skip` cannot mean what they
1989    // mean on a list view. Refuse them rather than accept and ignore them.
1990    if subscription.query.take.is_some()
1991        || subscription.query.skip.is_some()
1992        || subscription.query.snapshot_limit.is_some()
1993    {
1994        return Err(RejectedSubscription::into_error(SocketIssueMessage::protocol(
1995            Some(subscription_id),
1996            "invalid-subscription",
1997            format!(
1998                "take, skip and snapshotLimit are window options and do not apply to the replayable view {view_id}"
1999            ),
2000        )));
2001    }
2002
2003    // Reject a malformed cursor rather than silently replaying from the start,
2004    // which would look like success and duplicate everything.
2005    let cursor = match subscription.query.after.as_deref() {
2006        Some(raw) => match crate::journal::Cursor::parse(raw) {
2007            Some(cursor) => Some(cursor),
2008            None => {
2009                return Err(RejectedSubscription::into_error(
2010                    SocketIssueMessage::protocol(
2011                        Some(subscription_id),
2012                        "invalid-cursor",
2013                        format!(
2014                            "`after` must be an {{epoch}}:{{offset}} replay cursor for view {view_id}, got {raw:?}"
2015                        ),
2016                    ),
2017                ));
2018            }
2019        },
2020        None => None,
2021    };
2022
2023    // Subscribe before reading the journal so anything published during the
2024    // replay is still delivered; the offset filter below drops the overlap.
2025    let mut receiver = context.bus_manager.get_or_create_list_bus(&view_id).await;
2026
2027    let replayed = match journal.replay_after(&view_id, cursor.as_ref()).await {
2028        Ok(records) => records,
2029        Err(error) => {
2030            return Err(RejectedSubscription::into_error(
2031                SocketIssueMessage::replay_refused(Some(subscription_id), &error),
2032            ));
2033        }
2034    };
2035
2036    let frame = SubscribedFrame::new(
2037        subscription.subscription_id.clone(),
2038        subscription.query.clone(),
2039        view_spec.mode,
2040        extract_sort_config(&view_spec),
2041    )
2042    .with_replay_window(journal.window(&view_id).await);
2043    send_control_frame(context, &frame, &view_id)?;
2044
2045    // Everything past the acknowledgement runs on its own task. The replay can
2046    // be long and applies real backpressure, and `attach_client_to_bus` is
2047    // awaited directly on the connection's inbound loop — doing it there would
2048    // block unsubscribe, auth refresh and pong for the whole replay.
2049    let task_context = context.clone();
2050    let task_subscription_id = subscription.subscription_id.clone();
2051    let task_query = subscription.query.clone();
2052    let task_epoch = journal.epoch().await;
2053    let span_view = view_id.clone();
2054    tokio::spawn(
2055        async move {
2056            let mut last_sent = cursor.map(|cursor| cursor.offset);
2057            // Frames that published while the replay was still running. The
2058            // bus is a bounded broadcast, so it has to be drained as we go or
2059            // a busy view laps us before the replay finishes.
2060            let mut pending: VecDeque<Arc<BusMessage>> = VecDeque::new();
2061            let mut lagged: Option<u64> = None;
2062
2063            for record in replayed {
2064                if cancel_token.is_cancelled() {
2065                    return;
2066                }
2067                drain_available(&mut receiver, &mut pending, &mut lagged);
2068                if journal_record_matches(&task_query, &record)
2069                    && send_scoped_source_payload_async(
2070                        &task_context,
2071                        &task_subscription_id,
2072                        &span_view,
2073                        record.payload,
2074                    )
2075                    .await
2076                    .is_err()
2077                {
2078                    return;
2079                }
2080                // Advance past filtered records too: they were considered.
2081                last_sent = Some(record.offset);
2082            }
2083
2084            // Flush what arrived during the replay before going live, so the
2085            // handover keeps offset order.
2086            while let Some(envelope) = pending.pop_front() {
2087                if !forward_live_frame(
2088                    &task_context,
2089                    &task_subscription_id,
2090                    &span_view,
2091                    &task_query,
2092                    &envelope,
2093                    &mut last_sent,
2094                )
2095                .await
2096                {
2097                    return;
2098                }
2099            }
2100
2101            if let Some(skipped) = lagged {
2102                report_replay_gap(
2103                    &task_context,
2104                    &task_subscription_id,
2105                    &span_view,
2106                    &task_epoch,
2107                    skipped,
2108                    last_sent,
2109                );
2110                return;
2111            }
2112
2113            loop {
2114                tokio::select! {
2115                    _ = cancel_token.cancelled() => break,
2116                    received = receiver.recv() => {
2117                        let envelope = match received {
2118                            Ok(envelope) => envelope,
2119                            // A lagged tape is a gap. Report it with the last
2120                            // offset delivered *before* the gap and stop
2121                            // delivering on this subscription.
2122                            //
2123                            // Continuing would hand the consumer frames from
2124                            // after the gap, advancing its checkpoint past
2125                            // the skipped records so they could never be
2126                            // replayed. Stopping is not a silent stall: the
2127                            // consumer has an explicit error and a cursor
2128                            // that recovers exactly what it missed.
2129                            //
2130                            // The registration is deliberately left alone.
2131                            // Its lifecycle belongs to the connection loop,
2132                            // which holds the only handle to
2133                            // `active_subscriptions`; releasing half of it
2134                            // here would desynchronise unsubscribe, the
2135                            // duplicate-ID gate and close-time usage.
2136                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
2137                                report_replay_gap(
2138                                    &task_context,
2139                                    &task_subscription_id,
2140                                    &span_view,
2141                                    &task_epoch,
2142                                    skipped,
2143                                    last_sent,
2144                                );
2145                                break;
2146                            }
2147                            Err(broadcast::error::RecvError::Closed) => break,
2148                        };
2149
2150                        if !forward_live_frame(
2151                            &task_context,
2152                            &task_subscription_id,
2153                            &span_view,
2154                            &task_query,
2155                            &envelope,
2156                            &mut last_sent,
2157                        )
2158                        .await
2159                        {
2160                            break;
2161                        }
2162                    }
2163                }
2164            }
2165        }
2166        .instrument(info_span!(
2167            "ws.subscribe.replay",
2168            client_id = %context.client_id,
2169            view = %view_id
2170        )),
2171    );
2172    Ok(())
2173}
2174
2175/// Take whatever the bus already has without waiting, so a long replay cannot
2176/// be lapped by a busy view.
2177fn drain_available(
2178    receiver: &mut broadcast::Receiver<Arc<BusMessage>>,
2179    pending: &mut VecDeque<Arc<BusMessage>>,
2180    lagged: &mut Option<u64>,
2181) {
2182    // Once a gap is known, everything still on the bus is on the far side of
2183    // it. Buffering it would put post-gap frames in front of the lag report,
2184    // advancing `last_sent` past the hole and making `recoverFrom` point
2185    // after the very records it is supposed to recover.
2186    if lagged.is_some() {
2187        return;
2188    }
2189    // Bounded so a view publishing faster than the client drains cannot turn
2190    // the buffer into an unbounded queue. Filling it is not itself a gap:
2191    // nothing has been skipped at that instant, the buffer simply stopped
2192    // accepting. Stop buffering and let the bus report the loss, with the
2193    // count it actually measures, when delivery reaches it.
2194    const MAX_PENDING: usize = 8_192;
2195    loop {
2196        if pending.len() >= MAX_PENDING {
2197            return;
2198        }
2199        match receiver.try_recv() {
2200            Ok(envelope) => pending.push_back(envelope),
2201            Err(broadcast::error::TryRecvError::Empty)
2202            | Err(broadcast::error::TryRecvError::Closed) => return,
2203            Err(broadcast::error::TryRecvError::Lagged(skipped)) => {
2204                *lagged = Some(skipped);
2205                return;
2206            }
2207        }
2208    }
2209}
2210
2211/// Whether a live frame was already delivered by the replay that preceded it.
2212///
2213/// The bus is subscribed before the tape is read, so a record published in
2214/// between appears on both paths; without this the consumer sees it twice.
2215/// Advances the high-water mark as a side effect.
2216fn already_delivered(offset: Option<u64>, last_sent: &mut Option<u64>) -> bool {
2217    let Some(offset) = offset else {
2218        // A frame with no offset predates the tape, so it cannot have been
2219        // replayed and must not move the mark.
2220        return false;
2221    };
2222    if last_sent.is_some_and(|last| offset <= last) {
2223        return true;
2224    }
2225    *last_sent = Some(offset);
2226    false
2227}
2228
2229/// Deliver one live frame, skipping anything the replay already sent.
2230///
2231/// Returns false when the subscription should end.
2232async fn forward_live_frame(
2233    context: &SubscriptionContext,
2234    subscription_id: &str,
2235    view_id: &str,
2236    query: &SubscriptionQuery,
2237    envelope: &Arc<BusMessage>,
2238    last_sent: &mut Option<u64>,
2239) -> bool {
2240    let metadata = source_frame_metadata(&envelope.payload);
2241    if already_delivered(metadata.offset, last_sent) {
2242        return true;
2243    }
2244    if !live_frame_matches(query, &envelope.key, &envelope.payload) {
2245        return true;
2246    }
2247    send_scoped_source_payload(context, subscription_id, view_id, envelope.payload.clone()).is_ok()
2248}
2249
2250fn report_replay_gap(
2251    context: &SubscriptionContext,
2252    subscription_id: &str,
2253    view_id: &str,
2254    epoch: &crate::journal::JournalEpoch,
2255    skipped: u64,
2256    last_sent: Option<u64>,
2257) {
2258    warn!(
2259        "Replay subscription {} lagged past {} records; stopping with a recovery cursor",
2260        subscription_id, skipped
2261    );
2262    let recover_from = last_sent.map(|offset| crate::journal::Cursor {
2263        epoch: epoch.clone(),
2264        offset,
2265    });
2266    let _ = send_control_frame(
2267        context,
2268        &SocketIssueMessage::replay_lagged(
2269            Some(subscription_id.to_string()),
2270            skipped,
2271            recover_from,
2272        ),
2273        view_id,
2274    );
2275}
2276
2277/// Apply the subscription's `key`, `partition` and `filters` to a retained
2278/// record. A replay must honour the same predicates a live subscription does.
2279fn journal_record_matches(
2280    query: &SubscriptionQuery,
2281    record: &crate::journal::JournalRecord,
2282) -> bool {
2283    live_frame_matches(query, &record.key, &record.payload)
2284}
2285
2286fn live_frame_matches(query: &SubscriptionQuery, key: &str, payload: &[u8]) -> bool {
2287    if !query.matches_key(key) {
2288        return false;
2289    }
2290    if query.partition.is_none() && query.filters.is_empty() {
2291        return true;
2292    }
2293    let Ok(frame) = serde_json::from_slice::<Value>(payload) else {
2294        return false;
2295    };
2296    let Some(data) = frame.get("data") else {
2297        return false;
2298    };
2299    if let Some(partition) = &query.partition {
2300        if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2301            return false;
2302        }
2303    }
2304    query
2305        .filters
2306        .iter()
2307        .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2308}
2309
2310/// Awaiting variant of [`send_scoped_source_payload`], for replays that can
2311/// exceed the client's send queue.
2312async fn send_scoped_source_payload_async(
2313    context: &SubscriptionContext,
2314    subscription_id: &str,
2315    view_id: &str,
2316    payload: Arc<Bytes>,
2317) -> Result<()> {
2318    let mut value: Value = serde_json::from_slice(&payload)?;
2319    let object = value
2320        .as_object_mut()
2321        .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2322    object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2323    object.insert(
2324        "subscriptionId".to_string(),
2325        Value::String(subscription_id.to_string()),
2326    );
2327    let json = serde_json::to_vec(&value)?;
2328    let compressed = maybe_compress(&json);
2329    let bytes = compressed.as_bytes().len();
2330    context
2331        .client_manager
2332        .send_compressed_async(context.client_id, compressed)
2333        .await
2334        .map_err(|error| anyhow::anyhow!("failed to send replayed frame: {error}"))?;
2335    context.metrics.message_sent();
2336    emit_update_sent_for_client(
2337        &context.usage_emitter,
2338        &context.client_manager,
2339        context.client_id,
2340        view_id,
2341        bytes,
2342    );
2343    Ok(())
2344}
2345
2346#[derive(Default)]
2347struct SourceFrameMetadata {
2348    op: String,
2349    seq: Option<String>,
2350    offset: Option<u64>,
2351}
2352
2353fn source_frame_metadata(payload: &[u8]) -> SourceFrameMetadata {
2354    serde_json::from_slice::<Value>(payload)
2355        .ok()
2356        .map(|value| SourceFrameMetadata {
2357            op: value
2358                .get("op")
2359                .and_then(Value::as_str)
2360                .unwrap_or_default()
2361                .to_string(),
2362            seq: value.get("seq").and_then(Value::as_str).map(str::to_string),
2363            offset: value.get("offset").and_then(Value::as_u64),
2364        })
2365        .unwrap_or_default()
2366}
2367
2368/// The `data` of a source frame, as a client receives it.
2369fn source_frame_data(payload: &[u8]) -> Option<Value> {
2370    let mut frame: Value = serde_json::from_slice(payload).ok()?;
2371    frame.get_mut("data").map(Value::take)
2372}
2373
2374/// Tracks what a client holds after a frame forwarded to it as published. An
2375/// upsert hands it the entity whole, which a catch-up is then measured
2376/// against; a patch sets fields (see [`mark_fields`]).
2377fn record_forwarded(payload: &[u8], op: &str, synced: &mut Option<Value>, touched: &mut Value) {
2378    let data = source_frame_data(payload);
2379    if op == "upsert" {
2380        *synced = data;
2381        *touched = Value::Null;
2382    } else if let Some(data) = data {
2383        mark_fields(touched, &data);
2384    }
2385}
2386
2387/// Marks in `touched` the fields `patch` sets: objects field by field, any
2388/// other value whole (`true`). `null` marks nothing.
2389///
2390/// A client holds whatever a forwarded patch set, which the entity may have
2391/// set back before the next catch-up (or the client may have dropped the
2392/// patch), so a catch-up sends these fields even when they match the entity
2393/// the client last received whole.
2394fn mark_fields(touched: &mut Value, patch: &Value) {
2395    let Value::Object(fields) = patch else {
2396        *touched = Value::Bool(true);
2397        return;
2398    };
2399    if touched.is_null() {
2400        *touched = Value::Object(serde_json::Map::new());
2401    }
2402    let Value::Object(marked) = touched else {
2403        // Already marked whole.
2404        return;
2405    };
2406    for (field, value) in fields {
2407        mark_fields(marked.entry(field.clone()).or_insert(Value::Null), value);
2408    }
2409}
2410
2411/// The fields of `current` that differ from `base` or that `touched` marks
2412/// (see [`mark_fields`]), as a patch that turns what the client holds into
2413/// `current` when merged: objects compare field by field, any other value is
2414/// sent whole, and a field `current` no longer has is sent as `null`. `None`
2415/// when nothing needs sending.
2416fn changed_fields(base: &Value, current: &Value, touched: &Value) -> Option<Value> {
2417    match (base, current, touched) {
2418        (_, _, Value::Bool(true)) => Some(current.clone()),
2419        (Value::Object(base), Value::Object(current), _) => {
2420            let marked = touched.as_object();
2421            let mut patch = serde_json::Map::new();
2422            for (field, value) in current {
2423                let touched = marked
2424                    .and_then(|marked| marked.get(field))
2425                    .unwrap_or(&Value::Null);
2426                match base.get(field) {
2427                    Some(old) => {
2428                        if let Some(changed) = changed_fields(old, value, touched) {
2429                            patch.insert(field.clone(), changed);
2430                        }
2431                    }
2432                    None => {
2433                        patch.insert(field.clone(), value.clone());
2434                    }
2435                }
2436            }
2437            for field in base
2438                .keys()
2439                .chain(marked.into_iter().flat_map(|marked| marked.keys()))
2440            {
2441                if !current.contains_key(field) {
2442                    patch.insert(field.clone(), Value::Null);
2443                }
2444            }
2445            (!patch.is_empty()).then_some(Value::Object(patch))
2446        }
2447        _ => (base != current || !touched.is_null()).then(|| current.clone()),
2448    }
2449}
2450
2451fn send_scoped_source_payload(
2452    context: &SubscriptionContext,
2453    subscription_id: &str,
2454    view_id: &str,
2455    payload: Arc<Bytes>,
2456) -> Result<()> {
2457    let mut value: Value = serde_json::from_slice(&payload)?;
2458    let object = value
2459        .as_object_mut()
2460        .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2461    object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2462    object.insert(
2463        "subscriptionId".to_string(),
2464        Value::String(subscription_id.to_string()),
2465    );
2466    let encoded = Arc::new(Bytes::from(serde_json::to_vec(&value)?));
2467    let bytes = encoded.len();
2468    context
2469        .client_manager
2470        .send_to_client(context.client_id, encoded)
2471        .map_err(|error| anyhow::anyhow!("failed to send live frame: {error}"))?;
2472    context.metrics.message_sent();
2473    emit_update_sent_for_client(
2474        &context.usage_emitter,
2475        &context.client_manager,
2476        context.client_id,
2477        view_id,
2478        bytes,
2479    );
2480    Ok(())
2481}
2482
2483fn send_membership_frame(
2484    context: &SubscriptionContext,
2485    subscription_id: &str,
2486    view_spec: &ViewSpec,
2487    op: &str,
2488    key: &str,
2489    mut data: Value,
2490    seq: Option<String>,
2491) -> Result<()> {
2492    apply_wire_format(&mut data, &view_spec.wire_format);
2493    let frame = Frame::scoped(
2494        subscription_id,
2495        view_spec.mode,
2496        &view_spec.id,
2497        op,
2498        key,
2499        data,
2500        seq,
2501    );
2502    let encoded = Arc::new(Bytes::from(serde_json::to_vec(&frame)?));
2503    let bytes = encoded.len();
2504    context
2505        .client_manager
2506        .send_to_client(context.client_id, encoded)
2507        .map_err(|error| anyhow::anyhow!("failed to send membership frame: {error}"))?;
2508    context.metrics.message_sent();
2509    emit_update_sent_for_client(
2510        &context.usage_emitter,
2511        &context.client_manager,
2512        context.client_id,
2513        &view_spec.id,
2514        bytes,
2515    );
2516    Ok(())
2517}
2518
2519/// Send what one source mutation owes a subscriber, and keep `undelivered`
2520/// (window members the client holds no copy of) in step with what was sent.
2521#[allow(clippy::too_many_arguments)]
2522fn emit_collection_delta(
2523    context: &SubscriptionContext,
2524    subscription_id: &str,
2525    view_spec: &ViewSpec,
2526    current: &[(String, Value)],
2527    next: &[(String, Value)],
2528    envelope: &BusMessage,
2529    metadata: &SourceFrameMetadata,
2530    undelivered: &mut HashSet<String>,
2531) -> Result<()> {
2532    let current_keys: Vec<&str> = current.iter().map(|(key, _)| key.as_str()).collect();
2533    let next_keys: Vec<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2534    let next_set: HashSet<&str> = next_keys.iter().copied().collect();
2535
2536    for key in current_keys
2537        .iter()
2538        .copied()
2539        .filter(|key| !next_set.contains(key))
2540    {
2541        // A key the client was never sent has nothing to remove.
2542        if undelivered.remove(key) {
2543            continue;
2544        }
2545        let op = if metadata.op == "delete" && key == envelope.key {
2546            "delete"
2547        } else {
2548            "remove"
2549        };
2550        send_membership_frame(
2551            context,
2552            subscription_id,
2553            view_spec,
2554            op,
2555            key,
2556            Value::Null,
2557            metadata.seq.clone(),
2558        )?;
2559    }
2560
2561    for (key, data) in next.iter() {
2562        let in_window = current_keys.iter().any(|candidate| *candidate == key);
2563        match member_action(
2564            in_window,
2565            in_window && !undelivered.contains(key),
2566            key == &envelope.key,
2567            view_spec.is_derived(),
2568            &metadata.op,
2569        ) {
2570            MemberAction::Skip => {}
2571            MemberAction::ForwardPatch => send_scoped_source_payload(
2572                context,
2573                subscription_id,
2574                &view_spec.id,
2575                envelope.payload.clone(),
2576            )?,
2577            MemberAction::Upsert => {
2578                let seq = metadata
2579                    .seq
2580                    .clone()
2581                    .or_else(|| data.get("_seq").and_then(Value::as_str).map(str::to_string));
2582                send_membership_frame(
2583                    context,
2584                    subscription_id,
2585                    view_spec,
2586                    "upsert",
2587                    key,
2588                    data.clone(),
2589                    seq,
2590                )?;
2591                undelivered.remove(key);
2592            }
2593        }
2594    }
2595    Ok(())
2596}
2597
2598/// Emit the net effect of every source mutation seen during one coalescing
2599/// interval. Full entities are used for changed members because intermediate
2600/// sparse patches were intentionally discarded and can no longer be merged
2601/// safely by a client.
2602fn emit_coalesced_collection_delta(
2603    context: &SubscriptionContext,
2604    subscription_id: &str,
2605    view_spec: &ViewSpec,
2606    current: &[(String, Value)],
2607    next: &[(String, Value)],
2608    pending: &HashMap<String, Arc<BusMessage>>,
2609    undelivered: &mut HashSet<String>,
2610) -> Result<()> {
2611    let changes = plan_coalesced_collection_delta(current, next, pending, undelivered);
2612    // Every key that left the window or was just upserted is settled: the
2613    // client either never needs it or now holds all of it.
2614    let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2615    undelivered.retain(|key| next_keys.contains(key.as_str()));
2616    for change in changes {
2617        if change.op == "upsert" {
2618            undelivered.remove(&change.key);
2619        }
2620        send_membership_frame(
2621            context,
2622            subscription_id,
2623            view_spec,
2624            change.op,
2625            &change.key,
2626            change.data,
2627            change.seq,
2628        )?;
2629    }
2630    Ok(())
2631}
2632
2633#[derive(Debug, PartialEq)]
2634struct CollectionChange {
2635    op: &'static str,
2636    key: String,
2637    data: Value,
2638    seq: Option<String>,
2639}
2640
2641/// Plan one coalescing interval. `undelivered` holds window members the
2642/// client was never sent; they get no `remove`/`delete` when they leave, and
2643/// like every changed key they are sent whole when they change.
2644fn plan_coalesced_collection_delta(
2645    current: &[(String, Value)],
2646    next: &[(String, Value)],
2647    pending: &HashMap<String, Arc<BusMessage>>,
2648    undelivered: &HashSet<String>,
2649) -> Vec<CollectionChange> {
2650    let current_by_key: HashMap<&str, &Value> = current
2651        .iter()
2652        .map(|(key, data)| (key.as_str(), data))
2653        .collect();
2654    let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2655    let latest_seq = pending
2656        .values()
2657        .filter_map(|envelope| source_frame_metadata(&envelope.payload).seq)
2658        .max_by(|left, right| cmp_seq(left, right));
2659    let mut changes = Vec::new();
2660
2661    for (key, _) in current
2662        .iter()
2663        .filter(|(key, _)| !next_keys.contains(key.as_str()) && !undelivered.contains(key.as_str()))
2664    {
2665        let metadata = pending
2666            .get(key)
2667            .map(|envelope| source_frame_metadata(&envelope.payload));
2668        let op = if metadata
2669            .as_ref()
2670            .is_some_and(|metadata| metadata.op == "delete")
2671        {
2672            "delete"
2673        } else {
2674            "remove"
2675        };
2676        let seq = metadata
2677            .and_then(|metadata| metadata.seq)
2678            .or_else(|| latest_seq.clone());
2679        changes.push(CollectionChange {
2680            op,
2681            key: key.clone(),
2682            data: Value::Null,
2683            seq,
2684        });
2685    }
2686
2687    for (key, data) in next {
2688        let changed = current_by_key
2689            .get(key.as_str())
2690            .is_none_or(|previous| *previous != data);
2691        if !changed && !pending.contains_key(key) {
2692            continue;
2693        }
2694        let seq = data
2695            .get("_seq")
2696            .and_then(Value::as_str)
2697            .map(str::to_string)
2698            .or_else(|| {
2699                pending
2700                    .get(key)
2701                    .and_then(|envelope| source_frame_metadata(&envelope.payload).seq)
2702            })
2703            .or_else(|| latest_seq.clone());
2704        changes.push(CollectionChange {
2705            op: "upsert",
2706            key: key.clone(),
2707            data: data.clone(),
2708            seq,
2709        });
2710    }
2711    changes
2712}
2713
2714/// What one in-window key owes a subscriber after a source mutation.
2715#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2716enum MemberAction {
2717    /// Send nothing: the subscriber already holds this entity and its data did
2718    /// not change.
2719    Skip,
2720    /// Forward the source patch verbatim.
2721    ForwardPatch,
2722    /// Send the whole entity.
2723    Upsert,
2724}
2725
2726/// Decide what to send for one key in the query window.
2727///
2728/// `in_window` says the key was in the window before this mutation; `held`
2729/// says the client also has a copy of it. The two differ after a snapshot
2730/// that `snapshotLimit` truncated or that was disabled: those keys count for
2731/// the window's positions, but the client was never sent them.
2732///
2733/// Only the mutated key's data changed, so a key that was already in the
2734/// window has at most moved index — and index is not the server's to
2735/// communicate. `subscribed` announces the window's sort (see
2736/// [`extract_sort_config`], which is `_seq` descending for a plain list) and
2737/// every SDK re-sorts locally from it, so resending an unchanged entity to
2738/// convey its new position is pure waste. This matters because on a
2739/// `_seq`-ordered list the mutated entity jumps to the front on *every*
2740/// mutation: keying the decision on position change meant rebroadcasting the
2741/// whole window each time, and the mutated entity itself never kept its
2742/// position long enough to have its patch forwarded.
2743fn member_action(
2744    in_window: bool,
2745    held: bool,
2746    is_mutated_key: bool,
2747    is_derived: bool,
2748    op: &str,
2749) -> MemberAction {
2750    if !is_mutated_key {
2751        // A key entering the window has no local state to merge into, so it
2752        // needs the whole entity. One already in the window is unchanged:
2753        // if the client holds it there is nothing new, and if it doesn't the
2754        // snapshot limit deliberately left it out, so it waits for its own
2755        // next change.
2756        return if in_window {
2757            MemberAction::Skip
2758        } else {
2759            MemberAction::Upsert
2760        };
2761    }
2762    // The mutated entity rides its own patch through untouched, but only when
2763    // the subscriber already holds a copy; a patch for a key the client never
2764    // received is not an entity. Derived views still send whole entities: the
2765    // patch on the bus is scoped to the source view, not this one (see
2766    // A4-150).
2767    if held && !is_derived && op != "delete" {
2768        MemberAction::ForwardPatch
2769    } else {
2770        MemberAction::Upsert
2771    }
2772}
2773
2774fn to_wire_snapshot_entities(
2775    entities: Vec<(String, Value)>,
2776    view_spec: &ViewSpec,
2777) -> Vec<SnapshotEntity> {
2778    entities
2779        .into_iter()
2780        .map(|(key, mut data)| {
2781            apply_wire_format(&mut data, &view_spec.wire_format);
2782            SnapshotEntity { key, data }
2783        })
2784        .collect()
2785}
2786
2787async fn load_query_entities(
2788    entity_cache: &EntityCache,
2789    sorted_caches: Option<
2790        Arc<tokio::sync::RwLock<HashMap<String, crate::sorted_cache::SortedViewCache>>>,
2791    >,
2792    view_spec: &ViewSpec,
2793    query: &SubscriptionQuery,
2794    apply_snapshot_limit: bool,
2795) -> Vec<(String, Value)> {
2796    let (entities, preordered) = if let Some(sorted_caches) = sorted_caches {
2797        let mut caches = sorted_caches.write().await;
2798        let entities = caches
2799            .get_mut(&view_spec.id)
2800            .map(|cache| cache.get_all_ordered())
2801            .unwrap_or_default();
2802        (entities, true)
2803    } else if view_spec.mode == Mode::State {
2804        let entity = match query.key.as_deref() {
2805            Some(key) => entity_cache
2806                .get(&view_spec.id, key)
2807                .await
2808                .map(|data| vec![(key.to_string(), data)])
2809                .unwrap_or_default(),
2810            None => vec![],
2811        };
2812        (entity, true)
2813    } else {
2814        (entity_cache.get_all(&view_spec.id).await, false)
2815    };
2816    select_query_entities(entities, query, preordered, apply_snapshot_limit)
2817}
2818
2819fn select_query_entities(
2820    mut entities: Vec<(String, Value)>,
2821    query: &SubscriptionQuery,
2822    preordered: bool,
2823    apply_snapshot_limit: bool,
2824) -> Vec<(String, Value)> {
2825    entities.retain(|(key, data)| query_matches_entity(query, key, data));
2826    if !preordered {
2827        entities.sort_by(|left, right| {
2828            let left_seq = left.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2829            let right_seq = right.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2830            let order = if query.after.is_some() {
2831                cmp_seq(left_seq, right_seq)
2832            } else {
2833                cmp_seq(right_seq, left_seq)
2834            };
2835            order.then_with(|| left.0.cmp(&right.0))
2836        });
2837    }
2838
2839    let skip = query.skip.unwrap_or(0);
2840    let take = query.take.unwrap_or(usize::MAX);
2841    let mut selected: Vec<_> = entities.into_iter().skip(skip).take(take).collect();
2842    if apply_snapshot_limit {
2843        if let Some(limit) = query.snapshot_limit {
2844            selected.truncate(limit);
2845        }
2846    }
2847    selected
2848}
2849
2850fn query_matches_entity(query: &SubscriptionQuery, key: &str, data: &Value) -> bool {
2851    if !query.matches_key(key) {
2852        return false;
2853    }
2854    if let Some(partition) = &query.partition {
2855        if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2856            return false;
2857        }
2858    }
2859    if let Some(after) = &query.after {
2860        let Some(seq) = data.get("_seq").and_then(Value::as_str) else {
2861            return false;
2862        };
2863        if cmp_seq(seq, after) != std::cmp::Ordering::Greater {
2864            return false;
2865        }
2866    }
2867    query
2868        .filters
2869        .iter()
2870        .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2871}
2872
2873fn value_at_dot_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
2874    path.split('.')
2875        .try_fold(value, |current, segment| current.get(segment))
2876}
2877
2878#[cfg(test)]
2879mod tests {
2880    use super::*;
2881    use crate::cache::{EntityCacheConfig, PatchOrigin};
2882    use crate::view::{Delivery, Filters, Projection};
2883    use serde_json::json;
2884    use tokio::sync::oneshot;
2885
2886    #[test]
2887    fn a_catch_up_carries_only_the_fields_that_changed() {
2888        let base = json!({
2889            "id": 7,
2890            "state": {"deployed": [1, 2, 3], "total": 6},
2891            "results": {"square": null},
2892            "gone": true,
2893        });
2894        let current = json!({
2895            "id": 7,
2896            "state": {"deployed": [1, 2, 4], "total": 6},
2897            "results": {"square": 3},
2898            "fresh": "x",
2899        });
2900        assert_eq!(
2901            changed_fields(&base, &current, &Value::Null),
2902            Some(json!({
2903                // An array that changed is sent whole.
2904                "state": {"deployed": [1, 2, 4]},
2905                "results": {"square": 3},
2906                "fresh": "x",
2907                // A field the entity no longer has is cleared.
2908                "gone": null,
2909            }))
2910        );
2911        assert_eq!(changed_fields(&current, &current, &Value::Null), None);
2912        // Anything that is not an object on both sides is compared whole.
2913        assert_eq!(
2914            changed_fields(&json!(null), &current, &Value::Null),
2915            Some(current.clone())
2916        );
2917    }
2918
2919    #[test]
2920    fn a_catch_up_resends_the_fields_forwarded_patches_set() {
2921        let base = json!({"a": 1, "state": {"x": 1, "y": 1}, "list": [1]});
2922        let mut touched = Value::Null;
2923        mark_fields(&mut touched, &json!({"a": 2, "state": {"x": 2}}));
2924        mark_fields(&mut touched, &json!({"list": [2], "extra": true}));
2925        assert_eq!(
2926            touched,
2927            json!({"a": true, "state": {"x": true}, "list": true, "extra": true})
2928        );
2929        // The entity set `a` and `state.x` back and never kept `extra`: the
2930        // client still holds what the forwarded patches set, so a catch-up
2931        // compared with `base` alone would leave it stale.
2932        let current = json!({"a": 1, "state": {"x": 1, "y": 1}, "list": [1, 2]});
2933        assert_eq!(
2934            changed_fields(&base, &current, &Value::Null),
2935            Some(json!({"list": [1, 2]}))
2936        );
2937        assert_eq!(
2938            changed_fields(&base, &current, &touched),
2939            Some(json!({"a": 1, "state": {"x": 1}, "list": [1, 2], "extra": null}))
2940        );
2941        // A patch whose data is not an object marks everything.
2942        mark_fields(&mut touched, &json!(null));
2943        assert_eq!(
2944            changed_fields(&base, &current, &touched),
2945            Some(current.clone())
2946        );
2947    }
2948
2949    fn list_spec() -> ViewSpec {
2950        ViewSpec {
2951            id: "Thing/list".to_string(),
2952            export: "Thing".to_string(),
2953            mode: Mode::List,
2954            wire_format: Default::default(),
2955            projection: Projection::all(),
2956            filters: Filters::all(),
2957            delivery: Delivery::default(),
2958            pipeline: None,
2959            source_view: None,
2960        }
2961    }
2962
2963    /// Plan the whole window the way `emit_collection_delta` does, so a test
2964    /// can assert on what a subscriber is actually sent.
2965    fn plan_window(
2966        current_keys: &[&str],
2967        next_keys: &[&str],
2968        envelope_key: &str,
2969        is_derived: bool,
2970        op: &str,
2971    ) -> Vec<(String, MemberAction)> {
2972        plan_window_with_undelivered(current_keys, &[], next_keys, envelope_key, is_derived, op)
2973    }
2974
2975    /// [`plan_window`] for a subscriber that was never sent `undelivered`.
2976    fn plan_window_with_undelivered(
2977        current_keys: &[&str],
2978        undelivered: &[&str],
2979        next_keys: &[&str],
2980        envelope_key: &str,
2981        is_derived: bool,
2982        op: &str,
2983    ) -> Vec<(String, MemberAction)> {
2984        next_keys
2985            .iter()
2986            .map(|key| {
2987                let in_window = current_keys.contains(key);
2988                let held = in_window && !undelivered.contains(key);
2989                (
2990                    (*key).to_string(),
2991                    member_action(in_window, held, *key == envelope_key, is_derived, op),
2992                )
2993            })
2994            .collect()
2995    }
2996
2997    #[test]
2998    fn a_window_member_the_client_never_received_is_sent_whole() {
2999        // snapshotLimit cut the window [4, 3, 2, 1] to [4, 3]. "1" still
3000        // counts for positions, but a patch for it would reach a client with
3001        // nothing to merge into.
3002        let plan = plan_window_with_undelivered(
3003            &["4", "3", "2", "1"],
3004            &["2", "1"],
3005            &["1", "4", "3", "2"],
3006            "1",
3007            false,
3008            "patch",
3009        );
3010        assert_eq!(
3011            plan,
3012            vec![
3013                ("1".to_string(), MemberAction::Upsert),
3014                ("4".to_string(), MemberAction::Skip),
3015                ("3".to_string(), MemberAction::Skip),
3016                // Unchanged and never sent: it waits for its own change
3017                // rather than undoing the snapshot limit.
3018                ("2".to_string(), MemberAction::Skip),
3019            ]
3020        );
3021    }
3022
3023    #[test]
3024    fn a_held_member_still_rides_its_patch_after_a_truncated_snapshot() {
3025        let plan = plan_window_with_undelivered(
3026            &["4", "3", "2", "1"],
3027            &["2", "1"],
3028            &["4", "3", "2", "1"],
3029            "4",
3030            false,
3031            "patch",
3032        );
3033        assert_eq!(plan[0], ("4".to_string(), MemberAction::ForwardPatch));
3034    }
3035
3036    #[test]
3037    fn coalescing_sends_no_remove_for_a_key_the_client_never_received() {
3038        let current = vec![
3039            ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3040            (
3041                "unsent".to_string(),
3042                json!({"count": 1, "_seq": "10:000001"}),
3043            ),
3044            ("gone".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3045        ];
3046        let next = vec![
3047            ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3048            (
3049                "unsent".to_string(),
3050                json!({"count": 2, "_seq": "10:000003"}),
3051            ),
3052        ];
3053        let pending = HashMap::from([
3054            (
3055                "unsent".to_string(),
3056                list_message("unsent", "patch", "10:000003"),
3057            ),
3058            (
3059                "gone".to_string(),
3060                list_message("gone", "delete", "10:000002"),
3061            ),
3062        ]);
3063        let undelivered = HashSet::from(["unsent".to_string(), "gone".to_string()]);
3064
3065        assert_eq!(
3066            plan_coalesced_collection_delta(&current, &next, &pending, &undelivered),
3067            vec![CollectionChange {
3068                op: "upsert",
3069                key: "unsent".to_string(),
3070                data: json!({"count": 2, "_seq": "10:000003"}),
3071                seq: Some("10:000003".to_string()),
3072            }]
3073        );
3074    }
3075
3076    #[test]
3077    fn a_reordered_window_sends_only_the_mutated_entity() {
3078        // A `_seq`-descending list: mutating "1" moves it to the front and
3079        // shifts every other key down one. Only "1" changed, so only "1" is
3080        // sent — and it rides its own patch, not a full entity.
3081        let current = ["4", "3", "2", "1"];
3082        let next = ["1", "4", "3", "2"];
3083        let plan = plan_window(&current, &next, "1", false, "patch");
3084
3085        assert_eq!(
3086            plan,
3087            vec![
3088                ("1".to_string(), MemberAction::ForwardPatch),
3089                ("4".to_string(), MemberAction::Skip),
3090                ("3".to_string(), MemberAction::Skip),
3091                ("2".to_string(), MemberAction::Skip),
3092            ]
3093        );
3094    }
3095
3096    #[test]
3097    fn a_key_entering_the_window_gets_the_whole_entity() {
3098        // "5" has no local state for the subscriber to merge a patch into.
3099        let plan = plan_window(&["4", "3"], &["5", "4", "3"], "5", false, "patch");
3100        assert_eq!(
3101            plan,
3102            vec![
3103                ("5".to_string(), MemberAction::Upsert),
3104                ("4".to_string(), MemberAction::Skip),
3105                ("3".to_string(), MemberAction::Skip),
3106            ]
3107        );
3108    }
3109
3110    #[test]
3111    fn derived_views_still_send_whole_entities() {
3112        // The patch on the bus is scoped to the source view, so a derived
3113        // subscription cannot forward it verbatim (A4-150).
3114        let plan = plan_window(&["1", "2"], &["1", "2"], "1", true, "patch");
3115        assert_eq!(
3116            plan,
3117            vec![
3118                ("1".to_string(), MemberAction::Upsert),
3119                ("2".to_string(), MemberAction::Skip),
3120            ]
3121        );
3122    }
3123
3124    #[test]
3125    fn a_delete_envelope_never_forwards_a_patch() {
3126        // A surviving key on a delete envelope carries no mergeable patch.
3127        let plan = plan_window(&["1", "2"], &["1", "2"], "1", false, "delete");
3128        assert_eq!(
3129            plan,
3130            vec![
3131                ("1".to_string(), MemberAction::Upsert),
3132                ("2".to_string(), MemberAction::Skip),
3133            ]
3134        );
3135    }
3136
3137    #[test]
3138    fn an_unchanged_window_sends_one_frame_not_a_broadcast() {
3139        // The regression this guards: 500 members used to mean 500 full
3140        // entities on the wire for a single mutation.
3141        let keys: Vec<String> = (0..500).map(|index| index.to_string()).collect();
3142        let refs: Vec<&str> = keys.iter().map(String::as_str).collect();
3143        let plan = plan_window(&refs, &refs, "250", false, "patch");
3144
3145        let sent = plan
3146            .iter()
3147            .filter(|(_, action)| *action != MemberAction::Skip)
3148            .count();
3149        assert_eq!(sent, 1);
3150        assert_eq!(plan[250].1, MemberAction::ForwardPatch);
3151    }
3152
3153    fn list_message(key: &str, op: &str, seq: &str) -> Arc<BusMessage> {
3154        Arc::new(BusMessage {
3155            key: key.to_string(),
3156            entity: "Thing/list".to_string(),
3157            payload: Arc::new(Bytes::from(
3158                serde_json::to_vec(&json!({
3159                    "entity": "Thing/list",
3160                    "op": op,
3161                    "key": key,
3162                    "seq": seq,
3163                    "data": {},
3164                }))
3165                .unwrap(),
3166            )),
3167        })
3168    }
3169
3170    #[test]
3171    fn coalescing_emits_only_the_final_full_state_per_changed_key() {
3172        let current = vec![
3173            ("a".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3174            ("b".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3175            (
3176                "stable".to_string(),
3177                json!({"count": 1, "_seq": "10:000001"}),
3178            ),
3179        ];
3180        let next = vec![
3181            ("a".to_string(), json!({"count": 3, "_seq": "10:000004"})),
3182            ("c".to_string(), json!({"count": 1, "_seq": "10:000003"})),
3183            (
3184                "stable".to_string(),
3185                json!({"count": 1, "_seq": "10:000001"}),
3186            ),
3187        ];
3188        let pending = HashMap::from([
3189            ("a".to_string(), list_message("a", "patch", "10:000004")),
3190            ("b".to_string(), list_message("b", "delete", "10:000002")),
3191            ("c".to_string(), list_message("c", "patch", "10:000003")),
3192        ]);
3193
3194        assert_eq!(
3195            plan_coalesced_collection_delta(&current, &next, &pending, &HashSet::new()),
3196            vec![
3197                CollectionChange {
3198                    op: "delete",
3199                    key: "b".to_string(),
3200                    data: Value::Null,
3201                    seq: Some("10:000002".to_string()),
3202                },
3203                CollectionChange {
3204                    op: "upsert",
3205                    key: "a".to_string(),
3206                    data: json!({"count": 3, "_seq": "10:000004"}),
3207                    seq: Some("10:000004".to_string()),
3208                },
3209                CollectionChange {
3210                    op: "upsert",
3211                    key: "c".to_string(),
3212                    data: json!({"count": 1, "_seq": "10:000003"}),
3213                    seq: Some("10:000003".to_string()),
3214                },
3215            ]
3216        );
3217    }
3218
3219    #[test]
3220    fn a_delayed_delete_cannot_erase_a_newer_recreated_entity() {
3221        let recreated = json!({"balance": 2, "_seq": "10:000004"});
3222        assert!(source_delete_is_stale(Some(&recreated), Some("10:000002")));
3223        assert!(!source_delete_is_stale(Some(&recreated), Some("10:000004")));
3224        assert!(!source_delete_is_stale(Some(&recreated), None));
3225    }
3226
3227    #[test]
3228    fn a_lagged_snapshotless_subscription_gets_a_fatal_retryable_error() {
3229        let issue = SocketIssueMessage::subscription_lagged("balances".to_string(), 42);
3230        assert_eq!(issue.code, "subscription-lagged");
3231        assert!(issue.retryable);
3232        assert!(issue.fatal);
3233        assert!(issue.message.contains("42"));
3234        assert!(issue
3235            .suggested_action
3236            .unwrap()
3237            .contains("snapshots enabled"));
3238
3239        let append = SocketIssueMessage::append_subscription_lagged("trades".to_string(), 9);
3240        assert!(append.fatal);
3241        assert!(append.suggested_action.unwrap().contains("retained replay"));
3242    }
3243
3244    #[test]
3245    fn snapshot_batches_share_identity_and_completion() {
3246        let entities = ["one", "two", "three"].map(|key| SnapshotEntity {
3247            key: key.to_string(),
3248            data: json!({"key": key}),
3249        });
3250        let batches = create_snapshot_batches(
3251            &entities,
3252            SnapshotMetadata {
3253                subscription_id: "sub-1",
3254                snapshot_id: "snapshot-1",
3255                authoritative: true,
3256                mode: Mode::List,
3257                view_id: "Thing/list",
3258                key: None,
3259            },
3260            &SnapshotBatchConfig {
3261                initial_batch_size: 2,
3262                subsequent_batch_size: 1,
3263            },
3264        );
3265        assert_eq!(batches.len(), 2);
3266        assert!(batches.iter().all(|batch| batch.subscription_id == "sub-1"));
3267        assert!(batches
3268            .iter()
3269            .all(|batch| batch.snapshot_id == "snapshot-1"));
3270        assert!(!batches[0].complete);
3271        assert!(batches[1].complete);
3272        assert!(batches.iter().all(|batch| batch.authoritative));
3273    }
3274
3275    #[test]
3276    fn empty_incremental_snapshot_is_explicitly_non_authoritative() {
3277        let batches = create_snapshot_batches(
3278            &[],
3279            SnapshotMetadata {
3280                subscription_id: "sub-1",
3281                snapshot_id: "snapshot-1",
3282                authoritative: false,
3283                mode: Mode::State,
3284                view_id: "Thing/state",
3285                key: Some("missing"),
3286            },
3287            &SnapshotBatchConfig {
3288                initial_batch_size: 1,
3289                subsequent_batch_size: 1,
3290            },
3291        );
3292        assert_eq!(batches.len(), 1);
3293        assert!(!batches[0].authoritative);
3294        assert!(batches[0].complete);
3295        assert_eq!(batches[0].key.as_deref(), Some("missing"));
3296    }
3297
3298    #[test]
3299    fn lag_recovery_is_authoritative_for_an_after_query() {
3300        let subscription = Subscription {
3301            protocol_version: PROTOCOL_VERSION,
3302            subscription_id: "sub-1".to_string(),
3303            query: SubscriptionQuery {
3304                view: "Thing/list".to_string(),
3305                after: Some("40:000000000010".to_string()),
3306                ..Default::default()
3307            },
3308            snapshot: Default::default(),
3309        };
3310
3311        assert!(!SnapshotPurpose::Initial.authoritative(&subscription));
3312        assert!(SnapshotPurpose::Recovery.authoritative(&subscription));
3313    }
3314
3315    #[test]
3316    fn dot_path_filters_are_exact_and_type_sensitive() {
3317        let mut query = SubscriptionQuery {
3318            view: "Thing/list".to_string(),
3319            ..Default::default()
3320        };
3321        query
3322            .filters
3323            .insert("state.status".to_string(), json!("open"));
3324        query.filters.insert("metrics.count".to_string(), json!(2));
3325        assert!(query_matches_entity(
3326            &query,
3327            "one",
3328            &json!({"state": {"status": "open"}, "metrics": {"count": 2}}),
3329        ));
3330        assert!(!query_matches_entity(
3331            &query,
3332            "one",
3333            &json!({"state": {"status": "open"}, "metrics": {"count": "2"}}),
3334        ));
3335    }
3336
3337    #[test]
3338    fn take_and_skip_define_independent_deterministic_windows() {
3339        let entities: Vec<_> = (1..=6)
3340            .map(|id| {
3341                (
3342                    id.to_string(),
3343                    json!({"id": id, "_seq": format!("10:{id:012}")}),
3344                )
3345            })
3346            .collect();
3347        let first = SubscriptionQuery {
3348            view: "Thing/list".to_string(),
3349            take: Some(2),
3350            skip: Some(0),
3351            ..Default::default()
3352        };
3353        let second = SubscriptionQuery {
3354            skip: Some(2),
3355            ..first.clone()
3356        };
3357        let first_keys: Vec<_> = select_query_entities(entities.clone(), &first, false, false)
3358            .into_iter()
3359            .map(|(key, _)| key)
3360            .collect();
3361        let second_keys: Vec<_> = select_query_entities(entities, &second, false, false)
3362            .into_iter()
3363            .map(|(key, _)| key)
3364            .collect();
3365        assert_eq!(first_keys, ["6", "5"]);
3366        assert_eq!(second_keys, ["4", "3"]);
3367    }
3368
3369    #[tokio::test]
3370    async fn state_receiver_is_installed_before_snapshot_awaits() {
3371        let bus = BusManager::new();
3372        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3373        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3374        let bus_for_task = bus.clone();
3375        let task = tokio::spawn(async move {
3376            subscribe_state_then_snapshot(&bus_for_task, "Thing/state", "one", || async move {
3377                snapshot_started_tx.send(()).unwrap();
3378                release_snapshot_rx.await.unwrap();
3379            })
3380            .await
3381            .0
3382        });
3383        snapshot_started_rx.await.unwrap();
3384        bus.publish_state(
3385            "Thing/state",
3386            "one",
3387            Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3388        )
3389        .await;
3390        release_snapshot_tx.send(()).unwrap();
3391        let mut receiver = task.await.unwrap();
3392        receiver.changed().await.unwrap();
3393        assert!(!receiver.borrow().payload.is_empty());
3394    }
3395
3396    async fn assert_list_receiver_precedes_snapshot(view: &'static str) {
3397        let bus = BusManager::new();
3398        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3399        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3400        let bus_for_task = bus.clone();
3401        let task = tokio::spawn(async move {
3402            subscribe_list_then_snapshot(&bus_for_task, view, || async move {
3403                snapshot_started_tx.send(()).unwrap();
3404                release_snapshot_rx.await.unwrap();
3405            })
3406            .await
3407            .0
3408        });
3409        snapshot_started_rx.await.unwrap();
3410        bus.publish_list(
3411            view,
3412            Arc::new(BusMessage {
3413                key: "one".to_string(),
3414                entity: view.to_string(),
3415                payload: Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3416            }),
3417        )
3418        .await;
3419        release_snapshot_tx.send(()).unwrap();
3420        let mut receiver = task.await.unwrap();
3421        assert_eq!(receiver.recv().await.unwrap().key, "one");
3422    }
3423
3424    #[tokio::test]
3425    async fn list_receiver_is_installed_before_snapshot() {
3426        assert_list_receiver_precedes_snapshot("Thing/list").await;
3427    }
3428
3429    #[tokio::test]
3430    async fn append_receiver_is_installed_before_snapshot() {
3431        assert_list_receiver_precedes_snapshot("Thing/append").await;
3432    }
3433
3434    #[tokio::test]
3435    async fn derived_source_receiver_is_installed_before_snapshot() {
3436        assert_list_receiver_precedes_snapshot("Thing/list-source").await;
3437    }
3438
3439    #[tokio::test]
3440    async fn snapshot_limit_does_not_change_live_take_skip_membership() {
3441        let cache = EntityCache::with_config(EntityCacheConfig {
3442            max_entities_per_view: 10,
3443            ..Default::default()
3444        });
3445        for id in 1..=4 {
3446            cache
3447                .upsert(
3448                    "Thing/list",
3449                    &id.to_string(),
3450                    json!({"_seq": format!("10:{id:012}")}),
3451                )
3452                .await;
3453        }
3454        let query = SubscriptionQuery {
3455            view: "Thing/list".to_string(),
3456            take: Some(3),
3457            skip: Some(1),
3458            snapshot_limit: Some(1),
3459            ..Default::default()
3460        };
3461        let live = load_query_entities(&cache, None, &list_spec(), &query, false).await;
3462        let snapshot = load_query_entities(&cache, None, &list_spec(), &query, true).await;
3463        assert_eq!(live.len(), 3);
3464        assert_eq!(snapshot.len(), 1);
3465    }
3466
3467    #[test]
3468    fn fixture_manifest_covers_required_conformance_cases() {
3469        let manifest: Value = serde_json::from_str(include_str!(
3470            "../../../../tests/fixtures/websocket-v2/manifest.json"
3471        ))
3472        .unwrap();
3473        let names: HashSet<_> = manifest["fixtures"]
3474            .as_array()
3475            .unwrap()
3476            .iter()
3477            .filter_map(Value::as_str)
3478            .collect();
3479        for required in [
3480            "keyed-state.json",
3481            "list-windows.json",
3482            "filters.json",
3483            "multi-batch-authoritative.json",
3484            "empty-snapshot.json",
3485            "remove.json",
3486            "delete.json",
3487            "incremental-snapshot.json",
3488            "reconnect-replacement.json",
3489            "errors.json",
3490            "whole-entities.json",
3491        ] {
3492            assert!(names.contains(required), "missing fixture {required}");
3493        }
3494
3495        for document in [
3496            include_str!("../../../../tests/fixtures/websocket-v2/keyed-state.json"),
3497            include_str!("../../../../tests/fixtures/websocket-v2/list-windows.json"),
3498            include_str!("../../../../tests/fixtures/websocket-v2/filters.json"),
3499            include_str!("../../../../tests/fixtures/websocket-v2/multi-batch-authoritative.json"),
3500            include_str!("../../../../tests/fixtures/websocket-v2/empty-snapshot.json"),
3501            include_str!("../../../../tests/fixtures/websocket-v2/remove.json"),
3502            include_str!("../../../../tests/fixtures/websocket-v2/delete.json"),
3503            include_str!("../../../../tests/fixtures/websocket-v2/incremental-snapshot.json"),
3504            include_str!("../../../../tests/fixtures/websocket-v2/reconnect-replacement.json"),
3505            include_str!("../../../../tests/fixtures/websocket-v2/errors.json"),
3506            include_str!("../../../../tests/fixtures/websocket-v2/whole-entities.json"),
3507        ] {
3508            let fixture: Value = serde_json::from_str(document).unwrap();
3509            assert!(fixture["name"].is_string());
3510        }
3511    }
3512
3513    fn append_frame(key: &str, data: Value) -> Arc<Bytes> {
3514        let frame = json!({
3515            "entity": "Trade/append",
3516            "op": "patch",
3517            "key": key,
3518            "offset": 7,
3519            "data": data,
3520        });
3521        Arc::new(Bytes::from(serde_json::to_vec(&frame).unwrap()))
3522    }
3523
3524    /// A replayable subscription must honour the same predicates a live
3525    /// collection subscription does; otherwise a filtered consumer receives
3526    /// events outside the query it asked for.
3527    #[test]
3528    fn replay_delivery_applies_key_partition_and_filters() {
3529        let matching = append_frame("pool1", json!({"_partition": "us", "side": "buy"}));
3530        let other_partition = append_frame("pool1", json!({"_partition": "eu", "side": "buy"}));
3531        let other_side = append_frame("pool1", json!({"_partition": "us", "side": "sell"}));
3532
3533        let unfiltered = SubscriptionQuery {
3534            view: "Trade/append".to_string(),
3535            ..Default::default()
3536        };
3537        assert!(live_frame_matches(&unfiltered, "pool1", &matching));
3538        assert!(live_frame_matches(&unfiltered, "pool9", &matching));
3539
3540        let keyed = SubscriptionQuery {
3541            view: "Trade/append".to_string(),
3542            key: Some("pool1".to_string()),
3543            ..Default::default()
3544        };
3545        assert!(live_frame_matches(&keyed, "pool1", &matching));
3546        assert!(!live_frame_matches(&keyed, "pool2", &matching));
3547
3548        let partitioned = SubscriptionQuery {
3549            view: "Trade/append".to_string(),
3550            partition: Some("us".to_string()),
3551            ..Default::default()
3552        };
3553        assert!(live_frame_matches(&partitioned, "pool1", &matching));
3554        assert!(!live_frame_matches(&partitioned, "pool1", &other_partition));
3555
3556        let filtered = SubscriptionQuery {
3557            view: "Trade/append".to_string(),
3558            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3559            ..Default::default()
3560        };
3561        assert!(live_frame_matches(&filtered, "pool1", &matching));
3562        assert!(!live_frame_matches(&filtered, "pool1", &other_side));
3563    }
3564
3565    #[test]
3566    fn a_frame_without_decodable_data_does_not_satisfy_a_filter() {
3567        let filtered = SubscriptionQuery {
3568            view: "Trade/append".to_string(),
3569            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3570            ..Default::default()
3571        };
3572        let garbage = Arc::new(Bytes::from_static(b"not json"));
3573        assert!(!live_frame_matches(&filtered, "pool1", &garbage));
3574    }
3575
3576    /// The recovery cursor must be the last offset delivered *before* the
3577    /// gap. Reporting the newest offset seen would step the consumer over
3578    /// the skipped records permanently.
3579    #[test]
3580    fn replay_lagged_recovers_from_before_the_gap() {
3581        let epoch = crate::journal::JournalEpoch::new();
3582        let issue = SocketIssueMessage::replay_lagged(
3583            Some("trades".to_string()),
3584            37,
3585            Some(crate::journal::Cursor {
3586                epoch: epoch.clone(),
3587                offset: 4180,
3588            }),
3589        );
3590        assert_eq!(issue.code, "replay-lagged");
3591        assert_eq!(issue.recover_from, Some(format!("{epoch}:4180")));
3592        assert!(
3593            issue.suggested_action.unwrap().contains("4180"),
3594            "the consumer is told exactly which cursor recovers the gap"
3595        );
3596
3597        // Nothing delivered yet: there is no pre-gap offset, so the whole
3598        // retained window is the recovery.
3599        let from_scratch = SocketIssueMessage::replay_lagged(Some("trades".to_string()), 9, None);
3600        assert_eq!(from_scratch.recover_from, None);
3601        assert!(from_scratch
3602            .suggested_action
3603            .unwrap()
3604            .contains("without `after`"));
3605    }
3606
3607    fn bus_message(key: &str) -> Arc<BusMessage> {
3608        Arc::new(BusMessage {
3609            key: key.to_string(),
3610            entity: "Trade/append".to_string(),
3611            payload: Arc::new(Bytes::from_static(b"{}")),
3612        })
3613    }
3614
3615    /// A long replay must keep the bus drained. The broadcast buffer is
3616    /// bounded, so a busy view would otherwise lap the replay and the first
3617    /// live `recv` would return `Lagged` — telling the client to resubscribe,
3618    /// starting another long replay, which laps again.
3619    #[tokio::test]
3620    async fn draining_during_a_replay_keeps_a_busy_view_from_lapping_it() {
3621        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(16);
3622        let mut pending = VecDeque::new();
3623        let mut lagged = None;
3624
3625        // Publish more than the channel holds, draining as a replay would
3626        // between sends.
3627        for index in 0..48 {
3628            sender.send(bus_message(&format!("k{index}"))).unwrap();
3629            drain_available(&mut receiver, &mut pending, &mut lagged);
3630        }
3631
3632        assert_eq!(lagged, None, "draining as we go means nothing is dropped");
3633        assert_eq!(pending.len(), 48, "every published frame is buffered");
3634    }
3635
3636    #[tokio::test]
3637    async fn a_replay_that_never_drains_is_reported_as_a_gap() {
3638        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3639        for index in 0..32 {
3640            sender.send(bus_message(&format!("k{index}"))).unwrap();
3641        }
3642
3643        let mut pending = VecDeque::new();
3644        let mut lagged = None;
3645        drain_available(&mut receiver, &mut pending, &mut lagged);
3646
3647        assert!(
3648            lagged.is_some(),
3649            "overflowing the bus is a gap, not silent truncation"
3650        );
3651    }
3652
3653    /// Everything still on the bus after a gap is on the far side of it.
3654    /// Buffering it would put those frames in front of the lag report and
3655    /// advance the recovery cursor past the records it is meant to recover.
3656    #[tokio::test]
3657    async fn nothing_after_a_gap_is_buffered_ahead_of_the_report() {
3658        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3659        let mut pending = VecDeque::new();
3660        let mut lagged = None;
3661
3662        // Delivered and buffered normally.
3663        sender.send(bus_message("before")).unwrap();
3664        drain_available(&mut receiver, &mut pending, &mut lagged);
3665        assert_eq!(pending.len(), 1);
3666
3667        // Overflow the bus: everything published from here is past the gap.
3668        for index in 0..32 {
3669            sender.send(bus_message(&format!("lost{index}"))).unwrap();
3670        }
3671        drain_available(&mut receiver, &mut pending, &mut lagged);
3672        assert!(lagged.is_some());
3673
3674        // The bus still holds what survived the overflow, all of it past the
3675        // gap. A later iteration of the replay loop drains again.
3676        let buffered_at_gap = pending.len();
3677        drain_available(&mut receiver, &mut pending, &mut lagged);
3678        assert_eq!(
3679            pending.len(),
3680            buffered_at_gap,
3681            "post-gap frames must not join the pre-gap flush"
3682        );
3683        assert_eq!(pending.front().unwrap().key, "before");
3684    }
3685    /// The bus is subscribed before the tape is read, so a record published
3686    /// in that window arrives on both paths.
3687    #[test]
3688    fn the_seam_between_replay_and_live_neither_repeats_nor_skips() {
3689        let mut last_sent = Some(4211);
3690
3691        assert!(
3692            already_delivered(Some(4211), &mut last_sent),
3693            "the record the replay ended on must not be sent twice"
3694        );
3695        assert!(already_delivered(Some(4100), &mut last_sent));
3696        assert_eq!(last_sent, Some(4211), "a duplicate never moves the mark");
3697
3698        assert!(
3699            !already_delivered(Some(4212), &mut last_sent),
3700            "the next record is new"
3701        );
3702        assert_eq!(last_sent, Some(4212));
3703
3704        // A frame with no offset comes from a view with no tape; it cannot
3705        // have been replayed, and must not disturb the mark.
3706        assert!(!already_delivered(None, &mut last_sent));
3707        assert_eq!(last_sent, Some(4212));
3708    }
3709
3710    /// A subscription with no cursor has delivered nothing, so the first live
3711    /// frame is not a duplicate.
3712    #[test]
3713    fn a_fresh_subscription_delivers_its_first_live_frame() {
3714        let mut last_sent = None;
3715        assert!(!already_delivered(Some(0), &mut last_sent));
3716        assert_eq!(last_sent, Some(0));
3717    }
3718
3719    /// End-to-end over a real socket: the pieces above are unit-tested
3720    /// individually, but the thing a consumer actually does — reconnect with
3721    /// a stored cursor and keep reading — only exists once a subscription is
3722    /// attached to a connection.
3723    mod over_a_socket {
3724        use super::*;
3725        use crate::journal::{EventJournal, JournalConfig};
3726        use crate::projector::Projector;
3727        use crate::{MutationBatch, SlotContext};
3728        use arete_interpreter::Mutation;
3729        use futures_util::{SinkExt, StreamExt};
3730        use std::time::Duration;
3731        use tokio::net::{TcpListener, TcpStream};
3732        use tokio::sync::mpsc;
3733        use tokio_tungstenite::tungstenite::Message;
3734        use tokio_tungstenite::{client_async, WebSocketStream};
3735
3736        const RETAINED: u64 = 600;
3737
3738        fn append_index() -> ViewIndex {
3739            let mut index = ViewIndex::new();
3740            index.add_spec(ViewSpec {
3741                id: "Trade/append".to_string(),
3742                export: "Trade".to_string(),
3743                mode: Mode::Append,
3744                wire_format: Default::default(),
3745                projection: Projection::all(),
3746                filters: Filters::all(),
3747                delivery: Delivery::default(),
3748                pipeline: None,
3749                source_view: None,
3750            });
3751            index
3752        }
3753
3754        fn trade(index: u64) -> MutationBatch {
3755            MutationBatch::with_slot_context(
3756                vec![Mutation {
3757                    export: "Trade".to_string(),
3758                    key: json!(format!("pool{}", index % 4)),
3759                    patch: json!({"trade": index}),
3760                    append: vec![],
3761                }]
3762                .into_iter()
3763                .collect(),
3764                SlotContext::new(100 + index / 3, index % 3),
3765            )
3766        }
3767
3768        struct Harness {
3769            addr: SocketAddr,
3770            journal: Arc<EventJournal>,
3771            tx: mpsc::Sender<MutationBatch>,
3772        }
3773
3774        impl Harness {
3775            async fn start() -> Self {
3776                let view_index = Arc::new(append_index());
3777                let entity_cache = EntityCache::new();
3778                let bus_manager = BusManager::new();
3779                let journal = Arc::new(EventJournal::new(JournalConfig {
3780                    enabled: true,
3781                    max_bytes_per_view: u64::MAX,
3782                    max_records_per_view: 10_000,
3783                    max_age: Duration::from_secs(3_600),
3784                }));
3785
3786                let (tx, rx) = mpsc::channel::<MutationBatch>(256);
3787                tokio::spawn(
3788                    Projector::new(
3789                        view_index.clone(),
3790                        bus_manager.clone(),
3791                        entity_cache.clone(),
3792                        rx,
3793                        #[cfg(feature = "otel")]
3794                        None,
3795                    )
3796                    .with_journal(journal.clone())
3797                    .run(),
3798                );
3799
3800                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3801                let addr = listener.local_addr().unwrap();
3802                let server = WebSocketServer::new(
3803                    addr,
3804                    bus_manager,
3805                    entity_cache,
3806                    view_index,
3807                    #[cfg(feature = "otel")]
3808                    None,
3809                )
3810                .with_journal(journal.clone());
3811                let (acceptor, _cleanup) = server.into_acceptor();
3812                tokio::spawn(async move { acceptor.serve_listener(listener).await });
3813
3814                Self { addr, journal, tx }
3815            }
3816
3817            async fn publish(&self, range: std::ops::Range<u64>) {
3818                for index in range {
3819                    self.tx.send(trade(index)).await.unwrap();
3820                }
3821                let (ack, wait) = oneshot::channel();
3822                self.tx
3823                    .send(MutationBatch::flush_marker(ack))
3824                    .await
3825                    .unwrap();
3826                wait.await.unwrap();
3827            }
3828
3829            async fn connect(&self) -> WebSocketStream<TcpStream> {
3830                let stream = TcpStream::connect(self.addr).await.unwrap();
3831                client_async(format!("ws://{}/", self.addr), stream)
3832                    .await
3833                    .unwrap()
3834                    .0
3835            }
3836        }
3837
3838        async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
3839            loop {
3840                let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
3841                    .await
3842                    .expect("the server answers within the timeout")
3843                    .expect("the stream stays open")
3844                    .expect("a readable frame");
3845                // Control and data frames arrive as binary; issue frames as
3846                // text. Both are JSON.
3847                let bytes = match &message {
3848                    Message::Text(text) => text.as_bytes(),
3849                    Message::Binary(bytes) => bytes.as_ref(),
3850                    _ => continue,
3851                };
3852                return serde_json::from_slice(bytes).expect("frames are JSON");
3853            }
3854        }
3855
3856        /// Collect `count` event frames, ignoring anything else on the wire.
3857        async fn collect_trades(socket: &mut WebSocketStream<TcpStream>, count: usize) -> Vec<u64> {
3858            let mut offsets = Vec::with_capacity(count);
3859            while offsets.len() < count {
3860                let frame = next_frame(socket).await;
3861                assert_ne!(
3862                    frame["type"], "error",
3863                    "no error frame should interrupt delivery: {frame}"
3864                );
3865                if let Some(offset) = frame["offset"].as_u64() {
3866                    offsets.push(offset);
3867                }
3868            }
3869            offsets
3870        }
3871
3872        /// The headline claim: reconnecting with a stored cursor delivers every
3873        /// event published since it, in order, and then continues live without
3874        /// a duplicate or a hole at the seam.
3875        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3876        async fn a_reconnect_replays_from_a_cursor_and_continues_live() {
3877            let harness = Harness::start().await;
3878            harness.publish(0..RETAINED).await;
3879
3880            let cursor = harness.journal.window("Trade/append").await;
3881            let stored = format!("{}:{}", cursor.epoch, 99);
3882
3883            let mut socket = harness.connect().await;
3884            socket
3885                .send(Message::Text(
3886                    json!({
3887                        "type": "subscribe",
3888                        "protocolVersion": 2,
3889                        "subscriptionId": "trades",
3890                        "query": {"view": "Trade/append", "after": stored},
3891                    })
3892                    .to_string()
3893                    .into(),
3894                ))
3895                .await
3896                .unwrap();
3897
3898            let ack = next_frame(&mut socket).await;
3899            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3900            assert_eq!(ack["replayWindow"]["next"], json!(RETAINED));
3901
3902            // Well over the 500 a single read or buffer would cover.
3903            let replayed = collect_trades(&mut socket, (RETAINED - 100) as usize).await;
3904            assert_eq!(
3905                replayed,
3906                (100..RETAINED).collect::<Vec<_>>(),
3907                "every event after the cursor, in order, exactly once"
3908            );
3909
3910            // Published only now, so these can only arrive over the live path.
3911            harness.publish(RETAINED..RETAINED + 40).await;
3912            let live = collect_trades(&mut socket, 40).await;
3913            assert_eq!(
3914                live,
3915                (RETAINED..RETAINED + 40).collect::<Vec<_>>(),
3916                "the live stream resumes exactly where the replay stopped"
3917            );
3918
3919            socket.close(None).await.ok();
3920        }
3921
3922        /// Events published *during* the replay must still arrive. The replay
3923        /// and the live subscription are separate reads of the same tape, and
3924        /// the seam between them is where a naive implementation drops or
3925        /// repeats.
3926        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3927        async fn events_published_during_a_replay_are_not_lost() {
3928            let harness = Harness::start().await;
3929            harness.publish(0..RETAINED).await;
3930
3931            let epoch = harness.journal.window("Trade/append").await.epoch;
3932            let mut socket = harness.connect().await;
3933            socket
3934                .send(Message::Text(
3935                    json!({
3936                        "type": "subscribe",
3937                        "protocolVersion": 2,
3938                        "subscriptionId": "trades",
3939                        "query": {"view": "Trade/append", "after": format!("{epoch}:0")},
3940                    })
3941                    .to_string()
3942                    .into(),
3943                ))
3944                .await
3945                .unwrap();
3946            let ack = next_frame(&mut socket).await;
3947            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3948
3949            // Keep publishing while the replay is still draining.
3950            harness.publish(RETAINED..RETAINED + 200).await;
3951
3952            let total = (RETAINED + 200 - 1) as usize;
3953            let delivered = collect_trades(&mut socket, total).await;
3954            assert_eq!(
3955                delivered,
3956                (1..RETAINED + 200).collect::<Vec<_>>(),
3957                "replay and live output join without a gap or a repeat"
3958            );
3959
3960            socket.close(None).await.ok();
3961        }
3962
3963        /// A cursor from another tape lifetime is refused rather than served
3964        /// as a continuation, and the refusal releases the subscription id.
3965        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3966        async fn a_stale_epoch_is_refused_and_frees_the_subscription_id() {
3967            let harness = Harness::start().await;
3968            harness.publish(0..50).await;
3969
3970            let mut socket = harness.connect().await;
3971            socket
3972                .send(Message::Text(
3973                    json!({
3974                        "type": "subscribe",
3975                        "protocolVersion": 2,
3976                        "subscriptionId": "trades",
3977                        "query": {
3978                            "view": "Trade/append",
3979                            "after": format!("{}:10", crate::journal::JournalEpoch::new()),
3980                        },
3981                    })
3982                    .to_string()
3983                    .into(),
3984                ))
3985                .await
3986                .unwrap();
3987
3988            let error = next_frame(&mut socket).await;
3989            assert_eq!(error["type"], "error", "unexpected frame: {error}");
3990            assert_eq!(error["code"], "cursor-epoch-changed");
3991
3992            // The documented recovery is to resubscribe without a cursor. That
3993            // only works if the refused attempt released the id.
3994            socket
3995                .send(Message::Text(
3996                    json!({
3997                        "type": "subscribe",
3998                        "protocolVersion": 2,
3999                        "subscriptionId": "trades",
4000                        "query": {"view": "Trade/append"},
4001                    })
4002                    .to_string()
4003                    .into(),
4004                ))
4005                .await
4006                .unwrap();
4007            let ack = next_frame(&mut socket).await;
4008            assert_eq!(
4009                ack["op"], "subscribed",
4010                "the refused id must be reusable: {ack}"
4011            );
4012
4013            assert_eq!(collect_trades(&mut socket, 50).await.len(), 50);
4014            socket.close(None).await.ok();
4015        }
4016
4017        const ROUND: &str = "Round/state";
4018
4019        /// A state view with no projector: the test writes the cache and the
4020        /// bus itself, so it controls exactly when the subscriber can run.
4021        #[derive(Clone)]
4022        struct StateServer {
4023            addr: SocketAddr,
4024            bus_manager: BusManager,
4025            entity_cache: EntityCache,
4026        }
4027
4028        impl StateServer {
4029            async fn start() -> Self {
4030                let mut index = ViewIndex::new();
4031                index.add_spec(ViewSpec {
4032                    id: ROUND.to_string(),
4033                    export: "Round".to_string(),
4034                    mode: Mode::State,
4035                    wire_format: Default::default(),
4036                    projection: Projection::all(),
4037                    filters: Filters::all(),
4038                    delivery: Delivery::default(),
4039                    pipeline: None,
4040                    source_view: None,
4041                });
4042                let bus_manager = BusManager::new();
4043                let entity_cache = EntityCache::new();
4044                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4045                let addr = listener.local_addr().unwrap();
4046                let server = WebSocketServer::new(
4047                    addr,
4048                    bus_manager.clone(),
4049                    entity_cache.clone(),
4050                    Arc::new(index),
4051                    #[cfg(feature = "otel")]
4052                    None,
4053                );
4054                let (acceptor, _cleanup) = server.into_acceptor();
4055                tokio::spawn(async move { acceptor.serve_listener(listener).await });
4056                Self {
4057                    addr,
4058                    bus_manager,
4059                    entity_cache,
4060                }
4061            }
4062
4063            /// Publish a patch the way the projector does: stamp its `_seq`,
4064            /// write the cache, then the key's bus.
4065            async fn publish(&self, key: &str, mut patch: Value, seq: &str) {
4066                patch["_seq"] = Value::String(seq.to_string());
4067                self.entity_cache
4068                    .upsert_with_append(ROUND, key, patch.clone(), &[], PatchOrigin::Unknown)
4069                    .await;
4070                self.publish_frame(key, "patch", patch, seq).await;
4071            }
4072
4073            /// Publish a frame to the key's bus without touching the cache.
4074            async fn publish_frame(&self, key: &str, op: &str, data: Value, seq: &str) {
4075                let frame = json!({
4076                    "mode": "state",
4077                    "entity": ROUND,
4078                    "op": op,
4079                    "key": key,
4080                    "data": data,
4081                    "seq": seq,
4082                });
4083                self.bus_manager
4084                    .publish_state(ROUND, key, Arc::new(Bytes::from(frame.to_string())))
4085                    .await;
4086            }
4087
4088            /// Subscribe to one key and read through its snapshot.
4089            async fn subscribe(&self, key: &str) -> WebSocketStream<TcpStream> {
4090                let stream = TcpStream::connect(self.addr).await.unwrap();
4091                let mut socket = client_async(format!("ws://{}/", self.addr), stream)
4092                    .await
4093                    .unwrap()
4094                    .0;
4095                socket
4096                    .send(Message::Text(
4097                        json!({
4098                            "type": "subscribe",
4099                            "protocolVersion": 2,
4100                            "subscriptionId": "round",
4101                            "query": {"view": ROUND, "key": key},
4102                        })
4103                        .to_string()
4104                        .into(),
4105                    ))
4106                    .await
4107                    .unwrap();
4108                let ack = next_frame(&mut socket).await;
4109                assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
4110                loop {
4111                    let frame = next_frame(&mut socket).await;
4112                    assert_eq!(frame["op"], "snapshot", "unexpected frame: {frame}");
4113                    if frame["complete"] == true {
4114                        return socket;
4115                    }
4116                }
4117            }
4118        }
4119
4120        /// Two patches to one key, published before the subscriber reads the
4121        /// first, must both reach it. The bus holds only the latest frame, so
4122        /// forwarding it alone drops the fields of the one it replaced.
4123        // One worker, so tasks run one at a time. The handshake needs the
4124        // multi-threaded runtime.
4125        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4126        async fn a_state_subscriber_keeps_the_fields_of_an_overwritten_patch() {
4127            let server = StateServer::start().await;
4128            server
4129                .publish("7", json!({"id": 7}), "100:000000000001")
4130                .await;
4131            let mut socket = server.subscribe("7").await;
4132
4133            // A lone patch is forwarded as it was published. This one is an
4134            // account update, numbered by its write version.
4135            server
4136                .publish(
4137                    "7",
4138                    json!({"results": {"slot_hash": "abc"}}),
4139                    "140:003836292257",
4140                )
4141                .await;
4142            let frame = next_frame(&mut socket).await;
4143            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4144            assert_eq!(frame["seq"], "140:003836292257");
4145
4146            // Two instruction patches from later in the same slot, numbered by
4147            // transaction index. Both writes happen in one task that never
4148            // yields to the subscriber, so the second replaces the first on
4149            // the bus before the subscriber wakes.
4150            let writer = server.clone();
4151            tokio::spawn(async move {
4152                writer
4153                    .publish("7", json!({"entropy": {"seed": "def"}}), "140:000000001200")
4154                    .await;
4155                writer
4156                    .publish("7", json!({"total": 1}), "140:000000001200")
4157                    .await;
4158            })
4159            .await
4160            .unwrap();
4161            let frame = next_frame(&mut socket).await;
4162            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4163            // Everything that changed since the snapshot, the forwarded patch
4164            // included; `id` has not.
4165            assert_eq!(
4166                frame["data"],
4167                json!({
4168                    "results": {"slot_hash": "abc"},
4169                    "entropy": {"seed": "def"},
4170                    "total": 1,
4171                    "_seq": "140:000000001200",
4172                })
4173            );
4174            // The latest seq sorts below the one already delivered, and a
4175            // client would drop the frame as stale.
4176            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4177
4178            // Having caught up, the subscriber forwards patches again.
4179            server
4180                .publish("7", json!({"total": 2}), "141:000000000002")
4181                .await;
4182            let frame = next_frame(&mut socket).await;
4183            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4184            assert_eq!(
4185                frame["data"],
4186                json!({"total": 2, "_seq": "141:000000000002"})
4187            );
4188
4189            // The next catch-up is measured against the last one: only what
4190            // changed since it, not the fields that one already carried.
4191            let writer = server.clone();
4192            tokio::spawn(async move {
4193                writer
4194                    .publish("7", json!({"total": 3}), "142:000000000001")
4195                    .await;
4196                writer
4197                    .publish("7", json!({"flag": true}), "142:000000000001")
4198                    .await;
4199            })
4200            .await
4201            .unwrap();
4202            let frame = next_frame(&mut socket).await;
4203            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4204            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4205            assert_eq!(
4206                frame["data"],
4207                json!({"total": 3, "flag": true, "_seq": "142:000000000001"})
4208            );
4209
4210            socket.close(None).await.ok();
4211        }
4212
4213        /// A whole entity forwarded as it was published replaces the client's
4214        /// copy, so a later catch-up is measured against it.
4215        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4216        async fn a_forwarded_whole_entity_is_what_a_catch_up_is_measured_against() {
4217            let server = StateServer::start().await;
4218            server
4219                .publish("7", json!({"id": 7, "a": 1}), "100:000000000001")
4220                .await;
4221            let mut socket = server.subscribe("7").await;
4222
4223            let whole = json!({"id": 7, "a": 1, "b": 2, "_seq": "101:000000000001"});
4224            server
4225                .entity_cache
4226                .store_whole(ROUND, "7", whole.clone())
4227                .await;
4228            server
4229                .publish_frame("7", "upsert", whole.clone(), "101:000000000001")
4230                .await;
4231            let frame = next_frame(&mut socket).await;
4232            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4233            assert_eq!(frame["data"], whole);
4234
4235            let writer = server.clone();
4236            tokio::spawn(async move {
4237                writer
4238                    .publish("7", json!({"c": 3}), "102:000000000001")
4239                    .await;
4240                writer
4241                    .publish("7", json!({"d": 4}), "102:000000000002")
4242                    .await;
4243            })
4244            .await
4245            .unwrap();
4246            let frame = next_frame(&mut socket).await;
4247            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4248            assert_eq!(
4249                frame["data"],
4250                json!({"c": 3, "d": 4, "_seq": "102:000000000002"}),
4251                "`b` came with the whole entity, so it is not sent again"
4252            );
4253            socket.close(None).await.ok();
4254        }
4255
4256        /// A catch-up resends a field a forwarded patch set even when the
4257        /// entity has set it back to the value the client last received
4258        /// whole: the client still holds the forwarded value.
4259        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4260        async fn a_catch_up_restores_a_field_a_forwarded_patch_set() {
4261            let server = StateServer::start().await;
4262            server
4263                .publish("7", json!({"id": 7, "a": 1}), "100:000000000001")
4264                .await;
4265            let mut socket = server.subscribe("7").await;
4266
4267            server
4268                .publish("7", json!({"a": 2}), "101:000000000001")
4269                .await;
4270            let frame = next_frame(&mut socket).await;
4271            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4272            assert_eq!(frame["data"], json!({"a": 2, "_seq": "101:000000000001"}));
4273
4274            let writer = server.clone();
4275            tokio::spawn(async move {
4276                writer
4277                    .publish("7", json!({"a": 1}), "102:000000000001")
4278                    .await;
4279                writer
4280                    .publish("7", json!({"b": 1}), "102:000000000002")
4281                    .await;
4282            })
4283            .await
4284            .unwrap();
4285            let frame = next_frame(&mut socket).await;
4286            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4287            assert_eq!(
4288                frame["data"],
4289                json!({"a": 1, "b": 1, "_seq": "102:000000000002"}),
4290                "`a` matches the snapshot but the client holds 2"
4291            );
4292            socket.close(None).await.ok();
4293        }
4294
4295        /// A key deleted and created again before the subscriber reads
4296        /// either frame: the bus keeps only the new entity's. Merged into the
4297        /// client's copy it would leave the deleted entity's fields in place,
4298        /// so the client gets the new entity whole, replacing its copy.
4299        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4300        async fn a_key_deleted_and_recreated_between_reads_replaces_the_clients_copy() {
4301            let server = StateServer::start().await;
4302            server
4303                .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
4304                .await;
4305            let mut socket = server.subscribe("7").await;
4306
4307            let writer = server.clone();
4308            tokio::spawn(async move {
4309                writer.entity_cache.remove(ROUND, "7").await;
4310                writer
4311                    .publish_frame("7", "delete", Value::Null, "101:000000000001")
4312                    .await;
4313                writer
4314                    .publish("7", json!({"id": 7, "fresh": true}), "101:000000000002")
4315                    .await;
4316            })
4317            .await
4318            .unwrap();
4319            let frame = next_frame(&mut socket).await;
4320            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4321            assert_eq!(
4322                frame["data"],
4323                json!({"id": 7, "fresh": true, "_seq": "101:000000000002"})
4324            );
4325            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4326
4327            // Deleted, with a patch overwritten before it: the delete is final.
4328            let writer = server.clone();
4329            tokio::spawn(async move {
4330                writer
4331                    .publish("7", json!({"n": 1}), "102:000000000001")
4332                    .await;
4333                writer.entity_cache.remove(ROUND, "7").await;
4334                writer
4335                    .publish_frame("7", "delete", Value::Null, "102:000000000002")
4336                    .await;
4337            })
4338            .await
4339            .unwrap();
4340            let frame = next_frame(&mut socket).await;
4341            assert_eq!(frame["op"], "delete", "unexpected frame: {frame}");
4342
4343            // Created again: the client holds nothing, so it arrives whole.
4344            server
4345                .publish("7", json!({"id": 7, "third": true}), "103:000000000001")
4346                .await;
4347            let frame = next_frame(&mut socket).await;
4348            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4349            assert_eq!(frame["data"]["third"], true);
4350            assert!(frame["data"].get("fresh").is_none(), "{frame}");
4351            socket.close(None).await.ok();
4352        }
4353
4354        /// A missed delete, then a patch the cache refused (it lacks the new
4355        /// entity): the client's copy is of the deleted entity, so the patch
4356        /// is not merged into it. The new entity replaces it once it is whole.
4357        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4358        async fn a_copy_of_a_deleted_entity_takes_no_patch_until_it_is_replaced() {
4359            let server = StateServer::start().await;
4360            server
4361                .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
4362                .await;
4363            let mut socket = server.subscribe("7").await;
4364
4365            let writer = server.clone();
4366            tokio::spawn(async move {
4367                writer.entity_cache.remove(ROUND, "7").await;
4368                writer
4369                    .publish_frame("7", "delete", Value::Null, "101:000000000001")
4370                    .await;
4371                writer
4372                    .publish_frame("7", "patch", json!({"n": 1}), "101:000000000002")
4373                    .await;
4374            })
4375            .await
4376            .unwrap();
4377            // Let the subscriber read the patch, and withhold it, before the
4378            // whole entity arrives. Nothing is sent, so there is no frame to
4379            // wait for.
4380            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
4381
4382            let whole = json!({"id": 7, "n": 1, "_seq": "101:000000000002"});
4383            server
4384                .entity_cache
4385                .store_whole(ROUND, "7", whole.clone())
4386                .await;
4387            server
4388                .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
4389                .await;
4390            let frame = next_frame(&mut socket).await;
4391            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4392            assert_eq!(frame["data"], whole);
4393            socket.close(None).await.ok();
4394        }
4395
4396        /// The source replaced the entity whole (an `upsert`, dropping a field)
4397        /// and then patched it, both before the subscriber read either. The
4398        /// patch alone, or the entity merged in, would keep the dropped field.
4399        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4400        async fn a_whole_entity_among_overwritten_frames_replaces_the_clients_copy() {
4401            let server = StateServer::start().await;
4402            server
4403                .publish("7", json!({"id": 7, "dropped": 1}), "100:000000000001")
4404                .await;
4405            let mut socket = server.subscribe("7").await;
4406
4407            let writer = server.clone();
4408            tokio::spawn(async move {
4409                let whole = json!({"id": 7, "_seq": "101:000000000001"});
4410                writer
4411                    .entity_cache
4412                    .store_whole(ROUND, "7", whole.clone())
4413                    .await;
4414                writer
4415                    .publish_frame("7", "upsert", whole, "101:000000000001")
4416                    .await;
4417                writer
4418                    .publish("7", json!({"n": 1}), "101:000000000002")
4419                    .await;
4420            })
4421            .await
4422            .unwrap();
4423            let frame = next_frame(&mut socket).await;
4424            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4425            assert_eq!(
4426                frame["data"],
4427                json!({"id": 7, "n": 1, "_seq": "101:000000000002"})
4428            );
4429            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4430            socket.close(None).await.ok();
4431        }
4432
4433        /// Patches overwritten on the bus while the cache lacks the key (it
4434        /// evicted it, and refused them) leave nothing whole to send in their
4435        /// place: the holder gets the latest patch, as a holder of an evicted
4436        /// key does, and stays behind. The resend that brings the entity back
4437        /// carries the seq of the key's latest change, which the holder
4438        /// already has and would drop as stale, so it gets the cached entity
4439        /// without a seq instead: as an upsert, since the resend is the whole
4440        /// entity and replaces what the holder has.
4441        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4442        async fn a_holder_behind_on_an_evicted_key_catches_up_when_it_returns() {
4443            let server = StateServer::start().await;
4444            server
4445                .publish("7", json!({"id": 7}), "100:000000000001")
4446                .await;
4447            let mut socket = server.subscribe("7").await;
4448            server.entity_cache.remove(ROUND, "7").await;
4449
4450            let writer = server.clone();
4451            tokio::spawn(async move {
4452                writer
4453                    .publish_frame("7", "patch", json!({"a": 1}), "101:000000000001")
4454                    .await;
4455                writer
4456                    .publish_frame("7", "patch", json!({"b": 2}), "101:000000000002")
4457                    .await;
4458            })
4459            .await
4460            .unwrap();
4461            let frame = next_frame(&mut socket).await;
4462            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4463            assert_eq!(frame["data"], json!({"b": 2}));
4464            assert_eq!(frame["seq"], "101:000000000002");
4465
4466            // The resend: the whole entity, at the latest change's position.
4467            let whole = json!({"id": 7, "a": 1, "b": 2, "_seq": "101:000000000002"});
4468            server
4469                .entity_cache
4470                .store_whole(ROUND, "7", whole.clone())
4471                .await;
4472            server
4473                .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
4474                .await;
4475            let frame = next_frame(&mut socket).await;
4476            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4477            assert_eq!(frame["data"], whole);
4478            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4479
4480            // Caught up: patches are forwarded again.
4481            server
4482                .publish("7", json!({"c": 3}), "102:000000000001")
4483                .await;
4484            let frame = next_frame(&mut socket).await;
4485            assert_eq!(frame["data"], json!({"c": 3, "_seq": "102:000000000001"}));
4486            assert_eq!(frame["seq"], "102:000000000001");
4487            socket.close(None).await.ok();
4488        }
4489    }
4490
4491    /// A key only counts as held by the client once a frame actually carried
4492    /// it. A truncated or disabled snapshot leaves keys the client never saw,
4493    /// and their first change must arrive whole.
4494    mod partial_entities_over_a_socket {
4495        use super::*;
4496        use crate::projector::Projector;
4497        use crate::{MutationBatch, SlotContext};
4498        use arete_interpreter::Mutation;
4499        use futures_util::{SinkExt, StreamExt};
4500        use std::time::Duration;
4501        use tokio::net::{TcpListener, TcpStream};
4502        use tokio::sync::mpsc;
4503        use tokio_tungstenite::tungstenite::Message;
4504        use tokio_tungstenite::{client_async, WebSocketStream};
4505
4506        fn thing_index() -> ViewIndex {
4507            let mut index = ViewIndex::new();
4508            index.add_spec(list_spec());
4509            index.add_spec(ViewSpec {
4510                id: "Thing/state".to_string(),
4511                mode: Mode::State,
4512                ..list_spec()
4513            });
4514            index
4515        }
4516
4517        struct Harness {
4518            addr: SocketAddr,
4519            tx: mpsc::Sender<MutationBatch>,
4520            slot: std::sync::atomic::AtomicU64,
4521        }
4522
4523        impl Harness {
4524            async fn start(delivery: WebSocketDeliveryConfig) -> Self {
4525                Self::start_with_cache(delivery, EntityCache::new()).await
4526            }
4527
4528            async fn start_with_cache(
4529                delivery: WebSocketDeliveryConfig,
4530                entity_cache: EntityCache,
4531            ) -> Self {
4532                let view_index = Arc::new(thing_index());
4533                let bus_manager = BusManager::new();
4534                let (tx, rx) = mpsc::channel::<MutationBatch>(64);
4535                // Unconstrained, so Tokio's cooperative budget never makes the
4536                // projector yield partway through a batch: a test on one
4537                // worker then knows no subscriber runs until the whole batch
4538                // is published (see
4539                // `snapshot_rows_patches_and_catch_ups_carry_the_version`).
4540                tokio::spawn(tokio::task::unconstrained(
4541                    Projector::new(
4542                        view_index.clone(),
4543                        bus_manager.clone(),
4544                        entity_cache.clone(),
4545                        rx,
4546                        #[cfg(feature = "otel")]
4547                        None,
4548                    )
4549                    .run(),
4550                ));
4551                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4552                let addr = listener.local_addr().unwrap();
4553                let server = WebSocketServer::new(
4554                    addr,
4555                    bus_manager,
4556                    entity_cache,
4557                    view_index,
4558                    #[cfg(feature = "otel")]
4559                    None,
4560                )
4561                .with_delivery_config(delivery);
4562                let (acceptor, _cleanup) = server.into_acceptor();
4563                tokio::spawn(async move { acceptor.serve_listener(listener).await });
4564                Self {
4565                    addr,
4566                    tx,
4567                    slot: std::sync::atomic::AtomicU64::new(100),
4568                }
4569            }
4570
4571            /// Apply one source patch and wait until the projector published it.
4572            async fn patch(&self, key: &str, patch: Value) {
4573                self.mutate(Mutation {
4574                    export: "Thing".to_string(),
4575                    key: json!(key),
4576                    patch,
4577                    append: vec![],
4578                })
4579                .await;
4580            }
4581
4582            async fn mutate(&self, mutation: Mutation) {
4583                let slot = self.slot.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
4584                self.tx
4585                    .send(MutationBatch::with_slot_context(
4586                        vec![mutation].into_iter().collect(),
4587                        SlotContext::new(slot, 0),
4588                    ))
4589                    .await
4590                    .unwrap();
4591                let (ack, wait) = oneshot::channel();
4592                self.tx
4593                    .send(MutationBatch::flush_marker(ack))
4594                    .await
4595                    .unwrap();
4596                wait.await.unwrap();
4597            }
4598
4599            async fn subscribe(&self, query: Value, snapshot: bool) -> WebSocketStream<TcpStream> {
4600                let stream = TcpStream::connect(self.addr).await.unwrap();
4601                let mut socket = client_async(format!("ws://{}/", self.addr), stream)
4602                    .await
4603                    .unwrap()
4604                    .0;
4605                socket
4606                    .send(Message::Text(
4607                        json!({
4608                            "type": "subscribe",
4609                            "protocolVersion": 2,
4610                            "subscriptionId": "things",
4611                            "query": query,
4612                            "snapshot": {"enabled": snapshot},
4613                        })
4614                        .to_string()
4615                        .into(),
4616                    ))
4617                    .await
4618                    .unwrap();
4619                let ack = next_frame(&mut socket).await;
4620                assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
4621                assert_eq!(
4622                    ack["wholeEntities"],
4623                    json!(true),
4624                    "the ack advertises the whole-entity guarantee: {ack}"
4625                );
4626                socket
4627            }
4628        }
4629
4630        async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
4631            loop {
4632                let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
4633                    .await
4634                    .expect("the server answers within the timeout")
4635                    .expect("the stream stays open")
4636                    .expect("a readable frame");
4637                // Every frame here is far below the compression threshold.
4638                let bytes = match &message {
4639                    Message::Text(text) => text.as_bytes(),
4640                    Message::Binary(bytes) => bytes.as_ref(),
4641                    _ => continue,
4642                };
4643                return serde_json::from_slice(bytes).expect("frames are JSON");
4644            }
4645        }
4646
4647        /// The snapshot rows, in order, once the snapshot completes.
4648        async fn snapshot_keys(socket: &mut WebSocketStream<TcpStream>) -> Vec<String> {
4649            let mut keys = Vec::new();
4650            loop {
4651                let frame = next_frame(socket).await;
4652                assert_eq!(frame["op"], "snapshot", "expected a snapshot: {frame}");
4653                keys.extend(
4654                    frame["data"]
4655                        .as_array()
4656                        .unwrap()
4657                        .iter()
4658                        .map(|row| row["key"].as_str().unwrap().to_string()),
4659                );
4660                if frame["complete"] == json!(true) {
4661                    return keys;
4662                }
4663            }
4664        }
4665
4666        async fn seed(harness: &Harness) {
4667            for key in ["a", "b", "c"] {
4668                harness
4669                    .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
4670                    .await;
4671            }
4672        }
4673
4674        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4675        async fn a_change_to_a_key_cut_by_snapshot_limit_is_an_upsert() {
4676            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4677            seed(&harness).await;
4678
4679            let mut socket = harness
4680                .subscribe(json!({"view": "Thing/list", "snapshotLimit": 1}), true)
4681                .await;
4682            // Newest first: only "c" was sent.
4683            assert_eq!(snapshot_keys(&mut socket).await, ["c"]);
4684
4685            harness.patch("a", json!({"count": 2})).await;
4686            let frame = next_frame(&mut socket).await;
4687            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4688            assert_eq!(frame["op"], "upsert", "a key never sent arrives whole");
4689            assert_eq!(frame["data"]["name"], "thing-a");
4690            assert_eq!(frame["data"]["count"], 2);
4691
4692            // Now the client holds "a", so its next change is a patch again.
4693            harness.patch("a", json!({"count": 3})).await;
4694            let frame = next_frame(&mut socket).await;
4695            assert_eq!(frame["key"], "a");
4696            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4697            assert!(frame["data"].get("name").is_none());
4698
4699            // A key the snapshot did carry still rides its patch.
4700            harness.patch("c", json!({"count": 2})).await;
4701            let frame = next_frame(&mut socket).await;
4702            assert_eq!(frame["key"], "c");
4703            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4704            socket.close(None).await.ok();
4705        }
4706
4707        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4708        async fn a_list_without_a_snapshot_upserts_each_key_on_its_first_change() {
4709            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4710            seed(&harness).await;
4711
4712            let mut socket = harness
4713                .subscribe(json!({"view": "Thing/list"}), false)
4714                .await;
4715            harness.patch("b", json!({"count": 2})).await;
4716            let frame = next_frame(&mut socket).await;
4717            assert_eq!(frame["key"], "b");
4718            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4719            assert_eq!(frame["data"]["name"], "thing-b");
4720
4721            harness.patch("b", json!({"count": 3})).await;
4722            let frame = next_frame(&mut socket).await;
4723            assert_eq!(frame["key"], "b");
4724            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4725            socket.close(None).await.ok();
4726        }
4727
4728        /// Eviction from the server's bounded cache is not a change to the
4729        /// entity: a client holding it keeps getting its patches, never a
4730        /// `remove`, and the whole entity replaces its copy when it returns.
4731        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4732        async fn a_state_holder_keeps_its_patches_when_the_cache_evicts_the_key() {
4733            let harness = Harness::start_with_cache(
4734                WebSocketDeliveryConfig::default(),
4735                EntityCache::with_config(crate::cache::EntityCacheConfig {
4736                    max_entities_per_view: 2,
4737                    ..Default::default()
4738                }),
4739            )
4740            .await;
4741            harness
4742                .patch("a", json!({"name": "thing-a", "count": 1}))
4743                .await;
4744            let mut socket = harness
4745                .subscribe(json!({"view": "Thing/state", "key": "a"}), true)
4746                .await;
4747            assert_eq!(snapshot_keys(&mut socket).await, ["a"]);
4748
4749            // "b" and "c" push "a" out of the cache.
4750            for key in ["b", "c"] {
4751                harness
4752                    .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
4753                    .await;
4754            }
4755            harness.patch("a", json!({"count": 2})).await;
4756            let frame = next_frame(&mut socket).await;
4757            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4758            assert_eq!(frame["op"], "patch", "the holder merges it: {frame}");
4759            assert_eq!(frame["data"]["count"], 2);
4760
4761            // The whole entity comes back from the VM.
4762            let mut whole = Mutation {
4763                export: "Thing".to_string(),
4764                key: json!("a"),
4765                patch: json!({"name": "thing-a", "count": 2}),
4766                append: vec![],
4767            };
4768            whole.mark_whole_entity();
4769            harness.mutate(whole).await;
4770            let frame = next_frame(&mut socket).await;
4771            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4772            assert_eq!(frame["data"]["name"], "thing-a");
4773            assert_eq!(frame["data"]["count"], 2);
4774            assert!(frame["data"]
4775                .get(arete_interpreter::WHOLE_ENTITY_MARKER)
4776                .is_none());
4777
4778            harness.patch("a", json!({"count": 3})).await;
4779            let frame = next_frame(&mut socket).await;
4780            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4781            socket.close(None).await.ok();
4782        }
4783
4784        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4785        async fn a_state_view_without_a_snapshot_upserts_on_the_first_change() {
4786            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4787            seed(&harness).await;
4788
4789            let mut socket = harness
4790                .subscribe(json!({"view": "Thing/state", "key": "a"}), false)
4791                .await;
4792            harness.patch("a", json!({"count": 2})).await;
4793            let frame = next_frame(&mut socket).await;
4794            assert_eq!(frame["key"], "a");
4795            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4796            assert_eq!(frame["data"]["name"], "thing-a");
4797            assert_eq!(frame["data"]["count"], 2);
4798
4799            harness.patch("a", json!({"count": 3})).await;
4800            let frame = next_frame(&mut socket).await;
4801            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4802            socket.close(None).await.ok();
4803        }
4804
4805        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4806        async fn a_coalesced_list_sends_no_remove_for_a_key_it_never_sent() {
4807            let harness = Harness::start(WebSocketDeliveryConfig {
4808                collection_coalesce_ms: Some(10),
4809                ..Default::default()
4810            })
4811            .await;
4812            seed(&harness).await;
4813
4814            // A one-row window holding "c", which the disabled snapshot never
4815            // sent.
4816            let mut socket = harness
4817                .subscribe(json!({"view": "Thing/list", "take": 1}), false)
4818                .await;
4819            // "a" jumps to the front: it enters the window whole and "c",
4820            // never sent, leaves it without a `remove`.
4821            harness.patch("a", json!({"count": 2})).await;
4822            let frame = next_frame(&mut socket).await;
4823            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4824            assert_eq!(frame["op"], "upsert");
4825            assert_eq!(frame["data"]["name"], "thing-a");
4826
4827            // "c" re-enters whole; "a", held, leaves with a `remove`.
4828            harness.patch("c", json!({"count": 2})).await;
4829            let mut frames = [next_frame(&mut socket).await, next_frame(&mut socket).await];
4830            frames.sort_by_key(|frame| frame["key"].as_str().unwrap().to_string());
4831            assert_eq!(frames[0]["key"], "a");
4832            assert_eq!(frames[0]["op"], "remove");
4833            assert_eq!(frames[1]["key"], "c");
4834            assert_eq!(frames[1]["op"], "upsert");
4835            assert_eq!(frames[1]["data"]["name"], "thing-c");
4836            socket.close(None).await.ok();
4837        }
4838
4839        /// `(epoch, counter)` from an entity's `_version`.
4840        fn version_of(data: &Value) -> (String, u64) {
4841            let version = data["_version"]
4842                .as_str()
4843                .unwrap_or_else(|| panic!("no _version in {data}"));
4844            let (epoch, counter) = version.split_once(':').expect("epoch:counter");
4845            (epoch.to_string(), counter.parse().expect("decimal counter"))
4846        }
4847
4848        /// Every entity a subscriber receives carries the projector's
4849        /// `_version`: a snapshot row, a forwarded patch, and a state
4850        /// subscription's catch-up, which sends the cached entity after the
4851        /// bus overwrote a patch.
4852        // One worker, and a projector the cooperative budget cannot interrupt:
4853        // within a batch it awaits only locks no other task holds, so it
4854        // publishes both changes before the subscriber can run, and the
4855        // second overwrites the first on the state bus.
4856        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4857        async fn snapshot_rows_patches_and_catch_ups_carry_the_version() {
4858            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4859            harness.patch("7", json!({"name": "thing-7"})).await;
4860
4861            let mut list = harness.subscribe(json!({"view": "Thing/list"}), true).await;
4862            let snapshot = next_frame(&mut list).await;
4863            assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
4864            let (epoch, _) = version_of(&snapshot["data"][0]["data"]);
4865
4866            let mut state = harness
4867                .subscribe(json!({"view": "Thing/state", "key": "7"}), true)
4868                .await;
4869            let snapshot = next_frame(&mut state).await;
4870            assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
4871            let (state_epoch, seeded) = version_of(&snapshot["data"][0]["data"]);
4872            assert_eq!(state_epoch, epoch, "one projector, one epoch");
4873
4874            harness.patch("7", json!({"count": 1})).await;
4875            let patch = next_frame(&mut list).await;
4876            assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
4877            assert_eq!(version_of(&patch["data"]).0, epoch);
4878            let patch = next_frame(&mut state).await;
4879            assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
4880            let (_, patched) = version_of(&patch["data"]);
4881            assert!(patched > seeded, "{patched} after {seeded}");
4882
4883            // Two changes to the key in one batch.
4884            let slot = harness
4885                .slot
4886                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
4887            let mutation = |patch: Value| Mutation {
4888                export: "Thing".to_string(),
4889                key: json!("7"),
4890                patch,
4891                append: vec![],
4892            };
4893            harness
4894                .tx
4895                .send(MutationBatch::with_slot_context(
4896                    vec![
4897                        mutation(json!({"count": 2})),
4898                        mutation(json!({"flag": true})),
4899                    ]
4900                    .into_iter()
4901                    .collect(),
4902                    SlotContext::new(slot, 0),
4903                ))
4904                .await
4905                .unwrap();
4906
4907            let catch_up = next_frame(&mut state).await;
4908            assert_eq!(catch_up["op"], "patch", "unexpected frame: {catch_up}");
4909            assert!(catch_up.get("seq").is_none(), "a catch-up: {catch_up}");
4910            assert_eq!(catch_up["data"]["count"], 2);
4911            assert_eq!(catch_up["data"]["flag"], true);
4912            let (_, caught_up) = version_of(&catch_up["data"]);
4913            assert!(caught_up > patched, "{caught_up} after {patched}");
4914
4915            list.close(None).await.ok();
4916            state.close(None).await.ok();
4917        }
4918    }
4919
4920    /// Session tokens that expire while their socket is open.
4921    mod session_expiry {
4922        use super::*;
4923        use crate::websocket::auth::SignedSessionAuthPlugin;
4924        use arete_auth::{KeyClass, SessionClaims, SigningKey, TokenSigner, TokenVerifier};
4925        use futures_util::{SinkExt, StreamExt};
4926        use std::time::Duration;
4927        use tokio::net::{TcpListener, TcpStream};
4928        use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
4929        use tokio_tungstenite::tungstenite::protocol::CloseFrame;
4930        use tokio_tungstenite::tungstenite::Message;
4931        use tokio_tungstenite::{client_async, WebSocketStream};
4932
4933        struct Server {
4934            addr: SocketAddr,
4935            signer: TokenSigner,
4936        }
4937
4938        impl Server {
4939            async fn start() -> Self {
4940                let signing_key = SigningKey::generate();
4941                let verifier =
4942                    TokenVerifier::new(signing_key.verifying_key(), "test-issuer", "test-audience");
4943                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4944                let addr = listener.local_addr().unwrap();
4945                let server = WebSocketServer::new(
4946                    addr,
4947                    BusManager::new(),
4948                    EntityCache::new(),
4949                    Arc::new(ViewIndex::new()),
4950                    #[cfg(feature = "otel")]
4951                    None,
4952                )
4953                .with_auth_plugin(Arc::new(SignedSessionAuthPlugin::new(verifier)));
4954                let (acceptor, _cleanup) = server.into_acceptor();
4955                tokio::spawn(async move { acceptor.serve_listener(listener).await });
4956                Self {
4957                    addr,
4958                    signer: TokenSigner::new(signing_key, "test-issuer"),
4959                }
4960            }
4961
4962            fn token(&self, ttl_seconds: u64) -> String {
4963                let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
4964                    .with_scope("read")
4965                    .with_key_class(KeyClass::Secret)
4966                    .with_ttl(ttl_seconds)
4967                    .build();
4968                self.signer.sign(claims).unwrap()
4969            }
4970
4971            async fn connect(&self, token: &str) -> WebSocketStream<TcpStream> {
4972                let stream = TcpStream::connect(self.addr).await.unwrap();
4973                client_async(format!("ws://{}/?hs_token={token}", self.addr), stream)
4974                    .await
4975                    .unwrap()
4976                    .0
4977            }
4978        }
4979
4980        async fn send_json(socket: &mut WebSocketStream<TcpStream>, message: Value) {
4981            socket
4982                .send(Message::Text(message.to_string().into()))
4983                .await
4984                .unwrap();
4985        }
4986
4987        /// The close frame, if the server closes the socket within `wait`.
4988        async fn close_within(
4989            socket: &mut WebSocketStream<TcpStream>,
4990            wait: Duration,
4991        ) -> Option<Option<CloseFrame>> {
4992            tokio::time::timeout(wait, async {
4993                while let Some(Ok(message)) = socket.next().await {
4994                    if let Message::Close(frame) = message {
4995                        return Some(frame);
4996                    }
4997                }
4998                Some(None)
4999            })
5000            .await
5001            .ok()
5002            .flatten()
5003        }
5004
5005        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5006        async fn an_expired_session_is_closed_with_the_reason() {
5007            let server = Server::start().await;
5008            let mut socket = server.connect(&server.token(2)).await;
5009
5010            tokio::time::sleep(Duration::from_secs(3)).await;
5011            send_json(&mut socket, json!({"type": "ping"})).await;
5012
5013            let frame = close_within(&mut socket, Duration::from_secs(5))
5014                .await
5015                .expect("the server closes the expired session")
5016                .expect("the close frame carries a reason");
5017            assert_eq!(frame.code, CloseCode::Policy);
5018            assert_eq!(
5019                frame.reason.as_str(),
5020                "token-expired: Authentication token expired"
5021            );
5022        }
5023
5024        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5025        async fn a_session_refreshed_in_band_outlives_its_first_token() {
5026            let server = Server::start().await;
5027            let mut socket = server.connect(&server.token(2)).await;
5028
5029            send_json(
5030                &mut socket,
5031                json!({"type": "refresh_auth", "token": server.token(3_600)}),
5032            )
5033            .await;
5034            let reply = tokio::time::timeout(Duration::from_secs(5), async {
5035                while let Some(Ok(message)) = socket.next().await {
5036                    if let Message::Text(text) = message {
5037                        return serde_json::from_str::<Value>(text.as_str()).ok();
5038                    }
5039                }
5040                None
5041            })
5042            .await
5043            .expect("the server answers the refresh")
5044            .expect("the answer is JSON");
5045            assert_eq!(reply["success"], true, "refresh accepted: {reply}");
5046
5047            tokio::time::sleep(Duration::from_secs(3)).await;
5048            send_json(&mut socket, json!({"type": "ping"})).await;
5049            assert!(
5050                close_within(&mut socket, Duration::from_millis(1_500))
5051                    .await
5052                    .is_none(),
5053                "the socket stays open on the refreshed token"
5054            );
5055        }
5056    }
5057}