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