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::row_scope::{session_row_scopes, RowFilter};
5use crate::shared_entity::{lookup, EntityFields, SharedEntity};
6use crate::view::{ViewIndex, ViewSpec};
7use crate::websocket::admission::{WebSocketAdmissionProvider, WebSocketConnectionPermit};
8use crate::websocket::auth::{
9    AuthContext, AuthDecision, AuthDeny, ConnectionAuthRequest, WebSocketAuthPlugin,
10};
11use crate::websocket::client_manager::{ClientManager, RateLimitConfig};
12use crate::websocket::frame::{
13    apply_wire_format, Frame, Mode, SortConfig, SortOrder, SourceFields, SubscribedFrame,
14    UnsubscribedFrame, WireEntity, WireFormat,
15};
16use crate::websocket::subscription::{
17    ClientMessage, RefreshAuthRequest, RefreshAuthResponse, SocketIssueMessage, Subscription,
18    SubscriptionQuery, Unsubscription, PROTOCOL_VERSION,
19};
20use crate::websocket::usage::{UsageIdentity, WebSocketUsageEmitter, WebSocketUsageEvent};
21use crate::{WebSocketDeliveryConfig, WebSocketLimitsConfig};
22use anyhow::Result;
23use bytes::Bytes;
24use futures_util::StreamExt;
25use serde::Serialize;
26use serde_json::Value;
27use std::collections::{HashMap, HashSet, VecDeque};
28use std::future::Future;
29use std::net::SocketAddr;
30use std::sync::atomic::{AtomicUsize, Ordering};
31use std::sync::Arc;
32use std::time::{Duration, Instant};
33use tokio::net::{TcpListener, TcpStream};
34use tokio::sync::{broadcast, watch};
35use tokio_tungstenite::{
36    accept_hdr_async_with_config,
37    tungstenite::{
38        handshake::server::{ErrorResponse as HandshakeErrorResponse, Request, Response},
39        http::{header::CONTENT_TYPE, StatusCode},
40        protocol::WebSocketConfig,
41        Error as WsError,
42    },
43};
44use tokio_util::sync::CancellationToken;
45use tokio_util::task::TaskTracker;
46use tracing::{debug, error, info, info_span, warn, Instrument};
47use uuid::Uuid;
48
49#[cfg(feature = "otel")]
50use crate::metrics::Metrics;
51
52#[derive(Clone, Default)]
53struct WsMetrics {
54    #[cfg(feature = "otel")]
55    inner: Option<Arc<Metrics>>,
56    #[cfg(test)]
57    probe: Option<Arc<DeliveryProbe>>,
58}
59
60/// Test-only mirror of the delivery instruments recorded through
61/// [`WsMetrics`], so the real-socket load harness can report them without the
62/// `otel` feature. Each counter has the meaning of the metric named beside it.
63#[cfg(test)]
64#[derive(Debug, Default)]
65pub(crate) struct DeliveryProbe {
66    /// `arete.ws.connections.active`, by metering key
67    pub(crate) active_connections: std::sync::Mutex<std::collections::BTreeMap<String, i64>>,
68    /// `arete.ws.subscriptions.active`, by view and metering key
69    pub(crate) active_subscriptions: std::sync::Mutex<std::collections::BTreeMap<String, i64>>,
70    /// `arete.ws.messages.sent`
71    pub(crate) messages_sent: std::sync::atomic::AtomicU64,
72    /// `arete.ws.subscription.lagged`
73    pub(crate) lag_events: std::sync::atomic::AtomicU64,
74    /// `arete.ws.subscription.dropped_updates`
75    pub(crate) dropped_updates: std::sync::atomic::AtomicU64,
76    /// `arete.ws.subscription.resnapshots`
77    pub(crate) resnapshots: std::sync::atomic::AtomicU64,
78    /// `arete.ws.collection.coalesced_updates`
79    pub(crate) coalesced_updates: std::sync::atomic::AtomicU64,
80    /// `arete.ws.collection.coalesced_flushes`
81    pub(crate) coalesced_flushes: std::sync::atomic::AtomicU64,
82    /// `arete.ws.delivery.stopped`, by reason
83    pub(crate) delivery_stopped: std::sync::Mutex<std::collections::BTreeMap<&'static str, u64>>,
84}
85
86#[cfg(test)]
87impl DeliveryProbe {
88    fn add(counter: &std::sync::atomic::AtomicU64, value: u64) {
89        counter.fetch_add(value, std::sync::atomic::Ordering::Relaxed);
90    }
91
92    fn add_active(
93        counts: &std::sync::Mutex<std::collections::BTreeMap<String, i64>>,
94        key: String,
95        delta: i64,
96    ) {
97        *counts
98            .lock()
99            .expect("delivery probe lock poisoned")
100            .entry(key)
101            .or_default() += delta;
102    }
103}
104
105impl WsMetrics {
106    #[cfg(feature = "otel")]
107    fn new(inner: Option<Arc<Metrics>>) -> Self {
108        Self {
109            inner,
110            #[cfg(test)]
111            probe: None,
112        }
113    }
114
115    #[cfg(test)]
116    fn probe(&self, record: impl FnOnce(&DeliveryProbe)) {
117        if let Some(probe) = &self.probe {
118            record(probe);
119        }
120    }
121
122    fn connection_opened(&self, metering_key: Option<&str>) {
123        #[cfg(test)]
124        self.probe(|probe| {
125            DeliveryProbe::add_active(
126                &probe.active_connections,
127                metering_key.unwrap_or("<none>").to_string(),
128                1,
129            )
130        });
131        #[cfg(not(feature = "otel"))]
132        let _ = metering_key;
133        #[cfg(feature = "otel")]
134        if let Some(metrics) = &self.inner {
135            if let Some(metering_key) = metering_key {
136                metrics.record_ws_connection_with_metering(metering_key);
137            } else {
138                metrics.record_ws_connection();
139            }
140        }
141    }
142
143    fn connection_closed(&self, duration_secs: f64, metering_key: Option<&str>) {
144        #[cfg(test)]
145        self.probe(|probe| {
146            DeliveryProbe::add_active(
147                &probe.active_connections,
148                metering_key.unwrap_or("<none>").to_string(),
149                -1,
150            )
151        });
152        #[cfg(not(feature = "otel"))]
153        let _ = (duration_secs, metering_key);
154        #[cfg(feature = "otel")]
155        if let Some(metrics) = &self.inner {
156            if let Some(metering_key) = metering_key {
157                metrics.record_ws_disconnection_with_metering(duration_secs, metering_key);
158            } else {
159                metrics.record_ws_disconnection(duration_secs);
160            }
161        }
162    }
163
164    fn message_received(&self, metering_key: Option<&str>) {
165        #[cfg(not(feature = "otel"))]
166        let _ = metering_key;
167        #[cfg(feature = "otel")]
168        if let Some(metrics) = &self.inner {
169            if let Some(metering_key) = metering_key {
170                metrics.record_ws_message_received_with_metering(metering_key);
171            } else {
172                metrics.record_ws_message_received();
173            }
174        }
175    }
176
177    fn message_sent(&self) {
178        #[cfg(test)]
179        self.probe(|probe| DeliveryProbe::add(&probe.messages_sent, 1));
180        #[cfg(feature = "otel")]
181        if let Some(metrics) = &self.inner {
182            metrics.record_ws_message_sent();
183        }
184    }
185
186    fn subscription_created(&self, view: &str, metering_key: Option<&str>) {
187        #[cfg(test)]
188        self.probe(|probe| {
189            DeliveryProbe::add_active(
190                &probe.active_subscriptions,
191                format!("{view}|{}", metering_key.unwrap_or("<none>")),
192                1,
193            )
194        });
195        #[cfg(not(feature = "otel"))]
196        let _ = (view, metering_key);
197        #[cfg(feature = "otel")]
198        if let Some(metrics) = &self.inner {
199            if let Some(metering_key) = metering_key {
200                metrics.record_subscription_created_with_metering(view, metering_key);
201            } else {
202                metrics.record_subscription_created(view);
203            }
204        }
205    }
206
207    fn subscription_removed(&self, view: &str, metering_key: Option<&str>) {
208        #[cfg(test)]
209        self.probe(|probe| {
210            DeliveryProbe::add_active(
211                &probe.active_subscriptions,
212                format!("{view}|{}", metering_key.unwrap_or("<none>")),
213                -1,
214            )
215        });
216        #[cfg(not(feature = "otel"))]
217        let _ = (view, metering_key);
218        #[cfg(feature = "otel")]
219        if let Some(metrics) = &self.inner {
220            if let Some(metering_key) = metering_key {
221                metrics.record_subscription_removed_with_metering(view, metering_key);
222            } else {
223                metrics.record_subscription_removed(view);
224            }
225        }
226    }
227
228    fn protocol_error(&self, code: &str) {
229        #[cfg(not(feature = "otel"))]
230        let _ = code;
231        #[cfg(feature = "otel")]
232        if let Some(metrics) = &self.inner {
233            metrics.record_ws_protocol_error(code);
234        }
235    }
236
237    fn subscription_lagged(&self, view: &str, skipped: u64) {
238        #[cfg(test)]
239        self.probe(|probe| {
240            DeliveryProbe::add(&probe.lag_events, 1);
241            DeliveryProbe::add(&probe.dropped_updates, skipped);
242        });
243        #[cfg(not(feature = "otel"))]
244        let _ = (view, skipped);
245        #[cfg(feature = "otel")]
246        if let Some(metrics) = &self.inner {
247            metrics.record_ws_subscription_lagged(view, skipped);
248        }
249    }
250
251    fn subscription_resnapshot(&self, view: &str) {
252        #[cfg(test)]
253        self.probe(|probe| DeliveryProbe::add(&probe.resnapshots, 1));
254        #[cfg(not(feature = "otel"))]
255        let _ = view;
256        #[cfg(feature = "otel")]
257        if let Some(metrics) = &self.inner {
258            metrics.record_ws_subscription_resnapshot(view);
259        }
260    }
261
262    fn collection_coalesced(&self, view: &str, updates: u64) {
263        #[cfg(test)]
264        self.probe(|probe| {
265            DeliveryProbe::add(&probe.coalesced_updates, updates);
266            DeliveryProbe::add(&probe.coalesced_flushes, 1);
267        });
268        #[cfg(not(feature = "otel"))]
269        let _ = (view, updates);
270        #[cfg(feature = "otel")]
271        if let Some(metrics) = &self.inner {
272            metrics.record_ws_collection_coalesced(view, updates);
273        }
274    }
275
276    fn delivery_stopped(&self, view: &str, reason: &'static str) {
277        #[cfg(test)]
278        self.probe(|probe| {
279            *probe
280                .delivery_stopped
281                .lock()
282                .expect("delivery probe lock poisoned")
283                .entry(reason)
284                .or_default() += 1;
285        });
286        #[cfg(not(feature = "otel"))]
287        let _ = (view, reason);
288        #[cfg(feature = "otel")]
289        if let Some(metrics) = &self.inner {
290            metrics.record_ws_delivery_stopped(view, reason);
291        }
292    }
293}
294
295async fn handle_refresh_auth(
296    client_id: Uuid,
297    refresh_req: &RefreshAuthRequest,
298    client_manager: &ClientManager,
299    auth_plugin: &Arc<dyn WebSocketAuthPlugin>,
300) {
301    let refresh_result: Result<AuthContext, String> = if let Some(signed_plugin) = auth_plugin
302        .as_any()
303        .downcast_ref::<crate::websocket::auth::SignedSessionAuthPlugin>()
304    {
305        signed_plugin
306            .verify_refresh_token(&refresh_req.token)
307            .await
308            .map_err(|error| error.reason)
309    } else {
310        Err("In-band auth refresh not supported with current auth plugin".to_string())
311    };
312
313    let response = match refresh_result {
314        Ok(new_context) => {
315            let expires_at = new_context.expires_at;
316            match client_manager
317                .try_update_client_auth_async(client_id, new_context)
318                .await
319            {
320                Ok(true) => RefreshAuthResponse {
321                    success: true,
322                    error: None,
323                    expires_at: Some(expires_at),
324                },
325                Ok(false) => RefreshAuthResponse {
326                    success: false,
327                    error: Some("client-not-found".to_string()),
328                    expires_at: None,
329                },
330                Err(deny) => RefreshAuthResponse {
331                    success: false,
332                    error: Some(deny.code.as_str().to_string()),
333                    expires_at: None,
334                },
335            }
336        }
337        Err(error) => {
338            let code = if error.contains("expired") {
339                "token-expired"
340            } else if error.contains("signature") {
341                "token-invalid-signature"
342            } else if error.contains("issuer") {
343                "token-invalid-issuer"
344            } else if error.contains("audience") {
345                "token-invalid-audience"
346            } else {
347                "token-invalid"
348            };
349            RefreshAuthResponse {
350                success: false,
351                error: Some(code.to_string()),
352                expires_at: None,
353            }
354        }
355    };
356
357    if let Ok(json) = serde_json::to_string(&response) {
358        let _ = client_manager.send_text_to_client(client_id, json).await;
359    }
360}
361
362async fn send_socket_issue(
363    client_id: Uuid,
364    client_manager: &ClientManager,
365    deny: &AuthDeny,
366    fatal: bool,
367    subscription_id: Option<String>,
368) {
369    let message = SocketIssueMessage::from_auth_deny(deny, fatal, subscription_id);
370    if let Ok(json) = serde_json::to_string(&message) {
371        let _ = client_manager.send_text_to_client(client_id, json).await;
372    }
373}
374
375async fn send_protocol_issue(
376    client_id: Uuid,
377    client_manager: &ClientManager,
378    metrics: &WsMetrics,
379    subscription_id: Option<String>,
380    code: &str,
381    message: impl Into<String>,
382) {
383    metrics.protocol_error(code);
384    let issue = SocketIssueMessage::protocol(subscription_id, code, message);
385    if let Ok(json) = serde_json::to_string(&issue) {
386        let _ = client_manager.send_text_to_client(client_id, json).await;
387    }
388}
389
390/// Send an already-built issue, for refusals that carry structured detail.
391async fn send_prepared_issue(
392    client_id: Uuid,
393    client_manager: &ClientManager,
394    metrics: &WsMetrics,
395    issue: SocketIssueMessage,
396) {
397    metrics.protocol_error(&issue.code);
398    if let Ok(json) = serde_json::to_string(&issue) {
399        let _ = client_manager.send_text_to_client(client_id, json).await;
400    }
401}
402
403fn key_class_label(key_class: arete_auth::KeyClass) -> &'static str {
404    match key_class {
405        arete_auth::KeyClass::Secret => "secret",
406        arete_auth::KeyClass::Publishable => "publishable",
407    }
408}
409
410#[derive(Clone)]
411struct UsageEmitterHandle {
412    emitter: Arc<dyn WebSocketUsageEmitter>,
413    tasks: TaskTracker,
414}
415
416impl UsageEmitterHandle {
417    fn new(emitter: Arc<dyn WebSocketUsageEmitter>) -> Self {
418        Self {
419            emitter,
420            tasks: TaskTracker::new(),
421        }
422    }
423
424    fn emit(&self, event: WebSocketUsageEvent) {
425        let emitter = self.emitter.clone();
426        self.tasks.spawn(async move {
427            emitter.emit(event).await;
428        });
429    }
430
431    async fn close_and_wait(&self) {
432        self.tasks.close();
433        self.tasks.wait().await;
434    }
435}
436
437fn emit_usage_event(usage_emitter: &Option<UsageEmitterHandle>, event: WebSocketUsageEvent) {
438    if let Some(emitter) = usage_emitter {
439        emitter.emit(event);
440    }
441}
442
443fn usage_identity_from_context(
444    auth_context: Option<&AuthContext>,
445) -> (UsageIdentity, Option<String>) {
446    match auth_context {
447        Some(context) => {
448            let v2 = !context.is_legacy_policy();
449            (
450                UsageIdentity {
451                    metering_key: Some(context.metering_key.clone()),
452                    subject: Some(context.subject.clone()),
453                    key_class: Some(key_class_label(context.key_class).to_string()),
454                    actor_key: v2.then(|| context.actor_key.clone()).flatten(),
455                    account_key: v2.then(|| context.account_key.clone()).flatten(),
456                    consumer_key: v2.then(|| context.consumer_key.clone()).flatten(),
457                    plan_code: v2.then(|| context.plan.clone()).flatten(),
458                    policy_version: v2.then_some(context.policy_version).flatten(),
459                },
460                context.deployment_id.clone(),
461            )
462        }
463        None => (UsageIdentity::default(), None),
464    }
465}
466
467fn emit_update_sent_for_client(
468    usage_emitter: &Option<UsageEmitterHandle>,
469    client_manager: &ClientManager,
470    client_id: Uuid,
471    view_id: &str,
472    bytes: usize,
473) {
474    let auth_context = client_manager.get_auth_context(client_id);
475    let (identity, deployment_id) = usage_identity_from_context(auth_context.as_ref());
476    emit_usage_event(
477        usage_emitter,
478        WebSocketUsageEvent::UpdateSent {
479            client_id: client_id.to_string(),
480            deployment_id,
481            identity,
482            view_id: view_id.to_string(),
483            messages: 1,
484            bytes: bytes as u64,
485        },
486    );
487}
488
489#[derive(Clone)]
490struct SubscriptionContext {
491    client_id: Uuid,
492    client_manager: ClientManager,
493    bus_manager: BusManager,
494    entity_cache: EntityCache,
495    view_index: Arc<ViewIndex>,
496    usage_emitter: Option<UsageEmitterHandle>,
497    journal: Option<Arc<crate::journal::EventJournal>>,
498    metrics: WsMetrics,
499    delivery: WebSocketDeliveryConfig,
500    /// Cancelled when the server stops; every session ends through its normal
501    /// cleanup path rather than being dropped mid-flight.
502    shutdown: CancellationToken,
503    /// The connection's row-scope generation. A token refresh that changes
504    /// the session's row scopes bumps it under the write lock before it
505    /// resnapshots, which stops every subscription task still delivering on
506    /// the old scopes (see [`RowFilter::hold`]).
507    row_scope_lock: Arc<tokio::sync::RwLock<u64>>,
508    /// Set on one subscription's copy of the context when its view is
509    /// row-scoped: which rows that subscription may receive.
510    row_filter: Option<RowFilter>,
511    /// Set while re-attaching a subscription after its session's row scopes
512    /// changed: it sends no second `subscribed` frame, and its snapshot
513    /// replaces what the client holds.
514    resnapshot: bool,
515}
516
517pub struct WebSocketServer {
518    bind_addr: SocketAddr,
519    client_manager: ClientManager,
520    bus_manager: BusManager,
521    entity_cache: EntityCache,
522    view_index: Arc<ViewIndex>,
523    max_clients: usize,
524    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
525    usage_emitter: Option<UsageEmitterHandle>,
526    rate_limit_config: Option<RateLimitConfig>,
527    admission_provider: Option<Arc<dyn WebSocketAdmissionProvider>>,
528    journal: Option<Arc<crate::journal::EventJournal>>,
529    delivery: WebSocketDeliveryConfig,
530    limits: WebSocketLimitsConfig,
531    #[cfg(feature = "otel")]
532    metrics: Option<Arc<Metrics>>,
533}
534
535impl WebSocketServer {
536    #[cfg(feature = "otel")]
537    pub fn new(
538        bind_addr: SocketAddr,
539        bus_manager: BusManager,
540        entity_cache: EntityCache,
541        view_index: Arc<ViewIndex>,
542        metrics: Option<Arc<Metrics>>,
543    ) -> Self {
544        Self {
545            bind_addr,
546            client_manager: ClientManager::new(),
547            bus_manager,
548            entity_cache,
549            view_index,
550            max_clients: 10_000,
551            auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
552            usage_emitter: None,
553            rate_limit_config: None,
554            admission_provider: None,
555            journal: None,
556            delivery: WebSocketDeliveryConfig::default(),
557            limits: WebSocketLimitsConfig::default(),
558            metrics,
559        }
560    }
561
562    #[cfg(not(feature = "otel"))]
563    pub fn new(
564        bind_addr: SocketAddr,
565        bus_manager: BusManager,
566        entity_cache: EntityCache,
567        view_index: Arc<ViewIndex>,
568    ) -> Self {
569        Self {
570            bind_addr,
571            client_manager: ClientManager::new(),
572            bus_manager,
573            entity_cache,
574            view_index,
575            max_clients: 10_000,
576            auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
577            usage_emitter: None,
578            rate_limit_config: None,
579            admission_provider: None,
580            journal: None,
581            delivery: WebSocketDeliveryConfig::default(),
582            limits: WebSocketLimitsConfig::default(),
583        }
584    }
585
586    pub fn with_max_clients(mut self, max_clients: usize) -> Self {
587        self.max_clients = max_clients;
588        self
589    }
590
591    pub fn with_auth_plugin(mut self, auth_plugin: Arc<dyn WebSocketAuthPlugin>) -> Self {
592        self.auth_plugin = auth_plugin;
593        self
594    }
595
596    pub fn with_usage_emitter(mut self, usage_emitter: Arc<dyn WebSocketUsageEmitter>) -> Self {
597        self.usage_emitter = Some(UsageEmitterHandle::new(usage_emitter));
598        self
599    }
600
601    /// Serve replayable append subscriptions from the retained event journal.
602    pub fn with_journal(mut self, journal: Arc<crate::journal::EventJournal>) -> Self {
603        self.journal = Some(journal);
604        self
605    }
606
607    pub fn with_rate_limit_config(mut self, config: RateLimitConfig) -> Self {
608        self.rate_limit_config = Some(config);
609        self
610    }
611
612    /// Install an embedder-owned provider for cross-runtime admission.
613    pub fn with_admission_provider(
614        mut self,
615        provider: Arc<dyn WebSocketAdmissionProvider>,
616    ) -> Self {
617        self.admission_provider = Some(provider);
618        self
619    }
620
621    /// Message, frame and handshake limits; see [`WebSocketLimitsConfig`].
622    pub fn with_limits(mut self, limits: WebSocketLimitsConfig) -> Self {
623        self.limits = limits;
624        self
625    }
626
627    pub fn with_delivery_config(mut self, config: WebSocketDeliveryConfig) -> Self {
628        self.delivery = config;
629        self
630    }
631
632    /// Bind the configured address and serve connections until the task is
633    /// dropped. Equivalent to [`into_acceptor`](Self::into_acceptor) followed by
634    /// [`ConnectionAcceptor::serve_listener`].
635    pub async fn start(self) -> Result<()> {
636        info!(
637            "Starting WebSocket server on {} (max_clients: {})",
638            self.bind_addr, self.max_clients
639        );
640        let listener = TcpListener::bind(&self.bind_addr).await?;
641        let (acceptor, _cleanup) = self.into_acceptor();
642        acceptor.serve_listener(listener).await
643    }
644
645    /// Split this server into the part that serves connections and the
646    /// client-manager cleanup task, leaving the caller to own the listener.
647    ///
648    /// The cleanup handle is returned rather than detached so a caller that
649    /// stops serving can stop it too.
650    pub(crate) fn into_acceptor(self) -> (ConnectionAcceptor, tokio::task::JoinHandle<()>) {
651        let mut client_manager = self
652            .rate_limit_config
653            .map(ClientManager::with_config)
654            .unwrap_or(self.client_manager);
655        if let Some(provider) = self.admission_provider {
656            client_manager = client_manager.with_admission_provider(provider);
657        }
658        let cleanup = client_manager.start_cleanup_task();
659
660        #[cfg(feature = "otel")]
661        let metrics = WsMetrics::new(self.metrics.clone());
662        #[cfg(not(feature = "otel"))]
663        let metrics = WsMetrics::default();
664
665        let acceptor = ConnectionAcceptor {
666            client_manager,
667            bus_manager: self.bus_manager,
668            entity_cache: self.entity_cache,
669            view_index: self.view_index,
670            max_clients: self.max_clients,
671            auth_plugin: self.auth_plugin,
672            usage_emitter: self.usage_emitter,
673            journal: self.journal,
674            delivery: self.delivery,
675            limits: self.limits,
676            pending_handshakes: Arc::new(AtomicUsize::new(0)),
677            metrics,
678            shutdown: CancellationToken::new(),
679            sessions: TaskTracker::new(),
680        };
681        (acceptor, cleanup)
682    }
683}
684
685/// Serves already-accepted TCP connections against one server's buses, cache
686/// and views.
687///
688/// This is what [`WebSocketServer::start`] runs behind its listener, separated
689/// so that a caller that owns the listener (an application that terminates
690/// TLS itself, a test with an ephemeral port) can hand streams in without the
691/// server binding anything.
692#[derive(Clone)]
693pub(crate) struct ConnectionAcceptor {
694    client_manager: ClientManager,
695    bus_manager: BusManager,
696    entity_cache: EntityCache,
697    view_index: Arc<ViewIndex>,
698    max_clients: usize,
699    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
700    usage_emitter: Option<UsageEmitterHandle>,
701    journal: Option<Arc<crate::journal::EventJournal>>,
702    delivery: WebSocketDeliveryConfig,
703    limits: WebSocketLimitsConfig,
704    /// Connections accepted but not yet upgraded. They count against `max_clients` too, so
705    /// stalled handshakes cannot hold more sockets than established sessions could.
706    pending_handshakes: Arc<AtomicUsize>,
707    metrics: WsMetrics,
708    shutdown: CancellationToken,
709    /// Sessions spawned by [`serve_listener`](Self::serve_listener), so a
710    /// stop can wait for them. Sessions a caller serves on its own tasks are
711    /// the caller's to wait for.
712    sessions: TaskTracker,
713}
714
715impl ConnectionAcceptor {
716    /// Mirror this acceptor's delivery instruments into `probe`. Must be set
717    /// before serving: each session copies the metrics handle when it starts.
718    #[cfg(test)]
719    pub(crate) fn with_delivery_probe(mut self, probe: Arc<DeliveryProbe>) -> Self {
720        self.metrics.probe = Some(probe);
721        self
722    }
723
724    /// Number of clients currently connected to this server.
725    pub(crate) fn client_count(&self) -> usize {
726        self.client_manager.client_count()
727    }
728
729    /// End every session this acceptor is serving and stop accepting.
730    ///
731    /// Sessions notice on their next poll and leave through the same cleanup
732    /// as a client disconnect, so the client manager, buses and usage events
733    /// see an ordinary close.
734    pub(crate) fn shutdown(&self) {
735        self.shutdown.cancel();
736        self.sessions.close();
737    }
738
739    /// Resolves once every listener-spawned session has finished cleaning
740    /// up. Call after [`shutdown`](Self::shutdown).
741    pub(crate) async fn wait_for_sessions(&self) {
742        self.sessions.wait().await;
743    }
744
745    /// Stop accepting usage emissions and wait for every event spawned by a
746    /// completed session to enter the emitter queue.
747    pub(crate) async fn wait_for_usage_events(&self) {
748        if let Some(emitter) = &self.usage_emitter {
749            emitter.close_and_wait().await;
750        }
751    }
752
753    /// Serve one accepted connection: WebSocket handshake, authentication,
754    /// then the subscription session until the peer disconnects.
755    ///
756    /// Returns `Ok(())` without serving when the server is at its client
757    /// limit, exactly as the listener loop does.
758    pub(crate) async fn serve(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
759        let Some(pending) =
760            PendingHandshake::reserve(&self.pending_handshakes, self.limits.max_pending_handshakes)
761        else {
762            warn!(
763                "Rejecting connection from {}: too many pending handshakes",
764                remote_addr
765            );
766            return Ok(());
767        };
768        if self.client_manager.client_count() + pending.count() > self.max_clients {
769            warn!(
770                "Rejecting connection from {}: max clients reached",
771                remote_addr
772            );
773            return Ok(());
774        }
775
776        let context = SubscriptionContext {
777            client_id: Uuid::nil(),
778            client_manager: self.client_manager.clone(),
779            bus_manager: self.bus_manager.clone(),
780            entity_cache: self.entity_cache.clone(),
781            view_index: self.view_index.clone(),
782            usage_emitter: self.usage_emitter.clone(),
783            journal: self.journal.clone(),
784            delivery: self.delivery.clone(),
785            metrics: self.metrics.clone(),
786            shutdown: self.shutdown.clone(),
787            row_scope_lock: Arc::new(tokio::sync::RwLock::new(0)),
788            row_filter: None,
789            resnapshot: false,
790        };
791        handle_connection(
792            stream,
793            context,
794            remote_addr,
795            self.auth_plugin.clone(),
796            &self.limits,
797            pending,
798        )
799        .await
800    }
801
802    /// Accept from `listener` until [`shutdown`](Self::shutdown), serving each
803    /// connection on its own task.
804    pub(crate) async fn serve_listener(self, listener: TcpListener) -> Result<()> {
805        loop {
806            let accepted = tokio::select! {
807                _ = self.shutdown.cancelled() => return Ok(()),
808                accepted = listener.accept() => accepted,
809            };
810            match accepted {
811                Ok((stream, addr)) => {
812                    let acceptor = self.clone();
813                    self.sessions.spawn(
814                        async move {
815                            if let Err(error) = acceptor.serve(stream, addr).await {
816                                error!("WebSocket connection error: {}", error);
817                            }
818                        }
819                        .instrument(info_span!("ws.connection", %addr)),
820                    );
821                }
822                Err(error) => error!("Failed to accept connection: {}", error),
823            }
824        }
825    }
826}
827
828/// One connection's place among the pending handshakes, released when dropped.
829pub(crate) struct PendingHandshake {
830    counter: Arc<AtomicUsize>,
831    count: usize,
832}
833
834impl PendingHandshake {
835    fn reserve(counter: &Arc<AtomicUsize>, max: usize) -> Option<Self> {
836        let previous = counter
837            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |pending| {
838                (pending < max).then_some(pending + 1)
839            })
840            .ok()?;
841        Some(Self {
842            counter: counter.clone(),
843            count: previous + 1,
844        })
845    }
846
847    /// Pending handshakes, this one included, when it was reserved.
848    fn count(&self) -> usize {
849        self.count
850    }
851}
852
853impl Drop for PendingHandshake {
854    fn drop(&mut self) {
855        self.counter.fetch_sub(1, Ordering::AcqRel);
856    }
857}
858
859#[derive(Debug, Clone)]
860struct HandshakeReject {
861    status: StatusCode,
862    body: crate::websocket::auth::ErrorResponse,
863    error_code: String,
864    retry_after_secs: Option<u64>,
865}
866
867impl HandshakeReject {
868    fn from_deny(deny: &AuthDeny) -> Self {
869        let retry_after_secs = match deny.retry_policy {
870            crate::websocket::auth::RetryPolicy::RetryAfter(duration) => Some(duration.as_secs()),
871            _ => None,
872        };
873        Self {
874            status: StatusCode::from_u16(deny.http_status).unwrap_or(StatusCode::UNAUTHORIZED),
875            body: deny.to_error_response(),
876            error_code: deny.code.to_string(),
877            retry_after_secs,
878        }
879    }
880}
881
882fn build_handshake_error_response(
883    response: &Response,
884    reject: &HandshakeReject,
885) -> HandshakeErrorResponse {
886    let mut builder = Response::builder()
887        .status(reject.status)
888        .version(response.version())
889        .header(CONTENT_TYPE, "application/json; charset=utf-8")
890        .header("X-Error-Code", &reject.error_code)
891        .header("Cache-Control", "no-store");
892    if let Some(retry_after_secs) = reject.retry_after_secs {
893        builder = builder.header("Retry-After", retry_after_secs.to_string());
894    }
895    let body = serde_json::to_string(&reject.body).unwrap_or_else(|_| {
896        format!(
897            r#"{{"error":"{}","message":"{}","code":"{}","retryable":false}}"#,
898            reject.body.error, reject.body.message, reject.body.code
899        )
900    });
901    builder
902        .body(Some(body))
903        .expect("handshake rejection response should build")
904}
905
906#[allow(clippy::result_large_err)]
907async fn accept_authorized_connection(
908    stream: TcpStream,
909    remote_addr: SocketAddr,
910    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
911    client_manager: ClientManager,
912    limits: &WebSocketLimitsConfig,
913    deadline: tokio::time::Instant,
914) -> Result<
915    Option<(
916        tokio_tungstenite::WebSocketStream<TcpStream>,
917        AuthContext,
918        Option<Arc<dyn WebSocketConnectionPermit>>,
919    )>,
920> {
921    use std::sync::Mutex;
922
923    type Admission =
924        Result<(AuthContext, Option<Arc<dyn WebSocketConnectionPermit>>), HandshakeReject>;
925    let capture: Arc<Mutex<Option<Admission>>> = Arc::new(Mutex::new(None));
926    let capture_ref = capture.clone();
927    let auth_plugin_ref = auth_plugin.clone();
928    let manager_ref = client_manager.clone();
929
930    // A message or frame past these limits closes the connection rather than being buffered.
931    let config = WebSocketConfig::default()
932        .max_message_size(Some(limits.max_message_bytes))
933        .max_frame_size(Some(limits.max_frame_bytes));
934    let callback = move |request: &Request, response| {
935        let request = ConnectionAuthRequest::from_http_request(remote_addr, request);
936        let result = tokio::task::block_in_place(|| {
937            // This callback blocks its caller, so the handshake deadline
938            // outside cannot interrupt it: bound authentication and admission
939            // by the same deadline here.
940            tokio::runtime::Handle::current().block_on(async {
941                let admission = async {
942                    match auth_plugin_ref.authorize(&request).await {
943                        AuthDecision::Allow(context) => manager_ref
944                            .check_connection_allowed(remote_addr, &Some(context.clone()))
945                            .await
946                            .and_then(|()| manager_ref.reserve_connection_admission(&context))
947                            .map(|permit| (context, permit))
948                            .map_err(|deny| HandshakeReject::from_deny(&deny)),
949                        AuthDecision::Deny(deny) => Err(HandshakeReject::from_deny(&deny)),
950                    }
951                };
952                tokio::time::timeout_at(deadline, admission)
953                    .await
954                    .unwrap_or_else(|_| {
955                        Err(HandshakeReject::from_deny(&AuthDeny::new(
956                            crate::websocket::auth::AuthErrorCode::InternalError,
957                            "Authentication did not complete in time",
958                        )))
959                    })
960            })
961        });
962        *capture_ref.lock().expect("capture lock poisoned") = Some(result.clone());
963        match result {
964            Ok(_) => Ok(response),
965            Err(reject) => Err(build_handshake_error_response(&response, &reject)),
966        }
967    };
968    let handshake_result = accept_hdr_async_with_config(stream, callback, Some(config)).await;
969
970    let auth_result = capture.lock().expect("capture lock poisoned").take();
971    match handshake_result {
972        Ok(stream) => match auth_result {
973            Some(Ok((context, permit))) => Ok(Some((stream, context, permit))),
974            Some(Err(reject)) => Err(anyhow::anyhow!(
975                "handshake unexpectedly succeeded after rejection: {}",
976                reject.body.message
977            )),
978            None => Err(anyhow::anyhow!("no auth result captured during handshake")),
979        },
980        Err(WsError::Http(_)) => Ok(None),
981        Err(error) => Err(error.into()),
982    }
983}
984
985async fn handle_connection(
986    stream: TcpStream,
987    mut context: SubscriptionContext,
988    remote_addr: SocketAddr,
989    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
990    limits: &WebSocketLimitsConfig,
991    pending: PendingHandshake,
992) -> Result<()> {
993    // The handshake is raced against shutdown and a deadline: a peer that
994    // stalls it must not keep a socket and a task alive.
995    let deadline = tokio::time::Instant::now() + limits.handshake_timeout;
996    let accepted = tokio::select! {
997        _ = context.shutdown.cancelled() => return Ok(()),
998        accepted = tokio::time::timeout_at(
999            deadline,
1000            accept_authorized_connection(
1001                stream,
1002                remote_addr,
1003                auth_plugin.clone(),
1004                context.client_manager.clone(),
1005                limits,
1006                deadline,
1007            ),
1008        ) => match accepted {
1009            Ok(accepted) => accepted?,
1010            Err(_) => {
1011                debug!("WebSocket handshake from {} timed out", remote_addr);
1012                return Ok(());
1013            }
1014        },
1015    };
1016    let Some((ws_stream, auth_context, admission_permit)) = accepted else {
1017        return Ok(());
1018    };
1019
1020    let client_id = Uuid::new_v4();
1021    context.client_id = client_id;
1022    let connection_start = Instant::now();
1023    let (mut usage_identity, mut deployment_id) = usage_identity_from_context(Some(&auth_context));
1024    let mut metering_key = usage_identity.metering_key.clone();
1025    // Active gauges must be decremented with the same attributes used for
1026    // their increment. The signed usage identity may legitimately move to a
1027    // claimed owner's account during an in-band refresh, but that must not
1028    // leave the original account's active count stuck or make the new one
1029    // negative when this connection closes.
1030    let connection_metrics_metering_key = metering_key.clone();
1031    context
1032        .metrics
1033        .connection_opened(connection_metrics_metering_key.as_deref());
1034
1035    let (ws_sender, mut ws_receiver) = ws_stream.split();
1036    context.client_manager.add_client_with_admission(
1037        client_id,
1038        ws_sender,
1039        Some(auth_context),
1040        remote_addr,
1041        admission_permit,
1042    );
1043    // Registered: the connection now counts as a client, so its pending
1044    // place can go. Released earlier, a concurrent handshake could pass the
1045    // client limit in between.
1046    drop(pending);
1047    emit_usage_event(
1048        &context.usage_emitter,
1049        WebSocketUsageEvent::ConnectionEstablished {
1050            client_id: client_id.to_string(),
1051            remote_addr: remote_addr.to_string(),
1052            deployment_id: deployment_id.clone(),
1053            identity: usage_identity.clone(),
1054        },
1055    );
1056
1057    let mut active_subscriptions: HashMap<String, ActiveSubscription> = HashMap::new();
1058    loop {
1059        let message = tokio::select! {
1060            _ = context.shutdown.cancelled() => break,
1061            next = ws_receiver.next() => match next {
1062                Some(message) => message,
1063                None => break,
1064            },
1065        };
1066        let message = match message {
1067            Ok(message) => message,
1068            Err(error) => {
1069                warn!("WebSocket error for client {}: {}", client_id, error);
1070                break;
1071            }
1072        };
1073        if message.is_close() {
1074            break;
1075        }
1076        context.client_manager.update_client_last_seen(client_id);
1077        if !message.is_text() {
1078            continue;
1079        }
1080        if let Err(deny) = context
1081            .client_manager
1082            .check_inbound_message_allowed(client_id)
1083        {
1084            send_socket_issue(client_id, &context.client_manager, &deny, true, None).await;
1085            break;
1086        }
1087        context.metrics.message_received(metering_key.as_deref());
1088
1089        let text = match message.to_text() {
1090            Ok(text) => text,
1091            Err(_) => continue,
1092        };
1093        let client_message = match serde_json::from_str::<ClientMessage>(text) {
1094            Ok(message) => message,
1095            Err(parse_error) => {
1096                let subscription_id = extract_subscription_id(text);
1097                send_protocol_issue(
1098                    client_id,
1099                    &context.client_manager,
1100                    &context.metrics,
1101                    subscription_id,
1102                    "malformed-message",
1103                    format!("invalid protocol v2 message: {parse_error}"),
1104                )
1105                .await;
1106                continue;
1107            }
1108        };
1109
1110        match client_message {
1111            ClientMessage::Subscribe(subscription) => {
1112                let subscription_id = subscription.subscription_id.clone();
1113                if let Err(message) = subscription.validate() {
1114                    send_protocol_issue(
1115                        client_id,
1116                        &context.client_manager,
1117                        &context.metrics,
1118                        Some(subscription_id),
1119                        "invalid-subscription",
1120                        message,
1121                    )
1122                    .await;
1123                    continue;
1124                }
1125                if let Err(deny) = context
1126                    .client_manager
1127                    .check_subscription_allowed(client_id)
1128                    .await
1129                {
1130                    send_socket_issue(
1131                        client_id,
1132                        &context.client_manager,
1133                        &deny,
1134                        false,
1135                        Some(subscription_id),
1136                    )
1137                    .await;
1138                    continue;
1139                }
1140
1141                let cancel_token = CancellationToken::new();
1142                let added = match context
1143                    .client_manager
1144                    .try_add_client_subscription(
1145                        client_id,
1146                        subscription_id.clone(),
1147                        cancel_token.clone(),
1148                    )
1149                    .await
1150                {
1151                    Ok(added) => added,
1152                    Err(deny) => {
1153                        send_socket_issue(
1154                            client_id,
1155                            &context.client_manager,
1156                            &deny,
1157                            false,
1158                            Some(subscription_id),
1159                        )
1160                        .await;
1161                        continue;
1162                    }
1163                };
1164                if !added {
1165                    send_protocol_issue(
1166                        client_id,
1167                        &context.client_manager,
1168                        &context.metrics,
1169                        Some(subscription_id),
1170                        "duplicate-subscription-id",
1171                        "subscriptionId is already active on this connection",
1172                    )
1173                    .await;
1174                    continue;
1175                }
1176
1177                let view = subscription.query.view.clone();
1178                let original = subscription.clone();
1179                if let Err(error) = attach_client_to_bus(&context, subscription, cancel_token).await
1180                {
1181                    context
1182                        .client_manager
1183                        .remove_client_subscription(client_id, &subscription_id)
1184                        .await;
1185                    // A refusal that already knows what to tell the client
1186                    // (an expired cursor, a changed epoch) keeps its own
1187                    // frame; anything else is a generic rejection.
1188                    match error.downcast::<RejectedSubscription>() {
1189                        Ok(rejected) => {
1190                            send_prepared_issue(
1191                                client_id,
1192                                &context.client_manager,
1193                                &context.metrics,
1194                                rejected.0,
1195                            )
1196                            .await;
1197                        }
1198                        Err(error) => {
1199                            send_protocol_issue(
1200                                client_id,
1201                                &context.client_manager,
1202                                &context.metrics,
1203                                Some(subscription_id),
1204                                "subscription-rejected",
1205                                error.to_string(),
1206                            )
1207                            .await;
1208                        }
1209                    }
1210                    continue;
1211                }
1212
1213                active_subscriptions.insert(
1214                    subscription_id,
1215                    ActiveSubscription {
1216                        view: view.clone(),
1217                        metrics_metering_key: metering_key.clone(),
1218                        subscription: original,
1219                    },
1220                );
1221                context
1222                    .metrics
1223                    .subscription_created(&view, metering_key.as_deref());
1224                emit_usage_event(
1225                    &context.usage_emitter,
1226                    WebSocketUsageEvent::SubscriptionCreated {
1227                        client_id: client_id.to_string(),
1228                        deployment_id: deployment_id.clone(),
1229                        identity: usage_identity.clone(),
1230                        view_id: view,
1231                    },
1232                );
1233            }
1234            ClientMessage::Unsubscribe(unsubscription) => {
1235                handle_unsubscribe(
1236                    &context,
1237                    unsubscription,
1238                    &mut active_subscriptions,
1239                    &deployment_id,
1240                    &usage_identity,
1241                )
1242                .await;
1243            }
1244            ClientMessage::Ping => debug!("Received ping from client {}", client_id),
1245            ClientMessage::RefreshAuth(request) => {
1246                let row_scopes_before =
1247                    session_row_scopes(context.client_manager.get_auth_context(client_id).as_ref());
1248                handle_refresh_auth(client_id, &request, &context.client_manager, &auth_plugin)
1249                    .await;
1250                // A successful refresh replaces the verified context. Keep
1251                // every later usage event on the same current identity used
1252                // by snapshots and updates; on failure the manager still
1253                // returns the previous context.
1254                if let Some(current_auth) = context.client_manager.get_auth_context(client_id) {
1255                    (usage_identity, deployment_id) =
1256                        usage_identity_from_context(Some(&current_auth));
1257                    metering_key = usage_identity.metering_key.clone();
1258                    if session_row_scopes(Some(&current_auth)) != row_scopes_before {
1259                        rescope_subscriptions(&context, &active_subscriptions).await;
1260                    }
1261                }
1262            }
1263        }
1264    }
1265
1266    context
1267        .client_manager
1268        .cancel_all_client_subscriptions(client_id)
1269        .await;
1270    context.client_manager.remove_client(client_id);
1271    if let Some(rate_limiter) = context.client_manager.rate_limiter().cloned() {
1272        rate_limiter.remove_client_buckets(client_id).await;
1273    }
1274    for subscription in active_subscriptions.values() {
1275        context.metrics.subscription_removed(
1276            &subscription.view,
1277            subscription.metrics_metering_key.as_deref(),
1278        );
1279        emit_usage_event(
1280            &context.usage_emitter,
1281            WebSocketUsageEvent::SubscriptionRemoved {
1282                client_id: client_id.to_string(),
1283                deployment_id: deployment_id.clone(),
1284                identity: usage_identity.clone(),
1285                view_id: subscription.view.clone(),
1286            },
1287        );
1288    }
1289    let duration = connection_start.elapsed().as_secs_f64();
1290    context
1291        .metrics
1292        .connection_closed(duration, connection_metrics_metering_key.as_deref());
1293    emit_usage_event(
1294        &context.usage_emitter,
1295        WebSocketUsageEvent::ConnectionClosed {
1296            client_id: client_id.to_string(),
1297            deployment_id,
1298            identity: usage_identity,
1299            duration_secs: Some(duration),
1300            subscription_count: u32::try_from(active_subscriptions.len()).unwrap_or(u32::MAX),
1301        },
1302    );
1303    Ok(())
1304}
1305
1306struct ActiveSubscription {
1307    view: String,
1308    /// Attribute set used for the active-gauge increment. It is immutable
1309    /// even if a later token refresh moves usage to another billing account.
1310    metrics_metering_key: Option<String>,
1311    /// The subscription as requested, to re-attach it when the session's row
1312    /// scopes change.
1313    subscription: Subscription,
1314}
1315
1316/// Re-deliver every row-scoped subscription after a token refresh changed the
1317/// session's row scopes.
1318///
1319/// The generation bump comes first, under the write lock, so no task still
1320/// running on the old scopes can send once it is done: each send holds the
1321/// read lock and checks its generation. Each subscription then re-attaches
1322/// with the new scopes: its snapshot is authoritative, so the client drops
1323/// rows it may no longer see and receives the ones it now may. A replayable
1324/// tape cannot be resnapshotted without resending records the consumer
1325/// already has, so it stops with a retryable `subscription-rescoped` issue
1326/// and the consumer resubscribes from its cursor.
1327async fn rescope_subscriptions(
1328    context: &SubscriptionContext,
1329    active_subscriptions: &HashMap<String, ActiveSubscription>,
1330) {
1331    *context.row_scope_lock.write().await += 1;
1332    for (subscription_id, active) in active_subscriptions {
1333        if context
1334            .view_index
1335            .scope_field_for_view(&active.view)
1336            .is_none()
1337        {
1338            continue;
1339        }
1340        let cancel_token = CancellationToken::new();
1341        if !context
1342            .client_manager
1343            .replace_client_subscription_token(
1344                context.client_id,
1345                subscription_id,
1346                cancel_token.clone(),
1347            )
1348            .await
1349        {
1350            continue;
1351        }
1352        let replayable = context.journal.as_ref().is_some_and(|journal| {
1353            journal.is_enabled()
1354                && context
1355                    .view_index
1356                    .get_view(&active.view)
1357                    .is_some_and(|spec| spec.mode == Mode::Append)
1358        });
1359        if replayable {
1360            let _ = send_control_frame(
1361                context,
1362                &SocketIssueMessage::subscription_rescoped(subscription_id.clone()),
1363                &active.view,
1364            );
1365            continue;
1366        }
1367        let mut subscription = active.subscription.clone();
1368        subscription.snapshot.enabled = true;
1369        let mut resnapshot = context.clone();
1370        resnapshot.resnapshot = true;
1371        match attach_client_to_bus(&resnapshot, subscription, cancel_token).await {
1372            Ok(()) => context.metrics.subscription_resnapshot(&active.view),
1373            Err(error) => {
1374                warn!(
1375                    "Subscription {} could not be resnapshotted for new row scopes: {error:#}",
1376                    subscription_id
1377                );
1378                let _ = send_control_frame(
1379                    context,
1380                    &SocketIssueMessage::subscription_rescoped(subscription_id.clone()),
1381                    &active.view,
1382                );
1383            }
1384        }
1385    }
1386}
1387
1388async fn handle_unsubscribe(
1389    context: &SubscriptionContext,
1390    unsubscription: Unsubscription,
1391    active_subscriptions: &mut HashMap<String, ActiveSubscription>,
1392    deployment_id: &Option<String>,
1393    usage_identity: &UsageIdentity,
1394) {
1395    let subscription_id = unsubscription.subscription_id.clone();
1396    if let Err(message) = unsubscription.validate() {
1397        send_protocol_issue(
1398            context.client_id,
1399            &context.client_manager,
1400            &context.metrics,
1401            Some(subscription_id),
1402            "invalid-unsubscription",
1403            message,
1404        )
1405        .await;
1406        return;
1407    }
1408
1409    if !context
1410        .client_manager
1411        .remove_client_subscription(context.client_id, &subscription_id)
1412        .await
1413    {
1414        send_protocol_issue(
1415            context.client_id,
1416            &context.client_manager,
1417            &context.metrics,
1418            Some(subscription_id),
1419            "unknown-subscription-id",
1420            "subscriptionId is not active on this connection",
1421        )
1422        .await;
1423        return;
1424    }
1425
1426    let Some(subscription) = active_subscriptions.remove(&subscription_id) else {
1427        return;
1428    };
1429    let _ = send_control_frame(
1430        context,
1431        &UnsubscribedFrame::new(subscription_id),
1432        &subscription.view,
1433    );
1434    context.metrics.subscription_removed(
1435        &subscription.view,
1436        subscription.metrics_metering_key.as_deref(),
1437    );
1438    emit_usage_event(
1439        &context.usage_emitter,
1440        WebSocketUsageEvent::SubscriptionRemoved {
1441            client_id: context.client_id.to_string(),
1442            deployment_id: deployment_id.clone(),
1443            identity: usage_identity.clone(),
1444            view_id: subscription.view,
1445        },
1446    );
1447}
1448
1449fn extract_subscription_id(text: &str) -> Option<String> {
1450    serde_json::from_str::<Value>(text)
1451        .ok()?
1452        .get("subscriptionId")?
1453        .as_str()
1454        .map(str::to_string)
1455}
1456
1457struct SnapshotMetadata<'a> {
1458    subscription_id: &'a str,
1459    snapshot_id: &'a str,
1460    authoritative: bool,
1461    mode: Mode,
1462    view_id: &'a str,
1463    key: Option<&'a str>,
1464}
1465
1466#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1467enum SnapshotPurpose {
1468    Initial,
1469    Recovery,
1470}
1471
1472impl SnapshotPurpose {
1473    /// The purpose of a subscription's first snapshot: a resnapshot after
1474    /// its row scopes changed replaces what the client holds, like recovery.
1475    fn first(context: &SubscriptionContext) -> Self {
1476        if context.resnapshot {
1477            Self::Recovery
1478        } else {
1479            Self::Initial
1480        }
1481    }
1482
1483    fn authoritative(self, subscription: &Subscription) -> bool {
1484        match self {
1485            Self::Initial => subscription.query.after.is_none(),
1486            // Recovery replaces the exact query membership even when `after`
1487            // made the initial snapshot incremental. A merge-only snapshot
1488            // cannot remove members whose delete frames were skipped.
1489            Self::Recovery => true,
1490        }
1491    }
1492}
1493
1494/// One snapshot frame, serialized straight from the cached entities it
1495/// carries: the same JSON as a [`SnapshotFrame`] of
1496/// [`SnapshotEntity`](crate::websocket::frame::SnapshotEntity) rows with the
1497/// view's wire format applied, without building those rows.
1498///
1499/// [`SnapshotFrame`]: crate::websocket::frame::SnapshotFrame
1500#[derive(Serialize)]
1501#[serde(rename_all = "camelCase")]
1502struct SnapshotBatch<'a> {
1503    protocol_version: u8,
1504    subscription_id: &'a str,
1505    snapshot_id: &'a str,
1506    authoritative: bool,
1507    mode: Mode,
1508    #[serde(rename = "entity")]
1509    export: &'a str,
1510    op: &'static str,
1511    #[serde(skip_serializing_if = "Option::is_none")]
1512    key: Option<&'a str>,
1513    data: SnapshotRows<'a>,
1514    complete: bool,
1515}
1516
1517struct SnapshotRows<'a> {
1518    rows: &'a [(String, SharedEntity)],
1519    wire_format: &'a WireFormat,
1520}
1521
1522impl Serialize for SnapshotRows<'_> {
1523    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
1524        use serde::ser::SerializeSeq as _;
1525        let mut rows = serializer.serialize_seq(Some(self.rows.len()))?;
1526        for (key, entity) in self.rows {
1527            rows.serialize_element(&SnapshotRow {
1528                key,
1529                data: WireEntity {
1530                    entity,
1531                    wire_format: self.wire_format,
1532                },
1533            })?;
1534        }
1535        rows.end()
1536    }
1537}
1538
1539#[derive(Serialize)]
1540struct SnapshotRow<'a> {
1541    key: &'a str,
1542    data: WireEntity<'a>,
1543}
1544
1545/// Split `rows` into the snapshot's frames: the first batch for a fast first
1546/// render, then larger ones. An empty snapshot is one empty, complete frame.
1547fn create_snapshot_batches<'a>(
1548    rows: &'a [(String, SharedEntity)],
1549    wire_format: &'a WireFormat,
1550    metadata: SnapshotMetadata<'a>,
1551    batch_config: &SnapshotBatchConfig,
1552) -> Vec<SnapshotBatch<'a>> {
1553    let batch = |data: &'a [(String, SharedEntity)], complete: bool| SnapshotBatch {
1554        protocol_version: PROTOCOL_VERSION,
1555        subscription_id: metadata.subscription_id,
1556        snapshot_id: metadata.snapshot_id,
1557        authoritative: metadata.authoritative,
1558        mode: metadata.mode,
1559        export: metadata.view_id,
1560        op: "snapshot",
1561        key: metadata.key,
1562        data: SnapshotRows {
1563            rows: data,
1564            wire_format,
1565        },
1566        complete,
1567    };
1568    if rows.is_empty() {
1569        return vec![batch(&[], true)];
1570    }
1571
1572    let mut batches = Vec::new();
1573    let mut offset = 0;
1574    while offset < rows.len() {
1575        let configured_size = if offset == 0 {
1576            batch_config.initial_batch_size
1577        } else {
1578            batch_config.subsequent_batch_size
1579        };
1580        let end = (offset + configured_size.max(1)).min(rows.len());
1581        batches.push(batch(&rows[offset..end], end == rows.len()));
1582        offset = end;
1583    }
1584    batches
1585}
1586
1587/// Send `rows` as one snapshot, serializing each frame from the cached
1588/// entities as it goes: no row is copied.
1589async fn send_snapshot_batches(
1590    context: &SubscriptionContext,
1591    subscription: &Subscription,
1592    rows: &[(String, SharedEntity)],
1593    view_spec: &ViewSpec,
1594    purpose: SnapshotPurpose,
1595    batch_config: &SnapshotBatchConfig,
1596) -> Result<()> {
1597    let snapshot_id = Uuid::new_v4().to_string();
1598    let authoritative = purpose.authoritative(subscription);
1599    let mode = view_spec.mode;
1600    let frames = create_snapshot_batches(
1601        rows,
1602        &view_spec.wire_format,
1603        SnapshotMetadata {
1604            subscription_id: &subscription.subscription_id,
1605            snapshot_id: &snapshot_id,
1606            authoritative,
1607            mode,
1608            view_id: &subscription.query.view,
1609            key: subscription.query.key.as_deref(),
1610        },
1611        batch_config,
1612    );
1613
1614    for frame in frames {
1615        let rows = frame.data.rows.len() as u32;
1616        let json = serde_json::to_vec(&frame)?;
1617        let payload = maybe_compress(&json);
1618        let bytes = payload.as_bytes().len() as u64;
1619        let _held = hold_row_scopes_async(context).await?;
1620        context
1621            .client_manager
1622            .send_compressed_async(context.client_id, payload)
1623            .await
1624            .map_err(|error| anyhow::anyhow!("failed to send snapshot: {error}"))?;
1625        context.metrics.message_sent();
1626
1627        let auth_context = context.client_manager.get_auth_context(context.client_id);
1628        let (identity, deployment_id) = usage_identity_from_context(auth_context.as_ref());
1629        emit_usage_event(
1630            &context.usage_emitter,
1631            WebSocketUsageEvent::SnapshotSent {
1632                client_id: context.client_id.to_string(),
1633                deployment_id,
1634                identity,
1635                view_id: subscription.query.view.clone(),
1636                rows,
1637                messages: 1,
1638                bytes,
1639            },
1640        );
1641    }
1642    Ok(())
1643}
1644
1645fn extract_sort_config(view_spec: &ViewSpec) -> Option<SortConfig> {
1646    if let Some(sort) = view_spec
1647        .pipeline
1648        .as_ref()
1649        .and_then(|pipeline| pipeline.sort.as_ref())
1650    {
1651        return Some(SortConfig {
1652            field: sort.field_path.clone(),
1653            order: match sort.order {
1654                crate::materialized_view::SortOrder::Asc => SortOrder::Asc,
1655                crate::materialized_view::SortOrder::Desc => SortOrder::Desc,
1656            },
1657        });
1658    }
1659    (view_spec.mode == Mode::List).then(|| SortConfig {
1660        field: vec!["_seq".to_string()],
1661        order: SortOrder::Desc,
1662    })
1663}
1664
1665/// Hold a row-scoped subscription's right to send on the row scopes it
1666/// attached with, for one send: an error once a refresh changed them (see
1667/// [`RowFilter::try_hold`]). Nothing to hold for other subscriptions.
1668fn hold_row_scopes(
1669    context: &SubscriptionContext,
1670) -> Result<Option<tokio::sync::RwLockReadGuard<'_, u64>>> {
1671    match context.row_filter.as_ref() {
1672        None => Ok(None),
1673        Some(filter) => filter
1674            .try_hold()
1675            .map(Some)
1676            .ok_or_else(|| anyhow::anyhow!("the session's row scopes changed")),
1677    }
1678}
1679
1680/// [`hold_row_scopes`] for a send that awaits.
1681async fn hold_row_scopes_async(
1682    context: &SubscriptionContext,
1683) -> Result<Option<tokio::sync::RwLockReadGuard<'_, u64>>> {
1684    match context.row_filter.as_ref() {
1685        None => Ok(None),
1686        Some(filter) => filter
1687            .hold()
1688            .await
1689            .map(Some)
1690            .ok_or_else(|| anyhow::anyhow!("the session's row scopes changed")),
1691    }
1692}
1693
1694fn send_control_frame<T: Serialize>(
1695    context: &SubscriptionContext,
1696    frame: &T,
1697    view_id: &str,
1698) -> Result<()> {
1699    let json = serde_json::to_vec(frame)?;
1700    let bytes = json.len();
1701    let _held = hold_row_scopes(context)?;
1702    context
1703        .client_manager
1704        .send_to_client(context.client_id, Arc::new(Bytes::from(json)))
1705        .map_err(|error| anyhow::anyhow!("failed to send control frame: {error}"))?;
1706    context.metrics.message_sent();
1707    emit_update_sent_for_client(
1708        &context.usage_emitter,
1709        &context.client_manager,
1710        context.client_id,
1711        view_id,
1712        bytes,
1713    );
1714    Ok(())
1715}
1716
1717fn send_subscribed_frame(
1718    context: &SubscriptionContext,
1719    subscription: &Subscription,
1720    view_spec: &ViewSpec,
1721) -> Result<()> {
1722    let frame = SubscribedFrame::new(
1723        subscription.subscription_id.clone(),
1724        subscription.query.clone(),
1725        view_spec.mode,
1726        extract_sort_config(view_spec),
1727    );
1728    send_control_frame(context, &frame, &subscription.query.view)
1729}
1730
1731fn enforce_snapshot_limit(context: &SubscriptionContext, rows: usize) -> Result<()> {
1732    context
1733        .client_manager
1734        .check_snapshot_allowed(context.client_id, u32::try_from(rows).unwrap_or(u32::MAX))
1735        .map_err(|deny| anyhow::anyhow!(deny.reason))
1736}
1737
1738async fn subscribe_state_then_snapshot<F, Fut, T>(
1739    bus_manager: &BusManager,
1740    view_id: &str,
1741    key: &str,
1742    snapshot: F,
1743) -> (watch::Receiver<StateUpdate>, u64, T)
1744where
1745    F: FnOnce() -> Fut,
1746    Fut: Future<Output = T>,
1747{
1748    let mut receiver = bus_manager.get_or_create_state_bus(view_id, key).await;
1749    let published = receiver.borrow_and_update().published;
1750    let snapshot = snapshot().await;
1751    (receiver, published, snapshot)
1752}
1753
1754async fn subscribe_list_then_snapshot<F, Fut, T>(
1755    bus_manager: &BusManager,
1756    view_id: &str,
1757    snapshot: F,
1758) -> (broadcast::Receiver<Arc<BusMessage>>, T)
1759where
1760    F: FnOnce() -> Fut,
1761    Fut: Future<Output = T>,
1762{
1763    let receiver = bus_manager.get_or_create_list_bus(view_id).await;
1764    let snapshot = snapshot().await;
1765    (receiver, snapshot)
1766}
1767
1768async fn attach_client_to_bus(
1769    context: &SubscriptionContext,
1770    mut subscription: Subscription,
1771    cancel_token: CancellationToken,
1772) -> Result<()> {
1773    let view_spec = context
1774        .view_index
1775        .get_view(&subscription.query.view)
1776        .cloned()
1777        .ok_or_else(|| anyhow::anyhow!("unknown view: {}", subscription.query.view))?;
1778
1779    // A row-scoped view delivers only the rows this session's verified token
1780    // covers. The scopes are read here, once, from the token the connection
1781    // holds now; a refresh that changes them re-attaches the subscription.
1782    let mut scoped_context;
1783    let context = match context.view_index.scope_field_for_view(&view_spec.id) {
1784        Some(field) => {
1785            let generation = *context.row_scope_lock.read().await;
1786            scoped_context = context.clone();
1787            scoped_context.row_filter = Some(RowFilter::new(
1788                field.clone(),
1789                session_row_scopes(
1790                    context
1791                        .client_manager
1792                        .get_auth_context(context.client_id)
1793                        .as_ref(),
1794                ),
1795                context.row_scope_lock.clone(),
1796                generation,
1797            ));
1798            &scoped_context
1799        }
1800        None => context,
1801    };
1802
1803    if view_spec.mode == Mode::State && !view_spec.is_derived() && subscription.query.key.is_none()
1804    {
1805        return Err(anyhow::anyhow!("state subscriptions require query.key"));
1806    }
1807    if view_spec.is_derived() && subscription.query.take.is_none() {
1808        subscription.query.take = view_spec
1809            .pipeline
1810            .as_ref()
1811            .and_then(|pipeline| pipeline.limit);
1812    }
1813
1814    // A retained tape takes precedence for append views: it is the only
1815    // delivery that can honour a cursor. Without one, fall through to the
1816    // previous latest-state behaviour.
1817    let journal = context
1818        .journal
1819        .clone()
1820        .filter(|journal| journal.is_enabled() && view_spec.mode == Mode::Append);
1821    if let Some(journal) = journal {
1822        return attach_journal_subscription(
1823            context,
1824            subscription,
1825            view_spec,
1826            journal,
1827            cancel_token,
1828        )
1829        .await;
1830    }
1831
1832    if view_spec.mode == Mode::State && !view_spec.is_derived() {
1833        attach_state_subscription(context, subscription, view_spec, cancel_token).await
1834    } else {
1835        attach_collection_subscription(context, subscription, view_spec, cancel_token).await
1836    }
1837}
1838
1839async fn attach_state_subscription(
1840    context: &SubscriptionContext,
1841    subscription: Subscription,
1842    view_spec: ViewSpec,
1843    cancel_token: CancellationToken,
1844) -> Result<()> {
1845    let view_id = subscription.query.view.clone();
1846    let key = subscription.query.key.clone().unwrap_or_default();
1847    let query = subscription.query.clone();
1848    let cache = context.entity_cache.clone();
1849    let view_spec_for_snapshot = view_spec.clone();
1850    let row_filter = context.row_filter.clone();
1851    let (mut receiver, mut seen, initial) =
1852        subscribe_state_then_snapshot(&context.bus_manager, &view_id, &key, move || async move {
1853            load_query_entities(
1854                &cache,
1855                None,
1856                &view_spec_for_snapshot,
1857                &query,
1858                row_filter.as_ref(),
1859                false,
1860            )
1861            .await
1862        })
1863        .await;
1864
1865    let mut snapshot_entities = initial;
1866    if let Some(limit) = subscription.query.snapshot_limit {
1867        snapshot_entities.truncate(limit);
1868    }
1869    enforce_snapshot_limit(context, snapshot_entities.len())?;
1870    if !context.resnapshot {
1871        send_subscribed_frame(context, &subscription, &view_spec)?;
1872    }
1873    // The client holds the entity only if the snapshot actually carried it.
1874    // A disabled snapshot, or one `snapshotLimit` cut to nothing, leaves it
1875    // without a copy, so its first change must be a full `upsert`: a patch
1876    // would have nothing to merge into.
1877    let delivered = subscription.snapshot.enabled && !snapshot_entities.is_empty();
1878    // The entity as the client last received it whole, which a catch-up is
1879    // measured against.
1880    let synced = if delivered {
1881        snapshot_entities
1882            .first()
1883            .map(|(_, entity)| Synced::Entity(entity.clone()))
1884    } else {
1885        None
1886    };
1887    if subscription.snapshot.enabled {
1888        send_snapshot_batches(
1889            context,
1890            &subscription,
1891            &snapshot_entities,
1892            &view_spec,
1893            SnapshotPurpose::first(context),
1894            &context.entity_cache.snapshot_config(),
1895        )
1896        .await?;
1897    }
1898
1899    let task_context = context.clone();
1900    let subscription_id = subscription.subscription_id.clone();
1901    let query = subscription.query.clone();
1902    let view_spec_task = view_spec.clone();
1903    let span_view = view_id.clone();
1904    let span_key = key.clone();
1905    tokio::spawn(
1906        async move {
1907            // Whether the client holds this key, not whether it exists.
1908            let mut member = delivered;
1909            // The entity as the client last received it whole (its snapshot
1910            // row, an upsert or the last catch-up; see `Synced`). A catch-up
1911            // sends what changed since, and every field a patch forwarded in
1912            // between set (`touched`, see `mark_fields`): the entity may have
1913            // set it back, and a client may have dropped the patch.
1914            let mut synced = synced;
1915            let mut touched = Value::Null;
1916            // Whether the client lacks fields of frames the bus overwrote
1917            // before this task read them. The cached entity is sent in their
1918            // place, but while the cache lacks the key there is none, so the
1919            // debt waits for the next frame that finds it cached.
1920            let mut behind = false;
1921            // Whether, while behind, a frame the client missed replaced the
1922            // entity (a delete, or a whole upsert): its copy may then hold
1923            // fields the entity no longer has, which merging keeps, so the
1924            // catch-up must replace it.
1925            let mut replaced = false;
1926            loop {
1927                tokio::select! {
1928                    _ = cancel_token.cancelled() => break,
1929                    changed = receiver.changed() => {
1930                        if changed.is_err() {
1931                            break;
1932                        }
1933                        let (payload, published, replaced_at, row_scope) = {
1934                            let update = receiver.borrow_and_update();
1935                            (
1936                                update.payload.clone(),
1937                                update.published,
1938                                update.replaced,
1939                                update.row_scope.clone(),
1940                            )
1941                        };
1942                        // Authorise the frame against the scope it was
1943                        // published under. The cache may already hold a later
1944                        // state of the row in this session's scopes; a frame
1945                        // from before that is still not this session's to see.
1946                        let in_scope = frame_in_row_scopes(
1947                            task_context.row_filter.as_ref(),
1948                            row_scope.as_ref(),
1949                            &payload,
1950                        );
1951                        // The bus keeps only the latest frame. If more than one
1952                        // was published since the last read, the earlier ones
1953                        // were overwritten and the latest patch alone would
1954                        // drop their fields, so send the cached entity instead.
1955                        behind |= published > seen + 1;
1956                        replaced |= behind && replaced_at > seen;
1957                        seen = published;
1958                        let metadata = source_frame_metadata(&payload);
1959                        if metadata.op == "delete" {
1960                            if !task_context.entity_cache
1961                                .deletion_is_current(&query.view, &key)
1962                                .await
1963                            {
1964                                behind = true;
1965                                replaced = true;
1966                                continue;
1967                            }
1968                            if member && send_membership_frame(
1969                                &task_context,
1970                                &subscription_id,
1971                                &view_spec_task,
1972                                "delete",
1973                                &key,
1974                                Value::Null,
1975                                metadata.seq,
1976                            ).is_err() {
1977                                break;
1978                            }
1979                            member = false;
1980                            behind = false;
1981                            replaced = false;
1982                            synced = None;
1983                            touched = Value::Null;
1984                            continue;
1985                        }
1986
1987                        // A key missing from the cache was evicted there or
1988                        // never held whole, not deleted (deletes arrive
1989                        // above): the cache refused this patch and the whole
1990                        // entity is on its way (see `EntityResync`). The
1991                        // client's copy is still good, so a holder gets the
1992                        // patch and stays a holder; `remove` is only for an
1993                        // entity that stopped matching. If it is behind, it
1994                        // stays behind: the resend that brings the entity
1995                        // back carries the seq of the latest change, which it
1996                        // already has, so only the catch-up below reaches it.
1997                        let Some(cached) = task_context
1998                            .entity_cache
1999                            .get_shared(&view_spec_task.id, &key)
2000                            .await
2001                        else {
2002                            // A patch that moves the row out of this
2003                            // session's row scopes ends its membership.
2004                            if member
2005                                && task_context.row_filter.as_ref().is_some_and(|filter| {
2006                                    source_frame_data(&payload)
2007                                        .is_some_and(|data| filter.denies_frame_data(&data))
2008                                })
2009                            {
2010                                member = false;
2011                                synced = None;
2012                                touched = Value::Null;
2013                                if send_membership_frame(
2014                                    &task_context,
2015                                    &subscription_id,
2016                                    &view_spec_task,
2017                                    "remove",
2018                                    &key,
2019                                    Value::Null,
2020                                    metadata.seq,
2021                                ).is_err() {
2022                                    break;
2023                                }
2024                                continue;
2025                            }
2026                            // A frame outside this session's scopes is
2027                            // withheld; the client is then behind, and the
2028                            // next frame that finds the entity cached sends
2029                            // what it lacks from the cache.
2030                            if member && !in_scope {
2031                                behind = true;
2032                                continue;
2033                            }
2034                            // A copy of a replaced entity takes no patch: it
2035                            // waits for the whole entity to replace it.
2036                            if member && !replaced {
2037                                record_forwarded(&payload, &metadata.op, &mut synced, &mut touched);
2038                                if send_scoped_source_payload(
2039                                    &task_context,
2040                                    &subscription_id,
2041                                    &query.view,
2042                                    payload,
2043                                ).is_err() {
2044                                    break;
2045                                }
2046                            }
2047                            continue;
2048                        };
2049                        let selected = select_query_entities(
2050                            vec![(key.clone(), cached)],
2051                            &query,
2052                            task_context.row_filter.as_ref(),
2053                            true,
2054                            false,
2055                        );
2056                        let is_member = !selected.is_empty();
2057                        // A withheld frame is caught up from the cache, which
2058                        // `selected` has just shown to be in scope.
2059                        let caught_up = !std::mem::take(&mut behind) && in_scope;
2060                        let replace = std::mem::take(&mut replaced);
2061                        let result = match (member, selected.into_iter().next()) {
2062                            (true, Some(_)) if caught_up => {
2063                                record_forwarded(&payload, &metadata.op, &mut synced, &mut touched);
2064                                send_scoped_source_payload(
2065                                    &task_context,
2066                                    &subscription_id,
2067                                    &query.view,
2068                                    payload,
2069                                )
2070                            }
2071                            // The catch-up: a patch of the cached entity's
2072                            // fields that changed since the client last received
2073                            // it whole or that a forwarded patch set, which
2074                            // merges into what the client holds.
2075                            // If a missed frame replaced the entity, it is the
2076                            // whole entity as an upsert instead, which replaces
2077                            // the client's copy: merging would keep the fields of
2078                            // the entity that was deleted or replaced. Either
2079                            // carries no seq, because it is newer than anything
2080                            // this subscriber was sent and seqs are not ordered
2081                            // within a slot: account updates and instructions
2082                            // number themselves differently, so the latest
2083                            // patch's seq can sort below one already delivered.
2084                            (true, Some((entity_key, data))) => {
2085                                let wire_format = &view_spec_task.wire_format;
2086                                let base = synced
2087                                    .as_ref()
2088                                    .filter(|_| !replace)
2089                                    .and_then(|synced| synced.wire(wire_format));
2090                                let result = match base {
2091                                    Some(base) => match changed_fields(
2092                                        &base,
2093                                        &wire_value(&data, wire_format),
2094                                        &touched,
2095                                    ) {
2096                                        Some(body) => send_membership_frame(
2097                                            &task_context,
2098                                            &subscription_id,
2099                                            &view_spec_task,
2100                                            "patch",
2101                                            &entity_key,
2102                                            body,
2103                                            None,
2104                                        ),
2105                                        None => Ok(()),
2106                                    },
2107                                    // Nothing to measure against: the whole
2108                                    // entity, straight from the cache.
2109                                    None => send_entity_frame(
2110                                        &task_context,
2111                                        &subscription_id,
2112                                        &view_spec_task,
2113                                        if replace { "upsert" } else { "patch" },
2114                                        &entity_key,
2115                                        &data,
2116                                        None,
2117                                    ),
2118                                };
2119                                synced = Some(Synced::Entity(data));
2120                                touched = Value::Null;
2121                                result
2122                            }
2123                            (false, Some((entity_key, data))) => {
2124                                let result = send_entity_frame(
2125                                    &task_context,
2126                                    &subscription_id,
2127                                    &view_spec_task,
2128                                    "upsert",
2129                                    &entity_key,
2130                                    &data,
2131                                    metadata.seq.as_deref(),
2132                                );
2133                                synced = Some(Synced::Entity(data));
2134                                touched = Value::Null;
2135                                result
2136                            }
2137                            (true, None) => {
2138                                synced = None;
2139                                touched = Value::Null;
2140                                send_membership_frame(
2141                                &task_context,
2142                                &subscription_id,
2143                                &view_spec_task,
2144                                "remove",
2145                                &key,
2146                                Value::Null,
2147                                metadata.seq,
2148                                )
2149                            }
2150                            (false, None) => Ok(()),
2151                        };
2152                        if result.is_err() {
2153                            break;
2154                        }
2155                        member = is_member;
2156                    }
2157                }
2158            }
2159        }
2160        .instrument(info_span!("ws.subscribe.state", client_id = %context.client_id, view = %span_view, key = %span_key)),
2161    );
2162    Ok(())
2163}
2164
2165async fn attach_collection_subscription(
2166    context: &SubscriptionContext,
2167    subscription: Subscription,
2168    view_spec: ViewSpec,
2169    cancel_token: CancellationToken,
2170) -> Result<()> {
2171    let view_id = subscription.query.view.clone();
2172    let source_view_id = view_spec
2173        .source_view
2174        .clone()
2175        .unwrap_or_else(|| view_id.clone());
2176    let (mut receiver, initial_membership) = subscribe_collection_then_snapshot(
2177        context,
2178        &source_view_id,
2179        &view_spec,
2180        &subscription.query,
2181    )
2182    .await;
2183
2184    let mut snapshot_entities = initial_membership.clone();
2185    if let Some(limit) = subscription.query.snapshot_limit {
2186        snapshot_entities.truncate(limit);
2187    }
2188    enforce_snapshot_limit(context, snapshot_entities.len())?;
2189    if !context.resnapshot {
2190        send_subscribed_frame(context, &subscription, &view_spec)?;
2191    }
2192    // Window members the client was never sent: everything past a
2193    // `snapshotLimit` cut, or the whole window when the snapshot is disabled.
2194    // They stay in the window (take/skip positions still count them) but the
2195    // client holds no copy, so none of them may ever be sent a patch.
2196    let delivered_rows = if subscription.snapshot.enabled {
2197        snapshot_entities.len()
2198    } else {
2199        0
2200    };
2201    let mut undelivered: HashSet<String> = initial_membership
2202        .iter()
2203        .skip(delivered_rows)
2204        .map(|(key, _)| key.clone())
2205        .collect();
2206    let mut sent_whole = SentWhole::default();
2207    if subscription.snapshot.enabled {
2208        send_snapshot_batches(
2209            context,
2210            &subscription,
2211            &snapshot_entities,
2212            &view_spec,
2213            SnapshotPurpose::first(context),
2214            &context.entity_cache.snapshot_config(),
2215        )
2216        .await?;
2217    }
2218
2219    let task_context = context.clone();
2220    let task_subscription = subscription.clone();
2221    let subscription_id = subscription.subscription_id.clone();
2222    let query = subscription.query.clone();
2223    let view_spec_task = view_spec.clone();
2224    let span_view = view_id.clone();
2225    tokio::spawn(
2226        async move {
2227            let mut current = initial_membership;
2228            // Append views are event tapes: collapsing two records would lose
2229            // observable history. Coalescing is only valid for latest-state
2230            // list membership, where a full final entity preserves meaning.
2231            let coalesce_ms = (view_spec_task.mode == Mode::List)
2232                .then(|| {
2233                    view_spec_task
2234                        .delivery
2235                        .coalesce_ms
2236                        .or(task_context.delivery.collection_coalesce_ms)
2237                        .filter(|milliseconds| *milliseconds > 0)
2238                })
2239                .flatten();
2240            let mut flush_interval = coalesce_ms.map(|milliseconds| {
2241                let period = Duration::from_millis(milliseconds);
2242                let mut interval = tokio::time::interval_at(
2243                    tokio::time::Instant::now() + period,
2244                    period,
2245                );
2246                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
2247                interval
2248            });
2249            let mut pending = HashMap::<String, Arc<BusMessage>>::new();
2250            let mut pending_updates = 0_u64;
2251
2252            loop {
2253                tokio::select! {
2254                    _ = cancel_token.cancelled() => break,
2255                    _ = async {
2256                        flush_interval
2257                            .as_mut()
2258                            .expect("coalescing interval is guarded")
2259                            .tick()
2260                            .await;
2261                    }, if flush_interval.is_some() => {
2262                        if pending.is_empty() {
2263                            continue;
2264                        }
2265                        let sorted_caches = view_spec_task
2266                            .is_derived()
2267                            .then(|| task_context.view_index.sorted_caches());
2268                        let next = load_query_entities(
2269                            &task_context.entity_cache,
2270                            sorted_caches,
2271                            &view_spec_task,
2272                            &query,
2273                            task_context.row_filter.as_ref(),
2274                            false,
2275                        ).await;
2276                        if emit_coalesced_collection_delta(
2277                            &task_context,
2278                            &subscription_id,
2279                            &view_spec_task,
2280                            &current,
2281                            &next,
2282                            &pending,
2283                            &mut undelivered,
2284                        ).is_err() {
2285                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
2286                            break;
2287                        }
2288                        task_context
2289                            .metrics
2290                            .collection_coalesced(&view_id, pending_updates);
2291                        current = next;
2292                        pending.clear();
2293                        pending_updates = 0;
2294                    }
2295                    received = receiver.recv() => {
2296                        let envelope = match received {
2297                            Ok(envelope) => envelope,
2298                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
2299                                task_context.metrics.subscription_lagged(&view_id, skipped);
2300                                warn!(
2301                                    "Subscription {} lagged by {} updates",
2302                                    subscription_id, skipped
2303                                );
2304                                if view_spec_task.mode == Mode::Append {
2305                                    let _ = send_control_frame(
2306                                        &task_context,
2307                                        &SocketIssueMessage::append_subscription_lagged(
2308                                            subscription_id.clone(),
2309                                            skipped,
2310                                        ),
2311                                        &view_id,
2312                                    );
2313                                    task_context.metrics.delivery_stopped(&view_id, "append-lagged-without-replay");
2314                                    break;
2315                                }
2316                                if !task_subscription.snapshot.enabled {
2317                                    let _ = send_control_frame(
2318                                        &task_context,
2319                                        &SocketIssueMessage::subscription_lagged(
2320                                            subscription_id.clone(),
2321                                            skipped,
2322                                        ),
2323                                        &view_id,
2324                                    );
2325                                    task_context.metrics.delivery_stopped(&view_id, "lagged-without-snapshot");
2326                                    break;
2327                                }
2328                                info!(
2329                                    "Subscription {} is recovering from an authoritative snapshot",
2330                                    subscription_id
2331                                );
2332                                match recover_collection_subscription(
2333                                    &task_context,
2334                                    &task_subscription,
2335                                    &view_spec_task,
2336                                    &source_view_id,
2337                                ).await {
2338                                    Ok((next_receiver, recovered)) => {
2339                                        receiver = next_receiver;
2340                                        current = recovered;
2341                                        // Recovery replaces the whole live
2342                                        // membership, untruncated.
2343                                        undelivered.clear();
2344                                        sent_whole.clear();
2345                                        pending.clear();
2346                                        pending_updates = 0;
2347                                        task_context.metrics.subscription_resnapshot(&view_id);
2348                                        continue;
2349                                    }
2350                                    Err(error) => {
2351                                        warn!(
2352                                            "Subscription {} failed to recover from lag: {error:#}",
2353                                            subscription_id
2354                                        );
2355                                        task_context.metrics.delivery_stopped(&view_id, "resnapshot-failed");
2356                                        break;
2357                                    }
2358                                }
2359                            }
2360                            Err(broadcast::error::RecvError::Closed) => break,
2361                        };
2362
2363                        apply_collection_source_event(
2364                            &task_context,
2365                            &source_view_id,
2366                            &view_spec_task,
2367                            &query,
2368                            &envelope,
2369                        ).await;
2370
2371                        if flush_interval.is_some() {
2372                            pending_updates = pending_updates.saturating_add(1);
2373                            pending.insert(envelope.key.clone(), envelope);
2374                            continue;
2375                        }
2376
2377                        let metadata = source_frame_metadata(&envelope.payload);
2378                        let sorted_caches = view_spec_task
2379                            .is_derived()
2380                            .then(|| task_context.view_index.sorted_caches());
2381                        let next = load_query_entities(
2382                            &task_context.entity_cache,
2383                            sorted_caches,
2384                            &view_spec_task,
2385                            &query,
2386                            task_context.row_filter.as_ref(),
2387                            false,
2388                        ).await;
2389                        if emit_collection_delta(
2390                            &task_context,
2391                            &subscription_id,
2392                            &view_spec_task,
2393                            &current,
2394                            &next,
2395                            &envelope,
2396                            &metadata,
2397                            &mut undelivered,
2398                            &mut sent_whole,
2399                        ).is_err() {
2400                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
2401                            break;
2402                        }
2403                        current = next;
2404                    }
2405                }
2406            }
2407        }
2408        .instrument(info_span!("ws.subscribe.collection", client_id = %context.client_id, view = %span_view)),
2409    );
2410    Ok(())
2411}
2412
2413async fn subscribe_collection_then_snapshot(
2414    context: &SubscriptionContext,
2415    source_view_id: &str,
2416    view_spec: &ViewSpec,
2417    query: &SubscriptionQuery,
2418) -> (
2419    broadcast::Receiver<Arc<BusMessage>>,
2420    Vec<(String, SharedEntity)>,
2421) {
2422    let cache = context.entity_cache.clone();
2423    let sorted_caches = view_spec
2424        .is_derived()
2425        .then(|| context.view_index.sorted_caches());
2426    let view_spec = view_spec.clone();
2427    let query = query.clone();
2428    let row_filter = context.row_filter.clone();
2429    subscribe_list_then_snapshot(&context.bus_manager, source_view_id, move || async move {
2430        load_query_entities(
2431            &cache,
2432            sorted_caches,
2433            &view_spec,
2434            &query,
2435            row_filter.as_ref(),
2436            false,
2437        )
2438        .await
2439    })
2440    .await
2441}
2442
2443async fn recover_collection_subscription(
2444    context: &SubscriptionContext,
2445    subscription: &Subscription,
2446    view_spec: &ViewSpec,
2447    source_view_id: &str,
2448) -> Result<(
2449    broadcast::Receiver<Arc<BusMessage>>,
2450    Vec<(String, SharedEntity)>,
2451)> {
2452    let (receiver, membership) =
2453        subscribe_collection_then_snapshot(context, source_view_id, view_spec, &subscription.query)
2454            .await;
2455    // `snapshotLimit` caps only the initial transfer. Recovery has to replace
2456    // the complete live membership; truncating it would make the replacement
2457    // authoritative while immediately omitting members the server still
2458    // considers current.
2459    enforce_snapshot_limit(context, membership.len())?;
2460    send_snapshot_batches(
2461        context,
2462        subscription,
2463        &membership,
2464        view_spec,
2465        SnapshotPurpose::Recovery,
2466        &context.entity_cache.snapshot_config(),
2467    )
2468    .await?;
2469    Ok((receiver, membership))
2470}
2471
2472async fn apply_collection_source_event(
2473    context: &SubscriptionContext,
2474    source_view_id: &str,
2475    view_spec: &ViewSpec,
2476    query: &SubscriptionQuery,
2477    envelope: &BusMessage,
2478) {
2479    let metadata = source_frame_metadata(&envelope.payload);
2480    if metadata.op != "delete" {
2481        return;
2482    }
2483    // The projector changed the shared cache before publishing this frame.
2484    // A slow subscription must inspect that authoritative lifetime state, not
2485    // replay the delete against `_seq` and potentially erase a recreation.
2486    if !context
2487        .entity_cache
2488        .deletion_is_current(source_view_id, &envelope.key)
2489        .await
2490    {
2491        return;
2492    }
2493    if view_spec.is_derived() {
2494        let caches = context.view_index.sorted_caches();
2495        let mut guard = caches.write().await;
2496        if let Some(cache) = guard.get_mut(&query.view) {
2497            cache.remove_key(&envelope.key);
2498        }
2499    }
2500}
2501
2502/// A subscription refused for a reason the client needs spelled out.
2503///
2504/// Attach paths that return this get the registration released by the
2505/// connection loop, the same as any other failure, while the client still
2506/// receives the specific error rather than a generic `subscription-rejected`.
2507/// Sending the frame and returning `Ok` instead would leave a registered
2508/// subscription with nothing attached: it would hold a slot against the
2509/// client's limit and make the advertised "resubscribe" remediation fail with
2510/// `duplicate-subscription-id`.
2511#[derive(Debug)]
2512pub(crate) struct RejectedSubscription(pub SocketIssueMessage);
2513
2514impl RejectedSubscription {
2515    fn into_error(issue: SocketIssueMessage) -> anyhow::Error {
2516        anyhow::Error::new(Self(issue))
2517    }
2518}
2519
2520impl std::fmt::Display for RejectedSubscription {
2521    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2522        formatter.write_str(&self.0.message)
2523    }
2524}
2525
2526impl std::error::Error for RejectedSubscription {}
2527
2528/// Deliver an append view as an event tape: replay the retained records after
2529/// the cursor, then forward live frames.
2530///
2531/// This deliberately does not recompute membership from the entity cache the
2532/// way [`attach_collection_subscription`] does. The cache folds each patch
2533/// into the resident entity, so a membership diff cannot express "these three
2534/// events happened"; the retained records can.
2535///
2536/// For the same reason records are never promoted to a full `upsert` the
2537/// first time this subscription sees their key, the way a window entry is.
2538/// The cache only holds the latest state, not the state as of a record's
2539/// offset, so substituting it would hand a consumer state from after its
2540/// cursor. Records go out exactly as retained, `patch` included: a tape is
2541/// events, and the consumer that resumes from a cursor holds whatever came
2542/// before it. Clients therefore apply frames that carry an `offset` as events
2543/// and do not discard them as partial entities (see "Partial entities" in
2544/// `docs/websocket-v2-protocol.md`).
2545async fn attach_journal_subscription(
2546    context: &SubscriptionContext,
2547    subscription: Subscription,
2548    view_spec: ViewSpec,
2549    journal: Arc<crate::journal::EventJournal>,
2550    cancel_token: CancellationToken,
2551) -> Result<()> {
2552    let view_id = view_spec.id.clone();
2553    let subscription_id = subscription.subscription_id.clone();
2554
2555    // A tape has no membership window, so `take`/`skip` cannot mean what they
2556    // mean on a list view. Refuse them rather than accept and ignore them.
2557    if subscription.query.take.is_some()
2558        || subscription.query.skip.is_some()
2559        || subscription.query.snapshot_limit.is_some()
2560    {
2561        return Err(RejectedSubscription::into_error(SocketIssueMessage::protocol(
2562            Some(subscription_id),
2563            "invalid-subscription",
2564            format!(
2565                "take, skip and snapshotLimit are window options and do not apply to the replayable view {view_id}"
2566            ),
2567        )));
2568    }
2569
2570    // Reject a malformed cursor rather than silently replaying from the start,
2571    // which would look like success and duplicate everything.
2572    let cursor = match subscription.query.after.as_deref() {
2573        Some(raw) => match crate::journal::Cursor::parse(raw) {
2574            Some(cursor) => Some(cursor),
2575            None => {
2576                return Err(RejectedSubscription::into_error(
2577                    SocketIssueMessage::protocol(
2578                        Some(subscription_id),
2579                        "invalid-cursor",
2580                        format!(
2581                            "`after` must be an {{epoch}}:{{offset}} replay cursor for view {view_id}, got {raw:?}"
2582                        ),
2583                    ),
2584                ));
2585            }
2586        },
2587        None => None,
2588    };
2589
2590    // Subscribe before reading the journal so anything published during the
2591    // replay is still delivered; the offset filter below drops the overlap.
2592    let mut receiver = context.bus_manager.get_or_create_list_bus(&view_id).await;
2593
2594    let replayed = match journal.replay_after(&view_id, cursor.as_ref()).await {
2595        Ok(records) => records,
2596        Err(error) => {
2597            return Err(RejectedSubscription::into_error(
2598                SocketIssueMessage::replay_refused(Some(subscription_id), &error),
2599            ));
2600        }
2601    };
2602
2603    let frame = SubscribedFrame::new(
2604        subscription.subscription_id.clone(),
2605        subscription.query.clone(),
2606        view_spec.mode,
2607        extract_sort_config(&view_spec),
2608    )
2609    .with_replay_window(journal.window(&view_id).await);
2610    send_control_frame(context, &frame, &view_id)?;
2611
2612    // Everything past the acknowledgement runs on its own task. The replay can
2613    // be long and applies real backpressure, and `attach_client_to_bus` is
2614    // awaited directly on the connection's inbound loop — doing it there would
2615    // block unsubscribe, auth refresh and pong for the whole replay.
2616    let task_context = context.clone();
2617    let task_subscription_id = subscription.subscription_id.clone();
2618    let task_query = subscription.query.clone();
2619    let task_epoch = journal.epoch().await;
2620    let span_view = view_id.clone();
2621    tokio::spawn(
2622        async move {
2623            let mut last_sent = cursor.map(|cursor| cursor.offset);
2624            // Frames that published while the replay was still running. The
2625            // bus is a bounded broadcast, so it has to be drained as we go or
2626            // a busy view laps us before the replay finishes.
2627            let mut pending: VecDeque<Arc<BusMessage>> = VecDeque::new();
2628            let mut lagged: Option<u64> = None;
2629
2630            for record in replayed {
2631                if cancel_token.is_cancelled() {
2632                    return;
2633                }
2634                drain_available(&mut receiver, &mut pending, &mut lagged);
2635                if journal_record_matches(&task_query, &record)
2636                    && frame_in_row_scopes(
2637                        task_context.row_filter.as_ref(),
2638                        record.row_scope.as_ref(),
2639                        &record.payload,
2640                    )
2641                    && send_scoped_source_payload_async(
2642                        &task_context,
2643                        &task_subscription_id,
2644                        &span_view,
2645                        record.payload,
2646                    )
2647                    .await
2648                    .is_err()
2649                {
2650                    return;
2651                }
2652                // Advance past filtered records too: they were considered.
2653                last_sent = Some(record.offset);
2654            }
2655
2656            // Flush what arrived during the replay before going live, so the
2657            // handover keeps offset order.
2658            while let Some(envelope) = pending.pop_front() {
2659                if !forward_live_frame(
2660                    &task_context,
2661                    &task_subscription_id,
2662                    &span_view,
2663                    &task_query,
2664                    &envelope,
2665                    &mut last_sent,
2666                )
2667                .await
2668                {
2669                    return;
2670                }
2671            }
2672
2673            if let Some(skipped) = lagged {
2674                report_replay_gap(
2675                    &task_context,
2676                    &task_subscription_id,
2677                    &span_view,
2678                    &task_epoch,
2679                    skipped,
2680                    last_sent,
2681                );
2682                return;
2683            }
2684
2685            loop {
2686                tokio::select! {
2687                    _ = cancel_token.cancelled() => break,
2688                    received = receiver.recv() => {
2689                        let envelope = match received {
2690                            Ok(envelope) => envelope,
2691                            // A lagged tape is a gap. Report it with the last
2692                            // offset delivered *before* the gap and stop
2693                            // delivering on this subscription.
2694                            //
2695                            // Continuing would hand the consumer frames from
2696                            // after the gap, advancing its checkpoint past
2697                            // the skipped records so they could never be
2698                            // replayed. Stopping is not a silent stall: the
2699                            // consumer has an explicit error and a cursor
2700                            // that recovers exactly what it missed.
2701                            //
2702                            // The registration is deliberately left alone.
2703                            // Its lifecycle belongs to the connection loop,
2704                            // which holds the only handle to
2705                            // `active_subscriptions`; releasing half of it
2706                            // here would desynchronise unsubscribe, the
2707                            // duplicate-ID gate and close-time usage.
2708                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
2709                                report_replay_gap(
2710                                    &task_context,
2711                                    &task_subscription_id,
2712                                    &span_view,
2713                                    &task_epoch,
2714                                    skipped,
2715                                    last_sent,
2716                                );
2717                                break;
2718                            }
2719                            Err(broadcast::error::RecvError::Closed) => break,
2720                        };
2721
2722                        if !forward_live_frame(
2723                            &task_context,
2724                            &task_subscription_id,
2725                            &span_view,
2726                            &task_query,
2727                            &envelope,
2728                            &mut last_sent,
2729                        )
2730                        .await
2731                        {
2732                            break;
2733                        }
2734                    }
2735                }
2736            }
2737        }
2738        .instrument(info_span!(
2739            "ws.subscribe.replay",
2740            client_id = %context.client_id,
2741            view = %view_id
2742        )),
2743    );
2744    Ok(())
2745}
2746
2747/// Take whatever the bus already has without waiting, so a long replay cannot
2748/// be lapped by a busy view.
2749fn drain_available(
2750    receiver: &mut broadcast::Receiver<Arc<BusMessage>>,
2751    pending: &mut VecDeque<Arc<BusMessage>>,
2752    lagged: &mut Option<u64>,
2753) {
2754    // Once a gap is known, everything still on the bus is on the far side of
2755    // it. Buffering it would put post-gap frames in front of the lag report,
2756    // advancing `last_sent` past the hole and making `recoverFrom` point
2757    // after the very records it is supposed to recover.
2758    if lagged.is_some() {
2759        return;
2760    }
2761    // Bounded so a view publishing faster than the client drains cannot turn
2762    // the buffer into an unbounded queue. Filling it is not itself a gap:
2763    // nothing has been skipped at that instant, the buffer simply stopped
2764    // accepting. Stop buffering and let the bus report the loss, with the
2765    // count it actually measures, when delivery reaches it.
2766    const MAX_PENDING: usize = 8_192;
2767    loop {
2768        if pending.len() >= MAX_PENDING {
2769            return;
2770        }
2771        match receiver.try_recv() {
2772            Ok(envelope) => pending.push_back(envelope),
2773            Err(broadcast::error::TryRecvError::Empty)
2774            | Err(broadcast::error::TryRecvError::Closed) => return,
2775            Err(broadcast::error::TryRecvError::Lagged(skipped)) => {
2776                *lagged = Some(skipped);
2777                return;
2778            }
2779        }
2780    }
2781}
2782
2783/// Whether a live frame was already delivered by the replay that preceded it.
2784///
2785/// The bus is subscribed before the tape is read, so a record published in
2786/// between appears on both paths; without this the consumer sees it twice.
2787/// Advances the high-water mark as a side effect.
2788fn already_delivered(offset: Option<u64>, last_sent: &mut Option<u64>) -> bool {
2789    let Some(offset) = offset else {
2790        // A frame with no offset predates the tape, so it cannot have been
2791        // replayed and must not move the mark.
2792        return false;
2793    };
2794    if last_sent.is_some_and(|last| offset <= last) {
2795        return true;
2796    }
2797    *last_sent = Some(offset);
2798    false
2799}
2800
2801/// Deliver one live frame, skipping anything the replay already sent.
2802///
2803/// Returns false when the subscription should end.
2804async fn forward_live_frame(
2805    context: &SubscriptionContext,
2806    subscription_id: &str,
2807    view_id: &str,
2808    query: &SubscriptionQuery,
2809    envelope: &Arc<BusMessage>,
2810    last_sent: &mut Option<u64>,
2811) -> bool {
2812    let metadata = source_frame_metadata(&envelope.payload);
2813    if already_delivered(metadata.offset, last_sent) {
2814        return true;
2815    }
2816    if !live_frame_matches(query, &envelope.key, &envelope.payload) {
2817        return true;
2818    }
2819    if !frame_in_row_scopes(
2820        context.row_filter.as_ref(),
2821        envelope.row_scope.as_ref(),
2822        &envelope.payload,
2823    ) {
2824        return true;
2825    }
2826    send_scoped_source_payload(context, subscription_id, view_id, envelope.payload.clone()).is_ok()
2827}
2828
2829/// Check the scope frozen with a retained or queued frame. Legacy records
2830/// may prove scope from an explicit field, but never from today's row state.
2831fn frame_in_row_scopes(
2832    filter: Option<&RowFilter>,
2833    scope: Option<&crate::row_scope::RowScope>,
2834    payload: &[u8],
2835) -> bool {
2836    let Some(filter) = filter else {
2837        return true;
2838    };
2839    if let Some(scope) = scope {
2840        return filter.allows_scope(scope);
2841    }
2842    source_frame_data(payload)
2843        .as_ref()
2844        .and_then(|data| filter.field().scope_in(data))
2845        .is_some_and(|scope| filter.allows_scope(&scope))
2846}
2847
2848fn report_replay_gap(
2849    context: &SubscriptionContext,
2850    subscription_id: &str,
2851    view_id: &str,
2852    epoch: &crate::journal::JournalEpoch,
2853    skipped: u64,
2854    last_sent: Option<u64>,
2855) {
2856    warn!(
2857        "Replay subscription {} lagged past {} records; stopping with a recovery cursor",
2858        subscription_id, skipped
2859    );
2860    let recover_from = last_sent.map(|offset| crate::journal::Cursor {
2861        epoch: epoch.clone(),
2862        offset,
2863    });
2864    let _ = send_control_frame(
2865        context,
2866        &SocketIssueMessage::replay_lagged(
2867            Some(subscription_id.to_string()),
2868            skipped,
2869            recover_from,
2870        ),
2871        view_id,
2872    );
2873}
2874
2875/// Apply the subscription's `key`, `partition` and `filters` to a retained
2876/// record. A replay must honour the same predicates a live subscription does.
2877fn journal_record_matches(
2878    query: &SubscriptionQuery,
2879    record: &crate::journal::JournalRecord,
2880) -> bool {
2881    live_frame_matches(query, &record.key, &record.payload)
2882}
2883
2884fn live_frame_matches(query: &SubscriptionQuery, key: &str, payload: &[u8]) -> bool {
2885    if !query.matches_key(key) {
2886        return false;
2887    }
2888    if query.partition.is_none() && query.filters.is_empty() {
2889        return true;
2890    }
2891    let Some(data) = source_frame_data(payload) else {
2892        return false;
2893    };
2894    let data = &data;
2895    if let Some(partition) = &query.partition {
2896        if value_at_dot_path(data, "_partition").as_deref()
2897            != Some(&Value::String(partition.clone()))
2898        {
2899            return false;
2900        }
2901    }
2902    query
2903        .filters
2904        .iter()
2905        .all(|(path, expected)| value_at_dot_path(data, path).as_deref() == Some(expected))
2906}
2907
2908/// Awaiting variant of [`send_scoped_source_payload`], for replays that can
2909/// exceed the client's send queue.
2910async fn send_scoped_source_payload_async(
2911    context: &SubscriptionContext,
2912    subscription_id: &str,
2913    view_id: &str,
2914    payload: Arc<Bytes>,
2915) -> Result<()> {
2916    let json = SourceFields::scoped(&payload, subscription_id)?;
2917    let compressed = maybe_compress(&json);
2918    let bytes = compressed.as_bytes().len();
2919    let _held = hold_row_scopes_async(context).await?;
2920    context
2921        .client_manager
2922        .send_compressed_async(context.client_id, compressed)
2923        .await
2924        .map_err(|error| anyhow::anyhow!("failed to send replayed frame: {error}"))?;
2925    context.metrics.message_sent();
2926    emit_update_sent_for_client(
2927        &context.usage_emitter,
2928        &context.client_manager,
2929        context.client_id,
2930        view_id,
2931        bytes,
2932    );
2933    Ok(())
2934}
2935
2936#[derive(Default)]
2937struct SourceFrameMetadata {
2938    op: String,
2939    seq: Option<String>,
2940    offset: Option<u64>,
2941}
2942
2943/// A source frame's `op`, `seq` and `offset`, read without building its
2944/// `data`.
2945fn source_frame_metadata(payload: &[u8]) -> SourceFrameMetadata {
2946    SourceFields::parse(payload)
2947        .map(|fields| SourceFrameMetadata {
2948            op: fields
2949                .value("op")
2950                .as_ref()
2951                .and_then(Value::as_str)
2952                .unwrap_or_default()
2953                .to_string(),
2954            seq: fields
2955                .value("seq")
2956                .as_ref()
2957                .and_then(Value::as_str)
2958                .map(str::to_string),
2959            offset: fields.value("offset").as_ref().and_then(Value::as_u64),
2960        })
2961        .unwrap_or_default()
2962}
2963
2964/// The `data` of a source frame, as a client receives it.
2965fn source_frame_data(payload: &[u8]) -> Option<Value> {
2966    SourceFields::parse(payload)?.value("data")
2967}
2968
2969/// What a state subscriber last received whole, which a catch-up is
2970/// measured against. Both kinds are shared, with the entity cache or the bus,
2971/// so a subscriber keeps no copy of its own. A subscriber does keep the
2972/// version it was sent alive while the cache moves on; subscribers sent the
2973/// same version share it.
2974enum Synced {
2975    /// The cached entity, as sent in a snapshot row, an upsert or a catch-up.
2976    Entity(SharedEntity),
2977    /// A forwarded `upsert` source frame, whose `data` the client holds.
2978    Frame(Arc<Bytes>),
2979}
2980
2981impl Synced {
2982    /// What the client holds, as it received it: in wire format.
2983    fn wire(&self, wire_format: &WireFormat) -> Option<Value> {
2984        match self {
2985            Synced::Entity(entity) => Some(wire_value(entity, wire_format)),
2986            Synced::Frame(frame) => source_frame_data(frame),
2987        }
2988    }
2989}
2990
2991/// Tracks what a client holds after a frame forwarded to it as published. An
2992/// upsert hands it the entity whole, which a catch-up is then measured
2993/// against; a patch sets fields (see [`mark_fields`]).
2994fn record_forwarded(
2995    payload: &Arc<Bytes>,
2996    op: &str,
2997    synced: &mut Option<Synced>,
2998    touched: &mut Value,
2999) {
3000    if op == "upsert" {
3001        *synced = Some(Synced::Frame(payload.clone()));
3002        *touched = Value::Null;
3003    } else if let Some(data) = source_frame_data(payload) {
3004        mark_fields(touched, &data);
3005    }
3006}
3007
3008/// Marks in `touched` the fields `patch` sets: objects field by field, any
3009/// other value whole (`true`). `null` marks nothing.
3010///
3011/// A client holds whatever a forwarded patch set, which the entity may have
3012/// set back before the next catch-up (or the client may have dropped the
3013/// patch), so a catch-up sends these fields even when they match the entity
3014/// the client last received whole.
3015fn mark_fields(touched: &mut Value, patch: &Value) {
3016    let Value::Object(fields) = patch else {
3017        *touched = Value::Bool(true);
3018        return;
3019    };
3020    if touched.is_null() {
3021        *touched = Value::Object(serde_json::Map::new());
3022    }
3023    let Value::Object(marked) = touched else {
3024        // Already marked whole.
3025        return;
3026    };
3027    for (field, value) in fields {
3028        mark_fields(marked.entry(field.clone()).or_insert(Value::Null), value);
3029    }
3030}
3031
3032/// The fields of `current` that differ from `base` or that `touched` marks
3033/// (see [`mark_fields`]), as a patch that turns what the client holds into
3034/// `current` when merged: objects compare field by field, any other value is
3035/// sent whole, and a field `current` no longer has is sent as `null`. `None`
3036/// when nothing needs sending.
3037fn changed_fields(base: &Value, current: &Value, touched: &Value) -> Option<Value> {
3038    match (base, current, touched) {
3039        (_, _, Value::Bool(true)) => Some(current.clone()),
3040        (Value::Object(base), Value::Object(current), _) => {
3041            let marked = touched.as_object();
3042            let mut patch = serde_json::Map::new();
3043            for (field, value) in current {
3044                let touched = marked
3045                    .and_then(|marked| marked.get(field))
3046                    .unwrap_or(&Value::Null);
3047                match base.get(field) {
3048                    Some(old) => {
3049                        if let Some(changed) = changed_fields(old, value, touched) {
3050                            patch.insert(field.clone(), changed);
3051                        }
3052                    }
3053                    None => {
3054                        patch.insert(field.clone(), value.clone());
3055                    }
3056                }
3057            }
3058            for field in base
3059                .keys()
3060                .chain(marked.into_iter().flat_map(|marked| marked.keys()))
3061            {
3062                if !current.contains_key(field) {
3063                    patch.insert(field.clone(), Value::Null);
3064                }
3065            }
3066            (!patch.is_empty()).then_some(Value::Object(patch))
3067        }
3068        _ => (base != current || !touched.is_null()).then(|| current.clone()),
3069    }
3070}
3071
3072fn send_scoped_source_payload(
3073    context: &SubscriptionContext,
3074    subscription_id: &str,
3075    view_id: &str,
3076    payload: Arc<Bytes>,
3077) -> Result<()> {
3078    let encoded = Arc::new(Bytes::from(SourceFields::scoped(
3079        &payload,
3080        subscription_id,
3081    )?));
3082    let bytes = encoded.len();
3083    let _held = hold_row_scopes(context)?;
3084    context
3085        .client_manager
3086        .send_to_client(context.client_id, encoded)
3087        .map_err(|error| anyhow::anyhow!("failed to send live frame: {error}"))?;
3088    context.metrics.message_sent();
3089    emit_update_sent_for_client(
3090        &context.usage_emitter,
3091        &context.client_manager,
3092        context.client_id,
3093        view_id,
3094        bytes,
3095    );
3096    Ok(())
3097}
3098
3099fn send_membership_frame(
3100    context: &SubscriptionContext,
3101    subscription_id: &str,
3102    view_spec: &ViewSpec,
3103    op: &str,
3104    key: &str,
3105    mut data: Value,
3106    seq: Option<String>,
3107) -> Result<()> {
3108    apply_wire_format(&mut data, &view_spec.wire_format);
3109    let frame = Frame::scoped(
3110        subscription_id,
3111        view_spec.mode,
3112        &view_spec.id,
3113        op,
3114        key,
3115        data,
3116        seq,
3117    );
3118    send_membership_json(context, view_spec, serde_json::to_vec(&frame)?)
3119}
3120
3121/// A scoped frame whose data is a cached entity, serialized straight from
3122/// it: the same JSON as [`Frame::scoped`] with a formatted copy of the entity.
3123#[derive(Serialize)]
3124#[serde(rename_all = "camelCase")]
3125struct ScopedEntityFrame<'a> {
3126    protocol_version: u8,
3127    subscription_id: &'a str,
3128    mode: Mode,
3129    #[serde(rename = "entity")]
3130    export: &'a str,
3131    op: &'a str,
3132    key: &'a str,
3133    data: WireEntity<'a>,
3134    #[serde(skip_serializing_if = "Option::is_none")]
3135    seq: Option<&'a str>,
3136}
3137
3138/// [`send_membership_frame`] for a cached entity, without copying it.
3139fn send_entity_frame(
3140    context: &SubscriptionContext,
3141    subscription_id: &str,
3142    view_spec: &ViewSpec,
3143    op: &str,
3144    key: &str,
3145    entity: &SharedEntity,
3146    seq: Option<&str>,
3147) -> Result<()> {
3148    let frame = ScopedEntityFrame {
3149        protocol_version: PROTOCOL_VERSION,
3150        subscription_id,
3151        mode: view_spec.mode,
3152        export: &view_spec.id,
3153        op,
3154        key,
3155        data: WireEntity {
3156            entity,
3157            wire_format: &view_spec.wire_format,
3158        },
3159        seq,
3160    };
3161    send_membership_json(context, view_spec, serde_json::to_vec(&frame)?)
3162}
3163
3164fn send_membership_json(
3165    context: &SubscriptionContext,
3166    view_spec: &ViewSpec,
3167    json: Vec<u8>,
3168) -> Result<()> {
3169    let encoded = Arc::new(Bytes::from(json));
3170    let bytes = encoded.len();
3171    let _held = hold_row_scopes(context)?;
3172    context
3173        .client_manager
3174        .send_to_client(context.client_id, encoded)
3175        .map_err(|error| anyhow::anyhow!("failed to send membership frame: {error}"))?;
3176    context.metrics.message_sent();
3177    emit_update_sent_for_client(
3178        &context.usage_emitter,
3179        &context.client_manager,
3180        context.client_id,
3181        &view_spec.id,
3182        bytes,
3183    );
3184    Ok(())
3185}
3186
3187/// Send what one source mutation owes a subscriber, and keep `undelivered`
3188/// (window members the client holds no copy of) in step with what was sent.
3189#[allow(clippy::too_many_arguments)]
3190fn emit_collection_delta(
3191    context: &SubscriptionContext,
3192    subscription_id: &str,
3193    view_spec: &ViewSpec,
3194    current: &[(String, SharedEntity)],
3195    next: &[(String, SharedEntity)],
3196    envelope: &BusMessage,
3197    metadata: &SourceFrameMetadata,
3198    undelivered: &mut HashSet<String>,
3199    sent_whole: &mut SentWhole,
3200) -> Result<()> {
3201    let current_keys: HashSet<&str> = current.iter().map(|(key, _)| key.as_str()).collect();
3202    let next_set: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
3203
3204    // In window order, as the frames go out.
3205    for key in current
3206        .iter()
3207        .map(|(key, _)| key.as_str())
3208        .filter(|key| !next_set.contains(key))
3209    {
3210        sent_whole.forget(key);
3211        // A key the client was never sent has nothing to remove.
3212        if undelivered.remove(key) {
3213            continue;
3214        }
3215        let op = if metadata.op == "delete" && key == envelope.key {
3216            "delete"
3217        } else {
3218            "remove"
3219        };
3220        send_membership_frame(
3221            context,
3222            subscription_id,
3223            view_spec,
3224            op,
3225            key,
3226            Value::Null,
3227            metadata.seq.clone(),
3228        )?;
3229    }
3230
3231    for (key, data) in next.iter() {
3232        let in_window = current_keys.contains(key.as_str());
3233        let action = match member_action(
3234            in_window,
3235            in_window && !undelivered.contains(key),
3236            key == &envelope.key,
3237            view_spec.is_derived(),
3238            &metadata.op,
3239        ) {
3240            // The row is in this session's scopes now (it is in `next`), but
3241            // the patch was published under the scope it had then. If that
3242            // scope is not this session's, send the row as it is now instead.
3243            MemberAction::ForwardPatch
3244                if !frame_in_row_scopes(
3245                    context.row_filter.as_ref(),
3246                    envelope.row_scope.as_ref(),
3247                    &envelope.payload,
3248                ) =>
3249            {
3250                MemberAction::Upsert
3251            }
3252            action => action,
3253        };
3254        match action {
3255            MemberAction::Skip => {}
3256            MemberAction::ForwardPatch => {
3257                // The upsert that brought the key into the window may already
3258                // carry this patch: it was read from the cache, which folds
3259                // a mutation in before its frame reaches this task.
3260                if sent_whole.covers(key, || {
3261                    source_frame_data(&envelope.payload)
3262                        .and_then(|data| data.get("_version").cloned())
3263                }) {
3264                    continue;
3265                }
3266                send_scoped_source_payload(
3267                    context,
3268                    subscription_id,
3269                    &view_spec.id,
3270                    envelope.payload.clone(),
3271                )?
3272            }
3273            MemberAction::Upsert => {
3274                let held = in_window && !undelivered.contains(key);
3275                // A derived view sends the mutated key whole every time; the
3276                // client already holds this exact version.
3277                if held && sent_whole.holds(key, data.version()) {
3278                    continue;
3279                }
3280                sent_whole.sent(key, data.version());
3281                let seq = metadata
3282                    .seq
3283                    .as_deref()
3284                    .or_else(|| data.field("_seq").and_then(Value::as_str));
3285                send_entity_frame(
3286                    context,
3287                    subscription_id,
3288                    view_spec,
3289                    "upsert",
3290                    key,
3291                    data,
3292                    seq,
3293                )?;
3294                undelivered.remove(key);
3295            }
3296        }
3297    }
3298    Ok(())
3299}
3300
3301/// The `_version` of each entity a list subscriber was last sent whole (an
3302/// upsert), so a source patch the upsert already carried is not sent again.
3303///
3304/// A list subscriber sends a key entering its window as the cached entity,
3305/// read when the subscriber handles the key's first mutation. The projector
3306/// folds a mutation into the cache before the subscriber handles its frame,
3307/// so when one key changes twice in quick succession (an entity created by
3308/// one instruction and completed by the next in the same transaction), the
3309/// upsert already holds the second change, and forwarding the second patch
3310/// as well repeats it. Versions are `<epoch>:<counter>`; a patch whose
3311/// version is not newer than the one sent whole is skipped.
3312///
3313/// Forgetting an entry only means a redundant patch may be forwarded, which
3314/// clients ignore by version, so the map is capped rather than kept exact.
3315#[derive(Default)]
3316struct SentWhole(HashMap<String, Value>);
3317
3318impl SentWhole {
3319    /// Entries kept per subscription; beyond it the map starts over.
3320    const MAX_ENTRIES: usize = 4096;
3321
3322    fn sent(&mut self, key: &str, version: Option<&Value>) {
3323        match version {
3324            Some(version) => {
3325                if self.0.len() >= Self::MAX_ENTRIES && !self.0.contains_key(key) {
3326                    self.0.clear();
3327                }
3328                self.0.insert(key.to_string(), version.clone());
3329            }
3330            None => {
3331                self.0.remove(key);
3332            }
3333        }
3334    }
3335
3336    fn forget(&mut self, key: &str) {
3337        self.0.remove(key);
3338    }
3339
3340    fn clear(&mut self) {
3341        self.0.clear();
3342    }
3343
3344    /// Whether the client holds exactly `version` of `key`.
3345    fn holds(&self, key: &str, version: Option<&Value>) -> bool {
3346        version.is_some() && self.0.get(key) == version
3347    }
3348
3349    /// Whether the entity sent whole for `key` already carries a patch whose
3350    /// version `patch_version` reads. Reads the patch only when `key` was
3351    /// sent whole. A newer patch (or one without a comparable version) ends
3352    /// the entry: the client then holds more than what was sent whole.
3353    fn covers(&mut self, key: &str, patch_version: impl FnOnce() -> Option<Value>) -> bool {
3354        let Some(sent) = self.0.get(key) else {
3355            return false;
3356        };
3357        let covered = match (
3358            version_parts(sent),
3359            patch_version().as_ref().and_then(version_parts),
3360        ) {
3361            (Some((sent_epoch, sent_counter)), Some((epoch, counter))) => {
3362                epoch == sent_epoch && counter <= sent_counter
3363            }
3364            _ => false,
3365        };
3366        if !covered {
3367            self.0.remove(key);
3368        }
3369        covered
3370    }
3371}
3372
3373/// `(epoch, counter)` of a `_version` (`<epoch>:<counter>`).
3374fn version_parts(version: &Value) -> Option<(String, u64)> {
3375    let (epoch, counter) = version.as_str()?.rsplit_once(':')?;
3376    Some((epoch.to_string(), counter.parse().ok()?))
3377}
3378
3379/// Emit the net effect of every source mutation seen during one coalescing
3380/// interval. Full entities are used for changed members because intermediate
3381/// sparse patches were intentionally discarded and can no longer be merged
3382/// safely by a client.
3383fn emit_coalesced_collection_delta(
3384    context: &SubscriptionContext,
3385    subscription_id: &str,
3386    view_spec: &ViewSpec,
3387    current: &[(String, SharedEntity)],
3388    next: &[(String, SharedEntity)],
3389    pending: &HashMap<String, Arc<BusMessage>>,
3390    undelivered: &mut HashSet<String>,
3391) -> Result<()> {
3392    let changes = plan_coalesced_collection_delta(current, next, pending, undelivered);
3393    // Every key that left the window or was just upserted is settled: the
3394    // client either never needs it or now holds all of it.
3395    let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
3396    undelivered.retain(|key| next_keys.contains(key.as_str()));
3397    for change in changes {
3398        if change.op == "upsert" {
3399            undelivered.remove(&change.key);
3400        }
3401        match &change.data {
3402            Some(entity) => send_entity_frame(
3403                context,
3404                subscription_id,
3405                view_spec,
3406                change.op,
3407                &change.key,
3408                entity,
3409                change.seq.as_deref(),
3410            )?,
3411            None => send_membership_frame(
3412                context,
3413                subscription_id,
3414                view_spec,
3415                change.op,
3416                &change.key,
3417                Value::Null,
3418                change.seq,
3419            )?,
3420        }
3421    }
3422    Ok(())
3423}
3424
3425#[derive(Debug, PartialEq)]
3426struct CollectionChange {
3427    op: &'static str,
3428    key: String,
3429    /// The entity an `upsert` sends; `None` for a `remove` or `delete`.
3430    data: Option<SharedEntity>,
3431    seq: Option<String>,
3432}
3433
3434/// Plan one coalescing interval. `undelivered` holds window members the
3435/// client was never sent; they get no `remove`/`delete` when they leave, and
3436/// like every changed key they are sent whole when they change.
3437fn plan_coalesced_collection_delta(
3438    current: &[(String, SharedEntity)],
3439    next: &[(String, SharedEntity)],
3440    pending: &HashMap<String, Arc<BusMessage>>,
3441    undelivered: &HashSet<String>,
3442) -> Vec<CollectionChange> {
3443    let current_by_key: HashMap<&str, &SharedEntity> = current
3444        .iter()
3445        .map(|(key, data)| (key.as_str(), data))
3446        .collect();
3447    let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
3448    let latest_seq = pending
3449        .values()
3450        .filter_map(|envelope| source_frame_metadata(&envelope.payload).seq)
3451        .max_by(|left, right| cmp_seq(left, right));
3452    let mut changes = Vec::new();
3453
3454    for (key, _) in current
3455        .iter()
3456        .filter(|(key, _)| !next_keys.contains(key.as_str()) && !undelivered.contains(key.as_str()))
3457    {
3458        let metadata = pending
3459            .get(key)
3460            .map(|envelope| source_frame_metadata(&envelope.payload));
3461        let op = if metadata
3462            .as_ref()
3463            .is_some_and(|metadata| metadata.op == "delete")
3464        {
3465            "delete"
3466        } else {
3467            "remove"
3468        };
3469        let seq = metadata
3470            .and_then(|metadata| metadata.seq)
3471            .or_else(|| latest_seq.clone());
3472        changes.push(CollectionChange {
3473            op,
3474            key: key.clone(),
3475            data: None,
3476            seq,
3477        });
3478    }
3479
3480    for (key, data) in next {
3481        let changed = current_by_key
3482            .get(key.as_str())
3483            .is_none_or(|previous| *previous != data);
3484        if !changed && !pending.contains_key(key) {
3485            continue;
3486        }
3487        let seq = data
3488            .field("_seq")
3489            .and_then(Value::as_str)
3490            .map(str::to_string)
3491            .or_else(|| {
3492                pending
3493                    .get(key)
3494                    .and_then(|envelope| source_frame_metadata(&envelope.payload).seq)
3495            })
3496            .or_else(|| latest_seq.clone());
3497        changes.push(CollectionChange {
3498            op: "upsert",
3499            key: key.clone(),
3500            data: Some(data.clone()),
3501            seq,
3502        });
3503    }
3504    changes
3505}
3506
3507/// What one in-window key owes a subscriber after a source mutation.
3508#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3509enum MemberAction {
3510    /// Send nothing: the subscriber already holds this entity and its data did
3511    /// not change.
3512    Skip,
3513    /// Forward the source patch verbatim.
3514    ForwardPatch,
3515    /// Send the whole entity.
3516    Upsert,
3517}
3518
3519/// Decide what to send for one key in the query window.
3520///
3521/// `in_window` says the key was in the window before this mutation; `held`
3522/// says the client also has a copy of it. The two differ after a snapshot
3523/// that `snapshotLimit` truncated or that was disabled: those keys count for
3524/// the window's positions, but the client was never sent them.
3525///
3526/// Only the mutated key's data changed, so a key that was already in the
3527/// window has at most moved index — and index is not the server's to
3528/// communicate. `subscribed` announces the window's sort (see
3529/// [`extract_sort_config`], which is `_seq` descending for a plain list) and
3530/// every SDK re-sorts locally from it, so resending an unchanged entity to
3531/// convey its new position is pure waste. This matters because on a
3532/// `_seq`-ordered list the mutated entity jumps to the front on *every*
3533/// mutation: keying the decision on position change meant rebroadcasting the
3534/// whole window each time, and the mutated entity itself never kept its
3535/// position long enough to have its patch forwarded.
3536fn member_action(
3537    in_window: bool,
3538    held: bool,
3539    is_mutated_key: bool,
3540    is_derived: bool,
3541    op: &str,
3542) -> MemberAction {
3543    if !is_mutated_key {
3544        // A key entering the window has no local state to merge into, so it
3545        // needs the whole entity. One already in the window is unchanged:
3546        // if the client holds it there is nothing new, and if it doesn't the
3547        // snapshot limit deliberately left it out, so it waits for its own
3548        // next change.
3549        return if in_window {
3550            MemberAction::Skip
3551        } else {
3552            MemberAction::Upsert
3553        };
3554    }
3555    // The mutated entity rides its own patch through untouched, but only when
3556    // the subscriber already holds a copy; a patch for a key the client never
3557    // received is not an entity. Derived views still send whole entities: the
3558    // patch on the bus is scoped to the source view, not this one (see
3559    // A4-150).
3560    if held && !is_derived && op != "delete" {
3561        MemberAction::ForwardPatch
3562    } else {
3563        MemberAction::Upsert
3564    }
3565}
3566
3567/// A copy of `entity` with the view's wire format applied, for code that
3568/// needs it as a value (a catch-up diff) rather than serialized.
3569fn wire_value(entity: &SharedEntity, wire_format: &WireFormat) -> Value {
3570    let mut value = entity.to_value();
3571    apply_wire_format(&mut value, wire_format);
3572    value
3573}
3574
3575async fn load_query_entities(
3576    entity_cache: &EntityCache,
3577    sorted_caches: Option<
3578        Arc<tokio::sync::RwLock<HashMap<String, crate::sorted_cache::SortedViewCache>>>,
3579    >,
3580    view_spec: &ViewSpec,
3581    query: &SubscriptionQuery,
3582    row_filter: Option<&RowFilter>,
3583    apply_snapshot_limit: bool,
3584) -> Vec<(String, SharedEntity)> {
3585    // Rows share their fields with the caches; nothing here copies an entity.
3586    let mut query = query.clone();
3587    if let Some(limit) = view_spec
3588        .pipeline
3589        .as_ref()
3590        .and_then(|pipeline| pipeline.limit)
3591    {
3592        query.take = Some(query.take.unwrap_or(limit).min(limit));
3593    }
3594    if let Some(sorted_caches) = sorted_caches {
3595        // A sorted cache is read in order, and only as far as the window
3596        // reaches: rows past it are neither visited nor shared.
3597        let caches = sorted_caches.read().await;
3598        if let Some(cache) = caches.get(&view_spec.id) {
3599            // A row-scoped view keeps its rows per scope: read only the
3600            // partitions this session may see, so its window is as full as
3601            // its own rows allow.
3602            let no_scopes = std::collections::BTreeSet::new();
3603            let rows: Box<dyn Iterator<Item = (&str, &SharedEntity)> + '_> = match row_filter {
3604                Some(filter) => Box::new(cache.iter_ordered_in_scopes(filter.scopes())),
3605                None if cache.is_scope_partitioned() => {
3606                    Box::new(cache.iter_ordered_in_scopes(&no_scopes))
3607                }
3608                None => Box::new(cache.iter_ordered()),
3609            };
3610            return select_window(
3611                rows.filter(|(key, data)| query_matches_entity(&query, key, *data, row_filter))
3612                    .map(|(key, data)| (key.to_string(), data.clone())),
3613                &query,
3614                apply_snapshot_limit,
3615            );
3616        }
3617    }
3618    let (entities, preordered) = if view_spec.is_derived() {
3619        // Empty and filter-only pipelines have no sorted cache. Evaluate the
3620        // source rows, retaining the pipeline predicate before query selection.
3621        let mut entities = entity_cache
3622            .get_all_shared(view_spec.source_view.as_deref().unwrap_or(&view_spec.id))
3623            .await;
3624        if let Some(filter) = view_spec
3625            .pipeline
3626            .as_ref()
3627            .and_then(|pipeline| pipeline.filter.as_ref())
3628        {
3629            entities.retain(|(_, data)| filter.matches(data));
3630        }
3631        (entities, false)
3632    } else if view_spec.mode == Mode::State {
3633        let entity = match query.key.as_deref() {
3634            Some(key) => entity_cache
3635                .get_shared(&view_spec.id, key)
3636                .await
3637                .map(|data| vec![(key.to_string(), data)])
3638                .unwrap_or_default(),
3639            None => vec![],
3640        };
3641        (entity, true)
3642    } else {
3643        (entity_cache.get_all_shared(&view_spec.id).await, false)
3644    };
3645    select_query_entities(
3646        entities,
3647        &query,
3648        row_filter,
3649        preordered,
3650        apply_snapshot_limit,
3651    )
3652}
3653
3654fn select_query_entities<E: EntityFields>(
3655    mut entities: Vec<(String, E)>,
3656    query: &SubscriptionQuery,
3657    row_filter: Option<&RowFilter>,
3658    preordered: bool,
3659    apply_snapshot_limit: bool,
3660) -> Vec<(String, E)> {
3661    entities.retain(|(key, data)| query_matches_entity(query, key, data, row_filter));
3662    if !preordered {
3663        entities.sort_by(|left, right| {
3664            let left_seq = left.1.field("_seq").and_then(Value::as_str).unwrap_or("");
3665            let right_seq = right.1.field("_seq").and_then(Value::as_str).unwrap_or("");
3666            let order = if query.after.is_some() {
3667                cmp_seq(left_seq, right_seq)
3668            } else {
3669                cmp_seq(right_seq, left_seq)
3670            };
3671            order.then_with(|| left.0.cmp(&right.0))
3672        });
3673    }
3674
3675    select_window(entities.into_iter(), query, apply_snapshot_limit)
3676}
3677
3678/// The `[skip, skip + take)` window of rows already filtered and in order,
3679/// cut to the snapshot limit when it applies. Takes no row past the window.
3680fn select_window<E>(
3681    rows: impl Iterator<Item = (String, E)>,
3682    query: &SubscriptionQuery,
3683    apply_snapshot_limit: bool,
3684) -> Vec<(String, E)> {
3685    let skip = query.skip.unwrap_or(0);
3686    let mut take = query.take.unwrap_or(usize::MAX);
3687    if apply_snapshot_limit {
3688        if let Some(limit) = query.snapshot_limit {
3689            take = take.min(limit);
3690        }
3691    }
3692    rows.skip(skip).take(take).collect()
3693}
3694
3695/// Whether a whole row belongs to the subscription: in its row scopes (when
3696/// the view is row-scoped), then its key, `partition`, `after` and `filters`.
3697/// The client-supplied predicates only narrow; none of them can admit a row
3698/// the row filter refuses.
3699fn query_matches_entity<E: EntityFields + ?Sized>(
3700    query: &SubscriptionQuery,
3701    key: &str,
3702    data: &E,
3703    row_filter: Option<&RowFilter>,
3704) -> bool {
3705    if row_filter.is_some_and(|filter| !filter.allows(data)) {
3706        return false;
3707    }
3708    if !query.matches_key(key) {
3709        return false;
3710    }
3711    if let Some(partition) = &query.partition {
3712        if value_at_dot_path(data, "_partition").as_deref()
3713            != Some(&Value::String(partition.clone()))
3714        {
3715            return false;
3716        }
3717    }
3718    if let Some(after) = &query.after {
3719        let Some(seq) = data.field("_seq").and_then(Value::as_str) else {
3720            return false;
3721        };
3722        if cmp_seq(seq, after) != std::cmp::Ordering::Greater {
3723            return false;
3724        }
3725    }
3726    query
3727        .filters
3728        .iter()
3729        .all(|(path, expected)| value_at_dot_path(data, path).as_deref() == Some(expected))
3730}
3731
3732fn value_at_dot_path<'a, E: EntityFields + ?Sized>(
3733    entity: &'a E,
3734    path: &str,
3735) -> Option<std::borrow::Cow<'a, Value>> {
3736    lookup(entity, path.split('.'))
3737}
3738
3739#[cfg(test)]
3740mod tests {
3741    use super::*;
3742    use crate::cache::{EntityCacheConfig, PatchOrigin};
3743    use crate::view::{Delivery, Filters, Projection};
3744    use serde_json::json;
3745    use tokio::sync::oneshot;
3746
3747    #[derive(Default)]
3748    struct DelayedUsageEmitter {
3749        started: tokio::sync::Notify,
3750        release: tokio::sync::Notify,
3751        events: tokio::sync::Mutex<Vec<WebSocketUsageEvent>>,
3752    }
3753
3754    #[async_trait::async_trait]
3755    impl WebSocketUsageEmitter for DelayedUsageEmitter {
3756        async fn emit(&self, event: WebSocketUsageEvent) {
3757            self.started.notify_one();
3758            self.release.notified().await;
3759            self.events.lock().await.push(event);
3760        }
3761    }
3762
3763    #[tokio::test]
3764    async fn usage_emitter_handle_waits_for_spawned_emissions() {
3765        let emitter = Arc::new(DelayedUsageEmitter::default());
3766        let handle = UsageEmitterHandle::new(emitter.clone());
3767        handle.emit(WebSocketUsageEvent::ConnectionEstablished {
3768            client_id: "client-1".to_string(),
3769            remote_addr: "127.0.0.1:1234".to_string(),
3770            deployment_id: Some("1".to_string()),
3771            identity: UsageIdentity::default(),
3772        });
3773        emitter.started.notified().await;
3774
3775        let closing_handle = handle.clone();
3776        let closing = tokio::spawn(async move {
3777            closing_handle.close_and_wait().await;
3778        });
3779        tokio::task::yield_now().await;
3780        assert!(!closing.is_finished());
3781
3782        emitter.release.notify_one();
3783        closing.await.expect("usage task wait should finish");
3784        assert_eq!(emitter.events.lock().await.len(), 1);
3785    }
3786
3787    #[test]
3788    fn a_catch_up_carries_only_the_fields_that_changed() {
3789        let base = json!({
3790            "id": 7,
3791            "state": {"deployed": [1, 2, 3], "total": 6},
3792            "results": {"square": null},
3793            "gone": true,
3794        });
3795        let current = json!({
3796            "id": 7,
3797            "state": {"deployed": [1, 2, 4], "total": 6},
3798            "results": {"square": 3},
3799            "fresh": "x",
3800        });
3801        assert_eq!(
3802            changed_fields(&base, &current, &Value::Null),
3803            Some(json!({
3804                // An array that changed is sent whole.
3805                "state": {"deployed": [1, 2, 4]},
3806                "results": {"square": 3},
3807                "fresh": "x",
3808                // A field the entity no longer has is cleared.
3809                "gone": null,
3810            }))
3811        );
3812        assert_eq!(changed_fields(&current, &current, &Value::Null), None);
3813        // Anything that is not an object on both sides is compared whole.
3814        assert_eq!(
3815            changed_fields(&json!(null), &current, &Value::Null),
3816            Some(current.clone())
3817        );
3818    }
3819
3820    #[test]
3821    fn a_catch_up_resends_the_fields_forwarded_patches_set() {
3822        let base = json!({"a": 1, "state": {"x": 1, "y": 1}, "list": [1]});
3823        let mut touched = Value::Null;
3824        mark_fields(&mut touched, &json!({"a": 2, "state": {"x": 2}}));
3825        mark_fields(&mut touched, &json!({"list": [2], "extra": true}));
3826        assert_eq!(
3827            touched,
3828            json!({"a": true, "state": {"x": true}, "list": true, "extra": true})
3829        );
3830        // The entity set `a` and `state.x` back and never kept `extra`: the
3831        // client still holds what the forwarded patches set, so a catch-up
3832        // compared with `base` alone would leave it stale.
3833        let current = json!({"a": 1, "state": {"x": 1, "y": 1}, "list": [1, 2]});
3834        assert_eq!(
3835            changed_fields(&base, &current, &Value::Null),
3836            Some(json!({"list": [1, 2]}))
3837        );
3838        assert_eq!(
3839            changed_fields(&base, &current, &touched),
3840            Some(json!({"a": 1, "state": {"x": 1}, "list": [1, 2], "extra": null}))
3841        );
3842        // A patch whose data is not an object marks everything.
3843        mark_fields(&mut touched, &json!(null));
3844        assert_eq!(
3845            changed_fields(&base, &current, &touched),
3846            Some(current.clone())
3847        );
3848    }
3849
3850    fn list_spec() -> ViewSpec {
3851        ViewSpec {
3852            id: "Thing/list".to_string(),
3853            export: "Thing".to_string(),
3854            mode: Mode::List,
3855            wire_format: Default::default(),
3856            projection: Projection::all(),
3857            filters: Filters::all(),
3858            delivery: Delivery::default(),
3859            pipeline: None,
3860            source_view: None,
3861        }
3862    }
3863
3864    /// Plan the whole window the way `emit_collection_delta` does, so a test
3865    /// can assert on what a subscriber is actually sent.
3866    fn plan_window(
3867        current_keys: &[&str],
3868        next_keys: &[&str],
3869        envelope_key: &str,
3870        is_derived: bool,
3871        op: &str,
3872    ) -> Vec<(String, MemberAction)> {
3873        plan_window_with_undelivered(current_keys, &[], next_keys, envelope_key, is_derived, op)
3874    }
3875
3876    /// [`plan_window`] for a subscriber that was never sent `undelivered`.
3877    fn plan_window_with_undelivered(
3878        current_keys: &[&str],
3879        undelivered: &[&str],
3880        next_keys: &[&str],
3881        envelope_key: &str,
3882        is_derived: bool,
3883        op: &str,
3884    ) -> Vec<(String, MemberAction)> {
3885        next_keys
3886            .iter()
3887            .map(|key| {
3888                let in_window = current_keys.contains(key);
3889                let held = in_window && !undelivered.contains(key);
3890                (
3891                    (*key).to_string(),
3892                    member_action(in_window, held, *key == envelope_key, is_derived, op),
3893                )
3894            })
3895            .collect()
3896    }
3897
3898    #[test]
3899    fn a_window_member_the_client_never_received_is_sent_whole() {
3900        // snapshotLimit cut the window [4, 3, 2, 1] to [4, 3]. "1" still
3901        // counts for positions, but a patch for it would reach a client with
3902        // nothing to merge into.
3903        let plan = plan_window_with_undelivered(
3904            &["4", "3", "2", "1"],
3905            &["2", "1"],
3906            &["1", "4", "3", "2"],
3907            "1",
3908            false,
3909            "patch",
3910        );
3911        assert_eq!(
3912            plan,
3913            vec![
3914                ("1".to_string(), MemberAction::Upsert),
3915                ("4".to_string(), MemberAction::Skip),
3916                ("3".to_string(), MemberAction::Skip),
3917                // Unchanged and never sent: it waits for its own change
3918                // rather than undoing the snapshot limit.
3919                ("2".to_string(), MemberAction::Skip),
3920            ]
3921        );
3922    }
3923
3924    #[test]
3925    fn a_held_member_still_rides_its_patch_after_a_truncated_snapshot() {
3926        let plan = plan_window_with_undelivered(
3927            &["4", "3", "2", "1"],
3928            &["2", "1"],
3929            &["4", "3", "2", "1"],
3930            "4",
3931            false,
3932            "patch",
3933        );
3934        assert_eq!(plan[0], ("4".to_string(), MemberAction::ForwardPatch));
3935    }
3936
3937    #[test]
3938    fn coalescing_sends_no_remove_for_a_key_the_client_never_received() {
3939        let current = vec![
3940            ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3941            (
3942                "unsent".to_string(),
3943                json!({"count": 1, "_seq": "10:000001"}),
3944            ),
3945            ("gone".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3946        ];
3947        let next = vec![
3948            ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3949            (
3950                "unsent".to_string(),
3951                json!({"count": 2, "_seq": "10:000003"}),
3952            ),
3953        ];
3954        let pending = HashMap::from([
3955            (
3956                "unsent".to_string(),
3957                list_message("unsent", "patch", "10:000003"),
3958            ),
3959            (
3960                "gone".to_string(),
3961                list_message("gone", "delete", "10:000002"),
3962            ),
3963        ]);
3964        let undelivered = HashSet::from(["unsent".to_string(), "gone".to_string()]);
3965
3966        assert_eq!(
3967            plan_coalesced_collection_delta(
3968                &shared(&current),
3969                &shared(&next),
3970                &pending,
3971                &undelivered
3972            ),
3973            vec![CollectionChange {
3974                op: "upsert",
3975                key: "unsent".to_string(),
3976                data: Some(SharedEntity::new(json!({"count": 2, "_seq": "10:000003"}))),
3977                seq: Some("10:000003".to_string()),
3978            }]
3979        );
3980    }
3981
3982    #[test]
3983    fn a_reordered_window_sends_only_the_mutated_entity() {
3984        // A `_seq`-descending list: mutating "1" moves it to the front and
3985        // shifts every other key down one. Only "1" changed, so only "1" is
3986        // sent — and it rides its own patch, not a full entity.
3987        let current = ["4", "3", "2", "1"];
3988        let next = ["1", "4", "3", "2"];
3989        let plan = plan_window(&current, &next, "1", false, "patch");
3990
3991        assert_eq!(
3992            plan,
3993            vec![
3994                ("1".to_string(), MemberAction::ForwardPatch),
3995                ("4".to_string(), MemberAction::Skip),
3996                ("3".to_string(), MemberAction::Skip),
3997                ("2".to_string(), MemberAction::Skip),
3998            ]
3999        );
4000    }
4001
4002    #[test]
4003    fn a_patch_the_whole_entity_already_carried_is_not_sent_again() {
4004        let mut sent = SentWhole::default();
4005        sent.sent("k", Some(&json!("e1:5")));
4006        // Older or equal patches were folded into what was sent.
4007        assert!(sent.covers("k", || Some(json!("e1:4"))));
4008        assert!(sent.covers("k", || Some(json!("e1:5"))));
4009        assert!(sent.holds("k", Some(&json!("e1:5"))));
4010        // A newer one is forwarded and ends the entry.
4011        assert!(!sent.covers("k", || Some(json!("e1:6"))));
4012        assert!(!sent.covers("k", || Some(json!("e1:3"))));
4013        // Another epoch or no version is never taken as covered.
4014        sent.sent("k", Some(&json!("e1:5")));
4015        assert!(!sent.covers("k", || Some(json!("e2:1"))));
4016        sent.sent("k", Some(&json!("e1:5")));
4017        assert!(!sent.covers("k", || None));
4018        // A key never sent whole reads no patch.
4019        assert!(!sent.covers("other", || panic!("read only for keys sent whole")));
4020        // An unversioned upsert leaves nothing to compare against.
4021        sent.sent("k", Some(&json!("e1:5")));
4022        sent.sent("k", None);
4023        assert!(!sent.holds("k", None));
4024        assert!(!sent.covers("k", || Some(json!("e1:1"))));
4025    }
4026
4027    #[test]
4028    fn a_key_entering_the_window_gets_the_whole_entity() {
4029        // "5" has no local state for the subscriber to merge a patch into.
4030        let plan = plan_window(&["4", "3"], &["5", "4", "3"], "5", false, "patch");
4031        assert_eq!(
4032            plan,
4033            vec![
4034                ("5".to_string(), MemberAction::Upsert),
4035                ("4".to_string(), MemberAction::Skip),
4036                ("3".to_string(), MemberAction::Skip),
4037            ]
4038        );
4039    }
4040
4041    #[test]
4042    fn derived_views_still_send_whole_entities() {
4043        // The patch on the bus is scoped to the source view, so a derived
4044        // subscription cannot forward it verbatim (A4-150).
4045        let plan = plan_window(&["1", "2"], &["1", "2"], "1", true, "patch");
4046        assert_eq!(
4047            plan,
4048            vec![
4049                ("1".to_string(), MemberAction::Upsert),
4050                ("2".to_string(), MemberAction::Skip),
4051            ]
4052        );
4053    }
4054
4055    #[test]
4056    fn a_delete_envelope_never_forwards_a_patch() {
4057        // A surviving key on a delete envelope carries no mergeable patch.
4058        let plan = plan_window(&["1", "2"], &["1", "2"], "1", false, "delete");
4059        assert_eq!(
4060            plan,
4061            vec![
4062                ("1".to_string(), MemberAction::Upsert),
4063                ("2".to_string(), MemberAction::Skip),
4064            ]
4065        );
4066    }
4067
4068    #[test]
4069    fn an_unchanged_window_sends_one_frame_not_a_broadcast() {
4070        // The regression this guards: 500 members used to mean 500 full
4071        // entities on the wire for a single mutation.
4072        let keys: Vec<String> = (0..500).map(|index| index.to_string()).collect();
4073        let refs: Vec<&str> = keys.iter().map(String::as_str).collect();
4074        let plan = plan_window(&refs, &refs, "250", false, "patch");
4075
4076        let sent = plan
4077            .iter()
4078            .filter(|(_, action)| *action != MemberAction::Skip)
4079            .count();
4080        assert_eq!(sent, 1);
4081        assert_eq!(plan[250].1, MemberAction::ForwardPatch);
4082    }
4083
4084    fn shared(rows: &[(String, Value)]) -> Vec<(String, SharedEntity)> {
4085        rows.iter()
4086            .map(|(key, data)| (key.clone(), SharedEntity::new(data.clone())))
4087            .collect()
4088    }
4089
4090    fn list_message(key: &str, op: &str, seq: &str) -> Arc<BusMessage> {
4091        Arc::new(BusMessage {
4092            row_scope: None,
4093            key: key.to_string(),
4094            entity: "Thing/list".to_string(),
4095            payload: Arc::new(Bytes::from(
4096                serde_json::to_vec(&json!({
4097                    "entity": "Thing/list",
4098                    "op": op,
4099                    "key": key,
4100                    "seq": seq,
4101                    "data": {},
4102                }))
4103                .unwrap(),
4104            )),
4105        })
4106    }
4107
4108    #[test]
4109    fn coalescing_emits_only_the_final_full_state_per_changed_key() {
4110        let current = vec![
4111            ("a".to_string(), json!({"count": 1, "_seq": "10:000001"})),
4112            ("b".to_string(), json!({"count": 1, "_seq": "10:000001"})),
4113            (
4114                "stable".to_string(),
4115                json!({"count": 1, "_seq": "10:000001"}),
4116            ),
4117        ];
4118        let next = vec![
4119            ("a".to_string(), json!({"count": 3, "_seq": "10:000004"})),
4120            ("c".to_string(), json!({"count": 1, "_seq": "10:000003"})),
4121            (
4122                "stable".to_string(),
4123                json!({"count": 1, "_seq": "10:000001"}),
4124            ),
4125        ];
4126        let pending = HashMap::from([
4127            ("a".to_string(), list_message("a", "patch", "10:000004")),
4128            ("b".to_string(), list_message("b", "delete", "10:000002")),
4129            ("c".to_string(), list_message("c", "patch", "10:000003")),
4130        ]);
4131
4132        assert_eq!(
4133            plan_coalesced_collection_delta(
4134                &shared(&current),
4135                &shared(&next),
4136                &pending,
4137                &HashSet::new()
4138            ),
4139            vec![
4140                CollectionChange {
4141                    op: "delete",
4142                    key: "b".to_string(),
4143                    data: None,
4144                    seq: Some("10:000002".to_string()),
4145                },
4146                CollectionChange {
4147                    op: "upsert",
4148                    key: "a".to_string(),
4149                    data: Some(SharedEntity::new(json!({"count": 3, "_seq": "10:000004"}))),
4150                    seq: Some("10:000004".to_string()),
4151                },
4152                CollectionChange {
4153                    op: "upsert",
4154                    key: "c".to_string(),
4155                    data: Some(SharedEntity::new(json!({"count": 1, "_seq": "10:000003"}))),
4156                    seq: Some("10:000003".to_string()),
4157                },
4158            ]
4159        );
4160    }
4161
4162    #[test]
4163    fn a_lagged_snapshotless_subscription_gets_a_fatal_retryable_error() {
4164        let issue = SocketIssueMessage::subscription_lagged("balances".to_string(), 42);
4165        assert_eq!(issue.code, "subscription-lagged");
4166        assert!(issue.retryable);
4167        assert!(issue.fatal);
4168        assert!(issue.message.contains("42"));
4169        assert!(issue
4170            .suggested_action
4171            .unwrap()
4172            .contains("snapshots enabled"));
4173
4174        let append = SocketIssueMessage::append_subscription_lagged("trades".to_string(), 9);
4175        assert!(append.fatal);
4176        assert!(append.suggested_action.unwrap().contains("retained replay"));
4177    }
4178
4179    #[test]
4180    fn snapshot_batches_share_identity_and_completion() {
4181        let entities = ["one", "two", "three"]
4182            .map(|key| (key.to_string(), SharedEntity::new(json!({"key": key}))));
4183        let wire_format = WireFormat::default();
4184        let batches = create_snapshot_batches(
4185            &entities,
4186            &wire_format,
4187            SnapshotMetadata {
4188                subscription_id: "sub-1",
4189                snapshot_id: "snapshot-1",
4190                authoritative: true,
4191                mode: Mode::List,
4192                view_id: "Thing/list",
4193                key: None,
4194            },
4195            &SnapshotBatchConfig {
4196                initial_batch_size: 2,
4197                subsequent_batch_size: 1,
4198            },
4199        );
4200        assert_eq!(batches.len(), 2);
4201        assert!(batches.iter().all(|batch| batch.subscription_id == "sub-1"));
4202        assert!(batches
4203            .iter()
4204            .all(|batch| batch.snapshot_id == "snapshot-1"));
4205        assert!(!batches[0].complete);
4206        assert!(batches[1].complete);
4207        assert!(batches.iter().all(|batch| batch.authoritative));
4208    }
4209
4210    #[test]
4211    fn empty_incremental_snapshot_is_explicitly_non_authoritative() {
4212        let wire_format = WireFormat::default();
4213        let batches = create_snapshot_batches(
4214            &[],
4215            &wire_format,
4216            SnapshotMetadata {
4217                subscription_id: "sub-1",
4218                snapshot_id: "snapshot-1",
4219                authoritative: false,
4220                mode: Mode::State,
4221                view_id: "Thing/state",
4222                key: Some("missing"),
4223            },
4224            &SnapshotBatchConfig {
4225                initial_batch_size: 1,
4226                subsequent_batch_size: 1,
4227            },
4228        );
4229        assert_eq!(batches.len(), 1);
4230        assert!(!batches[0].authoritative);
4231        assert!(batches[0].complete);
4232        assert_eq!(batches[0].key, Some("missing"));
4233    }
4234
4235    fn wide_ints() -> WireFormat {
4236        WireFormat {
4237            wide_int_paths: vec![
4238                vec!["amount".to_string()],
4239                vec!["trades".to_string(), "price".to_string()],
4240            ],
4241        }
4242    }
4243
4244    fn sample_rows() -> Vec<(String, SharedEntity)> {
4245        (0..7)
4246            .map(|index| {
4247                (
4248                    index.to_string(),
4249                    SharedEntity::new(json!({
4250                        "_seq": format!("10:{index:012}"),
4251                        "_version": format!("e:{index}"),
4252                        "amount": u64::MAX - index,
4253                        "label": "Ï€ \"quoted\" \n",
4254                        "trades": [{"price": index, "side": "buy"}, {"price": -1}],
4255                    })),
4256                )
4257            })
4258            .collect()
4259    }
4260
4261    /// Snapshot frames serialized from shared rows are byte for byte the
4262    /// frames built from formatted copies of the rows.
4263    #[test]
4264    fn snapshot_batches_serialize_as_frames_of_formatted_copies() {
4265        use crate::websocket::frame::{SnapshotEntity, SnapshotFrame};
4266        let rows = sample_rows();
4267        for (rows, key, wire_format, (initial, subsequent)) in [
4268            (&rows[..], None, wide_ints(), (2, 3)),
4269            (&rows[..], Some("3"), WireFormat::default(), (50, 100)),
4270            (&rows[..1], Some("0"), wide_ints(), (1, 1)),
4271            (&rows[..0], None, wide_ints(), (2, 3)),
4272        ] {
4273            let metadata = SnapshotMetadata {
4274                subscription_id: "sub-1",
4275                snapshot_id: "snapshot-1",
4276                authoritative: true,
4277                mode: Mode::List,
4278                view_id: "Thing/list",
4279                key,
4280            };
4281            let config = SnapshotBatchConfig {
4282                initial_batch_size: initial,
4283                subsequent_batch_size: subsequent,
4284            };
4285            let batches = create_snapshot_batches(rows, &wire_format, metadata, &config);
4286            let mut offset = 0;
4287            for batch in &batches {
4288                let copies = rows[offset..offset + batch.data.rows.len()]
4289                    .iter()
4290                    .map(|(key, entity)| SnapshotEntity {
4291                        key: key.clone(),
4292                        data: wire_value(entity, &wire_format),
4293                    })
4294                    .collect();
4295                offset += batch.data.rows.len();
4296                let frame = SnapshotFrame {
4297                    protocol_version: PROTOCOL_VERSION,
4298                    subscription_id: "sub-1".to_string(),
4299                    snapshot_id: "snapshot-1".to_string(),
4300                    authoritative: true,
4301                    mode: Mode::List,
4302                    export: "Thing/list".to_string(),
4303                    op: "snapshot",
4304                    key: key.map(str::to_string),
4305                    data: copies,
4306                    complete: batch.complete,
4307                };
4308                assert_eq!(
4309                    serde_json::to_string(batch).unwrap(),
4310                    serde_json::to_string(&frame).unwrap()
4311                );
4312            }
4313            assert_eq!(offset, rows.len());
4314        }
4315    }
4316
4317    /// A membership frame serialized from a shared entity is byte for byte
4318    /// the frame built from a formatted copy of it.
4319    #[test]
4320    fn entity_frames_serialize_as_frames_of_formatted_copies() {
4321        for (_, entity) in sample_rows() {
4322            for (wire_format, seq) in [
4323                (wide_ints(), Some("10:000000000001")),
4324                (WireFormat::default(), None),
4325            ] {
4326                let frame = ScopedEntityFrame {
4327                    protocol_version: PROTOCOL_VERSION,
4328                    subscription_id: "sub-1",
4329                    mode: Mode::List,
4330                    export: "Thing/latest",
4331                    op: "upsert",
4332                    key: "7",
4333                    data: WireEntity {
4334                        entity: &entity,
4335                        wire_format: &wire_format,
4336                    },
4337                    seq,
4338                };
4339                let copy = Frame::scoped(
4340                    "sub-1",
4341                    Mode::List,
4342                    "Thing/latest",
4343                    "upsert",
4344                    "7",
4345                    wire_value(&entity, &wire_format),
4346                    seq.map(str::to_string),
4347                );
4348                assert_eq!(
4349                    serde_json::to_string(&frame).unwrap(),
4350                    serde_json::to_string(&copy).unwrap()
4351                );
4352            }
4353        }
4354    }
4355
4356    #[test]
4357    fn lag_recovery_is_authoritative_for_an_after_query() {
4358        let subscription = Subscription {
4359            protocol_version: PROTOCOL_VERSION,
4360            subscription_id: "sub-1".to_string(),
4361            query: SubscriptionQuery {
4362                view: "Thing/list".to_string(),
4363                after: Some("40:000000000010".to_string()),
4364                ..Default::default()
4365            },
4366            snapshot: Default::default(),
4367        };
4368
4369        assert!(!SnapshotPurpose::Initial.authoritative(&subscription));
4370        assert!(SnapshotPurpose::Recovery.authoritative(&subscription));
4371    }
4372
4373    #[test]
4374    fn dot_path_filters_are_exact_and_type_sensitive() {
4375        let mut query = SubscriptionQuery {
4376            view: "Thing/list".to_string(),
4377            ..Default::default()
4378        };
4379        query
4380            .filters
4381            .insert("state.status".to_string(), json!("open"));
4382        query.filters.insert("metrics.count".to_string(), json!(2));
4383        assert!(query_matches_entity(
4384            &query,
4385            "one",
4386            &json!({"state": {"status": "open"}, "metrics": {"count": 2}}),
4387            None,
4388        ));
4389        assert!(!query_matches_entity(
4390            &query,
4391            "one",
4392            &json!({"state": {"status": "open"}, "metrics": {"count": "2"}}),
4393            None,
4394        ));
4395    }
4396
4397    #[test]
4398    fn take_and_skip_define_independent_deterministic_windows() {
4399        let entities: Vec<_> = (1..=6)
4400            .map(|id| {
4401                (
4402                    id.to_string(),
4403                    json!({"id": id, "_seq": format!("10:{id:012}")}),
4404                )
4405            })
4406            .collect();
4407        let first = SubscriptionQuery {
4408            view: "Thing/list".to_string(),
4409            take: Some(2),
4410            skip: Some(0),
4411            ..Default::default()
4412        };
4413        let second = SubscriptionQuery {
4414            skip: Some(2),
4415            ..first.clone()
4416        };
4417        let first_keys: Vec<_> =
4418            select_query_entities(entities.clone(), &first, None, false, false)
4419                .into_iter()
4420                .map(|(key, _)| key)
4421                .collect();
4422        let second_keys: Vec<_> = select_query_entities(entities, &second, None, false, false)
4423            .into_iter()
4424            .map(|(key, _)| key)
4425            .collect();
4426        assert_eq!(first_keys, ["6", "5"]);
4427        assert_eq!(second_keys, ["4", "3"]);
4428    }
4429
4430    #[tokio::test]
4431    async fn state_receiver_is_installed_before_snapshot_awaits() {
4432        let bus = BusManager::new();
4433        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
4434        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
4435        let bus_for_task = bus.clone();
4436        let task = tokio::spawn(async move {
4437            subscribe_state_then_snapshot(&bus_for_task, "Thing/state", "one", || async move {
4438                snapshot_started_tx.send(()).unwrap();
4439                release_snapshot_rx.await.unwrap();
4440            })
4441            .await
4442            .0
4443        });
4444        snapshot_started_rx.await.unwrap();
4445        bus.publish_state(
4446            "Thing/state",
4447            "one",
4448            Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
4449            None,
4450        )
4451        .await;
4452        release_snapshot_tx.send(()).unwrap();
4453        let mut receiver = task.await.unwrap();
4454        receiver.changed().await.unwrap();
4455        assert!(!receiver.borrow().payload.is_empty());
4456    }
4457
4458    async fn assert_list_receiver_precedes_snapshot(view: &'static str) {
4459        let bus = BusManager::new();
4460        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
4461        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
4462        let bus_for_task = bus.clone();
4463        let task = tokio::spawn(async move {
4464            subscribe_list_then_snapshot(&bus_for_task, view, || async move {
4465                snapshot_started_tx.send(()).unwrap();
4466                release_snapshot_rx.await.unwrap();
4467            })
4468            .await
4469            .0
4470        });
4471        snapshot_started_rx.await.unwrap();
4472        bus.publish_list(
4473            view,
4474            Arc::new(BusMessage {
4475                row_scope: None,
4476                key: "one".to_string(),
4477                entity: view.to_string(),
4478                payload: Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
4479            }),
4480        )
4481        .await;
4482        release_snapshot_tx.send(()).unwrap();
4483        let mut receiver = task.await.unwrap();
4484        assert_eq!(receiver.recv().await.unwrap().key, "one");
4485    }
4486
4487    #[tokio::test]
4488    async fn list_receiver_is_installed_before_snapshot() {
4489        assert_list_receiver_precedes_snapshot("Thing/list").await;
4490    }
4491
4492    #[tokio::test]
4493    async fn append_receiver_is_installed_before_snapshot() {
4494        assert_list_receiver_precedes_snapshot("Thing/append").await;
4495    }
4496
4497    #[tokio::test]
4498    async fn derived_source_receiver_is_installed_before_snapshot() {
4499        assert_list_receiver_precedes_snapshot("Thing/list-source").await;
4500    }
4501
4502    /// A sorted derived view is read only as far as the window reaches, and
4503    /// gives exactly the rows filtering and windowing all of it gives.
4504    #[tokio::test]
4505    async fn sorted_derived_reads_take_the_window_of_the_whole_order() {
4506        use crate::materialized_view::{SortConfig, SortOrder, ViewPipeline};
4507        use crate::sorted_cache::{SortOrder as CacheOrder, SortedViewCache};
4508        let mut sorted =
4509            SortedViewCache::new("Thing/top".into(), vec!["score".into()], CacheOrder::Desc);
4510        for id in 0..40u64 {
4511            sorted.upsert(
4512                id.to_string(),
4513                json!({
4514                    "_seq": format!("{}:0", 100 + id * 7 % 40),
4515                    "score": (id * 13 % 17).to_string(),
4516                    "owner": if id % 3 == 0 { "alice" } else { "bob" },
4517                    "_partition": if id % 2 == 0 { "even" } else { "odd" },
4518                }),
4519            );
4520        }
4521        let whole = sorted.ordered_entities();
4522        let mut spec = list_spec();
4523        spec.id = "Thing/top".into();
4524        spec.source_view = Some("Thing/list".into());
4525        let caches = Arc::new(tokio::sync::RwLock::new(HashMap::from([(
4526            spec.id.clone(),
4527            sorted,
4528        )])));
4529        let cache = EntityCache::new();
4530        for limit in [None, Some(5)] {
4531            spec.pipeline = Some(ViewPipeline {
4532                filter: None,
4533                sort: Some(SortConfig {
4534                    field_path: vec!["score".into()],
4535                    order: SortOrder::Desc,
4536                }),
4537                limit,
4538            });
4539            for skip in [None, Some(0), Some(3), Some(50)] {
4540                for take in [None, Some(1), Some(4), Some(100)] {
4541                    for (owner, after, partition) in [
4542                        (None, None, None),
4543                        (Some("alice"), None, None),
4544                        (None, Some("120:0"), Some("even")),
4545                    ] {
4546                        let mut query = SubscriptionQuery {
4547                            view: spec.id.clone(),
4548                            skip,
4549                            take,
4550                            after: after.map(str::to_string),
4551                            partition: partition.map(str::to_string),
4552                            snapshot_limit: Some(2),
4553                            ..Default::default()
4554                        };
4555                        if let Some(owner) = owner {
4556                            query.filters.insert("owner".into(), json!(owner));
4557                        }
4558                        for apply_snapshot_limit in [false, true] {
4559                            let rows = load_query_entities(
4560                                &cache,
4561                                Some(caches.clone()),
4562                                &spec,
4563                                &query,
4564                                None,
4565                                apply_snapshot_limit,
4566                            )
4567                            .await;
4568                            let mut expected_query = query.clone();
4569                            if let Some(limit) = limit {
4570                                expected_query.take = Some(take.unwrap_or(limit).min(limit));
4571                            }
4572                            let expected = select_query_entities(
4573                                whole.clone(),
4574                                &expected_query,
4575                                None,
4576                                true,
4577                                apply_snapshot_limit,
4578                            );
4579                            assert_eq!(rows, expected, "{query:?} {limit:?}");
4580                        }
4581                    }
4582                }
4583            }
4584        }
4585    }
4586
4587    #[tokio::test]
4588    async fn unsorted_derived_reads_preserve_filters_order_and_snapshot_window() {
4589        use crate::materialized_view::{CompareOp, FilterConfig, ViewPipeline};
4590        let cache = EntityCache::new();
4591        for id in 1..=6 {
4592            cache
4593                .upsert(
4594                    "Thing/list",
4595                    &id.to_string(),
4596                    json!({
4597                        "_seq": format!("{}:0", id * 10), "active": id != 6,
4598                        "owner": "alice", "id": id
4599                    }),
4600                )
4601                .await;
4602        }
4603        let mut spec = list_spec();
4604        spec.id = "Thing/custom".into();
4605        spec.source_view = Some("Thing/list".into());
4606        spec.pipeline = Some(ViewPipeline::default());
4607        let caches = Arc::new(tokio::sync::RwLock::new(HashMap::new()));
4608        let mut query = SubscriptionQuery {
4609            view: spec.id.clone(),
4610            ..Default::default()
4611        };
4612        let rows =
4613            load_query_entities(&cache, Some(caches.clone()), &spec, &query, None, false).await;
4614        assert_eq!(
4615            rows.iter().map(|(key, _)| key.as_str()).collect::<Vec<_>>(),
4616            ["6", "5", "4", "3", "2", "1"]
4617        );
4618        spec.pipeline.as_mut().unwrap().filter = Some(FilterConfig {
4619            field_path: vec!["active".into()],
4620            op: CompareOp::Eq,
4621            value: json!(true),
4622        });
4623        query.filters.insert("owner".into(), json!("alice"));
4624        query.skip = Some(1);
4625        query.take = Some(3);
4626        query.snapshot_limit = Some(1);
4627        let rows =
4628            load_query_entities(&cache, Some(caches.clone()), &spec, &query, None, false).await;
4629        assert_eq!(
4630            rows.iter().map(|(key, _)| key.as_str()).collect::<Vec<_>>(),
4631            ["4", "3", "2"]
4632        );
4633        assert_eq!(
4634            load_query_entities(&cache, Some(caches.clone()), &spec, &query, None, true)
4635                .await
4636                .len(),
4637            1
4638        );
4639        query.skip = None;
4640        query.after = Some("20:0".into());
4641        let rows =
4642            load_query_entities(&cache, Some(caches.clone()), &spec, &query, None, false).await;
4643        assert_eq!(
4644            rows.iter().map(|(key, _)| key.as_str()).collect::<Vec<_>>(),
4645            ["3", "4", "5"]
4646        );
4647        query.filters.insert("owner".into(), json!("bob"));
4648        assert!(
4649            load_query_entities(&cache, Some(caches), &spec, &query, None, false)
4650                .await
4651                .is_empty()
4652        );
4653    }
4654
4655    #[tokio::test]
4656    async fn snapshot_limit_does_not_change_live_take_skip_membership() {
4657        let cache = EntityCache::with_config(EntityCacheConfig {
4658            max_entities_per_view: 10,
4659            ..Default::default()
4660        });
4661        for id in 1..=4 {
4662            cache
4663                .upsert(
4664                    "Thing/list",
4665                    &id.to_string(),
4666                    json!({"_seq": format!("10:{id:012}")}),
4667                )
4668                .await;
4669        }
4670        let query = SubscriptionQuery {
4671            view: "Thing/list".to_string(),
4672            take: Some(3),
4673            skip: Some(1),
4674            snapshot_limit: Some(1),
4675            ..Default::default()
4676        };
4677        let live = load_query_entities(&cache, None, &list_spec(), &query, None, false).await;
4678        let snapshot = load_query_entities(&cache, None, &list_spec(), &query, None, true).await;
4679        assert_eq!(live.len(), 3);
4680        assert_eq!(snapshot.len(), 1);
4681    }
4682
4683    #[test]
4684    fn fixture_manifest_covers_required_conformance_cases() {
4685        let manifest: Value = serde_json::from_str(include_str!(
4686            "../../../../tests/fixtures/websocket-v2/manifest.json"
4687        ))
4688        .unwrap();
4689        let names: HashSet<_> = manifest["fixtures"]
4690            .as_array()
4691            .unwrap()
4692            .iter()
4693            .filter_map(Value::as_str)
4694            .collect();
4695        for required in [
4696            "keyed-state.json",
4697            "list-windows.json",
4698            "filters.json",
4699            "multi-batch-authoritative.json",
4700            "empty-snapshot.json",
4701            "remove.json",
4702            "delete.json",
4703            "incremental-snapshot.json",
4704            "reconnect-replacement.json",
4705            "errors.json",
4706            "whole-entities.json",
4707        ] {
4708            assert!(names.contains(required), "missing fixture {required}");
4709        }
4710
4711        for document in [
4712            include_str!("../../../../tests/fixtures/websocket-v2/keyed-state.json"),
4713            include_str!("../../../../tests/fixtures/websocket-v2/list-windows.json"),
4714            include_str!("../../../../tests/fixtures/websocket-v2/filters.json"),
4715            include_str!("../../../../tests/fixtures/websocket-v2/multi-batch-authoritative.json"),
4716            include_str!("../../../../tests/fixtures/websocket-v2/empty-snapshot.json"),
4717            include_str!("../../../../tests/fixtures/websocket-v2/remove.json"),
4718            include_str!("../../../../tests/fixtures/websocket-v2/delete.json"),
4719            include_str!("../../../../tests/fixtures/websocket-v2/incremental-snapshot.json"),
4720            include_str!("../../../../tests/fixtures/websocket-v2/reconnect-replacement.json"),
4721            include_str!("../../../../tests/fixtures/websocket-v2/errors.json"),
4722            include_str!("../../../../tests/fixtures/websocket-v2/whole-entities.json"),
4723        ] {
4724            let fixture: Value = serde_json::from_str(document).unwrap();
4725            assert!(fixture["name"].is_string());
4726        }
4727    }
4728
4729    fn append_frame(key: &str, data: Value) -> Arc<Bytes> {
4730        let frame = json!({
4731            "entity": "Trade/append",
4732            "op": "patch",
4733            "key": key,
4734            "offset": 7,
4735            "data": data,
4736        });
4737        Arc::new(Bytes::from(serde_json::to_vec(&frame).unwrap()))
4738    }
4739
4740    /// A replayable subscription must honour the same predicates a live
4741    /// collection subscription does; otherwise a filtered consumer receives
4742    /// events outside the query it asked for.
4743    #[test]
4744    fn replay_delivery_applies_key_partition_and_filters() {
4745        let matching = append_frame("pool1", json!({"_partition": "us", "side": "buy"}));
4746        let other_partition = append_frame("pool1", json!({"_partition": "eu", "side": "buy"}));
4747        let other_side = append_frame("pool1", json!({"_partition": "us", "side": "sell"}));
4748
4749        let unfiltered = SubscriptionQuery {
4750            view: "Trade/append".to_string(),
4751            ..Default::default()
4752        };
4753        assert!(live_frame_matches(&unfiltered, "pool1", &matching));
4754        assert!(live_frame_matches(&unfiltered, "pool9", &matching));
4755
4756        let keyed = SubscriptionQuery {
4757            view: "Trade/append".to_string(),
4758            key: Some("pool1".to_string()),
4759            ..Default::default()
4760        };
4761        assert!(live_frame_matches(&keyed, "pool1", &matching));
4762        assert!(!live_frame_matches(&keyed, "pool2", &matching));
4763
4764        let partitioned = SubscriptionQuery {
4765            view: "Trade/append".to_string(),
4766            partition: Some("us".to_string()),
4767            ..Default::default()
4768        };
4769        assert!(live_frame_matches(&partitioned, "pool1", &matching));
4770        assert!(!live_frame_matches(&partitioned, "pool1", &other_partition));
4771
4772        let filtered = SubscriptionQuery {
4773            view: "Trade/append".to_string(),
4774            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
4775            ..Default::default()
4776        };
4777        assert!(live_frame_matches(&filtered, "pool1", &matching));
4778        assert!(!live_frame_matches(&filtered, "pool1", &other_side));
4779    }
4780
4781    #[test]
4782    fn a_frame_without_decodable_data_does_not_satisfy_a_filter() {
4783        let filtered = SubscriptionQuery {
4784            view: "Trade/append".to_string(),
4785            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
4786            ..Default::default()
4787        };
4788        let garbage = Arc::new(Bytes::from_static(b"not json"));
4789        assert!(!live_frame_matches(&filtered, "pool1", &garbage));
4790    }
4791
4792    /// The recovery cursor must be the last offset delivered *before* the
4793    /// gap. Reporting the newest offset seen would step the consumer over
4794    /// the skipped records permanently.
4795    #[test]
4796    fn replay_lagged_recovers_from_before_the_gap() {
4797        let epoch = crate::journal::JournalEpoch::new();
4798        let issue = SocketIssueMessage::replay_lagged(
4799            Some("trades".to_string()),
4800            37,
4801            Some(crate::journal::Cursor {
4802                epoch: epoch.clone(),
4803                offset: 4180,
4804            }),
4805        );
4806        assert_eq!(issue.code, "replay-lagged");
4807        assert_eq!(issue.recover_from, Some(format!("{epoch}:4180")));
4808        assert!(
4809            issue.suggested_action.unwrap().contains("4180"),
4810            "the consumer is told exactly which cursor recovers the gap"
4811        );
4812
4813        // Nothing delivered yet: there is no pre-gap offset, so the whole
4814        // retained window is the recovery.
4815        let from_scratch = SocketIssueMessage::replay_lagged(Some("trades".to_string()), 9, None);
4816        assert_eq!(from_scratch.recover_from, None);
4817        assert!(from_scratch
4818            .suggested_action
4819            .unwrap()
4820            .contains("without `after`"));
4821    }
4822
4823    fn bus_message(key: &str) -> Arc<BusMessage> {
4824        Arc::new(BusMessage {
4825            row_scope: None,
4826            key: key.to_string(),
4827            entity: "Trade/append".to_string(),
4828            payload: Arc::new(Bytes::from_static(b"{}")),
4829        })
4830    }
4831
4832    /// A long replay must keep the bus drained. The broadcast buffer is
4833    /// bounded, so a busy view would otherwise lap the replay and the first
4834    /// live `recv` would return `Lagged` — telling the client to resubscribe,
4835    /// starting another long replay, which laps again.
4836    #[tokio::test]
4837    async fn draining_during_a_replay_keeps_a_busy_view_from_lapping_it() {
4838        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(16);
4839        let mut pending = VecDeque::new();
4840        let mut lagged = None;
4841
4842        // Publish more than the channel holds, draining as a replay would
4843        // between sends.
4844        for index in 0..48 {
4845            sender.send(bus_message(&format!("k{index}"))).unwrap();
4846            drain_available(&mut receiver, &mut pending, &mut lagged);
4847        }
4848
4849        assert_eq!(lagged, None, "draining as we go means nothing is dropped");
4850        assert_eq!(pending.len(), 48, "every published frame is buffered");
4851    }
4852
4853    #[tokio::test]
4854    async fn a_replay_that_never_drains_is_reported_as_a_gap() {
4855        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
4856        for index in 0..32 {
4857            sender.send(bus_message(&format!("k{index}"))).unwrap();
4858        }
4859
4860        let mut pending = VecDeque::new();
4861        let mut lagged = None;
4862        drain_available(&mut receiver, &mut pending, &mut lagged);
4863
4864        assert!(
4865            lagged.is_some(),
4866            "overflowing the bus is a gap, not silent truncation"
4867        );
4868    }
4869
4870    /// Everything still on the bus after a gap is on the far side of it.
4871    /// Buffering it would put those frames in front of the lag report and
4872    /// advance the recovery cursor past the records it is meant to recover.
4873    #[tokio::test]
4874    async fn nothing_after_a_gap_is_buffered_ahead_of_the_report() {
4875        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
4876        let mut pending = VecDeque::new();
4877        let mut lagged = None;
4878
4879        // Delivered and buffered normally.
4880        sender.send(bus_message("before")).unwrap();
4881        drain_available(&mut receiver, &mut pending, &mut lagged);
4882        assert_eq!(pending.len(), 1);
4883
4884        // Overflow the bus: everything published from here is past the gap.
4885        for index in 0..32 {
4886            sender.send(bus_message(&format!("lost{index}"))).unwrap();
4887        }
4888        drain_available(&mut receiver, &mut pending, &mut lagged);
4889        assert!(lagged.is_some());
4890
4891        // The bus still holds what survived the overflow, all of it past the
4892        // gap. A later iteration of the replay loop drains again.
4893        let buffered_at_gap = pending.len();
4894        drain_available(&mut receiver, &mut pending, &mut lagged);
4895        assert_eq!(
4896            pending.len(),
4897            buffered_at_gap,
4898            "post-gap frames must not join the pre-gap flush"
4899        );
4900        assert_eq!(pending.front().unwrap().key, "before");
4901    }
4902    /// The bus is subscribed before the tape is read, so a record published
4903    /// in that window arrives on both paths.
4904    #[test]
4905    fn the_seam_between_replay_and_live_neither_repeats_nor_skips() {
4906        let mut last_sent = Some(4211);
4907
4908        assert!(
4909            already_delivered(Some(4211), &mut last_sent),
4910            "the record the replay ended on must not be sent twice"
4911        );
4912        assert!(already_delivered(Some(4100), &mut last_sent));
4913        assert_eq!(last_sent, Some(4211), "a duplicate never moves the mark");
4914
4915        assert!(
4916            !already_delivered(Some(4212), &mut last_sent),
4917            "the next record is new"
4918        );
4919        assert_eq!(last_sent, Some(4212));
4920
4921        // A frame with no offset comes from a view with no tape; it cannot
4922        // have been replayed, and must not disturb the mark.
4923        assert!(!already_delivered(None, &mut last_sent));
4924        assert_eq!(last_sent, Some(4212));
4925    }
4926
4927    /// A subscription with no cursor has delivered nothing, so the first live
4928    /// frame is not a duplicate.
4929    #[test]
4930    fn a_fresh_subscription_delivers_its_first_live_frame() {
4931        let mut last_sent = None;
4932        assert!(!already_delivered(Some(0), &mut last_sent));
4933        assert_eq!(last_sent, Some(0));
4934    }
4935
4936    /// End-to-end over a real socket: the pieces above are unit-tested
4937    /// individually, but the thing a consumer actually does — reconnect with
4938    /// a stored cursor and keep reading — only exists once a subscription is
4939    /// attached to a connection.
4940    mod over_a_socket {
4941        use super::*;
4942        use crate::journal::{EventJournal, JournalConfig};
4943        use crate::projector::Projector;
4944        use crate::{MutationBatch, SlotContext};
4945        use arete_interpreter::Mutation;
4946        use futures_util::{SinkExt, StreamExt};
4947        use std::time::Duration;
4948        use tokio::net::{TcpListener, TcpStream};
4949        use tokio::sync::mpsc;
4950        use tokio_tungstenite::tungstenite::Message;
4951        use tokio_tungstenite::{client_async, WebSocketStream};
4952
4953        const RETAINED: u64 = 600;
4954
4955        fn append_index() -> ViewIndex {
4956            let mut index = ViewIndex::new();
4957            index.add_spec(ViewSpec {
4958                id: "Trade/append".to_string(),
4959                export: "Trade".to_string(),
4960                mode: Mode::Append,
4961                wire_format: Default::default(),
4962                projection: Projection::all(),
4963                filters: Filters::all(),
4964                delivery: Delivery::default(),
4965                pipeline: None,
4966                source_view: None,
4967            });
4968            index
4969        }
4970
4971        fn trade(index: u64) -> MutationBatch {
4972            MutationBatch::with_slot_context(
4973                vec![Mutation {
4974                    export: "Trade".to_string(),
4975                    key: json!(format!("pool{}", index % 4)),
4976                    patch: json!({"trade": index}),
4977                    append: vec![],
4978                }]
4979                .into_iter()
4980                .collect(),
4981                SlotContext::new(100 + index / 3, index % 3),
4982            )
4983        }
4984
4985        struct Harness {
4986            addr: SocketAddr,
4987            journal: Arc<EventJournal>,
4988            tx: mpsc::Sender<MutationBatch>,
4989        }
4990
4991        impl Harness {
4992            async fn start() -> Self {
4993                let view_index = Arc::new(append_index());
4994                let entity_cache = EntityCache::new();
4995                let bus_manager = BusManager::new();
4996                let journal = Arc::new(EventJournal::new(JournalConfig {
4997                    enabled: true,
4998                    max_bytes_per_view: u64::MAX,
4999                    max_records_per_view: 10_000,
5000                    max_age: Duration::from_secs(3_600),
5001                }));
5002
5003                let (tx, rx) = mpsc::channel::<MutationBatch>(256);
5004                tokio::spawn(
5005                    Projector::new(
5006                        view_index.clone(),
5007                        bus_manager.clone(),
5008                        entity_cache.clone(),
5009                        rx,
5010                        #[cfg(feature = "otel")]
5011                        None,
5012                    )
5013                    .with_journal(journal.clone())
5014                    .run(),
5015                );
5016
5017                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5018                let addr = listener.local_addr().unwrap();
5019                let server = WebSocketServer::new(
5020                    addr,
5021                    bus_manager,
5022                    entity_cache,
5023                    view_index,
5024                    #[cfg(feature = "otel")]
5025                    None,
5026                )
5027                .with_journal(journal.clone());
5028                let (acceptor, _cleanup) = server.into_acceptor();
5029                tokio::spawn(async move { acceptor.serve_listener(listener).await });
5030
5031                Self { addr, journal, tx }
5032            }
5033
5034            async fn publish(&self, range: std::ops::Range<u64>) {
5035                for index in range {
5036                    self.tx.send(trade(index)).await.unwrap();
5037                }
5038                let (ack, wait) = oneshot::channel();
5039                self.tx
5040                    .send(MutationBatch::flush_marker(ack))
5041                    .await
5042                    .unwrap();
5043                wait.await.unwrap();
5044            }
5045
5046            async fn connect(&self) -> WebSocketStream<TcpStream> {
5047                let stream = TcpStream::connect(self.addr).await.unwrap();
5048                client_async(format!("ws://{}/", self.addr), stream)
5049                    .await
5050                    .unwrap()
5051                    .0
5052            }
5053        }
5054
5055        async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
5056            loop {
5057                let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
5058                    .await
5059                    .expect("the server answers within the timeout")
5060                    .expect("the stream stays open")
5061                    .expect("a readable frame");
5062                // Control and data frames arrive as binary; issue frames as
5063                // text. Both are JSON.
5064                let bytes = match &message {
5065                    Message::Text(text) => text.as_bytes(),
5066                    Message::Binary(bytes) => bytes.as_ref(),
5067                    _ => continue,
5068                };
5069                return serde_json::from_slice(bytes).expect("frames are JSON");
5070            }
5071        }
5072
5073        /// Collect `count` event frames, ignoring anything else on the wire.
5074        async fn collect_trades(socket: &mut WebSocketStream<TcpStream>, count: usize) -> Vec<u64> {
5075            let mut offsets = Vec::with_capacity(count);
5076            while offsets.len() < count {
5077                let frame = next_frame(socket).await;
5078                assert_ne!(
5079                    frame["type"], "error",
5080                    "no error frame should interrupt delivery: {frame}"
5081                );
5082                if let Some(offset) = frame["offset"].as_u64() {
5083                    offsets.push(offset);
5084                }
5085            }
5086            offsets
5087        }
5088
5089        /// The headline claim: reconnecting with a stored cursor delivers every
5090        /// event published since it, in order, and then continues live without
5091        /// a duplicate or a hole at the seam.
5092        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
5093        async fn a_reconnect_replays_from_a_cursor_and_continues_live() {
5094            let harness = Harness::start().await;
5095            harness.publish(0..RETAINED).await;
5096
5097            let cursor = harness.journal.window("Trade/append").await;
5098            let stored = format!("{}:{}", cursor.epoch, 99);
5099
5100            let mut socket = harness.connect().await;
5101            socket
5102                .send(Message::Text(
5103                    json!({
5104                        "type": "subscribe",
5105                        "protocolVersion": 2,
5106                        "subscriptionId": "trades",
5107                        "query": {"view": "Trade/append", "after": stored},
5108                    })
5109                    .to_string()
5110                    .into(),
5111                ))
5112                .await
5113                .unwrap();
5114
5115            let ack = next_frame(&mut socket).await;
5116            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
5117            assert_eq!(ack["replayWindow"]["next"], json!(RETAINED));
5118
5119            // Well over the 500 a single read or buffer would cover.
5120            let replayed = collect_trades(&mut socket, (RETAINED - 100) as usize).await;
5121            assert_eq!(
5122                replayed,
5123                (100..RETAINED).collect::<Vec<_>>(),
5124                "every event after the cursor, in order, exactly once"
5125            );
5126
5127            // Published only now, so these can only arrive over the live path.
5128            harness.publish(RETAINED..RETAINED + 40).await;
5129            let live = collect_trades(&mut socket, 40).await;
5130            assert_eq!(
5131                live,
5132                (RETAINED..RETAINED + 40).collect::<Vec<_>>(),
5133                "the live stream resumes exactly where the replay stopped"
5134            );
5135
5136            socket.close(None).await.ok();
5137        }
5138
5139        /// Events published *during* the replay must still arrive. The replay
5140        /// and the live subscription are separate reads of the same tape, and
5141        /// the seam between them is where a naive implementation drops or
5142        /// repeats.
5143        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
5144        async fn events_published_during_a_replay_are_not_lost() {
5145            let harness = Harness::start().await;
5146            harness.publish(0..RETAINED).await;
5147
5148            let epoch = harness.journal.window("Trade/append").await.epoch;
5149            let mut socket = harness.connect().await;
5150            socket
5151                .send(Message::Text(
5152                    json!({
5153                        "type": "subscribe",
5154                        "protocolVersion": 2,
5155                        "subscriptionId": "trades",
5156                        "query": {"view": "Trade/append", "after": format!("{epoch}:0")},
5157                    })
5158                    .to_string()
5159                    .into(),
5160                ))
5161                .await
5162                .unwrap();
5163            let ack = next_frame(&mut socket).await;
5164            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
5165
5166            // Keep publishing while the replay is still draining.
5167            harness.publish(RETAINED..RETAINED + 200).await;
5168
5169            let total = (RETAINED + 200 - 1) as usize;
5170            let delivered = collect_trades(&mut socket, total).await;
5171            assert_eq!(
5172                delivered,
5173                (1..RETAINED + 200).collect::<Vec<_>>(),
5174                "replay and live output join without a gap or a repeat"
5175            );
5176
5177            socket.close(None).await.ok();
5178        }
5179
5180        /// A cursor from another tape lifetime is refused rather than served
5181        /// as a continuation, and the refusal releases the subscription id.
5182        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5183        async fn a_stale_epoch_is_refused_and_frees_the_subscription_id() {
5184            let harness = Harness::start().await;
5185            harness.publish(0..50).await;
5186
5187            let mut socket = harness.connect().await;
5188            socket
5189                .send(Message::Text(
5190                    json!({
5191                        "type": "subscribe",
5192                        "protocolVersion": 2,
5193                        "subscriptionId": "trades",
5194                        "query": {
5195                            "view": "Trade/append",
5196                            "after": format!("{}:10", crate::journal::JournalEpoch::new()),
5197                        },
5198                    })
5199                    .to_string()
5200                    .into(),
5201                ))
5202                .await
5203                .unwrap();
5204
5205            let error = next_frame(&mut socket).await;
5206            assert_eq!(error["type"], "error", "unexpected frame: {error}");
5207            assert_eq!(error["code"], "cursor-epoch-changed");
5208
5209            // The documented recovery is to resubscribe without a cursor. That
5210            // only works if the refused attempt released the id.
5211            socket
5212                .send(Message::Text(
5213                    json!({
5214                        "type": "subscribe",
5215                        "protocolVersion": 2,
5216                        "subscriptionId": "trades",
5217                        "query": {"view": "Trade/append"},
5218                    })
5219                    .to_string()
5220                    .into(),
5221                ))
5222                .await
5223                .unwrap();
5224            let ack = next_frame(&mut socket).await;
5225            assert_eq!(
5226                ack["op"], "subscribed",
5227                "the refused id must be reusable: {ack}"
5228            );
5229
5230            assert_eq!(collect_trades(&mut socket, 50).await.len(), 50);
5231            socket.close(None).await.ok();
5232        }
5233
5234        const ROUND: &str = "Round/state";
5235
5236        /// A state view with no projector: the test writes the cache and the
5237        /// bus itself, so it controls exactly when the subscriber can run.
5238        #[derive(Clone)]
5239        struct StateServer {
5240            addr: SocketAddr,
5241            bus_manager: BusManager,
5242            entity_cache: EntityCache,
5243        }
5244
5245        impl StateServer {
5246            async fn start() -> Self {
5247                let mut index = ViewIndex::new();
5248                index.add_spec(ViewSpec {
5249                    id: ROUND.to_string(),
5250                    export: "Round".to_string(),
5251                    mode: Mode::State,
5252                    wire_format: Default::default(),
5253                    projection: Projection::all(),
5254                    filters: Filters::all(),
5255                    delivery: Delivery::default(),
5256                    pipeline: None,
5257                    source_view: None,
5258                });
5259                let bus_manager = BusManager::new();
5260                let entity_cache = EntityCache::new();
5261                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5262                let addr = listener.local_addr().unwrap();
5263                let server = WebSocketServer::new(
5264                    addr,
5265                    bus_manager.clone(),
5266                    entity_cache.clone(),
5267                    Arc::new(index),
5268                    #[cfg(feature = "otel")]
5269                    None,
5270                );
5271                let (acceptor, _cleanup) = server.into_acceptor();
5272                tokio::spawn(async move { acceptor.serve_listener(listener).await });
5273                Self {
5274                    addr,
5275                    bus_manager,
5276                    entity_cache,
5277                }
5278            }
5279
5280            /// Publish a patch the way the projector does: stamp its `_seq`,
5281            /// write the cache, then the key's bus.
5282            async fn publish(&self, key: &str, mut patch: Value, seq: &str) {
5283                patch["_seq"] = Value::String(seq.to_string());
5284                self.entity_cache
5285                    .upsert_with_append(ROUND, key, patch.clone(), &[], PatchOrigin::Unknown)
5286                    .await;
5287                self.publish_frame(key, "patch", patch, seq).await;
5288            }
5289
5290            /// A complete new lifetime after an explicit source deletion.
5291            async fn create(&self, key: &str, mut entity: Value, seq: &str) {
5292                entity["_seq"] = Value::String(seq.to_string());
5293                self.entity_cache
5294                    .upsert_with_append(ROUND, key, entity.clone(), &[], PatchOrigin::Creation)
5295                    .await;
5296                self.publish_frame(key, "upsert", entity, seq).await;
5297            }
5298
5299            /// Publish a frame to the key's bus without touching the cache.
5300            async fn publish_frame(&self, key: &str, op: &str, data: Value, seq: &str) {
5301                let frame = json!({
5302                    "mode": "state",
5303                    "entity": ROUND,
5304                    "op": op,
5305                    "key": key,
5306                    "data": data,
5307                    "seq": seq,
5308                });
5309                self.bus_manager
5310                    .publish_state(ROUND, key, Arc::new(Bytes::from(frame.to_string())), None)
5311                    .await;
5312            }
5313
5314            /// Subscribe to one key and read through its snapshot.
5315            async fn subscribe(&self, key: &str) -> WebSocketStream<TcpStream> {
5316                let stream = TcpStream::connect(self.addr).await.unwrap();
5317                let mut socket = client_async(format!("ws://{}/", self.addr), stream)
5318                    .await
5319                    .unwrap()
5320                    .0;
5321                socket
5322                    .send(Message::Text(
5323                        json!({
5324                            "type": "subscribe",
5325                            "protocolVersion": 2,
5326                            "subscriptionId": "round",
5327                            "query": {"view": ROUND, "key": key},
5328                        })
5329                        .to_string()
5330                        .into(),
5331                    ))
5332                    .await
5333                    .unwrap();
5334                let ack = next_frame(&mut socket).await;
5335                assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
5336                loop {
5337                    let frame = next_frame(&mut socket).await;
5338                    assert_eq!(frame["op"], "snapshot", "unexpected frame: {frame}");
5339                    if frame["complete"] == true {
5340                        return socket;
5341                    }
5342                }
5343            }
5344        }
5345
5346        /// Two patches to one key, published before the subscriber reads the
5347        /// first, must both reach it. The bus holds only the latest frame, so
5348        /// forwarding it alone drops the fields of the one it replaced.
5349        // One worker, so tasks run one at a time. The handshake needs the
5350        // multi-threaded runtime.
5351        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5352        async fn a_state_subscriber_keeps_the_fields_of_an_overwritten_patch() {
5353            let server = StateServer::start().await;
5354            server
5355                .publish("7", json!({"id": 7}), "100:000000000001")
5356                .await;
5357            let mut socket = server.subscribe("7").await;
5358
5359            // A lone patch is forwarded as it was published. This one is an
5360            // account update, numbered by its write version.
5361            server
5362                .publish(
5363                    "7",
5364                    json!({"results": {"slot_hash": "abc"}}),
5365                    "140:003836292257",
5366                )
5367                .await;
5368            let frame = next_frame(&mut socket).await;
5369            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5370            assert_eq!(frame["seq"], "140:003836292257");
5371
5372            // Two instruction patches from later in the same slot, numbered by
5373            // transaction index. Both writes happen in one task that never
5374            // yields to the subscriber, so the second replaces the first on
5375            // the bus before the subscriber wakes.
5376            let writer = server.clone();
5377            tokio::spawn(async move {
5378                writer
5379                    .publish("7", json!({"entropy": {"seed": "def"}}), "140:000000001200")
5380                    .await;
5381                writer
5382                    .publish("7", json!({"total": 1}), "140:000000001200")
5383                    .await;
5384            })
5385            .await
5386            .unwrap();
5387            let frame = next_frame(&mut socket).await;
5388            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5389            // Everything that changed since the snapshot, the forwarded patch
5390            // included; `id` has not.
5391            assert_eq!(
5392                frame["data"],
5393                json!({
5394                    "results": {"slot_hash": "abc"},
5395                    "entropy": {"seed": "def"},
5396                    "total": 1,
5397                    "_seq": "140:000000001200",
5398                })
5399            );
5400            // The latest seq sorts below the one already delivered, and a
5401            // client would drop the frame as stale.
5402            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
5403
5404            // Having caught up, the subscriber forwards patches again.
5405            server
5406                .publish("7", json!({"total": 2}), "141:000000000002")
5407                .await;
5408            let frame = next_frame(&mut socket).await;
5409            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5410            assert_eq!(
5411                frame["data"],
5412                json!({"total": 2, "_seq": "141:000000000002"})
5413            );
5414
5415            // The next catch-up is measured against the last one: only what
5416            // changed since it, not the fields that one already carried.
5417            let writer = server.clone();
5418            tokio::spawn(async move {
5419                writer
5420                    .publish("7", json!({"total": 3}), "142:000000000001")
5421                    .await;
5422                writer
5423                    .publish("7", json!({"flag": true}), "142:000000000001")
5424                    .await;
5425            })
5426            .await
5427            .unwrap();
5428            let frame = next_frame(&mut socket).await;
5429            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5430            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
5431            assert_eq!(
5432                frame["data"],
5433                json!({"total": 3, "flag": true, "_seq": "142:000000000001"})
5434            );
5435
5436            socket.close(None).await.ok();
5437        }
5438
5439        /// A whole entity forwarded as it was published replaces the client's
5440        /// copy, so a later catch-up is measured against it.
5441        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5442        async fn a_forwarded_whole_entity_is_what_a_catch_up_is_measured_against() {
5443            let server = StateServer::start().await;
5444            server
5445                .publish("7", json!({"id": 7, "a": 1}), "100:000000000001")
5446                .await;
5447            let mut socket = server.subscribe("7").await;
5448
5449            let whole = json!({"id": 7, "a": 1, "b": 2, "_seq": "101:000000000001"});
5450            server
5451                .entity_cache
5452                .store_whole(ROUND, "7", whole.clone())
5453                .await;
5454            server
5455                .publish_frame("7", "upsert", whole.clone(), "101:000000000001")
5456                .await;
5457            let frame = next_frame(&mut socket).await;
5458            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5459            assert_eq!(frame["data"], whole);
5460
5461            let writer = server.clone();
5462            tokio::spawn(async move {
5463                writer
5464                    .publish("7", json!({"c": 3}), "102:000000000001")
5465                    .await;
5466                writer
5467                    .publish("7", json!({"d": 4}), "102:000000000002")
5468                    .await;
5469            })
5470            .await
5471            .unwrap();
5472            let frame = next_frame(&mut socket).await;
5473            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5474            assert_eq!(
5475                frame["data"],
5476                json!({"c": 3, "d": 4, "_seq": "102:000000000002"}),
5477                "`b` came with the whole entity, so it is not sent again"
5478            );
5479            socket.close(None).await.ok();
5480        }
5481
5482        /// Holders of the cached entity's fields beyond the cache and the
5483        /// copy this reads.
5484        async fn holders(server: &StateServer, key: &str) -> usize {
5485            let entity = server.entity_cache.get_shared(ROUND, key).await.unwrap();
5486            Arc::strong_count(entity.fields()) - 2
5487        }
5488
5489        /// Wait until `held` reports `expected`: a subscription's task takes
5490        /// over what it holds just after its snapshot reaches the client.
5491        async fn settle(mut held: impl AsyncFnMut() -> usize, expected: usize) {
5492            for _ in 0..200 {
5493                if held().await == expected {
5494                    return;
5495                }
5496                tokio::time::sleep(Duration::from_millis(5)).await;
5497            }
5498            assert_eq!(held().await, expected);
5499        }
5500
5501        /// What a state subscriber keeps to measure catch-ups against is the
5502        /// cached entity it was sent, or the frame the bus forwarded: shared,
5503        /// never a copy of its own.
5504        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5505        async fn state_subscribers_share_what_they_were_sent() {
5506            let server = StateServer::start().await;
5507            server
5508                .create(
5509                    "7",
5510                    json!({"id": 7, "detail": {"a": 1}}),
5511                    "100:000000000001",
5512                )
5513                .await;
5514            let mut sockets = Vec::new();
5515            for _ in 0..3 {
5516                sockets.push(server.subscribe("7").await);
5517            }
5518            settle(|| holders(&server, "7"), 3).await;
5519
5520            // A whole entity forwarded from the bus: each subscriber keeps the
5521            // frame the bus holds, not the entity it carries.
5522            let sent = server.entity_cache.get_shared(ROUND, "7").await.unwrap();
5523            let whole = json!({"id": 7, "detail": {"b": 2}, "_seq": "101:000000000001"});
5524            server
5525                .entity_cache
5526                .store_whole(ROUND, "7", whole.clone())
5527                .await;
5528            let frame = Arc::new(Bytes::from(
5529                json!({
5530                    "mode": "state",
5531                    "entity": ROUND,
5532                    "op": "upsert",
5533                    "key": "7",
5534                    "data": whole,
5535                    "seq": "101:000000000001",
5536                })
5537                .to_string(),
5538            ));
5539            server
5540                .bus_manager
5541                .publish_state(ROUND, "7", frame.clone(), None)
5542                .await;
5543            for socket in &mut sockets {
5544                assert_eq!(next_frame(socket).await["op"], "upsert");
5545            }
5546            // This test's handle, the bus's, and one per subscriber.
5547            settle(async || Arc::strong_count(&frame), 2 + sockets.len()).await;
5548            // Only this test still holds the entity they were sent before.
5549            settle(async || Arc::strong_count(sent.fields()), 1).await;
5550            for mut socket in sockets {
5551                socket.close(None).await.ok();
5552            }
5553        }
5554
5555        /// A catch-up resends a field a forwarded patch set even when the
5556        /// entity has set it back to the value the client last received
5557        /// whole: the client still holds the forwarded value.
5558        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5559        async fn a_catch_up_restores_a_field_a_forwarded_patch_set() {
5560            let server = StateServer::start().await;
5561            server
5562                .publish("7", json!({"id": 7, "a": 1}), "100:000000000001")
5563                .await;
5564            let mut socket = server.subscribe("7").await;
5565
5566            server
5567                .publish("7", json!({"a": 2}), "101:000000000001")
5568                .await;
5569            let frame = next_frame(&mut socket).await;
5570            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5571            assert_eq!(frame["data"], json!({"a": 2, "_seq": "101:000000000001"}));
5572
5573            let writer = server.clone();
5574            tokio::spawn(async move {
5575                writer
5576                    .publish("7", json!({"a": 1}), "102:000000000001")
5577                    .await;
5578                writer
5579                    .publish("7", json!({"b": 1}), "102:000000000002")
5580                    .await;
5581            })
5582            .await
5583            .unwrap();
5584            let frame = next_frame(&mut socket).await;
5585            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5586            assert_eq!(
5587                frame["data"],
5588                json!({"a": 1, "b": 1, "_seq": "102:000000000002"}),
5589                "`a` matches the snapshot but the client holds 2"
5590            );
5591            socket.close(None).await.ok();
5592        }
5593
5594        /// A key deleted and created again before the subscriber reads
5595        /// either frame: the bus keeps only the new entity's. Merged into the
5596        /// client's copy it would leave the deleted entity's fields in place,
5597        /// so the client gets the new entity whole, replacing its copy.
5598        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5599        async fn a_key_deleted_and_recreated_between_reads_replaces_the_clients_copy() {
5600            let server = StateServer::start().await;
5601            server
5602                .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
5603                .await;
5604            let mut socket = server.subscribe("7").await;
5605
5606            let writer = server.clone();
5607            tokio::spawn(async move {
5608                writer.entity_cache.remove(ROUND, "7").await;
5609                writer
5610                    .publish_frame("7", "delete", Value::Null, "101:000000000001")
5611                    .await;
5612                writer
5613                    .publish("7", json!({"id": 7, "fresh": true}), "101:000000000002")
5614                    .await;
5615            })
5616            .await
5617            .unwrap();
5618            let frame = next_frame(&mut socket).await;
5619            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5620            assert_eq!(
5621                frame["data"],
5622                json!({"id": 7, "fresh": true, "_seq": "101:000000000002"})
5623            );
5624            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
5625
5626            // Deleted, with a patch overwritten before it: the delete is final.
5627            let writer = server.clone();
5628            tokio::spawn(async move {
5629                writer
5630                    .publish("7", json!({"n": 1}), "102:000000000001")
5631                    .await;
5632                writer.entity_cache.remove(ROUND, "7").await;
5633                writer
5634                    .publish_frame("7", "delete", Value::Null, "102:000000000002")
5635                    .await;
5636            })
5637            .await
5638            .unwrap();
5639            let frame = next_frame(&mut socket).await;
5640            assert_eq!(frame["op"], "delete", "unexpected frame: {frame}");
5641
5642            // Created again: the client holds nothing, so it arrives whole.
5643            server
5644                .create("7", json!({"id": 7, "third": true}), "103:000000000001")
5645                .await;
5646            let frame = next_frame(&mut socket).await;
5647            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5648            assert_eq!(frame["data"]["third"], true);
5649            assert!(frame["data"].get("fresh").is_none(), "{frame}");
5650            socket.close(None).await.ok();
5651        }
5652
5653        /// A missed delete, then a patch the cache refused (it lacks the new
5654        /// entity): the client's copy is of the deleted entity, so the patch
5655        /// is not merged into it. The new entity replaces it once it is whole.
5656        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5657        async fn a_copy_of_a_deleted_entity_takes_no_patch_until_it_is_replaced() {
5658            let server = StateServer::start().await;
5659            server
5660                .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
5661                .await;
5662            let mut socket = server.subscribe("7").await;
5663
5664            let writer = server.clone();
5665            tokio::spawn(async move {
5666                writer.entity_cache.remove(ROUND, "7").await;
5667                writer
5668                    .publish_frame("7", "delete", Value::Null, "101:000000000001")
5669                    .await;
5670                writer
5671                    .publish_frame("7", "patch", json!({"n": 1}), "101:000000000002")
5672                    .await;
5673            })
5674            .await
5675            .unwrap();
5676            // Let the subscriber read the patch, and withhold it, before the
5677            // whole entity arrives. Nothing is sent, so there is no frame to
5678            // wait for.
5679            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
5680
5681            let whole = json!({"id": 7, "n": 1, "_seq": "101:000000000002"});
5682            server
5683                .entity_cache
5684                .store_whole(ROUND, "7", whole.clone())
5685                .await;
5686            server
5687                .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
5688                .await;
5689            let frame = next_frame(&mut socket).await;
5690            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5691            assert_eq!(frame["data"], whole);
5692            socket.close(None).await.ok();
5693        }
5694
5695        /// The source replaced the entity whole (an `upsert`, dropping a field)
5696        /// and then patched it, both before the subscriber read either. The
5697        /// patch alone, or the entity merged in, would keep the dropped field.
5698        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5699        async fn a_whole_entity_among_overwritten_frames_replaces_the_clients_copy() {
5700            let server = StateServer::start().await;
5701            server
5702                .publish("7", json!({"id": 7, "dropped": 1}), "100:000000000001")
5703                .await;
5704            let mut socket = server.subscribe("7").await;
5705
5706            let writer = server.clone();
5707            tokio::spawn(async move {
5708                let whole = json!({"id": 7, "_seq": "101:000000000001"});
5709                writer
5710                    .entity_cache
5711                    .store_whole(ROUND, "7", whole.clone())
5712                    .await;
5713                writer
5714                    .publish_frame("7", "upsert", whole, "101:000000000001")
5715                    .await;
5716                writer
5717                    .publish("7", json!({"n": 1}), "101:000000000002")
5718                    .await;
5719            })
5720            .await
5721            .unwrap();
5722            let frame = next_frame(&mut socket).await;
5723            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5724            assert_eq!(
5725                frame["data"],
5726                json!({"id": 7, "n": 1, "_seq": "101:000000000002"})
5727            );
5728            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
5729            socket.close(None).await.ok();
5730        }
5731
5732        /// Patches overwritten on the bus while the cache lacks the key (it
5733        /// evicted it, and refused them) leave nothing whole to send in their
5734        /// place: the holder gets the latest patch, as a holder of an evicted
5735        /// key does, and stays behind. The resend that brings the entity back
5736        /// carries the seq of the key's latest change, which the holder
5737        /// already has and would drop as stale, so it gets the cached entity
5738        /// without a seq instead: as an upsert, since the resend is the whole
5739        /// entity and replaces what the holder has.
5740        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5741        async fn a_holder_behind_on_an_evicted_key_catches_up_when_it_returns() {
5742            let server = StateServer::start().await;
5743            server
5744                .publish("7", json!({"id": 7}), "100:000000000001")
5745                .await;
5746            let mut socket = server.subscribe("7").await;
5747            server.entity_cache.remove(ROUND, "7").await;
5748
5749            let writer = server.clone();
5750            tokio::spawn(async move {
5751                writer
5752                    .publish_frame("7", "patch", json!({"a": 1}), "101:000000000001")
5753                    .await;
5754                writer
5755                    .publish_frame("7", "patch", json!({"b": 2}), "101:000000000002")
5756                    .await;
5757            })
5758            .await
5759            .unwrap();
5760            let frame = next_frame(&mut socket).await;
5761            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5762            assert_eq!(frame["data"], json!({"b": 2}));
5763            assert_eq!(frame["seq"], "101:000000000002");
5764
5765            // The resend: the whole entity, at the latest change's position.
5766            let whole = json!({"id": 7, "a": 1, "b": 2, "_seq": "101:000000000002"});
5767            server
5768                .entity_cache
5769                .store_whole(ROUND, "7", whole.clone())
5770                .await;
5771            server
5772                .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
5773                .await;
5774            let frame = next_frame(&mut socket).await;
5775            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5776            assert_eq!(frame["data"], whole);
5777            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
5778
5779            // Caught up: patches are forwarded again.
5780            server
5781                .publish("7", json!({"c": 3}), "102:000000000001")
5782                .await;
5783            let frame = next_frame(&mut socket).await;
5784            assert_eq!(frame["data"], json!({"c": 3, "_seq": "102:000000000001"}));
5785            assert_eq!(frame["seq"], "102:000000000001");
5786            socket.close(None).await.ok();
5787        }
5788    }
5789
5790    /// A key only counts as held by the client once a frame actually carried
5791    /// it. A truncated or disabled snapshot leaves keys the client never saw,
5792    /// and their first change must arrive whole.
5793    mod partial_entities_over_a_socket {
5794        use super::*;
5795        use crate::projector::Projector;
5796        use crate::{MutationBatch, SlotContext};
5797        use arete_interpreter::{AccountPosition, Mutation};
5798        use futures_util::{SinkExt, StreamExt};
5799        use std::time::Duration;
5800        use tokio::net::{TcpListener, TcpStream};
5801        use tokio::sync::mpsc;
5802        use tokio_tungstenite::tungstenite::Message;
5803        use tokio_tungstenite::{client_async, WebSocketStream};
5804
5805        fn thing_index() -> ViewIndex {
5806            let mut index = ViewIndex::new();
5807            index.add_spec(list_spec());
5808            index.add_spec(ViewSpec {
5809                id: "Thing/state".to_string(),
5810                mode: Mode::State,
5811                ..list_spec()
5812            });
5813            index
5814        }
5815
5816        struct Harness {
5817            addr: SocketAddr,
5818            tx: mpsc::Sender<MutationBatch>,
5819            slot: std::sync::atomic::AtomicU64,
5820        }
5821
5822        impl Harness {
5823            async fn start(delivery: WebSocketDeliveryConfig) -> Self {
5824                Self::start_with_cache(delivery, EntityCache::new()).await
5825            }
5826
5827            async fn start_with_cache(
5828                delivery: WebSocketDeliveryConfig,
5829                entity_cache: EntityCache,
5830            ) -> Self {
5831                Self::start_with_index(delivery, entity_cache, thing_index()).await
5832            }
5833
5834            async fn start_with_index(
5835                delivery: WebSocketDeliveryConfig,
5836                entity_cache: EntityCache,
5837                index: ViewIndex,
5838            ) -> Self {
5839                let view_index = Arc::new(index);
5840                let bus_manager = BusManager::new();
5841                let (tx, rx) = mpsc::channel::<MutationBatch>(64);
5842                // Unconstrained, so Tokio's cooperative budget never makes the
5843                // projector yield partway through a batch: a test on one
5844                // worker then knows no subscriber runs until the whole batch
5845                // is published (see
5846                // `snapshot_rows_patches_and_catch_ups_carry_the_version`).
5847                tokio::spawn(tokio::task::unconstrained(
5848                    Projector::new(
5849                        view_index.clone(),
5850                        bus_manager.clone(),
5851                        entity_cache.clone(),
5852                        rx,
5853                        #[cfg(feature = "otel")]
5854                        None,
5855                    )
5856                    .run(),
5857                ));
5858                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5859                let addr = listener.local_addr().unwrap();
5860                let server = WebSocketServer::new(
5861                    addr,
5862                    bus_manager,
5863                    entity_cache,
5864                    view_index,
5865                    #[cfg(feature = "otel")]
5866                    None,
5867                )
5868                .with_delivery_config(delivery);
5869                let (acceptor, _cleanup) = server.into_acceptor();
5870                tokio::spawn(async move { acceptor.serve_listener(listener).await });
5871                Self {
5872                    addr,
5873                    tx,
5874                    slot: std::sync::atomic::AtomicU64::new(100),
5875                }
5876            }
5877
5878            /// Apply one source patch and wait until the projector published it.
5879            async fn patch(&self, key: &str, patch: Value) {
5880                self.mutate(Mutation {
5881                    export: "Thing".to_string(),
5882                    key: json!(key),
5883                    patch,
5884                    append: vec![],
5885                })
5886                .await;
5887            }
5888
5889            async fn mutate(&self, mutation: Mutation) {
5890                let slot = self.slot.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
5891                self.mutate_with_context(mutation, SlotContext::new(slot, 0))
5892                    .await;
5893            }
5894
5895            async fn mutate_with_context(&self, mutation: Mutation, context: SlotContext) {
5896                self.tx
5897                    .send(MutationBatch::with_slot_context(
5898                        vec![mutation].into_iter().collect(),
5899                        context,
5900                    ))
5901                    .await
5902                    .unwrap();
5903                let (ack, wait) = oneshot::channel();
5904                self.tx
5905                    .send(MutationBatch::flush_marker(ack))
5906                    .await
5907                    .unwrap();
5908                wait.await.unwrap();
5909            }
5910
5911            async fn subscribe(&self, query: Value, snapshot: bool) -> WebSocketStream<TcpStream> {
5912                let stream = TcpStream::connect(self.addr).await.unwrap();
5913                let mut socket = client_async(format!("ws://{}/", self.addr), stream)
5914                    .await
5915                    .unwrap()
5916                    .0;
5917                socket
5918                    .send(Message::Text(
5919                        json!({
5920                            "type": "subscribe",
5921                            "protocolVersion": 2,
5922                            "subscriptionId": "things",
5923                            "query": query,
5924                            "snapshot": {"enabled": snapshot},
5925                        })
5926                        .to_string()
5927                        .into(),
5928                    ))
5929                    .await
5930                    .unwrap();
5931                let ack = next_frame(&mut socket).await;
5932                assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
5933                assert_eq!(
5934                    ack["wholeEntities"],
5935                    json!(true),
5936                    "the ack advertises the whole-entity guarantee: {ack}"
5937                );
5938                socket
5939            }
5940        }
5941
5942        async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
5943            loop {
5944                let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
5945                    .await
5946                    .expect("the server answers within the timeout")
5947                    .expect("the stream stays open")
5948                    .expect("a readable frame");
5949                // Every frame here is far below the compression threshold.
5950                let bytes = match &message {
5951                    Message::Text(text) => text.as_bytes(),
5952                    Message::Binary(bytes) => bytes.as_ref(),
5953                    _ => continue,
5954                };
5955                return serde_json::from_slice(bytes).expect("frames are JSON");
5956            }
5957        }
5958
5959        /// The snapshot rows, in order, once the snapshot completes.
5960        async fn snapshot_keys(socket: &mut WebSocketStream<TcpStream>) -> Vec<String> {
5961            let mut keys = Vec::new();
5962            loop {
5963                let frame = next_frame(socket).await;
5964                assert_eq!(frame["op"], "snapshot", "expected a snapshot: {frame}");
5965                keys.extend(
5966                    frame["data"]
5967                        .as_array()
5968                        .unwrap()
5969                        .iter()
5970                        .map(|row| row["key"].as_str().unwrap().to_string()),
5971                );
5972                if frame["complete"] == json!(true) {
5973                    return keys;
5974                }
5975            }
5976        }
5977
5978        async fn seed(harness: &Harness) {
5979            for key in ["a", "b", "c"] {
5980                harness
5981                    .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
5982                    .await;
5983            }
5984        }
5985
5986        fn derived_index() -> ViewIndex {
5987            use crate::materialized_view::{
5988                CompareOp, FilterConfig, SortConfig, SortOrder, ViewPipeline,
5989            };
5990            let mut index = thing_index();
5991            for name in ["empty", "filtered", "sorted"] {
5992                let mut pipeline = ViewPipeline::default();
5993                if name != "empty" {
5994                    pipeline.filter = Some(FilterConfig {
5995                        field_path: vec!["owner".into()],
5996                        op: CompareOp::Eq,
5997                        value: json!("alice"),
5998                    });
5999                    pipeline.limit = Some(2);
6000                }
6001                if name == "sorted" {
6002                    pipeline.sort = Some(SortConfig {
6003                        field_path: vec!["amount".into()],
6004                        order: SortOrder::Desc,
6005                    });
6006                }
6007                index.add_spec(ViewSpec {
6008                    id: format!("Thing/{name}"),
6009                    source_view: Some("Thing/list".into()),
6010                    pipeline: Some(pipeline),
6011                    ..list_spec()
6012                });
6013            }
6014            index
6015        }
6016
6017        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6018        async fn unsorted_derived_idle_bootstrap_and_reconnect_are_authoritative() {
6019            let harness = Harness::start_with_index(
6020                WebSocketDeliveryConfig::default(),
6021                EntityCache::new(),
6022                derived_index(),
6023            )
6024            .await;
6025            harness
6026                .patch("a", json!({"owner":"alice", "amount":1}))
6027                .await;
6028            harness.patch("b", json!({"owner":"bob", "amount":2})).await;
6029            harness
6030                .patch("c", json!({"owner":"alice", "amount":3}))
6031                .await;
6032            harness
6033                .patch("d", json!({"owner":"alice", "amount":4}))
6034                .await;
6035            for _ in 0..2 {
6036                for (view, expected) in [
6037                    ("empty", vec!["d", "c", "b", "a"]),
6038                    ("filtered", vec!["d", "c"]),
6039                    ("sorted", vec!["d", "c"]),
6040                ] {
6041                    let mut socket = harness
6042                        .subscribe(json!({"view": format!("Thing/{view}"), "take":100}), true)
6043                        .await;
6044                    assert_eq!(snapshot_keys(&mut socket).await, expected);
6045                    socket.close(None).await.unwrap();
6046                }
6047            }
6048        }
6049
6050        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6051        async fn source_delete_and_recreation_clear_all_views_across_reconnect() {
6052            let harness = Harness::start_with_index(
6053                WebSocketDeliveryConfig::default(),
6054                EntityCache::new(),
6055                derived_index(),
6056            )
6057            .await;
6058            harness
6059                .patch("a", json!({"owner":"alice", "amount":1, "oldField":true}))
6060                .await;
6061            let mut list = harness
6062                .subscribe(json!({"view":"Thing/filtered"}), true)
6063                .await;
6064            assert_eq!(snapshot_keys(&mut list).await, ["a"]);
6065            let mut state = harness
6066                .subscribe(json!({"view":"Thing/state", "key":"a"}), true)
6067                .await;
6068            assert_eq!(snapshot_keys(&mut state).await, ["a"]);
6069            // Leaving a predicate is view eviction, not chain deletion.
6070            harness.patch("a", json!({"owner":"bob"})).await;
6071            assert_eq!(next_frame(&mut list).await["op"], "remove");
6072            harness.patch("a", json!({"owner":"alice"})).await;
6073            assert_eq!(next_frame(&mut list).await["op"], "upsert");
6074            let mut deletion = Mutation::delete("Thing", json!("a"));
6075            deletion.mark_account_position(AccountPosition::new(100, 1001));
6076            harness.mutate(deletion).await;
6077            assert_eq!(next_frame(&mut list).await["op"], "delete");
6078            // State patches may precede its deletion on the socket.
6079            loop {
6080                if next_frame(&mut state).await["op"] == "delete" {
6081                    break;
6082                }
6083            }
6084            for view in [
6085                "Thing/list",
6086                "Thing/empty",
6087                "Thing/filtered",
6088                "Thing/sorted",
6089                "Thing/state",
6090            ] {
6091                let mut socket = harness
6092                    .subscribe(json!({"view":view, "key":"a"}), true)
6093                    .await;
6094                assert!(snapshot_keys(&mut socket).await.is_empty(), "{view}");
6095                socket.close(None).await.unwrap();
6096            }
6097            // A whole resend or sparse change cannot recreate a deleted row.
6098            harness
6099                .patch("a", json!({"owner":"alice", "oldField":true}))
6100                .await;
6101            let mut stale = Mutation {
6102                export: "Thing".into(),
6103                key: json!("a"),
6104                patch: json!({"owner":"alice", "oldField":true}),
6105                append: vec![],
6106            };
6107            stale.mark_whole_entity();
6108            harness.mutate(stale).await;
6109            let mut socket = harness
6110                .subscribe(json!({"view":"Thing/filtered"}), true)
6111                .await;
6112            assert!(snapshot_keys(&mut socket).await.is_empty());
6113            socket.close(None).await.unwrap();
6114            let mut recreated = Mutation {
6115                export: "Thing".into(),
6116                key: json!("a"),
6117                patch: json!({"owner":"alice", "amount":9}),
6118                append: vec![],
6119            };
6120            recreated.mark_created();
6121            recreated.mark_account_position(AccountPosition::new(100, 1002));
6122            harness.mutate(recreated).await;
6123            for view in [
6124                "Thing/list",
6125                "Thing/empty",
6126                "Thing/filtered",
6127                "Thing/sorted",
6128                "Thing/state",
6129            ] {
6130                let mut socket = harness
6131                    .subscribe(json!({"view":view, "key":"a"}), true)
6132                    .await;
6133                let frame = next_frame(&mut socket).await;
6134                assert_eq!(frame["data"][0]["key"], "a");
6135                assert_eq!(frame["data"][0]["data"]["amount"], 9);
6136                assert!(
6137                    frame["data"][0]["data"].get("oldField").is_none(),
6138                    "{frame}"
6139                );
6140                socket.close(None).await.unwrap();
6141            }
6142        }
6143
6144        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6145        async fn resolver_offsets_do_not_suppress_later_instruction_updates() {
6146            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
6147            let mut created = Mutation {
6148                export: "Thing".into(),
6149                key: json!("a"),
6150                patch: json!({"value":"created"}),
6151                append: vec![],
6152            };
6153            created.mark_created();
6154            created.mark_account_position(AccountPosition::new(100, 9));
6155            harness
6156                .mutate_with_context(created, SlotContext::account(100, 9))
6157                .await;
6158
6159            harness
6160                .mutate_with_context(
6161                    Mutation {
6162                        export: "Thing".into(),
6163                        key: json!("a"),
6164                        patch: json!({"value":"resolver", "resolved":true}),
6165                        append: vec![],
6166                    },
6167                    SlotContext::resolver(100, 1_u64 << 63),
6168                )
6169                .await;
6170            harness
6171                .mutate_with_context(
6172                    Mutation {
6173                        export: "Thing".into(),
6174                        key: json!("a"),
6175                        patch: json!({"value":"instruction", "instruction":true}),
6176                        append: vec![],
6177                    },
6178                    SlotContext::instruction(100, 900),
6179                )
6180                .await;
6181
6182            let mut socket = harness
6183                .subscribe(json!({"view":"Thing/state", "key":"a"}), true)
6184                .await;
6185            let frame = next_frame(&mut socket).await;
6186            assert_eq!(frame["op"], "snapshot", "{frame}");
6187            let row = &frame["data"][0]["data"];
6188            assert_eq!(row["value"], "instruction", "{frame}");
6189            assert_eq!(row["resolved"], true, "{frame}");
6190            assert_eq!(row["instruction"], true, "{frame}");
6191
6192            // Recency still applies inside the instruction domain even after
6193            // activity from another source domain.
6194            harness
6195                .mutate_with_context(
6196                    Mutation {
6197                        export: "Thing".into(),
6198                        key: json!("a"),
6199                        patch: json!({"value":"stale", "stale":true}),
6200                        append: vec![],
6201                    },
6202                    SlotContext::instruction(100, 899),
6203                )
6204                .await;
6205            let mut reconnected = harness
6206                .subscribe(json!({"view":"Thing/state", "key":"a"}), true)
6207                .await;
6208            let frame = next_frame(&mut reconnected).await;
6209            let row = &frame["data"][0]["data"];
6210            assert_eq!(row["value"], "instruction", "{frame}");
6211            assert!(row.get("stale").is_none(), "{frame}");
6212        }
6213
6214        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
6215        async fn projected_delete_and_recreation_between_reads_replace_state() {
6216            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
6217            harness.patch("a", json!({"old":true})).await;
6218            let mut socket = harness
6219                .subscribe(json!({"view":"Thing/state", "key":"a"}), true)
6220                .await;
6221            assert_eq!(snapshot_keys(&mut socket).await, ["a"]);
6222            let mut creation = Mutation {
6223                export: "Thing".into(),
6224                key: json!("a"),
6225                patch: json!({"fresh":true}),
6226                append: vec![],
6227            };
6228            creation.mark_created();
6229            // Both ready batches run before the watch receiver on this worker.
6230            for mutation in [Mutation::delete("Thing", json!("a")), creation] {
6231                let slot = harness
6232                    .slot
6233                    .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
6234                harness
6235                    .tx
6236                    .try_send(MutationBatch::with_slot_context(
6237                        vec![mutation].into_iter().collect(),
6238                        SlotContext::new(slot, 0),
6239                    ))
6240                    .unwrap();
6241            }
6242            let (ack, wait) = oneshot::channel();
6243            harness
6244                .tx
6245                .try_send(MutationBatch::flush_marker(ack))
6246                .unwrap();
6247            wait.await.unwrap();
6248            // Usually both batches land before the socket reads and arrive as
6249            // one upsert. Under load the socket can read between them, so the
6250            // delete may arrive on its own first. Either way the recreated
6251            // entity must not carry the deleted one's fields.
6252            let mut frame = next_frame(&mut socket).await;
6253            if frame["op"] == "delete" {
6254                assert_eq!(frame["key"], "a", "{frame}");
6255                frame = next_frame(&mut socket).await;
6256            }
6257            assert_eq!(frame["op"], "upsert", "{frame}");
6258            assert_eq!(frame["data"]["fresh"], true);
6259            assert!(frame["data"].get("old").is_none(), "{frame}");
6260        }
6261
6262        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6263        async fn a_change_to_a_key_cut_by_snapshot_limit_is_an_upsert() {
6264            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
6265            seed(&harness).await;
6266
6267            let mut socket = harness
6268                .subscribe(json!({"view": "Thing/list", "snapshotLimit": 1}), true)
6269                .await;
6270            // Newest first: only "c" was sent.
6271            assert_eq!(snapshot_keys(&mut socket).await, ["c"]);
6272
6273            harness.patch("a", json!({"count": 2})).await;
6274            let frame = next_frame(&mut socket).await;
6275            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
6276            assert_eq!(frame["op"], "upsert", "a key never sent arrives whole");
6277            assert_eq!(frame["data"]["name"], "thing-a");
6278            assert_eq!(frame["data"]["count"], 2);
6279
6280            // Now the client holds "a", so its next change is a patch again.
6281            harness.patch("a", json!({"count": 3})).await;
6282            let frame = next_frame(&mut socket).await;
6283            assert_eq!(frame["key"], "a");
6284            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
6285            assert!(frame["data"].get("name").is_none());
6286
6287            // A key the snapshot did carry still rides its patch.
6288            harness.patch("c", json!({"count": 2})).await;
6289            let frame = next_frame(&mut socket).await;
6290            assert_eq!(frame["key"], "c");
6291            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
6292            socket.close(None).await.ok();
6293        }
6294
6295        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6296        async fn a_list_without_a_snapshot_upserts_each_key_on_its_first_change() {
6297            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
6298            seed(&harness).await;
6299
6300            let mut socket = harness
6301                .subscribe(json!({"view": "Thing/list"}), false)
6302                .await;
6303            harness.patch("b", json!({"count": 2})).await;
6304            let frame = next_frame(&mut socket).await;
6305            assert_eq!(frame["key"], "b");
6306            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
6307            assert_eq!(frame["data"]["name"], "thing-b");
6308
6309            harness.patch("b", json!({"count": 3})).await;
6310            let frame = next_frame(&mut socket).await;
6311            assert_eq!(frame["key"], "b");
6312            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
6313            socket.close(None).await.ok();
6314        }
6315
6316        /// A key created and completed by two mutations of one batch (an
6317        /// activity row and the instruction joined to it) enters a list as one
6318        /// upsert carrying both, read from the cache once both are folded in.
6319        /// The second mutation's patch is in that upsert, so it is not
6320        /// forwarded again; the key's next change is.
6321        // One worker, and a projector the cooperative budget cannot interrupt:
6322        // the whole batch is published before the subscriber runs.
6323        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
6324        async fn a_key_entering_a_list_is_not_sent_the_patch_its_upsert_carries() {
6325            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
6326            seed(&harness).await;
6327            let mut socket = harness
6328                .subscribe(json!({"view": "Thing/list"}), false)
6329                .await;
6330
6331            let slot = harness
6332                .slot
6333                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
6334            let mut creation = Mutation {
6335                export: "Thing".to_string(),
6336                key: json!("n"),
6337                patch: json!({"name": "thing-n"}),
6338                append: vec![],
6339            };
6340            creation.mark_created();
6341            let joined = Mutation {
6342                export: "Thing".to_string(),
6343                key: json!("n"),
6344                patch: json!({"flag": true}),
6345                append: vec![],
6346            };
6347            harness
6348                .tx
6349                .send(MutationBatch::with_slot_context(
6350                    vec![creation, joined].into_iter().collect(),
6351                    SlotContext::new(slot, 0),
6352                ))
6353                .await
6354                .unwrap();
6355
6356            let upsert = next_frame(&mut socket).await;
6357            assert_eq!(upsert["key"], "n");
6358            assert_eq!(upsert["op"], "upsert", "unexpected frame: {upsert}");
6359            assert_eq!(upsert["data"]["name"], "thing-n");
6360            assert_eq!(upsert["data"]["flag"], true);
6361            let (_, sent) = version_of(&upsert["data"]);
6362
6363            // The next frame is the key's next change, not the joined patch.
6364            harness.patch("n", json!({"count": 9})).await;
6365            let patch = next_frame(&mut socket).await;
6366            assert_eq!(patch["key"], "n");
6367            assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
6368            assert_eq!(patch["data"]["count"], 9, "unexpected frame: {patch}");
6369            let (_, patched) = version_of(&patch["data"]);
6370            assert!(patched > sent, "{patched} after {sent}");
6371            socket.close(None).await.ok();
6372        }
6373
6374        /// Eviction from the server's bounded cache is not a change to the
6375        /// entity: a client holding it keeps getting its patches, never a
6376        /// `remove`, and the whole entity replaces its copy when it returns.
6377        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6378        async fn a_state_holder_keeps_its_patches_when_the_cache_evicts_the_key() {
6379            let harness = Harness::start_with_cache(
6380                WebSocketDeliveryConfig::default(),
6381                EntityCache::with_config(crate::cache::EntityCacheConfig {
6382                    max_entities_per_view: 2,
6383                    ..Default::default()
6384                }),
6385            )
6386            .await;
6387            harness
6388                .patch("a", json!({"name": "thing-a", "count": 1}))
6389                .await;
6390            let mut socket = harness
6391                .subscribe(json!({"view": "Thing/state", "key": "a"}), true)
6392                .await;
6393            assert_eq!(snapshot_keys(&mut socket).await, ["a"]);
6394
6395            // "b" and "c" push "a" out of the cache.
6396            for key in ["b", "c"] {
6397                harness
6398                    .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
6399                    .await;
6400            }
6401            harness.patch("a", json!({"count": 2})).await;
6402            let frame = next_frame(&mut socket).await;
6403            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
6404            assert_eq!(frame["op"], "patch", "the holder merges it: {frame}");
6405            assert_eq!(frame["data"]["count"], 2);
6406
6407            // The whole entity comes back from the VM.
6408            let mut whole = Mutation {
6409                export: "Thing".to_string(),
6410                key: json!("a"),
6411                patch: json!({"name": "thing-a", "count": 2}),
6412                append: vec![],
6413            };
6414            whole.mark_whole_entity();
6415            harness.mutate(whole).await;
6416            let frame = next_frame(&mut socket).await;
6417            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
6418            assert_eq!(frame["data"]["name"], "thing-a");
6419            assert_eq!(frame["data"]["count"], 2);
6420            assert!(frame["data"]
6421                .get(arete_interpreter::WHOLE_ENTITY_MARKER)
6422                .is_none());
6423
6424            harness.patch("a", json!({"count": 3})).await;
6425            let frame = next_frame(&mut socket).await;
6426            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
6427            socket.close(None).await.ok();
6428        }
6429
6430        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6431        async fn a_state_view_without_a_snapshot_upserts_on_the_first_change() {
6432            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
6433            seed(&harness).await;
6434
6435            let mut socket = harness
6436                .subscribe(json!({"view": "Thing/state", "key": "a"}), false)
6437                .await;
6438            harness.patch("a", json!({"count": 2})).await;
6439            let frame = next_frame(&mut socket).await;
6440            assert_eq!(frame["key"], "a");
6441            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
6442            assert_eq!(frame["data"]["name"], "thing-a");
6443            assert_eq!(frame["data"]["count"], 2);
6444
6445            harness.patch("a", json!({"count": 3})).await;
6446            let frame = next_frame(&mut socket).await;
6447            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
6448            socket.close(None).await.ok();
6449        }
6450
6451        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6452        async fn a_coalesced_list_sends_no_remove_for_a_key_it_never_sent() {
6453            let harness = Harness::start(WebSocketDeliveryConfig {
6454                collection_coalesce_ms: Some(10),
6455                ..Default::default()
6456            })
6457            .await;
6458            seed(&harness).await;
6459
6460            // A one-row window holding "c", which the disabled snapshot never
6461            // sent.
6462            let mut socket = harness
6463                .subscribe(json!({"view": "Thing/list", "take": 1}), false)
6464                .await;
6465            // "a" jumps to the front: it enters the window whole and "c",
6466            // never sent, leaves it without a `remove`.
6467            harness.patch("a", json!({"count": 2})).await;
6468            let frame = next_frame(&mut socket).await;
6469            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
6470            assert_eq!(frame["op"], "upsert");
6471            assert_eq!(frame["data"]["name"], "thing-a");
6472
6473            // "c" re-enters whole; "a", held, leaves with a `remove`.
6474            harness.patch("c", json!({"count": 2})).await;
6475            let mut frames = [next_frame(&mut socket).await, next_frame(&mut socket).await];
6476            frames.sort_by_key(|frame| frame["key"].as_str().unwrap().to_string());
6477            assert_eq!(frames[0]["key"], "a");
6478            assert_eq!(frames[0]["op"], "remove");
6479            assert_eq!(frames[1]["key"], "c");
6480            assert_eq!(frames[1]["op"], "upsert");
6481            assert_eq!(frames[1]["data"]["name"], "thing-c");
6482            socket.close(None).await.ok();
6483        }
6484
6485        /// `(epoch, counter)` from an entity's `_version`.
6486        fn version_of(data: &Value) -> (String, u64) {
6487            let version = data["_version"]
6488                .as_str()
6489                .unwrap_or_else(|| panic!("no _version in {data}"));
6490            let (epoch, counter) = version.split_once(':').expect("epoch:counter");
6491            (epoch.to_string(), counter.parse().expect("decimal counter"))
6492        }
6493
6494        /// Every entity a subscriber receives carries the projector's
6495        /// `_version`: a snapshot row, a forwarded patch, and a state
6496        /// subscription's catch-up, which sends the cached entity after the
6497        /// bus overwrote a patch.
6498        // One worker, and a projector the cooperative budget cannot interrupt:
6499        // within a batch it awaits only locks no other task holds, so it
6500        // publishes both changes before the subscriber can run, and the
6501        // second overwrites the first on the state bus.
6502        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
6503        async fn snapshot_rows_patches_and_catch_ups_carry_the_version() {
6504            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
6505            harness.patch("7", json!({"name": "thing-7"})).await;
6506
6507            let mut list = harness.subscribe(json!({"view": "Thing/list"}), true).await;
6508            let snapshot = next_frame(&mut list).await;
6509            assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
6510            let (epoch, _) = version_of(&snapshot["data"][0]["data"]);
6511
6512            let mut state = harness
6513                .subscribe(json!({"view": "Thing/state", "key": "7"}), true)
6514                .await;
6515            let snapshot = next_frame(&mut state).await;
6516            assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
6517            let (state_epoch, seeded) = version_of(&snapshot["data"][0]["data"]);
6518            assert_eq!(state_epoch, epoch, "one projector, one epoch");
6519
6520            harness.patch("7", json!({"count": 1})).await;
6521            let patch = next_frame(&mut list).await;
6522            assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
6523            assert_eq!(version_of(&patch["data"]).0, epoch);
6524            let patch = next_frame(&mut state).await;
6525            assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
6526            let (_, patched) = version_of(&patch["data"]);
6527            assert!(patched > seeded, "{patched} after {seeded}");
6528
6529            // Two changes to the key in one batch.
6530            let slot = harness
6531                .slot
6532                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
6533            let mutation = |patch: Value| Mutation {
6534                export: "Thing".to_string(),
6535                key: json!("7"),
6536                patch,
6537                append: vec![],
6538            };
6539            harness
6540                .tx
6541                .send(MutationBatch::with_slot_context(
6542                    vec![
6543                        mutation(json!({"count": 2})),
6544                        mutation(json!({"flag": true})),
6545                    ]
6546                    .into_iter()
6547                    .collect(),
6548                    SlotContext::new(slot, 0),
6549                ))
6550                .await
6551                .unwrap();
6552
6553            let catch_up = next_frame(&mut state).await;
6554            assert_eq!(catch_up["op"], "patch", "unexpected frame: {catch_up}");
6555            assert!(catch_up.get("seq").is_none(), "a catch-up: {catch_up}");
6556            assert_eq!(catch_up["data"]["count"], 2);
6557            assert_eq!(catch_up["data"]["flag"], true);
6558            let (_, caught_up) = version_of(&catch_up["data"]);
6559            assert!(caught_up > patched, "{caught_up} after {patched}");
6560
6561            list.close(None).await.ok();
6562            state.close(None).await.ok();
6563        }
6564    }
6565
6566    /// Session tokens that expire while their socket is open.
6567    mod session_expiry {
6568        use super::*;
6569        use crate::websocket::auth::SignedSessionAuthPlugin;
6570        use arete_auth::{KeyClass, SessionClaims, SigningKey, TokenSigner, TokenVerifier};
6571        use futures_util::{SinkExt, StreamExt};
6572        use std::time::Duration;
6573        use tokio::net::{TcpListener, TcpStream};
6574        use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
6575        use tokio_tungstenite::tungstenite::protocol::CloseFrame;
6576        use tokio_tungstenite::tungstenite::Message;
6577        use tokio_tungstenite::{client_async, WebSocketStream};
6578
6579        #[derive(Default)]
6580        struct RecordingUsageEmitter {
6581            events: tokio::sync::Mutex<Vec<WebSocketUsageEvent>>,
6582        }
6583
6584        #[async_trait::async_trait]
6585        impl WebSocketUsageEmitter for RecordingUsageEmitter {
6586            async fn emit(&self, event: WebSocketUsageEvent) {
6587                self.events.lock().await.push(event);
6588            }
6589        }
6590
6591        struct Server {
6592            addr: SocketAddr,
6593            signer: TokenSigner,
6594            usage: Arc<RecordingUsageEmitter>,
6595            metrics: Arc<DeliveryProbe>,
6596        }
6597
6598        impl Server {
6599            async fn start() -> Self {
6600                let signing_key = SigningKey::generate();
6601                let verifier =
6602                    TokenVerifier::new(signing_key.verifying_key(), "test-issuer", "test-audience");
6603                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6604                let addr = listener.local_addr().unwrap();
6605                let usage = Arc::new(RecordingUsageEmitter::default());
6606                let server = WebSocketServer::new(
6607                    addr,
6608                    BusManager::new(),
6609                    EntityCache::new(),
6610                    Arc::new(ViewIndex::new()),
6611                    #[cfg(feature = "otel")]
6612                    None,
6613                )
6614                .with_auth_plugin(Arc::new(SignedSessionAuthPlugin::new(verifier)))
6615                .with_usage_emitter(usage.clone());
6616                let metrics = Arc::new(DeliveryProbe::default());
6617                let (acceptor, _cleanup) = server.into_acceptor();
6618                let acceptor = acceptor.with_delivery_probe(metrics.clone());
6619                tokio::spawn(async move { acceptor.serve_listener(listener).await });
6620                Self {
6621                    addr,
6622                    signer: TokenSigner::new(signing_key, "test-issuer"),
6623                    usage,
6624                    metrics,
6625                }
6626            }
6627
6628            fn token(&self, ttl_seconds: u64) -> String {
6629                let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
6630                    .with_scope("read")
6631                    .with_key_class(KeyClass::Secret)
6632                    .with_ttl(ttl_seconds)
6633                    .build();
6634                self.signer.sign(claims).unwrap()
6635            }
6636
6637            fn account_token(&self, ttl_seconds: u64, account: &str) -> String {
6638                let claims = SessionClaims::builder("test-issuer", account, "test-audience")
6639                    .with_scope("read")
6640                    .with_key_class(KeyClass::Secret)
6641                    .with_metering_key(account)
6642                    .with_plan("agent_trial")
6643                    .with_actor_key(account)
6644                    .with_account_key(account)
6645                    .with_consumer_key(format!("consumer:{account}"))
6646                    .with_policy_version(3)
6647                    .with_account_limits(Default::default())
6648                    .with_ttl(ttl_seconds)
6649                    .build();
6650                self.signer.sign(claims).unwrap()
6651            }
6652
6653            async fn connect(&self, token: &str) -> WebSocketStream<TcpStream> {
6654                let stream = TcpStream::connect(self.addr).await.unwrap();
6655                client_async(format!("ws://{}/?hs_token={token}", self.addr), stream)
6656                    .await
6657                    .unwrap()
6658                    .0
6659            }
6660        }
6661
6662        async fn send_json(socket: &mut WebSocketStream<TcpStream>, message: Value) {
6663            socket
6664                .send(Message::Text(message.to_string().into()))
6665                .await
6666                .unwrap();
6667        }
6668
6669        /// The close frame, if the server closes the socket within `wait`.
6670        async fn close_within(
6671            socket: &mut WebSocketStream<TcpStream>,
6672            wait: Duration,
6673        ) -> Option<Option<CloseFrame>> {
6674            tokio::time::timeout(wait, async {
6675                while let Some(Ok(message)) = socket.next().await {
6676                    if let Message::Close(frame) = message {
6677                        return Some(frame);
6678                    }
6679                }
6680                Some(None)
6681            })
6682            .await
6683            .ok()
6684            .flatten()
6685        }
6686
6687        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6688        async fn an_expired_session_is_closed_with_the_reason() {
6689            let server = Server::start().await;
6690            let mut socket = server.connect(&server.token(2)).await;
6691
6692            tokio::time::sleep(Duration::from_secs(3)).await;
6693            send_json(&mut socket, json!({"type": "ping"})).await;
6694
6695            let frame = close_within(&mut socket, Duration::from_secs(5))
6696                .await
6697                .expect("the server closes the expired session")
6698                .expect("the close frame carries a reason");
6699            assert_eq!(frame.code, CloseCode::Policy);
6700            assert_eq!(
6701                frame.reason.as_str(),
6702                "token-expired: Authentication token expired"
6703            );
6704        }
6705
6706        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6707        async fn a_session_refreshed_in_band_outlives_its_first_token() {
6708            let server = Server::start().await;
6709            let mut socket = server.connect(&server.token(2)).await;
6710
6711            send_json(
6712                &mut socket,
6713                json!({"type": "refresh_auth", "token": server.token(3_600)}),
6714            )
6715            .await;
6716            let reply = tokio::time::timeout(Duration::from_secs(5), async {
6717                while let Some(Ok(message)) = socket.next().await {
6718                    if let Message::Text(text) = message {
6719                        return serde_json::from_str::<Value>(text.as_str()).ok();
6720                    }
6721                }
6722                None
6723            })
6724            .await
6725            .expect("the server answers the refresh")
6726            .expect("the answer is JSON");
6727            assert_eq!(reply["success"], true, "refresh accepted: {reply}");
6728
6729            tokio::time::sleep(Duration::from_secs(3)).await;
6730            send_json(&mut socket, json!({"type": "ping"})).await;
6731            assert!(
6732                close_within(&mut socket, Duration::from_millis(1_500))
6733                    .await
6734                    .is_none(),
6735                "the socket stays open on the refreshed token"
6736            );
6737        }
6738
6739        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6740        async fn usage_after_auth_refresh_uses_the_refreshed_account() {
6741            let server = Server::start().await;
6742            let mut socket = server
6743                .connect(&server.account_token(3_600, "account:1"))
6744                .await;
6745
6746            send_json(
6747                &mut socket,
6748                json!({
6749                    "type": "refresh_auth",
6750                    "token": server.account_token(3_600, "account:2")
6751                }),
6752            )
6753            .await;
6754            let reply = tokio::time::timeout(Duration::from_secs(5), async {
6755                while let Some(Ok(message)) = socket.next().await {
6756                    if let Message::Text(text) = message {
6757                        return serde_json::from_str::<Value>(text.as_str()).ok();
6758                    }
6759                }
6760                None
6761            })
6762            .await
6763            .expect("the server answers the refresh")
6764            .expect("the answer is JSON");
6765            assert_eq!(reply["success"], true, "refresh accepted: {reply}");
6766            socket.close(None).await.unwrap();
6767
6768            let account = tokio::time::timeout(Duration::from_secs(5), async {
6769                loop {
6770                    let events = server.usage.events.lock().await;
6771                    if let Some(account) = events.iter().find_map(|event| match event {
6772                        WebSocketUsageEvent::ConnectionClosed { identity, .. } => {
6773                            identity.account_key.clone()
6774                        }
6775                        _ => None,
6776                    }) {
6777                        return account;
6778                    }
6779                    drop(events);
6780                    tokio::task::yield_now().await;
6781                }
6782            })
6783            .await
6784            .expect("connection close usage event");
6785            assert_eq!(account, "account:2");
6786            let active = server
6787                .metrics
6788                .active_connections
6789                .lock()
6790                .expect("delivery probe lock poisoned");
6791            assert_eq!(active.get("account:1"), Some(&0));
6792            assert_eq!(active.get("account:2"), None);
6793        }
6794    }
6795}
6796
6797#[cfg(test)]
6798mod retained_scope_regressions {
6799    use super::*;
6800    use crate::row_scope::{RowScope, ScopeField};
6801
6802    #[test]
6803    fn sparse_private_frames_never_inherit_a_later_public_scope() {
6804        let public = RowFilter::new(
6805            ScopeField::new("meta.scope"),
6806            Default::default(),
6807            Arc::new(tokio::sync::RwLock::new(0)),
6808            0,
6809        );
6810        let entitled = RowFilter::new(
6811            ScopeField::new("meta.scope"),
6812            ["B".to_string()].into_iter().collect(),
6813            Arc::new(tokio::sync::RwLock::new(0)),
6814            0,
6815        );
6816        let patch = br#"{"data":{"secret":"private-at-the-time"}}"#;
6817        let original_scope = RowScope::Scoped("B".into());
6818        assert!(!frame_in_row_scopes(
6819            Some(&public),
6820            Some(&original_scope),
6821            patch
6822        ));
6823        assert!(frame_in_row_scopes(
6824            Some(&entitled),
6825            Some(&original_scope),
6826            patch
6827        ));
6828        // The public transition itself is delivered, without relabeling old frames.
6829        assert!(frame_in_row_scopes(
6830            Some(&public),
6831            Some(&RowScope::Public),
6832            br#"{"data":{"meta":{"scope":null}}}"#
6833        ));
6834        // Legacy sparse records have no proof of membership: fail closed.
6835        assert!(!frame_in_row_scopes(Some(&public), None, patch));
6836        assert!(!frame_in_row_scopes(Some(&entitled), None, patch));
6837        assert!(frame_in_row_scopes(None, None, patch));
6838    }
6839}
6840
6841#[cfg(test)]
6842mod limits_tests {
6843    use super::*;
6844    use futures_util::SinkExt;
6845    use tokio::io::AsyncReadExt;
6846    use tokio_tungstenite::client_async;
6847    use tokio_tungstenite::tungstenite::Message;
6848
6849    fn limits() -> WebSocketLimitsConfig {
6850        WebSocketLimitsConfig {
6851            max_message_bytes: 16 * 1024,
6852            max_frame_bytes: 16 * 1024,
6853            handshake_timeout: Duration::from_millis(300),
6854            max_pending_handshakes: 1,
6855        }
6856    }
6857
6858    async fn start(limits: WebSocketLimitsConfig) -> SocketAddr {
6859        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6860        let addr = listener.local_addr().unwrap();
6861        let server = WebSocketServer::new(
6862            addr,
6863            BusManager::new(),
6864            EntityCache::new(),
6865            Arc::new(ViewIndex::new()),
6866            #[cfg(feature = "otel")]
6867            None,
6868        )
6869        .with_limits(limits);
6870        let (acceptor, _cleanup) = server.into_acceptor();
6871        tokio::spawn(async move { acceptor.serve_listener(listener).await });
6872        addr
6873    }
6874
6875    /// Resolves once the server closes `stream`, or fails the test after `within`.
6876    async fn closed_within(stream: &mut TcpStream, within: Duration) {
6877        let mut buffer = [0; 1024];
6878        tokio::time::timeout(within, async {
6879            loop {
6880                match stream.read(&mut buffer).await {
6881                    Ok(0) | Err(_) => return,
6882                    Ok(_) => continue,
6883                }
6884            }
6885        })
6886        .await
6887        .expect("the server closes the connection");
6888    }
6889
6890    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6891    async fn a_connection_that_never_upgrades_is_dropped_at_the_deadline() {
6892        let addr = start(limits()).await;
6893        let mut stream = TcpStream::connect(addr).await.unwrap();
6894        let started = Instant::now();
6895        closed_within(&mut stream, Duration::from_secs(5)).await;
6896        assert!(started.elapsed() >= Duration::from_millis(250));
6897    }
6898
6899    /// Allows every connection, but only after `delay`.
6900    struct SlowAuthPlugin {
6901        delay: Duration,
6902    }
6903
6904    #[async_trait::async_trait]
6905    impl WebSocketAuthPlugin for SlowAuthPlugin {
6906        async fn authorize(&self, request: &ConnectionAuthRequest) -> AuthDecision {
6907            tokio::time::sleep(self.delay).await;
6908            crate::websocket::auth::AllowAllAuthPlugin
6909                .authorize(request)
6910                .await
6911        }
6912
6913        fn as_any(&self) -> &dyn std::any::Any {
6914            self
6915        }
6916    }
6917
6918    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6919    async fn authentication_slower_than_the_deadline_is_cut_off() {
6920        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6921        let addr = listener.local_addr().unwrap();
6922        let server = WebSocketServer::new(
6923            addr,
6924            BusManager::new(),
6925            EntityCache::new(),
6926            Arc::new(ViewIndex::new()),
6927            #[cfg(feature = "otel")]
6928            None,
6929        )
6930        .with_auth_plugin(Arc::new(SlowAuthPlugin {
6931            delay: Duration::from_secs(30),
6932        }))
6933        .with_limits(limits());
6934        let (acceptor, _cleanup) = server.into_acceptor();
6935        tokio::spawn(async move { acceptor.serve_listener(listener).await });
6936
6937        let started = Instant::now();
6938        let stream = TcpStream::connect(addr).await.unwrap();
6939        let handshake = tokio::time::timeout(
6940            Duration::from_secs(5),
6941            client_async(format!("ws://{addr}/"), stream),
6942        )
6943        .await
6944        .expect("the server ends the handshake before the test deadline");
6945        assert!(handshake.is_err(), "the upgrade is refused");
6946        assert!(started.elapsed() < Duration::from_secs(5));
6947    }
6948
6949    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6950    async fn pending_handshakes_past_the_cap_are_refused() {
6951        let addr = start(WebSocketLimitsConfig {
6952            handshake_timeout: Duration::from_secs(30),
6953            ..limits()
6954        })
6955        .await;
6956        let mut stalled = TcpStream::connect(addr).await.unwrap();
6957        // Let the server register the first connection as pending.
6958        tokio::time::sleep(Duration::from_millis(100)).await;
6959        let mut refused = TcpStream::connect(addr).await.unwrap();
6960        closed_within(&mut refused, Duration::from_secs(5)).await;
6961
6962        // The first is still held open within its own deadline.
6963        let mut buffer = [0; 16];
6964        assert!(
6965            tokio::time::timeout(Duration::from_millis(200), stalled.read(&mut buffer))
6966                .await
6967                .is_err()
6968        );
6969    }
6970
6971    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6972    async fn a_message_over_the_limit_closes_the_socket() {
6973        let addr = start(limits()).await;
6974        let stream = TcpStream::connect(addr).await.unwrap();
6975        let (mut socket, _) = client_async(format!("ws://{addr}/"), stream).await.unwrap();
6976        socket
6977            .send(Message::Text("x".repeat(64 * 1024).into()))
6978            .await
6979            .unwrap();
6980        let ended = tokio::time::timeout(Duration::from_secs(5), async {
6981            loop {
6982                match socket.next().await {
6983                    None | Some(Err(_)) | Some(Ok(Message::Close(_))) => return,
6984                    Some(Ok(_)) => continue,
6985                }
6986            }
6987        })
6988        .await;
6989        assert!(ended.is_ok(), "the server ends the session");
6990    }
6991}