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    if subscription.snapshot.enabled {
1395        send_snapshot_batches(
1396            context,
1397            &subscription,
1398            &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1399            view_spec.mode,
1400            SnapshotPurpose::Initial,
1401            &context.entity_cache.snapshot_config(),
1402        )
1403        .await?;
1404    }
1405
1406    let task_context = context.clone();
1407    let subscription_id = subscription.subscription_id.clone();
1408    let query = subscription.query.clone();
1409    let view_spec_task = view_spec.clone();
1410    let span_view = view_id.clone();
1411    let span_key = key.clone();
1412    tokio::spawn(
1413        async move {
1414            // Whether the client holds this key, not whether it exists.
1415            let mut member = delivered;
1416            // Whether the client lacks fields of frames the bus overwrote
1417            // before this task read them. The cached entity is sent in their
1418            // place, but while the cache lacks the key there is none, so the
1419            // debt waits for the next frame that finds it cached.
1420            let mut behind = false;
1421            // Whether, while behind, a frame the client missed replaced the
1422            // entity (a delete, or a whole upsert): its copy may then hold
1423            // fields the entity no longer has, which merging keeps, so the
1424            // catch-up must replace it.
1425            let mut replaced = false;
1426            loop {
1427                tokio::select! {
1428                    _ = cancel_token.cancelled() => break,
1429                    changed = receiver.changed() => {
1430                        if changed.is_err() {
1431                            break;
1432                        }
1433                        let (payload, published, replaced_at) = {
1434                            let update = receiver.borrow_and_update();
1435                            (update.payload.clone(), update.published, update.replaced)
1436                        };
1437                        // The bus keeps only the latest frame. If more than one
1438                        // was published since the last read, the earlier ones
1439                        // were overwritten and the latest patch alone would
1440                        // drop their fields, so send the cached entity instead.
1441                        behind |= published > seen + 1;
1442                        replaced |= behind && replaced_at > seen;
1443                        seen = published;
1444                        let metadata = source_frame_metadata(&payload);
1445                        if metadata.op == "delete" {
1446                            task_context.entity_cache.remove(&query.view, &key).await;
1447                            if member && send_membership_frame(
1448                                &task_context,
1449                                &subscription_id,
1450                                &view_spec_task,
1451                                "delete",
1452                                &key,
1453                                Value::Null,
1454                                metadata.seq,
1455                            ).is_err() {
1456                                break;
1457                            }
1458                            member = false;
1459                            behind = false;
1460                            replaced = false;
1461                            continue;
1462                        }
1463
1464                        // A key missing from the cache was evicted there or
1465                        // never held whole, not deleted (deletes arrive
1466                        // above): the cache refused this patch and the whole
1467                        // entity is on its way (see `EntityResync`). The
1468                        // client's copy is still good, so a holder gets the
1469                        // patch and stays a holder; `remove` is only for an
1470                        // entity that stopped matching. If it is behind, it
1471                        // stays behind: the resend that brings the entity
1472                        // back carries the seq of the latest change, which it
1473                        // already has, so only the catch-up below reaches it.
1474                        let Some(cached) = task_context
1475                            .entity_cache
1476                            .get(&view_spec_task.id, &key)
1477                            .await
1478                        else {
1479                            // A copy of a replaced entity takes no patch: it
1480                            // waits for the whole entity to replace it.
1481                            if member && !replaced && send_scoped_source_payload(
1482                                &task_context,
1483                                &subscription_id,
1484                                &query.view,
1485                                payload,
1486                            ).is_err() {
1487                                break;
1488                            }
1489                            continue;
1490                        };
1491                        let selected =
1492                            select_query_entities(vec![(key.clone(), cached)], &query, true, false);
1493                        let is_member = !selected.is_empty();
1494                        let caught_up = !std::mem::take(&mut behind);
1495                        let replace = std::mem::take(&mut replaced);
1496                        let result = match (member, selected.into_iter().next()) {
1497                            (true, Some(_)) if caught_up => send_scoped_source_payload(
1498                                &task_context,
1499                                &subscription_id,
1500                                &query.view,
1501                                payload,
1502                            ),
1503                            // The whole cached entity, as a patch: it merges into
1504                            // what the client holds, so a field the cache lost
1505                            // to eviction is kept, not replaced. If a missed frame
1506                            // replaced the entity, it is an upsert instead, which
1507                            // replaces the client's copy: merging would keep the
1508                            // fields of the entity that was deleted or replaced.
1509                            // Either carries no seq, because it is newer than
1510                            // anything this subscriber was sent and seqs are not
1511                            // ordered within a slot: account updates and
1512                            // instructions number themselves differently, so the
1513                            // latest patch's seq can sort below one already
1514                            // delivered.
1515                            (true, Some((entity_key, data))) => send_membership_frame(
1516                                &task_context,
1517                                &subscription_id,
1518                                &view_spec_task,
1519                                if replace { "upsert" } else { "patch" },
1520                                &entity_key,
1521                                data,
1522                                None,
1523                            ),
1524                            (false, Some((entity_key, data))) => send_membership_frame(
1525                                &task_context,
1526                                &subscription_id,
1527                                &view_spec_task,
1528                                "upsert",
1529                                &entity_key,
1530                                data,
1531                                metadata.seq,
1532                            ),
1533                            (true, None) => send_membership_frame(
1534                                &task_context,
1535                                &subscription_id,
1536                                &view_spec_task,
1537                                "remove",
1538                                &key,
1539                                Value::Null,
1540                                metadata.seq,
1541                            ),
1542                            (false, None) => Ok(()),
1543                        };
1544                        if result.is_err() {
1545                            break;
1546                        }
1547                        member = is_member;
1548                    }
1549                }
1550            }
1551        }
1552        .instrument(info_span!("ws.subscribe.state", client_id = %context.client_id, view = %span_view, key = %span_key)),
1553    );
1554    Ok(())
1555}
1556
1557async fn attach_collection_subscription(
1558    context: &SubscriptionContext,
1559    subscription: Subscription,
1560    view_spec: ViewSpec,
1561    cancel_token: CancellationToken,
1562) -> Result<()> {
1563    let view_id = subscription.query.view.clone();
1564    let source_view_id = view_spec
1565        .source_view
1566        .clone()
1567        .unwrap_or_else(|| view_id.clone());
1568    let (mut receiver, initial_membership) = subscribe_collection_then_snapshot(
1569        context,
1570        &source_view_id,
1571        &view_spec,
1572        &subscription.query,
1573    )
1574    .await;
1575
1576    let mut snapshot_entities = initial_membership.clone();
1577    if let Some(limit) = subscription.query.snapshot_limit {
1578        snapshot_entities.truncate(limit);
1579    }
1580    enforce_snapshot_limit(context, snapshot_entities.len())?;
1581    send_subscribed_frame(context, &subscription, &view_spec)?;
1582    // Window members the client was never sent: everything past a
1583    // `snapshotLimit` cut, or the whole window when the snapshot is disabled.
1584    // They stay in the window (take/skip positions still count them) but the
1585    // client holds no copy, so none of them may ever be sent a patch.
1586    let delivered_rows = if subscription.snapshot.enabled {
1587        snapshot_entities.len()
1588    } else {
1589        0
1590    };
1591    let mut undelivered: HashSet<String> = initial_membership
1592        .iter()
1593        .skip(delivered_rows)
1594        .map(|(key, _)| key.clone())
1595        .collect();
1596    if subscription.snapshot.enabled {
1597        send_snapshot_batches(
1598            context,
1599            &subscription,
1600            &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1601            view_spec.mode,
1602            SnapshotPurpose::Initial,
1603            &context.entity_cache.snapshot_config(),
1604        )
1605        .await?;
1606    }
1607
1608    let task_context = context.clone();
1609    let task_subscription = subscription.clone();
1610    let subscription_id = subscription.subscription_id.clone();
1611    let query = subscription.query.clone();
1612    let view_spec_task = view_spec.clone();
1613    let span_view = view_id.clone();
1614    tokio::spawn(
1615        async move {
1616            let mut current = initial_membership;
1617            // Append views are event tapes: collapsing two records would lose
1618            // observable history. Coalescing is only valid for latest-state
1619            // list membership, where a full final entity preserves meaning.
1620            let coalesce_ms = (view_spec_task.mode == Mode::List)
1621                .then(|| {
1622                    view_spec_task
1623                        .delivery
1624                        .coalesce_ms
1625                        .or(task_context.delivery.collection_coalesce_ms)
1626                        .filter(|milliseconds| *milliseconds > 0)
1627                })
1628                .flatten();
1629            let mut flush_interval = coalesce_ms.map(|milliseconds| {
1630                let period = Duration::from_millis(milliseconds);
1631                let mut interval = tokio::time::interval_at(
1632                    tokio::time::Instant::now() + period,
1633                    period,
1634                );
1635                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
1636                interval
1637            });
1638            let mut pending = HashMap::<String, Arc<BusMessage>>::new();
1639            let mut pending_updates = 0_u64;
1640
1641            loop {
1642                tokio::select! {
1643                    _ = cancel_token.cancelled() => break,
1644                    _ = async {
1645                        flush_interval
1646                            .as_mut()
1647                            .expect("coalescing interval is guarded")
1648                            .tick()
1649                            .await;
1650                    }, if flush_interval.is_some() => {
1651                        if pending.is_empty() {
1652                            continue;
1653                        }
1654                        let sorted_caches = view_spec_task
1655                            .is_derived()
1656                            .then(|| task_context.view_index.sorted_caches());
1657                        let next = load_query_entities(
1658                            &task_context.entity_cache,
1659                            sorted_caches,
1660                            &view_spec_task,
1661                            &query,
1662                            false,
1663                        ).await;
1664                        if emit_coalesced_collection_delta(
1665                            &task_context,
1666                            &subscription_id,
1667                            &view_spec_task,
1668                            &current,
1669                            &next,
1670                            &pending,
1671                            &mut undelivered,
1672                        ).is_err() {
1673                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
1674                            break;
1675                        }
1676                        task_context
1677                            .metrics
1678                            .collection_coalesced(&view_id, pending_updates);
1679                        current = next;
1680                        pending.clear();
1681                        pending_updates = 0;
1682                    }
1683                    received = receiver.recv() => {
1684                        let envelope = match received {
1685                            Ok(envelope) => envelope,
1686                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
1687                                task_context.metrics.subscription_lagged(&view_id, skipped);
1688                                warn!(
1689                                    "Subscription {} lagged by {} updates",
1690                                    subscription_id, skipped
1691                                );
1692                                if view_spec_task.mode == Mode::Append {
1693                                    let _ = send_control_frame(
1694                                        &task_context,
1695                                        &SocketIssueMessage::append_subscription_lagged(
1696                                            subscription_id.clone(),
1697                                            skipped,
1698                                        ),
1699                                        &view_id,
1700                                    );
1701                                    task_context.metrics.delivery_stopped(&view_id, "append-lagged-without-replay");
1702                                    break;
1703                                }
1704                                if !task_subscription.snapshot.enabled {
1705                                    let _ = send_control_frame(
1706                                        &task_context,
1707                                        &SocketIssueMessage::subscription_lagged(
1708                                            subscription_id.clone(),
1709                                            skipped,
1710                                        ),
1711                                        &view_id,
1712                                    );
1713                                    task_context.metrics.delivery_stopped(&view_id, "lagged-without-snapshot");
1714                                    break;
1715                                }
1716                                info!(
1717                                    "Subscription {} is recovering from an authoritative snapshot",
1718                                    subscription_id
1719                                );
1720                                match recover_collection_subscription(
1721                                    &task_context,
1722                                    &task_subscription,
1723                                    &view_spec_task,
1724                                    &source_view_id,
1725                                ).await {
1726                                    Ok((next_receiver, recovered)) => {
1727                                        receiver = next_receiver;
1728                                        current = recovered;
1729                                        // Recovery replaces the whole live
1730                                        // membership, untruncated.
1731                                        undelivered.clear();
1732                                        pending.clear();
1733                                        pending_updates = 0;
1734                                        task_context.metrics.subscription_resnapshot(&view_id);
1735                                        continue;
1736                                    }
1737                                    Err(error) => {
1738                                        warn!(
1739                                            "Subscription {} failed to recover from lag: {error:#}",
1740                                            subscription_id
1741                                        );
1742                                        task_context.metrics.delivery_stopped(&view_id, "resnapshot-failed");
1743                                        break;
1744                                    }
1745                                }
1746                            }
1747                            Err(broadcast::error::RecvError::Closed) => break,
1748                        };
1749
1750                        apply_collection_source_event(
1751                            &task_context,
1752                            &source_view_id,
1753                            &view_spec_task,
1754                            &query,
1755                            &envelope,
1756                        ).await;
1757
1758                        if flush_interval.is_some() {
1759                            pending_updates = pending_updates.saturating_add(1);
1760                            pending.insert(envelope.key.clone(), envelope);
1761                            continue;
1762                        }
1763
1764                        let metadata = source_frame_metadata(&envelope.payload);
1765                        let sorted_caches = view_spec_task
1766                            .is_derived()
1767                            .then(|| task_context.view_index.sorted_caches());
1768                        let next = load_query_entities(
1769                            &task_context.entity_cache,
1770                            sorted_caches,
1771                            &view_spec_task,
1772                            &query,
1773                            false,
1774                        ).await;
1775                        if emit_collection_delta(
1776                            &task_context,
1777                            &subscription_id,
1778                            &view_spec_task,
1779                            &current,
1780                            &next,
1781                            &envelope,
1782                            &metadata,
1783                            &mut undelivered,
1784                        ).is_err() {
1785                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
1786                            break;
1787                        }
1788                        current = next;
1789                    }
1790                }
1791            }
1792        }
1793        .instrument(info_span!("ws.subscribe.collection", client_id = %context.client_id, view = %span_view)),
1794    );
1795    Ok(())
1796}
1797
1798async fn subscribe_collection_then_snapshot(
1799    context: &SubscriptionContext,
1800    source_view_id: &str,
1801    view_spec: &ViewSpec,
1802    query: &SubscriptionQuery,
1803) -> (broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>) {
1804    let cache = context.entity_cache.clone();
1805    let sorted_caches = view_spec
1806        .is_derived()
1807        .then(|| context.view_index.sorted_caches());
1808    let view_spec = view_spec.clone();
1809    let query = query.clone();
1810    subscribe_list_then_snapshot(&context.bus_manager, source_view_id, move || async move {
1811        load_query_entities(&cache, sorted_caches, &view_spec, &query, false).await
1812    })
1813    .await
1814}
1815
1816async fn recover_collection_subscription(
1817    context: &SubscriptionContext,
1818    subscription: &Subscription,
1819    view_spec: &ViewSpec,
1820    source_view_id: &str,
1821) -> Result<(broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>)> {
1822    let (receiver, membership) =
1823        subscribe_collection_then_snapshot(context, source_view_id, view_spec, &subscription.query)
1824            .await;
1825    // `snapshotLimit` caps only the initial transfer. Recovery has to replace
1826    // the complete live membership; truncating it would make the replacement
1827    // authoritative while immediately omitting members the server still
1828    // considers current.
1829    enforce_snapshot_limit(context, membership.len())?;
1830    send_snapshot_batches(
1831        context,
1832        subscription,
1833        &to_wire_snapshot_entities(membership.clone(), view_spec),
1834        view_spec.mode,
1835        SnapshotPurpose::Recovery,
1836        &context.entity_cache.snapshot_config(),
1837    )
1838    .await?;
1839    Ok((receiver, membership))
1840}
1841
1842async fn apply_collection_source_event(
1843    context: &SubscriptionContext,
1844    source_view_id: &str,
1845    view_spec: &ViewSpec,
1846    query: &SubscriptionQuery,
1847    envelope: &BusMessage,
1848) {
1849    let metadata = source_frame_metadata(&envelope.payload);
1850    if metadata.op != "delete" {
1851        return;
1852    }
1853    // A slow subscription can observe an old delete after the projector has
1854    // already recreated the key. Never let subscriber-local lag erase newer
1855    // shared cache state.
1856    let current = context
1857        .entity_cache
1858        .get(source_view_id, &envelope.key)
1859        .await;
1860    if source_delete_is_stale(current.as_ref(), metadata.seq.as_deref()) {
1861        return;
1862    }
1863    context
1864        .entity_cache
1865        .remove(source_view_id, &envelope.key)
1866        .await;
1867    if view_spec.is_derived() {
1868        let caches = context.view_index.sorted_caches();
1869        let mut guard = caches.write().await;
1870        if let Some(cache) = guard.get_mut(&query.view) {
1871            cache.remove(&envelope.key);
1872        }
1873    }
1874}
1875
1876fn source_delete_is_stale(current: Option<&Value>, delete_seq: Option<&str>) -> bool {
1877    match (current, delete_seq) {
1878        (Some(current), Some(delete_seq)) => current
1879            .get("_seq")
1880            .and_then(Value::as_str)
1881            .is_some_and(|current_seq| {
1882                cmp_seq(current_seq, delete_seq) == std::cmp::Ordering::Greater
1883            }),
1884        _ => false,
1885    }
1886}
1887
1888/// A subscription refused for a reason the client needs spelled out.
1889///
1890/// Attach paths that return this get the registration released by the
1891/// connection loop, the same as any other failure, while the client still
1892/// receives the specific error rather than a generic `subscription-rejected`.
1893/// Sending the frame and returning `Ok` instead would leave a registered
1894/// subscription with nothing attached: it would hold a slot against the
1895/// client's limit and make the advertised "resubscribe" remediation fail with
1896/// `duplicate-subscription-id`.
1897#[derive(Debug)]
1898pub(crate) struct RejectedSubscription(pub SocketIssueMessage);
1899
1900impl RejectedSubscription {
1901    fn into_error(issue: SocketIssueMessage) -> anyhow::Error {
1902        anyhow::Error::new(Self(issue))
1903    }
1904}
1905
1906impl std::fmt::Display for RejectedSubscription {
1907    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1908        formatter.write_str(&self.0.message)
1909    }
1910}
1911
1912impl std::error::Error for RejectedSubscription {}
1913
1914/// Deliver an append view as an event tape: replay the retained records after
1915/// the cursor, then forward live frames.
1916///
1917/// This deliberately does not recompute membership from the entity cache the
1918/// way [`attach_collection_subscription`] does. The cache folds each patch
1919/// into the resident entity, so a membership diff cannot express "these three
1920/// events happened"; the retained records can.
1921///
1922/// For the same reason records are never promoted to a full `upsert` the
1923/// first time this subscription sees their key, the way a window entry is.
1924/// The cache only holds the latest state, not the state as of a record's
1925/// offset, so substituting it would hand a consumer state from after its
1926/// cursor. Records go out exactly as retained, `patch` included: a tape is
1927/// events, and the consumer that resumes from a cursor holds whatever came
1928/// before it. Clients therefore apply frames that carry an `offset` as events
1929/// and do not discard them as partial entities (see "Partial entities" in
1930/// `docs/websocket-v2-protocol.md`).
1931async fn attach_journal_subscription(
1932    context: &SubscriptionContext,
1933    subscription: Subscription,
1934    view_spec: ViewSpec,
1935    journal: Arc<crate::journal::EventJournal>,
1936    cancel_token: CancellationToken,
1937) -> Result<()> {
1938    let view_id = view_spec.id.clone();
1939    let subscription_id = subscription.subscription_id.clone();
1940
1941    // A tape has no membership window, so `take`/`skip` cannot mean what they
1942    // mean on a list view. Refuse them rather than accept and ignore them.
1943    if subscription.query.take.is_some()
1944        || subscription.query.skip.is_some()
1945        || subscription.query.snapshot_limit.is_some()
1946    {
1947        return Err(RejectedSubscription::into_error(SocketIssueMessage::protocol(
1948            Some(subscription_id),
1949            "invalid-subscription",
1950            format!(
1951                "take, skip and snapshotLimit are window options and do not apply to the replayable view {view_id}"
1952            ),
1953        )));
1954    }
1955
1956    // Reject a malformed cursor rather than silently replaying from the start,
1957    // which would look like success and duplicate everything.
1958    let cursor = match subscription.query.after.as_deref() {
1959        Some(raw) => match crate::journal::Cursor::parse(raw) {
1960            Some(cursor) => Some(cursor),
1961            None => {
1962                return Err(RejectedSubscription::into_error(
1963                    SocketIssueMessage::protocol(
1964                        Some(subscription_id),
1965                        "invalid-cursor",
1966                        format!(
1967                            "`after` must be an {{epoch}}:{{offset}} replay cursor for view {view_id}, got {raw:?}"
1968                        ),
1969                    ),
1970                ));
1971            }
1972        },
1973        None => None,
1974    };
1975
1976    // Subscribe before reading the journal so anything published during the
1977    // replay is still delivered; the offset filter below drops the overlap.
1978    let mut receiver = context.bus_manager.get_or_create_list_bus(&view_id).await;
1979
1980    let replayed = match journal.replay_after(&view_id, cursor.as_ref()).await {
1981        Ok(records) => records,
1982        Err(error) => {
1983            return Err(RejectedSubscription::into_error(
1984                SocketIssueMessage::replay_refused(Some(subscription_id), &error),
1985            ));
1986        }
1987    };
1988
1989    let frame = SubscribedFrame::new(
1990        subscription.subscription_id.clone(),
1991        subscription.query.clone(),
1992        view_spec.mode,
1993        extract_sort_config(&view_spec),
1994    )
1995    .with_replay_window(journal.window(&view_id).await);
1996    send_control_frame(context, &frame, &view_id)?;
1997
1998    // Everything past the acknowledgement runs on its own task. The replay can
1999    // be long and applies real backpressure, and `attach_client_to_bus` is
2000    // awaited directly on the connection's inbound loop — doing it there would
2001    // block unsubscribe, auth refresh and pong for the whole replay.
2002    let task_context = context.clone();
2003    let task_subscription_id = subscription.subscription_id.clone();
2004    let task_query = subscription.query.clone();
2005    let task_epoch = journal.epoch().await;
2006    let span_view = view_id.clone();
2007    tokio::spawn(
2008        async move {
2009            let mut last_sent = cursor.map(|cursor| cursor.offset);
2010            // Frames that published while the replay was still running. The
2011            // bus is a bounded broadcast, so it has to be drained as we go or
2012            // a busy view laps us before the replay finishes.
2013            let mut pending: VecDeque<Arc<BusMessage>> = VecDeque::new();
2014            let mut lagged: Option<u64> = None;
2015
2016            for record in replayed {
2017                if cancel_token.is_cancelled() {
2018                    return;
2019                }
2020                drain_available(&mut receiver, &mut pending, &mut lagged);
2021                if journal_record_matches(&task_query, &record)
2022                    && send_scoped_source_payload_async(
2023                        &task_context,
2024                        &task_subscription_id,
2025                        &span_view,
2026                        record.payload,
2027                    )
2028                    .await
2029                    .is_err()
2030                {
2031                    return;
2032                }
2033                // Advance past filtered records too: they were considered.
2034                last_sent = Some(record.offset);
2035            }
2036
2037            // Flush what arrived during the replay before going live, so the
2038            // handover keeps offset order.
2039            while let Some(envelope) = pending.pop_front() {
2040                if !forward_live_frame(
2041                    &task_context,
2042                    &task_subscription_id,
2043                    &span_view,
2044                    &task_query,
2045                    &envelope,
2046                    &mut last_sent,
2047                )
2048                .await
2049                {
2050                    return;
2051                }
2052            }
2053
2054            if let Some(skipped) = lagged {
2055                report_replay_gap(
2056                    &task_context,
2057                    &task_subscription_id,
2058                    &span_view,
2059                    &task_epoch,
2060                    skipped,
2061                    last_sent,
2062                );
2063                return;
2064            }
2065
2066            loop {
2067                tokio::select! {
2068                    _ = cancel_token.cancelled() => break,
2069                    received = receiver.recv() => {
2070                        let envelope = match received {
2071                            Ok(envelope) => envelope,
2072                            // A lagged tape is a gap. Report it with the last
2073                            // offset delivered *before* the gap and stop
2074                            // delivering on this subscription.
2075                            //
2076                            // Continuing would hand the consumer frames from
2077                            // after the gap, advancing its checkpoint past
2078                            // the skipped records so they could never be
2079                            // replayed. Stopping is not a silent stall: the
2080                            // consumer has an explicit error and a cursor
2081                            // that recovers exactly what it missed.
2082                            //
2083                            // The registration is deliberately left alone.
2084                            // Its lifecycle belongs to the connection loop,
2085                            // which holds the only handle to
2086                            // `active_subscriptions`; releasing half of it
2087                            // here would desynchronise unsubscribe, the
2088                            // duplicate-ID gate and close-time usage.
2089                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
2090                                report_replay_gap(
2091                                    &task_context,
2092                                    &task_subscription_id,
2093                                    &span_view,
2094                                    &task_epoch,
2095                                    skipped,
2096                                    last_sent,
2097                                );
2098                                break;
2099                            }
2100                            Err(broadcast::error::RecvError::Closed) => break,
2101                        };
2102
2103                        if !forward_live_frame(
2104                            &task_context,
2105                            &task_subscription_id,
2106                            &span_view,
2107                            &task_query,
2108                            &envelope,
2109                            &mut last_sent,
2110                        )
2111                        .await
2112                        {
2113                            break;
2114                        }
2115                    }
2116                }
2117            }
2118        }
2119        .instrument(info_span!(
2120            "ws.subscribe.replay",
2121            client_id = %context.client_id,
2122            view = %view_id
2123        )),
2124    );
2125    Ok(())
2126}
2127
2128/// Take whatever the bus already has without waiting, so a long replay cannot
2129/// be lapped by a busy view.
2130fn drain_available(
2131    receiver: &mut broadcast::Receiver<Arc<BusMessage>>,
2132    pending: &mut VecDeque<Arc<BusMessage>>,
2133    lagged: &mut Option<u64>,
2134) {
2135    // Once a gap is known, everything still on the bus is on the far side of
2136    // it. Buffering it would put post-gap frames in front of the lag report,
2137    // advancing `last_sent` past the hole and making `recoverFrom` point
2138    // after the very records it is supposed to recover.
2139    if lagged.is_some() {
2140        return;
2141    }
2142    // Bounded so a view publishing faster than the client drains cannot turn
2143    // the buffer into an unbounded queue. Filling it is not itself a gap:
2144    // nothing has been skipped at that instant, the buffer simply stopped
2145    // accepting. Stop buffering and let the bus report the loss, with the
2146    // count it actually measures, when delivery reaches it.
2147    const MAX_PENDING: usize = 8_192;
2148    loop {
2149        if pending.len() >= MAX_PENDING {
2150            return;
2151        }
2152        match receiver.try_recv() {
2153            Ok(envelope) => pending.push_back(envelope),
2154            Err(broadcast::error::TryRecvError::Empty)
2155            | Err(broadcast::error::TryRecvError::Closed) => return,
2156            Err(broadcast::error::TryRecvError::Lagged(skipped)) => {
2157                *lagged = Some(skipped);
2158                return;
2159            }
2160        }
2161    }
2162}
2163
2164/// Whether a live frame was already delivered by the replay that preceded it.
2165///
2166/// The bus is subscribed before the tape is read, so a record published in
2167/// between appears on both paths; without this the consumer sees it twice.
2168/// Advances the high-water mark as a side effect.
2169fn already_delivered(offset: Option<u64>, last_sent: &mut Option<u64>) -> bool {
2170    let Some(offset) = offset else {
2171        // A frame with no offset predates the tape, so it cannot have been
2172        // replayed and must not move the mark.
2173        return false;
2174    };
2175    if last_sent.is_some_and(|last| offset <= last) {
2176        return true;
2177    }
2178    *last_sent = Some(offset);
2179    false
2180}
2181
2182/// Deliver one live frame, skipping anything the replay already sent.
2183///
2184/// Returns false when the subscription should end.
2185async fn forward_live_frame(
2186    context: &SubscriptionContext,
2187    subscription_id: &str,
2188    view_id: &str,
2189    query: &SubscriptionQuery,
2190    envelope: &Arc<BusMessage>,
2191    last_sent: &mut Option<u64>,
2192) -> bool {
2193    let metadata = source_frame_metadata(&envelope.payload);
2194    if already_delivered(metadata.offset, last_sent) {
2195        return true;
2196    }
2197    if !live_frame_matches(query, &envelope.key, &envelope.payload) {
2198        return true;
2199    }
2200    send_scoped_source_payload(context, subscription_id, view_id, envelope.payload.clone()).is_ok()
2201}
2202
2203fn report_replay_gap(
2204    context: &SubscriptionContext,
2205    subscription_id: &str,
2206    view_id: &str,
2207    epoch: &crate::journal::JournalEpoch,
2208    skipped: u64,
2209    last_sent: Option<u64>,
2210) {
2211    warn!(
2212        "Replay subscription {} lagged past {} records; stopping with a recovery cursor",
2213        subscription_id, skipped
2214    );
2215    let recover_from = last_sent.map(|offset| crate::journal::Cursor {
2216        epoch: epoch.clone(),
2217        offset,
2218    });
2219    let _ = send_control_frame(
2220        context,
2221        &SocketIssueMessage::replay_lagged(
2222            Some(subscription_id.to_string()),
2223            skipped,
2224            recover_from,
2225        ),
2226        view_id,
2227    );
2228}
2229
2230/// Apply the subscription's `key`, `partition` and `filters` to a retained
2231/// record. A replay must honour the same predicates a live subscription does.
2232fn journal_record_matches(
2233    query: &SubscriptionQuery,
2234    record: &crate::journal::JournalRecord,
2235) -> bool {
2236    live_frame_matches(query, &record.key, &record.payload)
2237}
2238
2239fn live_frame_matches(query: &SubscriptionQuery, key: &str, payload: &[u8]) -> bool {
2240    if !query.matches_key(key) {
2241        return false;
2242    }
2243    if query.partition.is_none() && query.filters.is_empty() {
2244        return true;
2245    }
2246    let Ok(frame) = serde_json::from_slice::<Value>(payload) else {
2247        return false;
2248    };
2249    let Some(data) = frame.get("data") else {
2250        return false;
2251    };
2252    if let Some(partition) = &query.partition {
2253        if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2254            return false;
2255        }
2256    }
2257    query
2258        .filters
2259        .iter()
2260        .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2261}
2262
2263/// Awaiting variant of [`send_scoped_source_payload`], for replays that can
2264/// exceed the client's send queue.
2265async fn send_scoped_source_payload_async(
2266    context: &SubscriptionContext,
2267    subscription_id: &str,
2268    view_id: &str,
2269    payload: Arc<Bytes>,
2270) -> Result<()> {
2271    let mut value: Value = serde_json::from_slice(&payload)?;
2272    let object = value
2273        .as_object_mut()
2274        .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2275    object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2276    object.insert(
2277        "subscriptionId".to_string(),
2278        Value::String(subscription_id.to_string()),
2279    );
2280    let json = serde_json::to_vec(&value)?;
2281    let compressed = maybe_compress(&json);
2282    let bytes = compressed.as_bytes().len();
2283    context
2284        .client_manager
2285        .send_compressed_async(context.client_id, compressed)
2286        .await
2287        .map_err(|error| anyhow::anyhow!("failed to send replayed frame: {error}"))?;
2288    context.metrics.message_sent();
2289    emit_update_sent_for_client(
2290        &context.usage_emitter,
2291        &context.client_manager,
2292        context.client_id,
2293        view_id,
2294        bytes,
2295    );
2296    Ok(())
2297}
2298
2299#[derive(Default)]
2300struct SourceFrameMetadata {
2301    op: String,
2302    seq: Option<String>,
2303    offset: Option<u64>,
2304}
2305
2306fn source_frame_metadata(payload: &[u8]) -> SourceFrameMetadata {
2307    serde_json::from_slice::<Value>(payload)
2308        .ok()
2309        .map(|value| SourceFrameMetadata {
2310            op: value
2311                .get("op")
2312                .and_then(Value::as_str)
2313                .unwrap_or_default()
2314                .to_string(),
2315            seq: value.get("seq").and_then(Value::as_str).map(str::to_string),
2316            offset: value.get("offset").and_then(Value::as_u64),
2317        })
2318        .unwrap_or_default()
2319}
2320
2321fn send_scoped_source_payload(
2322    context: &SubscriptionContext,
2323    subscription_id: &str,
2324    view_id: &str,
2325    payload: Arc<Bytes>,
2326) -> Result<()> {
2327    let mut value: Value = serde_json::from_slice(&payload)?;
2328    let object = value
2329        .as_object_mut()
2330        .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2331    object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2332    object.insert(
2333        "subscriptionId".to_string(),
2334        Value::String(subscription_id.to_string()),
2335    );
2336    let encoded = Arc::new(Bytes::from(serde_json::to_vec(&value)?));
2337    let bytes = encoded.len();
2338    context
2339        .client_manager
2340        .send_to_client(context.client_id, encoded)
2341        .map_err(|error| anyhow::anyhow!("failed to send live frame: {error}"))?;
2342    context.metrics.message_sent();
2343    emit_update_sent_for_client(
2344        &context.usage_emitter,
2345        &context.client_manager,
2346        context.client_id,
2347        view_id,
2348        bytes,
2349    );
2350    Ok(())
2351}
2352
2353fn send_membership_frame(
2354    context: &SubscriptionContext,
2355    subscription_id: &str,
2356    view_spec: &ViewSpec,
2357    op: &str,
2358    key: &str,
2359    mut data: Value,
2360    seq: Option<String>,
2361) -> Result<()> {
2362    apply_wire_format(&mut data, &view_spec.wire_format);
2363    let frame = Frame::scoped(
2364        subscription_id,
2365        view_spec.mode,
2366        &view_spec.id,
2367        op,
2368        key,
2369        data,
2370        seq,
2371    );
2372    let encoded = Arc::new(Bytes::from(serde_json::to_vec(&frame)?));
2373    let bytes = encoded.len();
2374    context
2375        .client_manager
2376        .send_to_client(context.client_id, encoded)
2377        .map_err(|error| anyhow::anyhow!("failed to send membership frame: {error}"))?;
2378    context.metrics.message_sent();
2379    emit_update_sent_for_client(
2380        &context.usage_emitter,
2381        &context.client_manager,
2382        context.client_id,
2383        &view_spec.id,
2384        bytes,
2385    );
2386    Ok(())
2387}
2388
2389/// Send what one source mutation owes a subscriber, and keep `undelivered`
2390/// (window members the client holds no copy of) in step with what was sent.
2391#[allow(clippy::too_many_arguments)]
2392fn emit_collection_delta(
2393    context: &SubscriptionContext,
2394    subscription_id: &str,
2395    view_spec: &ViewSpec,
2396    current: &[(String, Value)],
2397    next: &[(String, Value)],
2398    envelope: &BusMessage,
2399    metadata: &SourceFrameMetadata,
2400    undelivered: &mut HashSet<String>,
2401) -> Result<()> {
2402    let current_keys: Vec<&str> = current.iter().map(|(key, _)| key.as_str()).collect();
2403    let next_keys: Vec<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2404    let next_set: HashSet<&str> = next_keys.iter().copied().collect();
2405
2406    for key in current_keys
2407        .iter()
2408        .copied()
2409        .filter(|key| !next_set.contains(key))
2410    {
2411        // A key the client was never sent has nothing to remove.
2412        if undelivered.remove(key) {
2413            continue;
2414        }
2415        let op = if metadata.op == "delete" && key == envelope.key {
2416            "delete"
2417        } else {
2418            "remove"
2419        };
2420        send_membership_frame(
2421            context,
2422            subscription_id,
2423            view_spec,
2424            op,
2425            key,
2426            Value::Null,
2427            metadata.seq.clone(),
2428        )?;
2429    }
2430
2431    for (key, data) in next.iter() {
2432        let in_window = current_keys.iter().any(|candidate| *candidate == key);
2433        match member_action(
2434            in_window,
2435            in_window && !undelivered.contains(key),
2436            key == &envelope.key,
2437            view_spec.is_derived(),
2438            &metadata.op,
2439        ) {
2440            MemberAction::Skip => {}
2441            MemberAction::ForwardPatch => send_scoped_source_payload(
2442                context,
2443                subscription_id,
2444                &view_spec.id,
2445                envelope.payload.clone(),
2446            )?,
2447            MemberAction::Upsert => {
2448                let seq = metadata
2449                    .seq
2450                    .clone()
2451                    .or_else(|| data.get("_seq").and_then(Value::as_str).map(str::to_string));
2452                send_membership_frame(
2453                    context,
2454                    subscription_id,
2455                    view_spec,
2456                    "upsert",
2457                    key,
2458                    data.clone(),
2459                    seq,
2460                )?;
2461                undelivered.remove(key);
2462            }
2463        }
2464    }
2465    Ok(())
2466}
2467
2468/// Emit the net effect of every source mutation seen during one coalescing
2469/// interval. Full entities are used for changed members because intermediate
2470/// sparse patches were intentionally discarded and can no longer be merged
2471/// safely by a client.
2472fn emit_coalesced_collection_delta(
2473    context: &SubscriptionContext,
2474    subscription_id: &str,
2475    view_spec: &ViewSpec,
2476    current: &[(String, Value)],
2477    next: &[(String, Value)],
2478    pending: &HashMap<String, Arc<BusMessage>>,
2479    undelivered: &mut HashSet<String>,
2480) -> Result<()> {
2481    let changes = plan_coalesced_collection_delta(current, next, pending, undelivered);
2482    // Every key that left the window or was just upserted is settled: the
2483    // client either never needs it or now holds all of it.
2484    let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2485    undelivered.retain(|key| next_keys.contains(key.as_str()));
2486    for change in changes {
2487        if change.op == "upsert" {
2488            undelivered.remove(&change.key);
2489        }
2490        send_membership_frame(
2491            context,
2492            subscription_id,
2493            view_spec,
2494            change.op,
2495            &change.key,
2496            change.data,
2497            change.seq,
2498        )?;
2499    }
2500    Ok(())
2501}
2502
2503#[derive(Debug, PartialEq)]
2504struct CollectionChange {
2505    op: &'static str,
2506    key: String,
2507    data: Value,
2508    seq: Option<String>,
2509}
2510
2511/// Plan one coalescing interval. `undelivered` holds window members the
2512/// client was never sent; they get no `remove`/`delete` when they leave, and
2513/// like every changed key they are sent whole when they change.
2514fn plan_coalesced_collection_delta(
2515    current: &[(String, Value)],
2516    next: &[(String, Value)],
2517    pending: &HashMap<String, Arc<BusMessage>>,
2518    undelivered: &HashSet<String>,
2519) -> Vec<CollectionChange> {
2520    let current_by_key: HashMap<&str, &Value> = current
2521        .iter()
2522        .map(|(key, data)| (key.as_str(), data))
2523        .collect();
2524    let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2525    let latest_seq = pending
2526        .values()
2527        .filter_map(|envelope| source_frame_metadata(&envelope.payload).seq)
2528        .max_by(|left, right| cmp_seq(left, right));
2529    let mut changes = Vec::new();
2530
2531    for (key, _) in current
2532        .iter()
2533        .filter(|(key, _)| !next_keys.contains(key.as_str()) && !undelivered.contains(key.as_str()))
2534    {
2535        let metadata = pending
2536            .get(key)
2537            .map(|envelope| source_frame_metadata(&envelope.payload));
2538        let op = if metadata
2539            .as_ref()
2540            .is_some_and(|metadata| metadata.op == "delete")
2541        {
2542            "delete"
2543        } else {
2544            "remove"
2545        };
2546        let seq = metadata
2547            .and_then(|metadata| metadata.seq)
2548            .or_else(|| latest_seq.clone());
2549        changes.push(CollectionChange {
2550            op,
2551            key: key.clone(),
2552            data: Value::Null,
2553            seq,
2554        });
2555    }
2556
2557    for (key, data) in next {
2558        let changed = current_by_key
2559            .get(key.as_str())
2560            .is_none_or(|previous| *previous != data);
2561        if !changed && !pending.contains_key(key) {
2562            continue;
2563        }
2564        let seq = data
2565            .get("_seq")
2566            .and_then(Value::as_str)
2567            .map(str::to_string)
2568            .or_else(|| {
2569                pending
2570                    .get(key)
2571                    .and_then(|envelope| source_frame_metadata(&envelope.payload).seq)
2572            })
2573            .or_else(|| latest_seq.clone());
2574        changes.push(CollectionChange {
2575            op: "upsert",
2576            key: key.clone(),
2577            data: data.clone(),
2578            seq,
2579        });
2580    }
2581    changes
2582}
2583
2584/// What one in-window key owes a subscriber after a source mutation.
2585#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2586enum MemberAction {
2587    /// Send nothing: the subscriber already holds this entity and its data did
2588    /// not change.
2589    Skip,
2590    /// Forward the source patch verbatim.
2591    ForwardPatch,
2592    /// Send the whole entity.
2593    Upsert,
2594}
2595
2596/// Decide what to send for one key in the query window.
2597///
2598/// `in_window` says the key was in the window before this mutation; `held`
2599/// says the client also has a copy of it. The two differ after a snapshot
2600/// that `snapshotLimit` truncated or that was disabled: those keys count for
2601/// the window's positions, but the client was never sent them.
2602///
2603/// Only the mutated key's data changed, so a key that was already in the
2604/// window has at most moved index — and index is not the server's to
2605/// communicate. `subscribed` announces the window's sort (see
2606/// [`extract_sort_config`], which is `_seq` descending for a plain list) and
2607/// every SDK re-sorts locally from it, so resending an unchanged entity to
2608/// convey its new position is pure waste. This matters because on a
2609/// `_seq`-ordered list the mutated entity jumps to the front on *every*
2610/// mutation: keying the decision on position change meant rebroadcasting the
2611/// whole window each time, and the mutated entity itself never kept its
2612/// position long enough to have its patch forwarded.
2613fn member_action(
2614    in_window: bool,
2615    held: bool,
2616    is_mutated_key: bool,
2617    is_derived: bool,
2618    op: &str,
2619) -> MemberAction {
2620    if !is_mutated_key {
2621        // A key entering the window has no local state to merge into, so it
2622        // needs the whole entity. One already in the window is unchanged:
2623        // if the client holds it there is nothing new, and if it doesn't the
2624        // snapshot limit deliberately left it out, so it waits for its own
2625        // next change.
2626        return if in_window {
2627            MemberAction::Skip
2628        } else {
2629            MemberAction::Upsert
2630        };
2631    }
2632    // The mutated entity rides its own patch through untouched, but only when
2633    // the subscriber already holds a copy; a patch for a key the client never
2634    // received is not an entity. Derived views still send whole entities: the
2635    // patch on the bus is scoped to the source view, not this one (see
2636    // A4-150).
2637    if held && !is_derived && op != "delete" {
2638        MemberAction::ForwardPatch
2639    } else {
2640        MemberAction::Upsert
2641    }
2642}
2643
2644fn to_wire_snapshot_entities(
2645    entities: Vec<(String, Value)>,
2646    view_spec: &ViewSpec,
2647) -> Vec<SnapshotEntity> {
2648    entities
2649        .into_iter()
2650        .map(|(key, mut data)| {
2651            apply_wire_format(&mut data, &view_spec.wire_format);
2652            SnapshotEntity { key, data }
2653        })
2654        .collect()
2655}
2656
2657async fn load_query_entities(
2658    entity_cache: &EntityCache,
2659    sorted_caches: Option<
2660        Arc<tokio::sync::RwLock<HashMap<String, crate::sorted_cache::SortedViewCache>>>,
2661    >,
2662    view_spec: &ViewSpec,
2663    query: &SubscriptionQuery,
2664    apply_snapshot_limit: bool,
2665) -> Vec<(String, Value)> {
2666    let (entities, preordered) = if let Some(sorted_caches) = sorted_caches {
2667        let mut caches = sorted_caches.write().await;
2668        let entities = caches
2669            .get_mut(&view_spec.id)
2670            .map(|cache| cache.get_all_ordered())
2671            .unwrap_or_default();
2672        (entities, true)
2673    } else if view_spec.mode == Mode::State {
2674        let entity = match query.key.as_deref() {
2675            Some(key) => entity_cache
2676                .get(&view_spec.id, key)
2677                .await
2678                .map(|data| vec![(key.to_string(), data)])
2679                .unwrap_or_default(),
2680            None => vec![],
2681        };
2682        (entity, true)
2683    } else {
2684        (entity_cache.get_all(&view_spec.id).await, false)
2685    };
2686    select_query_entities(entities, query, preordered, apply_snapshot_limit)
2687}
2688
2689fn select_query_entities(
2690    mut entities: Vec<(String, Value)>,
2691    query: &SubscriptionQuery,
2692    preordered: bool,
2693    apply_snapshot_limit: bool,
2694) -> Vec<(String, Value)> {
2695    entities.retain(|(key, data)| query_matches_entity(query, key, data));
2696    if !preordered {
2697        entities.sort_by(|left, right| {
2698            let left_seq = left.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2699            let right_seq = right.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2700            let order = if query.after.is_some() {
2701                cmp_seq(left_seq, right_seq)
2702            } else {
2703                cmp_seq(right_seq, left_seq)
2704            };
2705            order.then_with(|| left.0.cmp(&right.0))
2706        });
2707    }
2708
2709    let skip = query.skip.unwrap_or(0);
2710    let take = query.take.unwrap_or(usize::MAX);
2711    let mut selected: Vec<_> = entities.into_iter().skip(skip).take(take).collect();
2712    if apply_snapshot_limit {
2713        if let Some(limit) = query.snapshot_limit {
2714            selected.truncate(limit);
2715        }
2716    }
2717    selected
2718}
2719
2720fn query_matches_entity(query: &SubscriptionQuery, key: &str, data: &Value) -> bool {
2721    if !query.matches_key(key) {
2722        return false;
2723    }
2724    if let Some(partition) = &query.partition {
2725        if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2726            return false;
2727        }
2728    }
2729    if let Some(after) = &query.after {
2730        let Some(seq) = data.get("_seq").and_then(Value::as_str) else {
2731            return false;
2732        };
2733        if cmp_seq(seq, after) != std::cmp::Ordering::Greater {
2734            return false;
2735        }
2736    }
2737    query
2738        .filters
2739        .iter()
2740        .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2741}
2742
2743fn value_at_dot_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
2744    path.split('.')
2745        .try_fold(value, |current, segment| current.get(segment))
2746}
2747
2748#[cfg(test)]
2749mod tests {
2750    use super::*;
2751    use crate::cache::{EntityCacheConfig, PatchOrigin};
2752    use crate::view::{Delivery, Filters, Projection};
2753    use serde_json::json;
2754    use tokio::sync::oneshot;
2755
2756    fn list_spec() -> ViewSpec {
2757        ViewSpec {
2758            id: "Thing/list".to_string(),
2759            export: "Thing".to_string(),
2760            mode: Mode::List,
2761            wire_format: Default::default(),
2762            projection: Projection::all(),
2763            filters: Filters::all(),
2764            delivery: Delivery::default(),
2765            pipeline: None,
2766            source_view: None,
2767        }
2768    }
2769
2770    /// Plan the whole window the way `emit_collection_delta` does, so a test
2771    /// can assert on what a subscriber is actually sent.
2772    fn plan_window(
2773        current_keys: &[&str],
2774        next_keys: &[&str],
2775        envelope_key: &str,
2776        is_derived: bool,
2777        op: &str,
2778    ) -> Vec<(String, MemberAction)> {
2779        plan_window_with_undelivered(current_keys, &[], next_keys, envelope_key, is_derived, op)
2780    }
2781
2782    /// [`plan_window`] for a subscriber that was never sent `undelivered`.
2783    fn plan_window_with_undelivered(
2784        current_keys: &[&str],
2785        undelivered: &[&str],
2786        next_keys: &[&str],
2787        envelope_key: &str,
2788        is_derived: bool,
2789        op: &str,
2790    ) -> Vec<(String, MemberAction)> {
2791        next_keys
2792            .iter()
2793            .map(|key| {
2794                let in_window = current_keys.contains(key);
2795                let held = in_window && !undelivered.contains(key);
2796                (
2797                    (*key).to_string(),
2798                    member_action(in_window, held, *key == envelope_key, is_derived, op),
2799                )
2800            })
2801            .collect()
2802    }
2803
2804    #[test]
2805    fn a_window_member_the_client_never_received_is_sent_whole() {
2806        // snapshotLimit cut the window [4, 3, 2, 1] to [4, 3]. "1" still
2807        // counts for positions, but a patch for it would reach a client with
2808        // nothing to merge into.
2809        let plan = plan_window_with_undelivered(
2810            &["4", "3", "2", "1"],
2811            &["2", "1"],
2812            &["1", "4", "3", "2"],
2813            "1",
2814            false,
2815            "patch",
2816        );
2817        assert_eq!(
2818            plan,
2819            vec![
2820                ("1".to_string(), MemberAction::Upsert),
2821                ("4".to_string(), MemberAction::Skip),
2822                ("3".to_string(), MemberAction::Skip),
2823                // Unchanged and never sent: it waits for its own change
2824                // rather than undoing the snapshot limit.
2825                ("2".to_string(), MemberAction::Skip),
2826            ]
2827        );
2828    }
2829
2830    #[test]
2831    fn a_held_member_still_rides_its_patch_after_a_truncated_snapshot() {
2832        let plan = plan_window_with_undelivered(
2833            &["4", "3", "2", "1"],
2834            &["2", "1"],
2835            &["4", "3", "2", "1"],
2836            "4",
2837            false,
2838            "patch",
2839        );
2840        assert_eq!(plan[0], ("4".to_string(), MemberAction::ForwardPatch));
2841    }
2842
2843    #[test]
2844    fn coalescing_sends_no_remove_for_a_key_the_client_never_received() {
2845        let current = vec![
2846            ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2847            (
2848                "unsent".to_string(),
2849                json!({"count": 1, "_seq": "10:000001"}),
2850            ),
2851            ("gone".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2852        ];
2853        let next = vec![
2854            ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2855            (
2856                "unsent".to_string(),
2857                json!({"count": 2, "_seq": "10:000003"}),
2858            ),
2859        ];
2860        let pending = HashMap::from([
2861            (
2862                "unsent".to_string(),
2863                list_message("unsent", "patch", "10:000003"),
2864            ),
2865            (
2866                "gone".to_string(),
2867                list_message("gone", "delete", "10:000002"),
2868            ),
2869        ]);
2870        let undelivered = HashSet::from(["unsent".to_string(), "gone".to_string()]);
2871
2872        assert_eq!(
2873            plan_coalesced_collection_delta(&current, &next, &pending, &undelivered),
2874            vec![CollectionChange {
2875                op: "upsert",
2876                key: "unsent".to_string(),
2877                data: json!({"count": 2, "_seq": "10:000003"}),
2878                seq: Some("10:000003".to_string()),
2879            }]
2880        );
2881    }
2882
2883    #[test]
2884    fn a_reordered_window_sends_only_the_mutated_entity() {
2885        // A `_seq`-descending list: mutating "1" moves it to the front and
2886        // shifts every other key down one. Only "1" changed, so only "1" is
2887        // sent — and it rides its own patch, not a full entity.
2888        let current = ["4", "3", "2", "1"];
2889        let next = ["1", "4", "3", "2"];
2890        let plan = plan_window(&current, &next, "1", false, "patch");
2891
2892        assert_eq!(
2893            plan,
2894            vec![
2895                ("1".to_string(), MemberAction::ForwardPatch),
2896                ("4".to_string(), MemberAction::Skip),
2897                ("3".to_string(), MemberAction::Skip),
2898                ("2".to_string(), MemberAction::Skip),
2899            ]
2900        );
2901    }
2902
2903    #[test]
2904    fn a_key_entering_the_window_gets_the_whole_entity() {
2905        // "5" has no local state for the subscriber to merge a patch into.
2906        let plan = plan_window(&["4", "3"], &["5", "4", "3"], "5", false, "patch");
2907        assert_eq!(
2908            plan,
2909            vec![
2910                ("5".to_string(), MemberAction::Upsert),
2911                ("4".to_string(), MemberAction::Skip),
2912                ("3".to_string(), MemberAction::Skip),
2913            ]
2914        );
2915    }
2916
2917    #[test]
2918    fn derived_views_still_send_whole_entities() {
2919        // The patch on the bus is scoped to the source view, so a derived
2920        // subscription cannot forward it verbatim (A4-150).
2921        let plan = plan_window(&["1", "2"], &["1", "2"], "1", true, "patch");
2922        assert_eq!(
2923            plan,
2924            vec![
2925                ("1".to_string(), MemberAction::Upsert),
2926                ("2".to_string(), MemberAction::Skip),
2927            ]
2928        );
2929    }
2930
2931    #[test]
2932    fn a_delete_envelope_never_forwards_a_patch() {
2933        // A surviving key on a delete envelope carries no mergeable patch.
2934        let plan = plan_window(&["1", "2"], &["1", "2"], "1", false, "delete");
2935        assert_eq!(
2936            plan,
2937            vec![
2938                ("1".to_string(), MemberAction::Upsert),
2939                ("2".to_string(), MemberAction::Skip),
2940            ]
2941        );
2942    }
2943
2944    #[test]
2945    fn an_unchanged_window_sends_one_frame_not_a_broadcast() {
2946        // The regression this guards: 500 members used to mean 500 full
2947        // entities on the wire for a single mutation.
2948        let keys: Vec<String> = (0..500).map(|index| index.to_string()).collect();
2949        let refs: Vec<&str> = keys.iter().map(String::as_str).collect();
2950        let plan = plan_window(&refs, &refs, "250", false, "patch");
2951
2952        let sent = plan
2953            .iter()
2954            .filter(|(_, action)| *action != MemberAction::Skip)
2955            .count();
2956        assert_eq!(sent, 1);
2957        assert_eq!(plan[250].1, MemberAction::ForwardPatch);
2958    }
2959
2960    fn list_message(key: &str, op: &str, seq: &str) -> Arc<BusMessage> {
2961        Arc::new(BusMessage {
2962            key: key.to_string(),
2963            entity: "Thing/list".to_string(),
2964            payload: Arc::new(Bytes::from(
2965                serde_json::to_vec(&json!({
2966                    "entity": "Thing/list",
2967                    "op": op,
2968                    "key": key,
2969                    "seq": seq,
2970                    "data": {},
2971                }))
2972                .unwrap(),
2973            )),
2974        })
2975    }
2976
2977    #[test]
2978    fn coalescing_emits_only_the_final_full_state_per_changed_key() {
2979        let current = vec![
2980            ("a".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2981            ("b".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2982            (
2983                "stable".to_string(),
2984                json!({"count": 1, "_seq": "10:000001"}),
2985            ),
2986        ];
2987        let next = vec![
2988            ("a".to_string(), json!({"count": 3, "_seq": "10:000004"})),
2989            ("c".to_string(), json!({"count": 1, "_seq": "10:000003"})),
2990            (
2991                "stable".to_string(),
2992                json!({"count": 1, "_seq": "10:000001"}),
2993            ),
2994        ];
2995        let pending = HashMap::from([
2996            ("a".to_string(), list_message("a", "patch", "10:000004")),
2997            ("b".to_string(), list_message("b", "delete", "10:000002")),
2998            ("c".to_string(), list_message("c", "patch", "10:000003")),
2999        ]);
3000
3001        assert_eq!(
3002            plan_coalesced_collection_delta(&current, &next, &pending, &HashSet::new()),
3003            vec![
3004                CollectionChange {
3005                    op: "delete",
3006                    key: "b".to_string(),
3007                    data: Value::Null,
3008                    seq: Some("10:000002".to_string()),
3009                },
3010                CollectionChange {
3011                    op: "upsert",
3012                    key: "a".to_string(),
3013                    data: json!({"count": 3, "_seq": "10:000004"}),
3014                    seq: Some("10:000004".to_string()),
3015                },
3016                CollectionChange {
3017                    op: "upsert",
3018                    key: "c".to_string(),
3019                    data: json!({"count": 1, "_seq": "10:000003"}),
3020                    seq: Some("10:000003".to_string()),
3021                },
3022            ]
3023        );
3024    }
3025
3026    #[test]
3027    fn a_delayed_delete_cannot_erase_a_newer_recreated_entity() {
3028        let recreated = json!({"balance": 2, "_seq": "10:000004"});
3029        assert!(source_delete_is_stale(Some(&recreated), Some("10:000002")));
3030        assert!(!source_delete_is_stale(Some(&recreated), Some("10:000004")));
3031        assert!(!source_delete_is_stale(Some(&recreated), None));
3032    }
3033
3034    #[test]
3035    fn a_lagged_snapshotless_subscription_gets_a_fatal_retryable_error() {
3036        let issue = SocketIssueMessage::subscription_lagged("balances".to_string(), 42);
3037        assert_eq!(issue.code, "subscription-lagged");
3038        assert!(issue.retryable);
3039        assert!(issue.fatal);
3040        assert!(issue.message.contains("42"));
3041        assert!(issue
3042            .suggested_action
3043            .unwrap()
3044            .contains("snapshots enabled"));
3045
3046        let append = SocketIssueMessage::append_subscription_lagged("trades".to_string(), 9);
3047        assert!(append.fatal);
3048        assert!(append.suggested_action.unwrap().contains("retained replay"));
3049    }
3050
3051    #[test]
3052    fn snapshot_batches_share_identity_and_completion() {
3053        let entities = ["one", "two", "three"].map(|key| SnapshotEntity {
3054            key: key.to_string(),
3055            data: json!({"key": key}),
3056        });
3057        let batches = create_snapshot_batches(
3058            &entities,
3059            SnapshotMetadata {
3060                subscription_id: "sub-1",
3061                snapshot_id: "snapshot-1",
3062                authoritative: true,
3063                mode: Mode::List,
3064                view_id: "Thing/list",
3065                key: None,
3066            },
3067            &SnapshotBatchConfig {
3068                initial_batch_size: 2,
3069                subsequent_batch_size: 1,
3070            },
3071        );
3072        assert_eq!(batches.len(), 2);
3073        assert!(batches.iter().all(|batch| batch.subscription_id == "sub-1"));
3074        assert!(batches
3075            .iter()
3076            .all(|batch| batch.snapshot_id == "snapshot-1"));
3077        assert!(!batches[0].complete);
3078        assert!(batches[1].complete);
3079        assert!(batches.iter().all(|batch| batch.authoritative));
3080    }
3081
3082    #[test]
3083    fn empty_incremental_snapshot_is_explicitly_non_authoritative() {
3084        let batches = create_snapshot_batches(
3085            &[],
3086            SnapshotMetadata {
3087                subscription_id: "sub-1",
3088                snapshot_id: "snapshot-1",
3089                authoritative: false,
3090                mode: Mode::State,
3091                view_id: "Thing/state",
3092                key: Some("missing"),
3093            },
3094            &SnapshotBatchConfig {
3095                initial_batch_size: 1,
3096                subsequent_batch_size: 1,
3097            },
3098        );
3099        assert_eq!(batches.len(), 1);
3100        assert!(!batches[0].authoritative);
3101        assert!(batches[0].complete);
3102        assert_eq!(batches[0].key.as_deref(), Some("missing"));
3103    }
3104
3105    #[test]
3106    fn lag_recovery_is_authoritative_for_an_after_query() {
3107        let subscription = Subscription {
3108            protocol_version: PROTOCOL_VERSION,
3109            subscription_id: "sub-1".to_string(),
3110            query: SubscriptionQuery {
3111                view: "Thing/list".to_string(),
3112                after: Some("40:000000000010".to_string()),
3113                ..Default::default()
3114            },
3115            snapshot: Default::default(),
3116        };
3117
3118        assert!(!SnapshotPurpose::Initial.authoritative(&subscription));
3119        assert!(SnapshotPurpose::Recovery.authoritative(&subscription));
3120    }
3121
3122    #[test]
3123    fn dot_path_filters_are_exact_and_type_sensitive() {
3124        let mut query = SubscriptionQuery {
3125            view: "Thing/list".to_string(),
3126            ..Default::default()
3127        };
3128        query
3129            .filters
3130            .insert("state.status".to_string(), json!("open"));
3131        query.filters.insert("metrics.count".to_string(), json!(2));
3132        assert!(query_matches_entity(
3133            &query,
3134            "one",
3135            &json!({"state": {"status": "open"}, "metrics": {"count": 2}}),
3136        ));
3137        assert!(!query_matches_entity(
3138            &query,
3139            "one",
3140            &json!({"state": {"status": "open"}, "metrics": {"count": "2"}}),
3141        ));
3142    }
3143
3144    #[test]
3145    fn take_and_skip_define_independent_deterministic_windows() {
3146        let entities: Vec<_> = (1..=6)
3147            .map(|id| {
3148                (
3149                    id.to_string(),
3150                    json!({"id": id, "_seq": format!("10:{id:012}")}),
3151                )
3152            })
3153            .collect();
3154        let first = SubscriptionQuery {
3155            view: "Thing/list".to_string(),
3156            take: Some(2),
3157            skip: Some(0),
3158            ..Default::default()
3159        };
3160        let second = SubscriptionQuery {
3161            skip: Some(2),
3162            ..first.clone()
3163        };
3164        let first_keys: Vec<_> = select_query_entities(entities.clone(), &first, false, false)
3165            .into_iter()
3166            .map(|(key, _)| key)
3167            .collect();
3168        let second_keys: Vec<_> = select_query_entities(entities, &second, false, false)
3169            .into_iter()
3170            .map(|(key, _)| key)
3171            .collect();
3172        assert_eq!(first_keys, ["6", "5"]);
3173        assert_eq!(second_keys, ["4", "3"]);
3174    }
3175
3176    #[tokio::test]
3177    async fn state_receiver_is_installed_before_snapshot_awaits() {
3178        let bus = BusManager::new();
3179        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3180        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3181        let bus_for_task = bus.clone();
3182        let task = tokio::spawn(async move {
3183            subscribe_state_then_snapshot(&bus_for_task, "Thing/state", "one", || async move {
3184                snapshot_started_tx.send(()).unwrap();
3185                release_snapshot_rx.await.unwrap();
3186            })
3187            .await
3188            .0
3189        });
3190        snapshot_started_rx.await.unwrap();
3191        bus.publish_state(
3192            "Thing/state",
3193            "one",
3194            Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3195        )
3196        .await;
3197        release_snapshot_tx.send(()).unwrap();
3198        let mut receiver = task.await.unwrap();
3199        receiver.changed().await.unwrap();
3200        assert!(!receiver.borrow().payload.is_empty());
3201    }
3202
3203    async fn assert_list_receiver_precedes_snapshot(view: &'static str) {
3204        let bus = BusManager::new();
3205        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3206        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3207        let bus_for_task = bus.clone();
3208        let task = tokio::spawn(async move {
3209            subscribe_list_then_snapshot(&bus_for_task, view, || async move {
3210                snapshot_started_tx.send(()).unwrap();
3211                release_snapshot_rx.await.unwrap();
3212            })
3213            .await
3214            .0
3215        });
3216        snapshot_started_rx.await.unwrap();
3217        bus.publish_list(
3218            view,
3219            Arc::new(BusMessage {
3220                key: "one".to_string(),
3221                entity: view.to_string(),
3222                payload: Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3223            }),
3224        )
3225        .await;
3226        release_snapshot_tx.send(()).unwrap();
3227        let mut receiver = task.await.unwrap();
3228        assert_eq!(receiver.recv().await.unwrap().key, "one");
3229    }
3230
3231    #[tokio::test]
3232    async fn list_receiver_is_installed_before_snapshot() {
3233        assert_list_receiver_precedes_snapshot("Thing/list").await;
3234    }
3235
3236    #[tokio::test]
3237    async fn append_receiver_is_installed_before_snapshot() {
3238        assert_list_receiver_precedes_snapshot("Thing/append").await;
3239    }
3240
3241    #[tokio::test]
3242    async fn derived_source_receiver_is_installed_before_snapshot() {
3243        assert_list_receiver_precedes_snapshot("Thing/list-source").await;
3244    }
3245
3246    #[tokio::test]
3247    async fn snapshot_limit_does_not_change_live_take_skip_membership() {
3248        let cache = EntityCache::with_config(EntityCacheConfig {
3249            max_entities_per_view: 10,
3250            ..Default::default()
3251        });
3252        for id in 1..=4 {
3253            cache
3254                .upsert(
3255                    "Thing/list",
3256                    &id.to_string(),
3257                    json!({"_seq": format!("10:{id:012}")}),
3258                )
3259                .await;
3260        }
3261        let query = SubscriptionQuery {
3262            view: "Thing/list".to_string(),
3263            take: Some(3),
3264            skip: Some(1),
3265            snapshot_limit: Some(1),
3266            ..Default::default()
3267        };
3268        let live = load_query_entities(&cache, None, &list_spec(), &query, false).await;
3269        let snapshot = load_query_entities(&cache, None, &list_spec(), &query, true).await;
3270        assert_eq!(live.len(), 3);
3271        assert_eq!(snapshot.len(), 1);
3272    }
3273
3274    #[test]
3275    fn fixture_manifest_covers_required_conformance_cases() {
3276        let manifest: Value = serde_json::from_str(include_str!(
3277            "../../../../tests/fixtures/websocket-v2/manifest.json"
3278        ))
3279        .unwrap();
3280        let names: HashSet<_> = manifest["fixtures"]
3281            .as_array()
3282            .unwrap()
3283            .iter()
3284            .filter_map(Value::as_str)
3285            .collect();
3286        for required in [
3287            "keyed-state.json",
3288            "list-windows.json",
3289            "filters.json",
3290            "multi-batch-authoritative.json",
3291            "empty-snapshot.json",
3292            "remove.json",
3293            "delete.json",
3294            "incremental-snapshot.json",
3295            "reconnect-replacement.json",
3296            "errors.json",
3297            "whole-entities.json",
3298        ] {
3299            assert!(names.contains(required), "missing fixture {required}");
3300        }
3301
3302        for document in [
3303            include_str!("../../../../tests/fixtures/websocket-v2/keyed-state.json"),
3304            include_str!("../../../../tests/fixtures/websocket-v2/list-windows.json"),
3305            include_str!("../../../../tests/fixtures/websocket-v2/filters.json"),
3306            include_str!("../../../../tests/fixtures/websocket-v2/multi-batch-authoritative.json"),
3307            include_str!("../../../../tests/fixtures/websocket-v2/empty-snapshot.json"),
3308            include_str!("../../../../tests/fixtures/websocket-v2/remove.json"),
3309            include_str!("../../../../tests/fixtures/websocket-v2/delete.json"),
3310            include_str!("../../../../tests/fixtures/websocket-v2/incremental-snapshot.json"),
3311            include_str!("../../../../tests/fixtures/websocket-v2/reconnect-replacement.json"),
3312            include_str!("../../../../tests/fixtures/websocket-v2/errors.json"),
3313            include_str!("../../../../tests/fixtures/websocket-v2/whole-entities.json"),
3314        ] {
3315            let fixture: Value = serde_json::from_str(document).unwrap();
3316            assert!(fixture["name"].is_string());
3317        }
3318    }
3319
3320    fn append_frame(key: &str, data: Value) -> Arc<Bytes> {
3321        let frame = json!({
3322            "entity": "Trade/append",
3323            "op": "patch",
3324            "key": key,
3325            "offset": 7,
3326            "data": data,
3327        });
3328        Arc::new(Bytes::from(serde_json::to_vec(&frame).unwrap()))
3329    }
3330
3331    /// A replayable subscription must honour the same predicates a live
3332    /// collection subscription does; otherwise a filtered consumer receives
3333    /// events outside the query it asked for.
3334    #[test]
3335    fn replay_delivery_applies_key_partition_and_filters() {
3336        let matching = append_frame("pool1", json!({"_partition": "us", "side": "buy"}));
3337        let other_partition = append_frame("pool1", json!({"_partition": "eu", "side": "buy"}));
3338        let other_side = append_frame("pool1", json!({"_partition": "us", "side": "sell"}));
3339
3340        let unfiltered = SubscriptionQuery {
3341            view: "Trade/append".to_string(),
3342            ..Default::default()
3343        };
3344        assert!(live_frame_matches(&unfiltered, "pool1", &matching));
3345        assert!(live_frame_matches(&unfiltered, "pool9", &matching));
3346
3347        let keyed = SubscriptionQuery {
3348            view: "Trade/append".to_string(),
3349            key: Some("pool1".to_string()),
3350            ..Default::default()
3351        };
3352        assert!(live_frame_matches(&keyed, "pool1", &matching));
3353        assert!(!live_frame_matches(&keyed, "pool2", &matching));
3354
3355        let partitioned = SubscriptionQuery {
3356            view: "Trade/append".to_string(),
3357            partition: Some("us".to_string()),
3358            ..Default::default()
3359        };
3360        assert!(live_frame_matches(&partitioned, "pool1", &matching));
3361        assert!(!live_frame_matches(&partitioned, "pool1", &other_partition));
3362
3363        let filtered = SubscriptionQuery {
3364            view: "Trade/append".to_string(),
3365            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3366            ..Default::default()
3367        };
3368        assert!(live_frame_matches(&filtered, "pool1", &matching));
3369        assert!(!live_frame_matches(&filtered, "pool1", &other_side));
3370    }
3371
3372    #[test]
3373    fn a_frame_without_decodable_data_does_not_satisfy_a_filter() {
3374        let filtered = SubscriptionQuery {
3375            view: "Trade/append".to_string(),
3376            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3377            ..Default::default()
3378        };
3379        let garbage = Arc::new(Bytes::from_static(b"not json"));
3380        assert!(!live_frame_matches(&filtered, "pool1", &garbage));
3381    }
3382
3383    /// The recovery cursor must be the last offset delivered *before* the
3384    /// gap. Reporting the newest offset seen would step the consumer over
3385    /// the skipped records permanently.
3386    #[test]
3387    fn replay_lagged_recovers_from_before_the_gap() {
3388        let epoch = crate::journal::JournalEpoch::new();
3389        let issue = SocketIssueMessage::replay_lagged(
3390            Some("trades".to_string()),
3391            37,
3392            Some(crate::journal::Cursor {
3393                epoch: epoch.clone(),
3394                offset: 4180,
3395            }),
3396        );
3397        assert_eq!(issue.code, "replay-lagged");
3398        assert_eq!(issue.recover_from, Some(format!("{epoch}:4180")));
3399        assert!(
3400            issue.suggested_action.unwrap().contains("4180"),
3401            "the consumer is told exactly which cursor recovers the gap"
3402        );
3403
3404        // Nothing delivered yet: there is no pre-gap offset, so the whole
3405        // retained window is the recovery.
3406        let from_scratch = SocketIssueMessage::replay_lagged(Some("trades".to_string()), 9, None);
3407        assert_eq!(from_scratch.recover_from, None);
3408        assert!(from_scratch
3409            .suggested_action
3410            .unwrap()
3411            .contains("without `after`"));
3412    }
3413
3414    fn bus_message(key: &str) -> Arc<BusMessage> {
3415        Arc::new(BusMessage {
3416            key: key.to_string(),
3417            entity: "Trade/append".to_string(),
3418            payload: Arc::new(Bytes::from_static(b"{}")),
3419        })
3420    }
3421
3422    /// A long replay must keep the bus drained. The broadcast buffer is
3423    /// bounded, so a busy view would otherwise lap the replay and the first
3424    /// live `recv` would return `Lagged` — telling the client to resubscribe,
3425    /// starting another long replay, which laps again.
3426    #[tokio::test]
3427    async fn draining_during_a_replay_keeps_a_busy_view_from_lapping_it() {
3428        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(16);
3429        let mut pending = VecDeque::new();
3430        let mut lagged = None;
3431
3432        // Publish more than the channel holds, draining as a replay would
3433        // between sends.
3434        for index in 0..48 {
3435            sender.send(bus_message(&format!("k{index}"))).unwrap();
3436            drain_available(&mut receiver, &mut pending, &mut lagged);
3437        }
3438
3439        assert_eq!(lagged, None, "draining as we go means nothing is dropped");
3440        assert_eq!(pending.len(), 48, "every published frame is buffered");
3441    }
3442
3443    #[tokio::test]
3444    async fn a_replay_that_never_drains_is_reported_as_a_gap() {
3445        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3446        for index in 0..32 {
3447            sender.send(bus_message(&format!("k{index}"))).unwrap();
3448        }
3449
3450        let mut pending = VecDeque::new();
3451        let mut lagged = None;
3452        drain_available(&mut receiver, &mut pending, &mut lagged);
3453
3454        assert!(
3455            lagged.is_some(),
3456            "overflowing the bus is a gap, not silent truncation"
3457        );
3458    }
3459
3460    /// Everything still on the bus after a gap is on the far side of it.
3461    /// Buffering it would put those frames in front of the lag report and
3462    /// advance the recovery cursor past the records it is meant to recover.
3463    #[tokio::test]
3464    async fn nothing_after_a_gap_is_buffered_ahead_of_the_report() {
3465        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3466        let mut pending = VecDeque::new();
3467        let mut lagged = None;
3468
3469        // Delivered and buffered normally.
3470        sender.send(bus_message("before")).unwrap();
3471        drain_available(&mut receiver, &mut pending, &mut lagged);
3472        assert_eq!(pending.len(), 1);
3473
3474        // Overflow the bus: everything published from here is past the gap.
3475        for index in 0..32 {
3476            sender.send(bus_message(&format!("lost{index}"))).unwrap();
3477        }
3478        drain_available(&mut receiver, &mut pending, &mut lagged);
3479        assert!(lagged.is_some());
3480
3481        // The bus still holds what survived the overflow, all of it past the
3482        // gap. A later iteration of the replay loop drains again.
3483        let buffered_at_gap = pending.len();
3484        drain_available(&mut receiver, &mut pending, &mut lagged);
3485        assert_eq!(
3486            pending.len(),
3487            buffered_at_gap,
3488            "post-gap frames must not join the pre-gap flush"
3489        );
3490        assert_eq!(pending.front().unwrap().key, "before");
3491    }
3492    /// The bus is subscribed before the tape is read, so a record published
3493    /// in that window arrives on both paths.
3494    #[test]
3495    fn the_seam_between_replay_and_live_neither_repeats_nor_skips() {
3496        let mut last_sent = Some(4211);
3497
3498        assert!(
3499            already_delivered(Some(4211), &mut last_sent),
3500            "the record the replay ended on must not be sent twice"
3501        );
3502        assert!(already_delivered(Some(4100), &mut last_sent));
3503        assert_eq!(last_sent, Some(4211), "a duplicate never moves the mark");
3504
3505        assert!(
3506            !already_delivered(Some(4212), &mut last_sent),
3507            "the next record is new"
3508        );
3509        assert_eq!(last_sent, Some(4212));
3510
3511        // A frame with no offset comes from a view with no tape; it cannot
3512        // have been replayed, and must not disturb the mark.
3513        assert!(!already_delivered(None, &mut last_sent));
3514        assert_eq!(last_sent, Some(4212));
3515    }
3516
3517    /// A subscription with no cursor has delivered nothing, so the first live
3518    /// frame is not a duplicate.
3519    #[test]
3520    fn a_fresh_subscription_delivers_its_first_live_frame() {
3521        let mut last_sent = None;
3522        assert!(!already_delivered(Some(0), &mut last_sent));
3523        assert_eq!(last_sent, Some(0));
3524    }
3525
3526    /// End-to-end over a real socket: the pieces above are unit-tested
3527    /// individually, but the thing a consumer actually does — reconnect with
3528    /// a stored cursor and keep reading — only exists once a subscription is
3529    /// attached to a connection.
3530    mod over_a_socket {
3531        use super::*;
3532        use crate::journal::{EventJournal, JournalConfig};
3533        use crate::projector::Projector;
3534        use crate::{MutationBatch, SlotContext};
3535        use arete_interpreter::Mutation;
3536        use futures_util::{SinkExt, StreamExt};
3537        use std::time::Duration;
3538        use tokio::net::{TcpListener, TcpStream};
3539        use tokio::sync::mpsc;
3540        use tokio_tungstenite::tungstenite::Message;
3541        use tokio_tungstenite::{client_async, WebSocketStream};
3542
3543        const RETAINED: u64 = 600;
3544
3545        fn append_index() -> ViewIndex {
3546            let mut index = ViewIndex::new();
3547            index.add_spec(ViewSpec {
3548                id: "Trade/append".to_string(),
3549                export: "Trade".to_string(),
3550                mode: Mode::Append,
3551                wire_format: Default::default(),
3552                projection: Projection::all(),
3553                filters: Filters::all(),
3554                delivery: Delivery::default(),
3555                pipeline: None,
3556                source_view: None,
3557            });
3558            index
3559        }
3560
3561        fn trade(index: u64) -> MutationBatch {
3562            MutationBatch::with_slot_context(
3563                vec![Mutation {
3564                    export: "Trade".to_string(),
3565                    key: json!(format!("pool{}", index % 4)),
3566                    patch: json!({"trade": index}),
3567                    append: vec![],
3568                }]
3569                .into_iter()
3570                .collect(),
3571                SlotContext::new(100 + index / 3, index % 3),
3572            )
3573        }
3574
3575        struct Harness {
3576            addr: SocketAddr,
3577            journal: Arc<EventJournal>,
3578            tx: mpsc::Sender<MutationBatch>,
3579        }
3580
3581        impl Harness {
3582            async fn start() -> Self {
3583                let view_index = Arc::new(append_index());
3584                let entity_cache = EntityCache::new();
3585                let bus_manager = BusManager::new();
3586                let journal = Arc::new(EventJournal::new(JournalConfig {
3587                    enabled: true,
3588                    max_bytes_per_view: u64::MAX,
3589                    max_records_per_view: 10_000,
3590                    max_age: Duration::from_secs(3_600),
3591                }));
3592
3593                let (tx, rx) = mpsc::channel::<MutationBatch>(256);
3594                tokio::spawn(
3595                    Projector::new(
3596                        view_index.clone(),
3597                        bus_manager.clone(),
3598                        entity_cache.clone(),
3599                        rx,
3600                        #[cfg(feature = "otel")]
3601                        None,
3602                    )
3603                    .with_journal(journal.clone())
3604                    .run(),
3605                );
3606
3607                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3608                let addr = listener.local_addr().unwrap();
3609                let server = WebSocketServer::new(
3610                    addr,
3611                    bus_manager,
3612                    entity_cache,
3613                    view_index,
3614                    #[cfg(feature = "otel")]
3615                    None,
3616                )
3617                .with_journal(journal.clone());
3618                let (acceptor, _cleanup) = server.into_acceptor();
3619                tokio::spawn(async move { acceptor.serve_listener(listener).await });
3620
3621                Self { addr, journal, tx }
3622            }
3623
3624            async fn publish(&self, range: std::ops::Range<u64>) {
3625                for index in range {
3626                    self.tx.send(trade(index)).await.unwrap();
3627                }
3628                let (ack, wait) = oneshot::channel();
3629                self.tx
3630                    .send(MutationBatch::flush_marker(ack))
3631                    .await
3632                    .unwrap();
3633                wait.await.unwrap();
3634            }
3635
3636            async fn connect(&self) -> WebSocketStream<TcpStream> {
3637                let stream = TcpStream::connect(self.addr).await.unwrap();
3638                client_async(format!("ws://{}/", self.addr), stream)
3639                    .await
3640                    .unwrap()
3641                    .0
3642            }
3643        }
3644
3645        async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
3646            loop {
3647                let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
3648                    .await
3649                    .expect("the server answers within the timeout")
3650                    .expect("the stream stays open")
3651                    .expect("a readable frame");
3652                // Control and data frames arrive as binary; issue frames as
3653                // text. Both are JSON.
3654                let bytes = match &message {
3655                    Message::Text(text) => text.as_bytes(),
3656                    Message::Binary(bytes) => bytes.as_ref(),
3657                    _ => continue,
3658                };
3659                return serde_json::from_slice(bytes).expect("frames are JSON");
3660            }
3661        }
3662
3663        /// Collect `count` event frames, ignoring anything else on the wire.
3664        async fn collect_trades(socket: &mut WebSocketStream<TcpStream>, count: usize) -> Vec<u64> {
3665            let mut offsets = Vec::with_capacity(count);
3666            while offsets.len() < count {
3667                let frame = next_frame(socket).await;
3668                assert_ne!(
3669                    frame["type"], "error",
3670                    "no error frame should interrupt delivery: {frame}"
3671                );
3672                if let Some(offset) = frame["offset"].as_u64() {
3673                    offsets.push(offset);
3674                }
3675            }
3676            offsets
3677        }
3678
3679        /// The headline claim: reconnecting with a stored cursor delivers every
3680        /// event published since it, in order, and then continues live without
3681        /// a duplicate or a hole at the seam.
3682        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3683        async fn a_reconnect_replays_from_a_cursor_and_continues_live() {
3684            let harness = Harness::start().await;
3685            harness.publish(0..RETAINED).await;
3686
3687            let cursor = harness.journal.window("Trade/append").await;
3688            let stored = format!("{}:{}", cursor.epoch, 99);
3689
3690            let mut socket = harness.connect().await;
3691            socket
3692                .send(Message::Text(
3693                    json!({
3694                        "type": "subscribe",
3695                        "protocolVersion": 2,
3696                        "subscriptionId": "trades",
3697                        "query": {"view": "Trade/append", "after": stored},
3698                    })
3699                    .to_string()
3700                    .into(),
3701                ))
3702                .await
3703                .unwrap();
3704
3705            let ack = next_frame(&mut socket).await;
3706            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3707            assert_eq!(ack["replayWindow"]["next"], json!(RETAINED));
3708
3709            // Well over the 500 a single read or buffer would cover.
3710            let replayed = collect_trades(&mut socket, (RETAINED - 100) as usize).await;
3711            assert_eq!(
3712                replayed,
3713                (100..RETAINED).collect::<Vec<_>>(),
3714                "every event after the cursor, in order, exactly once"
3715            );
3716
3717            // Published only now, so these can only arrive over the live path.
3718            harness.publish(RETAINED..RETAINED + 40).await;
3719            let live = collect_trades(&mut socket, 40).await;
3720            assert_eq!(
3721                live,
3722                (RETAINED..RETAINED + 40).collect::<Vec<_>>(),
3723                "the live stream resumes exactly where the replay stopped"
3724            );
3725
3726            socket.close(None).await.ok();
3727        }
3728
3729        /// Events published *during* the replay must still arrive. The replay
3730        /// and the live subscription are separate reads of the same tape, and
3731        /// the seam between them is where a naive implementation drops or
3732        /// repeats.
3733        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3734        async fn events_published_during_a_replay_are_not_lost() {
3735            let harness = Harness::start().await;
3736            harness.publish(0..RETAINED).await;
3737
3738            let epoch = harness.journal.window("Trade/append").await.epoch;
3739            let mut socket = harness.connect().await;
3740            socket
3741                .send(Message::Text(
3742                    json!({
3743                        "type": "subscribe",
3744                        "protocolVersion": 2,
3745                        "subscriptionId": "trades",
3746                        "query": {"view": "Trade/append", "after": format!("{epoch}:0")},
3747                    })
3748                    .to_string()
3749                    .into(),
3750                ))
3751                .await
3752                .unwrap();
3753            let ack = next_frame(&mut socket).await;
3754            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3755
3756            // Keep publishing while the replay is still draining.
3757            harness.publish(RETAINED..RETAINED + 200).await;
3758
3759            let total = (RETAINED + 200 - 1) as usize;
3760            let delivered = collect_trades(&mut socket, total).await;
3761            assert_eq!(
3762                delivered,
3763                (1..RETAINED + 200).collect::<Vec<_>>(),
3764                "replay and live output join without a gap or a repeat"
3765            );
3766
3767            socket.close(None).await.ok();
3768        }
3769
3770        /// A cursor from another tape lifetime is refused rather than served
3771        /// as a continuation, and the refusal releases the subscription id.
3772        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3773        async fn a_stale_epoch_is_refused_and_frees_the_subscription_id() {
3774            let harness = Harness::start().await;
3775            harness.publish(0..50).await;
3776
3777            let mut socket = harness.connect().await;
3778            socket
3779                .send(Message::Text(
3780                    json!({
3781                        "type": "subscribe",
3782                        "protocolVersion": 2,
3783                        "subscriptionId": "trades",
3784                        "query": {
3785                            "view": "Trade/append",
3786                            "after": format!("{}:10", crate::journal::JournalEpoch::new()),
3787                        },
3788                    })
3789                    .to_string()
3790                    .into(),
3791                ))
3792                .await
3793                .unwrap();
3794
3795            let error = next_frame(&mut socket).await;
3796            assert_eq!(error["type"], "error", "unexpected frame: {error}");
3797            assert_eq!(error["code"], "cursor-epoch-changed");
3798
3799            // The documented recovery is to resubscribe without a cursor. That
3800            // only works if the refused attempt released the id.
3801            socket
3802                .send(Message::Text(
3803                    json!({
3804                        "type": "subscribe",
3805                        "protocolVersion": 2,
3806                        "subscriptionId": "trades",
3807                        "query": {"view": "Trade/append"},
3808                    })
3809                    .to_string()
3810                    .into(),
3811                ))
3812                .await
3813                .unwrap();
3814            let ack = next_frame(&mut socket).await;
3815            assert_eq!(
3816                ack["op"], "subscribed",
3817                "the refused id must be reusable: {ack}"
3818            );
3819
3820            assert_eq!(collect_trades(&mut socket, 50).await.len(), 50);
3821            socket.close(None).await.ok();
3822        }
3823
3824        const ROUND: &str = "Round/state";
3825
3826        /// A state view with no projector: the test writes the cache and the
3827        /// bus itself, so it controls exactly when the subscriber can run.
3828        #[derive(Clone)]
3829        struct StateServer {
3830            addr: SocketAddr,
3831            bus_manager: BusManager,
3832            entity_cache: EntityCache,
3833        }
3834
3835        impl StateServer {
3836            async fn start() -> Self {
3837                let mut index = ViewIndex::new();
3838                index.add_spec(ViewSpec {
3839                    id: ROUND.to_string(),
3840                    export: "Round".to_string(),
3841                    mode: Mode::State,
3842                    wire_format: Default::default(),
3843                    projection: Projection::all(),
3844                    filters: Filters::all(),
3845                    delivery: Delivery::default(),
3846                    pipeline: None,
3847                    source_view: None,
3848                });
3849                let bus_manager = BusManager::new();
3850                let entity_cache = EntityCache::new();
3851                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3852                let addr = listener.local_addr().unwrap();
3853                let server = WebSocketServer::new(
3854                    addr,
3855                    bus_manager.clone(),
3856                    entity_cache.clone(),
3857                    Arc::new(index),
3858                    #[cfg(feature = "otel")]
3859                    None,
3860                );
3861                let (acceptor, _cleanup) = server.into_acceptor();
3862                tokio::spawn(async move { acceptor.serve_listener(listener).await });
3863                Self {
3864                    addr,
3865                    bus_manager,
3866                    entity_cache,
3867                }
3868            }
3869
3870            /// Publish a patch the way the projector does: stamp its `_seq`,
3871            /// write the cache, then the key's bus.
3872            async fn publish(&self, key: &str, mut patch: Value, seq: &str) {
3873                patch["_seq"] = Value::String(seq.to_string());
3874                self.entity_cache
3875                    .upsert_with_append(ROUND, key, patch.clone(), &[], PatchOrigin::Unknown)
3876                    .await;
3877                self.publish_frame(key, "patch", patch, seq).await;
3878            }
3879
3880            /// Publish a frame to the key's bus without touching the cache.
3881            async fn publish_frame(&self, key: &str, op: &str, data: Value, seq: &str) {
3882                let frame = json!({
3883                    "mode": "state",
3884                    "entity": ROUND,
3885                    "op": op,
3886                    "key": key,
3887                    "data": data,
3888                    "seq": seq,
3889                });
3890                self.bus_manager
3891                    .publish_state(ROUND, key, Arc::new(Bytes::from(frame.to_string())))
3892                    .await;
3893            }
3894
3895            /// Subscribe to one key and read through its snapshot.
3896            async fn subscribe(&self, key: &str) -> WebSocketStream<TcpStream> {
3897                let stream = TcpStream::connect(self.addr).await.unwrap();
3898                let mut socket = client_async(format!("ws://{}/", self.addr), stream)
3899                    .await
3900                    .unwrap()
3901                    .0;
3902                socket
3903                    .send(Message::Text(
3904                        json!({
3905                            "type": "subscribe",
3906                            "protocolVersion": 2,
3907                            "subscriptionId": "round",
3908                            "query": {"view": ROUND, "key": key},
3909                        })
3910                        .to_string()
3911                        .into(),
3912                    ))
3913                    .await
3914                    .unwrap();
3915                let ack = next_frame(&mut socket).await;
3916                assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3917                loop {
3918                    let frame = next_frame(&mut socket).await;
3919                    assert_eq!(frame["op"], "snapshot", "unexpected frame: {frame}");
3920                    if frame["complete"] == true {
3921                        return socket;
3922                    }
3923                }
3924            }
3925        }
3926
3927        /// Two patches to one key, published before the subscriber reads the
3928        /// first, must both reach it. The bus holds only the latest frame, so
3929        /// forwarding it alone drops the fields of the one it replaced.
3930        // One worker, so tasks run one at a time. The handshake needs the
3931        // multi-threaded runtime.
3932        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
3933        async fn a_state_subscriber_keeps_the_fields_of_an_overwritten_patch() {
3934            let server = StateServer::start().await;
3935            server
3936                .publish("7", json!({"id": 7}), "100:000000000001")
3937                .await;
3938            let mut socket = server.subscribe("7").await;
3939
3940            // A lone patch is forwarded as it was published. This one is an
3941            // account update, numbered by its write version.
3942            server
3943                .publish(
3944                    "7",
3945                    json!({"results": {"slot_hash": "abc"}}),
3946                    "140:003836292257",
3947                )
3948                .await;
3949            let frame = next_frame(&mut socket).await;
3950            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3951            assert_eq!(frame["seq"], "140:003836292257");
3952
3953            // Two instruction patches from later in the same slot, numbered by
3954            // transaction index. Both writes happen in one task that never
3955            // yields to the subscriber, so the second replaces the first on
3956            // the bus before the subscriber wakes.
3957            let writer = server.clone();
3958            tokio::spawn(async move {
3959                writer
3960                    .publish("7", json!({"entropy": {"seed": "def"}}), "140:000000001200")
3961                    .await;
3962                writer
3963                    .publish("7", json!({"total": 1}), "140:000000001200")
3964                    .await;
3965            })
3966            .await
3967            .unwrap();
3968            let frame = next_frame(&mut socket).await;
3969            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3970            assert_eq!(
3971                frame["data"],
3972                json!({
3973                    "id": 7,
3974                    "results": {"slot_hash": "abc"},
3975                    "entropy": {"seed": "def"},
3976                    "total": 1,
3977                    "_seq": "140:000000001200",
3978                })
3979            );
3980            // The latest seq sorts below the one already delivered, and a
3981            // client would drop the frame as stale.
3982            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
3983
3984            // Having caught up, the subscriber forwards patches again.
3985            server
3986                .publish("7", json!({"total": 2}), "141:000000000002")
3987                .await;
3988            let frame = next_frame(&mut socket).await;
3989            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3990            assert_eq!(
3991                frame["data"],
3992                json!({"total": 2, "_seq": "141:000000000002"})
3993            );
3994
3995            socket.close(None).await.ok();
3996        }
3997
3998        /// A key deleted and created again before the subscriber reads
3999        /// either frame: the bus keeps only the new entity's. Merged into the
4000        /// client's copy it would leave the deleted entity's fields in place,
4001        /// so the client gets the new entity whole, replacing its copy.
4002        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4003        async fn a_key_deleted_and_recreated_between_reads_replaces_the_clients_copy() {
4004            let server = StateServer::start().await;
4005            server
4006                .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
4007                .await;
4008            let mut socket = server.subscribe("7").await;
4009
4010            let writer = server.clone();
4011            tokio::spawn(async move {
4012                writer.entity_cache.remove(ROUND, "7").await;
4013                writer
4014                    .publish_frame("7", "delete", Value::Null, "101:000000000001")
4015                    .await;
4016                writer
4017                    .publish("7", json!({"id": 7, "fresh": true}), "101:000000000002")
4018                    .await;
4019            })
4020            .await
4021            .unwrap();
4022            let frame = next_frame(&mut socket).await;
4023            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4024            assert_eq!(
4025                frame["data"],
4026                json!({"id": 7, "fresh": true, "_seq": "101:000000000002"})
4027            );
4028            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4029
4030            // Deleted, with a patch overwritten before it: the delete is final.
4031            let writer = server.clone();
4032            tokio::spawn(async move {
4033                writer
4034                    .publish("7", json!({"n": 1}), "102:000000000001")
4035                    .await;
4036                writer.entity_cache.remove(ROUND, "7").await;
4037                writer
4038                    .publish_frame("7", "delete", Value::Null, "102:000000000002")
4039                    .await;
4040            })
4041            .await
4042            .unwrap();
4043            let frame = next_frame(&mut socket).await;
4044            assert_eq!(frame["op"], "delete", "unexpected frame: {frame}");
4045
4046            // Created again: the client holds nothing, so it arrives whole.
4047            server
4048                .publish("7", json!({"id": 7, "third": true}), "103:000000000001")
4049                .await;
4050            let frame = next_frame(&mut socket).await;
4051            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4052            assert_eq!(frame["data"]["third"], true);
4053            assert!(frame["data"].get("fresh").is_none(), "{frame}");
4054            socket.close(None).await.ok();
4055        }
4056
4057        /// A missed delete, then a patch the cache refused (it lacks the new
4058        /// entity): the client's copy is of the deleted entity, so the patch
4059        /// is not merged into it. The new entity replaces it once it is whole.
4060        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4061        async fn a_copy_of_a_deleted_entity_takes_no_patch_until_it_is_replaced() {
4062            let server = StateServer::start().await;
4063            server
4064                .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
4065                .await;
4066            let mut socket = server.subscribe("7").await;
4067
4068            let writer = server.clone();
4069            tokio::spawn(async move {
4070                writer.entity_cache.remove(ROUND, "7").await;
4071                writer
4072                    .publish_frame("7", "delete", Value::Null, "101:000000000001")
4073                    .await;
4074                writer
4075                    .publish_frame("7", "patch", json!({"n": 1}), "101:000000000002")
4076                    .await;
4077            })
4078            .await
4079            .unwrap();
4080            // Let the subscriber read the patch, and withhold it, before the
4081            // whole entity arrives. Nothing is sent, so there is no frame to
4082            // wait for.
4083            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
4084
4085            let whole = json!({"id": 7, "n": 1, "_seq": "101:000000000002"});
4086            server
4087                .entity_cache
4088                .store_whole(ROUND, "7", whole.clone())
4089                .await;
4090            server
4091                .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
4092                .await;
4093            let frame = next_frame(&mut socket).await;
4094            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4095            assert_eq!(frame["data"], whole);
4096            socket.close(None).await.ok();
4097        }
4098
4099        /// The source replaced the entity whole (an `upsert`, dropping a field)
4100        /// and then patched it, both before the subscriber read either. The
4101        /// patch alone, or the entity merged in, would keep the dropped field.
4102        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4103        async fn a_whole_entity_among_overwritten_frames_replaces_the_clients_copy() {
4104            let server = StateServer::start().await;
4105            server
4106                .publish("7", json!({"id": 7, "dropped": 1}), "100:000000000001")
4107                .await;
4108            let mut socket = server.subscribe("7").await;
4109
4110            let writer = server.clone();
4111            tokio::spawn(async move {
4112                let whole = json!({"id": 7, "_seq": "101:000000000001"});
4113                writer
4114                    .entity_cache
4115                    .store_whole(ROUND, "7", whole.clone())
4116                    .await;
4117                writer
4118                    .publish_frame("7", "upsert", whole, "101:000000000001")
4119                    .await;
4120                writer
4121                    .publish("7", json!({"n": 1}), "101:000000000002")
4122                    .await;
4123            })
4124            .await
4125            .unwrap();
4126            let frame = next_frame(&mut socket).await;
4127            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4128            assert_eq!(
4129                frame["data"],
4130                json!({"id": 7, "n": 1, "_seq": "101:000000000002"})
4131            );
4132            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4133            socket.close(None).await.ok();
4134        }
4135
4136        /// Patches overwritten on the bus while the cache lacks the key (it
4137        /// evicted it, and refused them) leave nothing whole to send in their
4138        /// place: the holder gets the latest patch, as a holder of an evicted
4139        /// key does, and stays behind. The resend that brings the entity back
4140        /// carries the seq of the key's latest change, which the holder
4141        /// already has and would drop as stale, so it gets the cached entity
4142        /// without a seq instead: as an upsert, since the resend is the whole
4143        /// entity and replaces what the holder has.
4144        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4145        async fn a_holder_behind_on_an_evicted_key_catches_up_when_it_returns() {
4146            let server = StateServer::start().await;
4147            server
4148                .publish("7", json!({"id": 7}), "100:000000000001")
4149                .await;
4150            let mut socket = server.subscribe("7").await;
4151            server.entity_cache.remove(ROUND, "7").await;
4152
4153            let writer = server.clone();
4154            tokio::spawn(async move {
4155                writer
4156                    .publish_frame("7", "patch", json!({"a": 1}), "101:000000000001")
4157                    .await;
4158                writer
4159                    .publish_frame("7", "patch", json!({"b": 2}), "101:000000000002")
4160                    .await;
4161            })
4162            .await
4163            .unwrap();
4164            let frame = next_frame(&mut socket).await;
4165            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4166            assert_eq!(frame["data"], json!({"b": 2}));
4167            assert_eq!(frame["seq"], "101:000000000002");
4168
4169            // The resend: the whole entity, at the latest change's position.
4170            let whole = json!({"id": 7, "a": 1, "b": 2, "_seq": "101:000000000002"});
4171            server
4172                .entity_cache
4173                .store_whole(ROUND, "7", whole.clone())
4174                .await;
4175            server
4176                .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
4177                .await;
4178            let frame = next_frame(&mut socket).await;
4179            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4180            assert_eq!(frame["data"], whole);
4181            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4182
4183            // Caught up: patches are forwarded again.
4184            server
4185                .publish("7", json!({"c": 3}), "102:000000000001")
4186                .await;
4187            let frame = next_frame(&mut socket).await;
4188            assert_eq!(frame["data"], json!({"c": 3, "_seq": "102:000000000001"}));
4189            assert_eq!(frame["seq"], "102:000000000001");
4190            socket.close(None).await.ok();
4191        }
4192    }
4193
4194    /// A key only counts as held by the client once a frame actually carried
4195    /// it. A truncated or disabled snapshot leaves keys the client never saw,
4196    /// and their first change must arrive whole.
4197    mod partial_entities_over_a_socket {
4198        use super::*;
4199        use crate::projector::Projector;
4200        use crate::{MutationBatch, SlotContext};
4201        use arete_interpreter::Mutation;
4202        use futures_util::{SinkExt, StreamExt};
4203        use std::time::Duration;
4204        use tokio::net::{TcpListener, TcpStream};
4205        use tokio::sync::mpsc;
4206        use tokio_tungstenite::tungstenite::Message;
4207        use tokio_tungstenite::{client_async, WebSocketStream};
4208
4209        fn thing_index() -> ViewIndex {
4210            let mut index = ViewIndex::new();
4211            index.add_spec(list_spec());
4212            index.add_spec(ViewSpec {
4213                id: "Thing/state".to_string(),
4214                mode: Mode::State,
4215                ..list_spec()
4216            });
4217            index
4218        }
4219
4220        struct Harness {
4221            addr: SocketAddr,
4222            tx: mpsc::Sender<MutationBatch>,
4223            slot: std::sync::atomic::AtomicU64,
4224        }
4225
4226        impl Harness {
4227            async fn start(delivery: WebSocketDeliveryConfig) -> Self {
4228                Self::start_with_cache(delivery, EntityCache::new()).await
4229            }
4230
4231            async fn start_with_cache(
4232                delivery: WebSocketDeliveryConfig,
4233                entity_cache: EntityCache,
4234            ) -> Self {
4235                let view_index = Arc::new(thing_index());
4236                let bus_manager = BusManager::new();
4237                let (tx, rx) = mpsc::channel::<MutationBatch>(64);
4238                // Unconstrained, so Tokio's cooperative budget never makes the
4239                // projector yield partway through a batch: a test on one
4240                // worker then knows no subscriber runs until the whole batch
4241                // is published (see
4242                // `snapshot_rows_patches_and_catch_ups_carry_the_version`).
4243                tokio::spawn(tokio::task::unconstrained(
4244                    Projector::new(
4245                        view_index.clone(),
4246                        bus_manager.clone(),
4247                        entity_cache.clone(),
4248                        rx,
4249                        #[cfg(feature = "otel")]
4250                        None,
4251                    )
4252                    .run(),
4253                ));
4254                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4255                let addr = listener.local_addr().unwrap();
4256                let server = WebSocketServer::new(
4257                    addr,
4258                    bus_manager,
4259                    entity_cache,
4260                    view_index,
4261                    #[cfg(feature = "otel")]
4262                    None,
4263                )
4264                .with_delivery_config(delivery);
4265                let (acceptor, _cleanup) = server.into_acceptor();
4266                tokio::spawn(async move { acceptor.serve_listener(listener).await });
4267                Self {
4268                    addr,
4269                    tx,
4270                    slot: std::sync::atomic::AtomicU64::new(100),
4271                }
4272            }
4273
4274            /// Apply one source patch and wait until the projector published it.
4275            async fn patch(&self, key: &str, patch: Value) {
4276                self.mutate(Mutation {
4277                    export: "Thing".to_string(),
4278                    key: json!(key),
4279                    patch,
4280                    append: vec![],
4281                })
4282                .await;
4283            }
4284
4285            async fn mutate(&self, mutation: Mutation) {
4286                let slot = self.slot.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
4287                self.tx
4288                    .send(MutationBatch::with_slot_context(
4289                        vec![mutation].into_iter().collect(),
4290                        SlotContext::new(slot, 0),
4291                    ))
4292                    .await
4293                    .unwrap();
4294                let (ack, wait) = oneshot::channel();
4295                self.tx
4296                    .send(MutationBatch::flush_marker(ack))
4297                    .await
4298                    .unwrap();
4299                wait.await.unwrap();
4300            }
4301
4302            async fn subscribe(&self, query: Value, snapshot: bool) -> WebSocketStream<TcpStream> {
4303                let stream = TcpStream::connect(self.addr).await.unwrap();
4304                let mut socket = client_async(format!("ws://{}/", self.addr), stream)
4305                    .await
4306                    .unwrap()
4307                    .0;
4308                socket
4309                    .send(Message::Text(
4310                        json!({
4311                            "type": "subscribe",
4312                            "protocolVersion": 2,
4313                            "subscriptionId": "things",
4314                            "query": query,
4315                            "snapshot": {"enabled": snapshot},
4316                        })
4317                        .to_string()
4318                        .into(),
4319                    ))
4320                    .await
4321                    .unwrap();
4322                let ack = next_frame(&mut socket).await;
4323                assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
4324                assert_eq!(
4325                    ack["wholeEntities"],
4326                    json!(true),
4327                    "the ack advertises the whole-entity guarantee: {ack}"
4328                );
4329                socket
4330            }
4331        }
4332
4333        async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
4334            loop {
4335                let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
4336                    .await
4337                    .expect("the server answers within the timeout")
4338                    .expect("the stream stays open")
4339                    .expect("a readable frame");
4340                // Every frame here is far below the compression threshold.
4341                let bytes = match &message {
4342                    Message::Text(text) => text.as_bytes(),
4343                    Message::Binary(bytes) => bytes.as_ref(),
4344                    _ => continue,
4345                };
4346                return serde_json::from_slice(bytes).expect("frames are JSON");
4347            }
4348        }
4349
4350        /// The snapshot rows, in order, once the snapshot completes.
4351        async fn snapshot_keys(socket: &mut WebSocketStream<TcpStream>) -> Vec<String> {
4352            let mut keys = Vec::new();
4353            loop {
4354                let frame = next_frame(socket).await;
4355                assert_eq!(frame["op"], "snapshot", "expected a snapshot: {frame}");
4356                keys.extend(
4357                    frame["data"]
4358                        .as_array()
4359                        .unwrap()
4360                        .iter()
4361                        .map(|row| row["key"].as_str().unwrap().to_string()),
4362                );
4363                if frame["complete"] == json!(true) {
4364                    return keys;
4365                }
4366            }
4367        }
4368
4369        async fn seed(harness: &Harness) {
4370            for key in ["a", "b", "c"] {
4371                harness
4372                    .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
4373                    .await;
4374            }
4375        }
4376
4377        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4378        async fn a_change_to_a_key_cut_by_snapshot_limit_is_an_upsert() {
4379            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4380            seed(&harness).await;
4381
4382            let mut socket = harness
4383                .subscribe(json!({"view": "Thing/list", "snapshotLimit": 1}), true)
4384                .await;
4385            // Newest first: only "c" was sent.
4386            assert_eq!(snapshot_keys(&mut socket).await, ["c"]);
4387
4388            harness.patch("a", json!({"count": 2})).await;
4389            let frame = next_frame(&mut socket).await;
4390            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4391            assert_eq!(frame["op"], "upsert", "a key never sent arrives whole");
4392            assert_eq!(frame["data"]["name"], "thing-a");
4393            assert_eq!(frame["data"]["count"], 2);
4394
4395            // Now the client holds "a", so its next change is a patch again.
4396            harness.patch("a", json!({"count": 3})).await;
4397            let frame = next_frame(&mut socket).await;
4398            assert_eq!(frame["key"], "a");
4399            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4400            assert!(frame["data"].get("name").is_none());
4401
4402            // A key the snapshot did carry still rides its patch.
4403            harness.patch("c", json!({"count": 2})).await;
4404            let frame = next_frame(&mut socket).await;
4405            assert_eq!(frame["key"], "c");
4406            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4407            socket.close(None).await.ok();
4408        }
4409
4410        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4411        async fn a_list_without_a_snapshot_upserts_each_key_on_its_first_change() {
4412            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4413            seed(&harness).await;
4414
4415            let mut socket = harness
4416                .subscribe(json!({"view": "Thing/list"}), false)
4417                .await;
4418            harness.patch("b", json!({"count": 2})).await;
4419            let frame = next_frame(&mut socket).await;
4420            assert_eq!(frame["key"], "b");
4421            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4422            assert_eq!(frame["data"]["name"], "thing-b");
4423
4424            harness.patch("b", json!({"count": 3})).await;
4425            let frame = next_frame(&mut socket).await;
4426            assert_eq!(frame["key"], "b");
4427            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4428            socket.close(None).await.ok();
4429        }
4430
4431        /// Eviction from the server's bounded cache is not a change to the
4432        /// entity: a client holding it keeps getting its patches, never a
4433        /// `remove`, and the whole entity replaces its copy when it returns.
4434        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4435        async fn a_state_holder_keeps_its_patches_when_the_cache_evicts_the_key() {
4436            let harness = Harness::start_with_cache(
4437                WebSocketDeliveryConfig::default(),
4438                EntityCache::with_config(crate::cache::EntityCacheConfig {
4439                    max_entities_per_view: 2,
4440                    ..Default::default()
4441                }),
4442            )
4443            .await;
4444            harness
4445                .patch("a", json!({"name": "thing-a", "count": 1}))
4446                .await;
4447            let mut socket = harness
4448                .subscribe(json!({"view": "Thing/state", "key": "a"}), true)
4449                .await;
4450            assert_eq!(snapshot_keys(&mut socket).await, ["a"]);
4451
4452            // "b" and "c" push "a" out of the cache.
4453            for key in ["b", "c"] {
4454                harness
4455                    .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
4456                    .await;
4457            }
4458            harness.patch("a", json!({"count": 2})).await;
4459            let frame = next_frame(&mut socket).await;
4460            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4461            assert_eq!(frame["op"], "patch", "the holder merges it: {frame}");
4462            assert_eq!(frame["data"]["count"], 2);
4463
4464            // The whole entity comes back from the VM.
4465            let mut whole = Mutation {
4466                export: "Thing".to_string(),
4467                key: json!("a"),
4468                patch: json!({"name": "thing-a", "count": 2}),
4469                append: vec![],
4470            };
4471            whole.mark_whole_entity();
4472            harness.mutate(whole).await;
4473            let frame = next_frame(&mut socket).await;
4474            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4475            assert_eq!(frame["data"]["name"], "thing-a");
4476            assert_eq!(frame["data"]["count"], 2);
4477            assert!(frame["data"]
4478                .get(arete_interpreter::WHOLE_ENTITY_MARKER)
4479                .is_none());
4480
4481            harness.patch("a", json!({"count": 3})).await;
4482            let frame = next_frame(&mut socket).await;
4483            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4484            socket.close(None).await.ok();
4485        }
4486
4487        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4488        async fn a_state_view_without_a_snapshot_upserts_on_the_first_change() {
4489            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4490            seed(&harness).await;
4491
4492            let mut socket = harness
4493                .subscribe(json!({"view": "Thing/state", "key": "a"}), false)
4494                .await;
4495            harness.patch("a", json!({"count": 2})).await;
4496            let frame = next_frame(&mut socket).await;
4497            assert_eq!(frame["key"], "a");
4498            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4499            assert_eq!(frame["data"]["name"], "thing-a");
4500            assert_eq!(frame["data"]["count"], 2);
4501
4502            harness.patch("a", json!({"count": 3})).await;
4503            let frame = next_frame(&mut socket).await;
4504            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4505            socket.close(None).await.ok();
4506        }
4507
4508        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4509        async fn a_coalesced_list_sends_no_remove_for_a_key_it_never_sent() {
4510            let harness = Harness::start(WebSocketDeliveryConfig {
4511                collection_coalesce_ms: Some(10),
4512                ..Default::default()
4513            })
4514            .await;
4515            seed(&harness).await;
4516
4517            // A one-row window holding "c", which the disabled snapshot never
4518            // sent.
4519            let mut socket = harness
4520                .subscribe(json!({"view": "Thing/list", "take": 1}), false)
4521                .await;
4522            // "a" jumps to the front: it enters the window whole and "c",
4523            // never sent, leaves it without a `remove`.
4524            harness.patch("a", json!({"count": 2})).await;
4525            let frame = next_frame(&mut socket).await;
4526            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4527            assert_eq!(frame["op"], "upsert");
4528            assert_eq!(frame["data"]["name"], "thing-a");
4529
4530            // "c" re-enters whole; "a", held, leaves with a `remove`.
4531            harness.patch("c", json!({"count": 2})).await;
4532            let mut frames = [next_frame(&mut socket).await, next_frame(&mut socket).await];
4533            frames.sort_by_key(|frame| frame["key"].as_str().unwrap().to_string());
4534            assert_eq!(frames[0]["key"], "a");
4535            assert_eq!(frames[0]["op"], "remove");
4536            assert_eq!(frames[1]["key"], "c");
4537            assert_eq!(frames[1]["op"], "upsert");
4538            assert_eq!(frames[1]["data"]["name"], "thing-c");
4539            socket.close(None).await.ok();
4540        }
4541
4542        /// `(epoch, counter)` from an entity's `_version`.
4543        fn version_of(data: &Value) -> (String, u64) {
4544            let version = data["_version"]
4545                .as_str()
4546                .unwrap_or_else(|| panic!("no _version in {data}"));
4547            let (epoch, counter) = version.split_once(':').expect("epoch:counter");
4548            (epoch.to_string(), counter.parse().expect("decimal counter"))
4549        }
4550
4551        /// Every entity a subscriber receives carries the projector's
4552        /// `_version`: a snapshot row, a forwarded patch, and a state
4553        /// subscription's catch-up, which sends the cached entity after the
4554        /// bus overwrote a patch.
4555        // One worker, and a projector the cooperative budget cannot interrupt:
4556        // within a batch it awaits only locks no other task holds, so it
4557        // publishes both changes before the subscriber can run, and the
4558        // second overwrites the first on the state bus.
4559        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4560        async fn snapshot_rows_patches_and_catch_ups_carry_the_version() {
4561            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4562            harness.patch("7", json!({"name": "thing-7"})).await;
4563
4564            let mut list = harness.subscribe(json!({"view": "Thing/list"}), true).await;
4565            let snapshot = next_frame(&mut list).await;
4566            assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
4567            let (epoch, _) = version_of(&snapshot["data"][0]["data"]);
4568
4569            let mut state = harness
4570                .subscribe(json!({"view": "Thing/state", "key": "7"}), true)
4571                .await;
4572            let snapshot = next_frame(&mut state).await;
4573            assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
4574            let (state_epoch, seeded) = version_of(&snapshot["data"][0]["data"]);
4575            assert_eq!(state_epoch, epoch, "one projector, one epoch");
4576
4577            harness.patch("7", json!({"count": 1})).await;
4578            let patch = next_frame(&mut list).await;
4579            assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
4580            assert_eq!(version_of(&patch["data"]).0, epoch);
4581            let patch = next_frame(&mut state).await;
4582            assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
4583            let (_, patched) = version_of(&patch["data"]);
4584            assert!(patched > seeded, "{patched} after {seeded}");
4585
4586            // Two changes to the key in one batch.
4587            let slot = harness
4588                .slot
4589                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
4590            let mutation = |patch: Value| Mutation {
4591                export: "Thing".to_string(),
4592                key: json!("7"),
4593                patch,
4594                append: vec![],
4595            };
4596            harness
4597                .tx
4598                .send(MutationBatch::with_slot_context(
4599                    vec![
4600                        mutation(json!({"count": 2})),
4601                        mutation(json!({"flag": true})),
4602                    ]
4603                    .into_iter()
4604                    .collect(),
4605                    SlotContext::new(slot, 0),
4606                ))
4607                .await
4608                .unwrap();
4609
4610            let catch_up = next_frame(&mut state).await;
4611            assert_eq!(catch_up["op"], "patch", "unexpected frame: {catch_up}");
4612            assert!(catch_up.get("seq").is_none(), "a catch-up: {catch_up}");
4613            assert_eq!(catch_up["data"]["count"], 2);
4614            assert_eq!(catch_up["data"]["flag"], true);
4615            let (_, caught_up) = version_of(&catch_up["data"]);
4616            assert!(caught_up > patched, "{caught_up} after {patched}");
4617
4618            list.close(None).await.ok();
4619            state.close(None).await.ok();
4620        }
4621    }
4622
4623    /// Session tokens that expire while their socket is open.
4624    mod session_expiry {
4625        use super::*;
4626        use crate::websocket::auth::SignedSessionAuthPlugin;
4627        use arete_auth::{KeyClass, SessionClaims, SigningKey, TokenSigner, TokenVerifier};
4628        use futures_util::{SinkExt, StreamExt};
4629        use std::time::Duration;
4630        use tokio::net::{TcpListener, TcpStream};
4631        use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
4632        use tokio_tungstenite::tungstenite::protocol::CloseFrame;
4633        use tokio_tungstenite::tungstenite::Message;
4634        use tokio_tungstenite::{client_async, WebSocketStream};
4635
4636        struct Server {
4637            addr: SocketAddr,
4638            signer: TokenSigner,
4639        }
4640
4641        impl Server {
4642            async fn start() -> Self {
4643                let signing_key = SigningKey::generate();
4644                let verifier =
4645                    TokenVerifier::new(signing_key.verifying_key(), "test-issuer", "test-audience");
4646                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4647                let addr = listener.local_addr().unwrap();
4648                let server = WebSocketServer::new(
4649                    addr,
4650                    BusManager::new(),
4651                    EntityCache::new(),
4652                    Arc::new(ViewIndex::new()),
4653                    #[cfg(feature = "otel")]
4654                    None,
4655                )
4656                .with_auth_plugin(Arc::new(SignedSessionAuthPlugin::new(verifier)));
4657                let (acceptor, _cleanup) = server.into_acceptor();
4658                tokio::spawn(async move { acceptor.serve_listener(listener).await });
4659                Self {
4660                    addr,
4661                    signer: TokenSigner::new(signing_key, "test-issuer"),
4662                }
4663            }
4664
4665            fn token(&self, ttl_seconds: u64) -> String {
4666                let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
4667                    .with_scope("read")
4668                    .with_key_class(KeyClass::Secret)
4669                    .with_ttl(ttl_seconds)
4670                    .build();
4671                self.signer.sign(claims).unwrap()
4672            }
4673
4674            async fn connect(&self, token: &str) -> WebSocketStream<TcpStream> {
4675                let stream = TcpStream::connect(self.addr).await.unwrap();
4676                client_async(format!("ws://{}/?hs_token={token}", self.addr), stream)
4677                    .await
4678                    .unwrap()
4679                    .0
4680            }
4681        }
4682
4683        async fn send_json(socket: &mut WebSocketStream<TcpStream>, message: Value) {
4684            socket
4685                .send(Message::Text(message.to_string().into()))
4686                .await
4687                .unwrap();
4688        }
4689
4690        /// The close frame, if the server closes the socket within `wait`.
4691        async fn close_within(
4692            socket: &mut WebSocketStream<TcpStream>,
4693            wait: Duration,
4694        ) -> Option<Option<CloseFrame>> {
4695            tokio::time::timeout(wait, async {
4696                while let Some(Ok(message)) = socket.next().await {
4697                    if let Message::Close(frame) = message {
4698                        return Some(frame);
4699                    }
4700                }
4701                Some(None)
4702            })
4703            .await
4704            .ok()
4705            .flatten()
4706        }
4707
4708        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4709        async fn an_expired_session_is_closed_with_the_reason() {
4710            let server = Server::start().await;
4711            let mut socket = server.connect(&server.token(2)).await;
4712
4713            tokio::time::sleep(Duration::from_secs(3)).await;
4714            send_json(&mut socket, json!({"type": "ping"})).await;
4715
4716            let frame = close_within(&mut socket, Duration::from_secs(5))
4717                .await
4718                .expect("the server closes the expired session")
4719                .expect("the close frame carries a reason");
4720            assert_eq!(frame.code, CloseCode::Policy);
4721            assert_eq!(
4722                frame.reason.as_str(),
4723                "token-expired: Authentication token expired"
4724            );
4725        }
4726
4727        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4728        async fn a_session_refreshed_in_band_outlives_its_first_token() {
4729            let server = Server::start().await;
4730            let mut socket = server.connect(&server.token(2)).await;
4731
4732            send_json(
4733                &mut socket,
4734                json!({"type": "refresh_auth", "token": server.token(3_600)}),
4735            )
4736            .await;
4737            let reply = tokio::time::timeout(Duration::from_secs(5), async {
4738                while let Some(Ok(message)) = socket.next().await {
4739                    if let Message::Text(text) = message {
4740                        return serde_json::from_str::<Value>(text.as_str()).ok();
4741                    }
4742                }
4743                None
4744            })
4745            .await
4746            .expect("the server answers the refresh")
4747            .expect("the answer is JSON");
4748            assert_eq!(reply["success"], true, "refresh accepted: {reply}");
4749
4750            tokio::time::sleep(Duration::from_secs(3)).await;
4751            send_json(&mut socket, json!({"type": "ping"})).await;
4752            assert!(
4753                close_within(&mut socket, Duration::from_millis(1_500))
4754                    .await
4755                    .is_none(),
4756                "the socket stays open on the refreshed token"
4757            );
4758        }
4759    }
4760}