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, SortConfig, SortOrder, SourceFields, SubscribedFrame,
13    UnsubscribedFrame, WireEntity, WireFormat,
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
1284/// One snapshot frame, serialized straight from the cached entities it
1285/// carries: the same JSON as a [`SnapshotFrame`] of
1286/// [`SnapshotEntity`](crate::websocket::frame::SnapshotEntity) rows with the
1287/// view's wire format applied, without building those rows.
1288///
1289/// [`SnapshotFrame`]: crate::websocket::frame::SnapshotFrame
1290#[derive(Serialize)]
1291#[serde(rename_all = "camelCase")]
1292struct SnapshotBatch<'a> {
1293    protocol_version: u8,
1294    subscription_id: &'a str,
1295    snapshot_id: &'a str,
1296    authoritative: bool,
1297    mode: Mode,
1298    #[serde(rename = "entity")]
1299    export: &'a str,
1300    op: &'static str,
1301    #[serde(skip_serializing_if = "Option::is_none")]
1302    key: Option<&'a str>,
1303    data: SnapshotRows<'a>,
1304    complete: bool,
1305}
1306
1307struct SnapshotRows<'a> {
1308    rows: &'a [(String, SharedEntity)],
1309    wire_format: &'a WireFormat,
1310}
1311
1312impl Serialize for SnapshotRows<'_> {
1313    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
1314        use serde::ser::SerializeSeq as _;
1315        let mut rows = serializer.serialize_seq(Some(self.rows.len()))?;
1316        for (key, entity) in self.rows {
1317            rows.serialize_element(&SnapshotRow {
1318                key,
1319                data: WireEntity {
1320                    entity,
1321                    wire_format: self.wire_format,
1322                },
1323            })?;
1324        }
1325        rows.end()
1326    }
1327}
1328
1329#[derive(Serialize)]
1330struct SnapshotRow<'a> {
1331    key: &'a str,
1332    data: WireEntity<'a>,
1333}
1334
1335/// Split `rows` into the snapshot's frames: the first batch for a fast first
1336/// render, then larger ones. An empty snapshot is one empty, complete frame.
1337fn create_snapshot_batches<'a>(
1338    rows: &'a [(String, SharedEntity)],
1339    wire_format: &'a WireFormat,
1340    metadata: SnapshotMetadata<'a>,
1341    batch_config: &SnapshotBatchConfig,
1342) -> Vec<SnapshotBatch<'a>> {
1343    let batch = |data: &'a [(String, SharedEntity)], complete: bool| SnapshotBatch {
1344        protocol_version: PROTOCOL_VERSION,
1345        subscription_id: metadata.subscription_id,
1346        snapshot_id: metadata.snapshot_id,
1347        authoritative: metadata.authoritative,
1348        mode: metadata.mode,
1349        export: metadata.view_id,
1350        op: "snapshot",
1351        key: metadata.key,
1352        data: SnapshotRows {
1353            rows: data,
1354            wire_format,
1355        },
1356        complete,
1357    };
1358    if rows.is_empty() {
1359        return vec![batch(&[], true)];
1360    }
1361
1362    let mut batches = Vec::new();
1363    let mut offset = 0;
1364    while offset < rows.len() {
1365        let configured_size = if offset == 0 {
1366            batch_config.initial_batch_size
1367        } else {
1368            batch_config.subsequent_batch_size
1369        };
1370        let end = (offset + configured_size.max(1)).min(rows.len());
1371        batches.push(batch(&rows[offset..end], end == rows.len()));
1372        offset = end;
1373    }
1374    batches
1375}
1376
1377/// Send `rows` as one snapshot, serializing each frame from the cached
1378/// entities as it goes: no row is copied.
1379async fn send_snapshot_batches(
1380    context: &SubscriptionContext,
1381    subscription: &Subscription,
1382    rows: &[(String, SharedEntity)],
1383    view_spec: &ViewSpec,
1384    purpose: SnapshotPurpose,
1385    batch_config: &SnapshotBatchConfig,
1386) -> Result<()> {
1387    let snapshot_id = Uuid::new_v4().to_string();
1388    let authoritative = purpose.authoritative(subscription);
1389    let mode = view_spec.mode;
1390    let frames = create_snapshot_batches(
1391        rows,
1392        &view_spec.wire_format,
1393        SnapshotMetadata {
1394            subscription_id: &subscription.subscription_id,
1395            snapshot_id: &snapshot_id,
1396            authoritative,
1397            mode,
1398            view_id: &subscription.query.view,
1399            key: subscription.query.key.as_deref(),
1400        },
1401        batch_config,
1402    );
1403
1404    for frame in frames {
1405        let rows = frame.data.rows.len() as u32;
1406        let json = serde_json::to_vec(&frame)?;
1407        let payload = maybe_compress(&json);
1408        let bytes = payload.as_bytes().len() as u64;
1409        context
1410            .client_manager
1411            .send_compressed_async(context.client_id, payload)
1412            .await
1413            .map_err(|error| anyhow::anyhow!("failed to send snapshot: {error}"))?;
1414        context.metrics.message_sent();
1415
1416        let auth_context = context.client_manager.get_auth_context(context.client_id);
1417        let (identity, deployment_id) = usage_identity_from_context(auth_context.as_ref());
1418        emit_usage_event(
1419            &context.usage_emitter,
1420            WebSocketUsageEvent::SnapshotSent {
1421                client_id: context.client_id.to_string(),
1422                deployment_id,
1423                identity,
1424                view_id: subscription.query.view.clone(),
1425                rows,
1426                messages: 1,
1427                bytes,
1428            },
1429        );
1430    }
1431    Ok(())
1432}
1433
1434fn extract_sort_config(view_spec: &ViewSpec) -> Option<SortConfig> {
1435    if let Some(sort) = view_spec
1436        .pipeline
1437        .as_ref()
1438        .and_then(|pipeline| pipeline.sort.as_ref())
1439    {
1440        return Some(SortConfig {
1441            field: sort.field_path.clone(),
1442            order: match sort.order {
1443                crate::materialized_view::SortOrder::Asc => SortOrder::Asc,
1444                crate::materialized_view::SortOrder::Desc => SortOrder::Desc,
1445            },
1446        });
1447    }
1448    (view_spec.mode == Mode::List).then(|| SortConfig {
1449        field: vec!["_seq".to_string()],
1450        order: SortOrder::Desc,
1451    })
1452}
1453
1454fn send_control_frame<T: Serialize>(
1455    context: &SubscriptionContext,
1456    frame: &T,
1457    view_id: &str,
1458) -> Result<()> {
1459    let json = serde_json::to_vec(frame)?;
1460    let bytes = json.len();
1461    context
1462        .client_manager
1463        .send_to_client(context.client_id, Arc::new(Bytes::from(json)))
1464        .map_err(|error| anyhow::anyhow!("failed to send control frame: {error}"))?;
1465    context.metrics.message_sent();
1466    emit_update_sent_for_client(
1467        &context.usage_emitter,
1468        &context.client_manager,
1469        context.client_id,
1470        view_id,
1471        bytes,
1472    );
1473    Ok(())
1474}
1475
1476fn send_subscribed_frame(
1477    context: &SubscriptionContext,
1478    subscription: &Subscription,
1479    view_spec: &ViewSpec,
1480) -> Result<()> {
1481    let frame = SubscribedFrame::new(
1482        subscription.subscription_id.clone(),
1483        subscription.query.clone(),
1484        view_spec.mode,
1485        extract_sort_config(view_spec),
1486    );
1487    send_control_frame(context, &frame, &subscription.query.view)
1488}
1489
1490fn enforce_snapshot_limit(context: &SubscriptionContext, rows: usize) -> Result<()> {
1491    context
1492        .client_manager
1493        .check_snapshot_allowed(context.client_id, u32::try_from(rows).unwrap_or(u32::MAX))
1494        .map_err(|deny| anyhow::anyhow!(deny.reason))
1495}
1496
1497async fn subscribe_state_then_snapshot<F, Fut, T>(
1498    bus_manager: &BusManager,
1499    view_id: &str,
1500    key: &str,
1501    snapshot: F,
1502) -> (watch::Receiver<StateUpdate>, u64, T)
1503where
1504    F: FnOnce() -> Fut,
1505    Fut: Future<Output = T>,
1506{
1507    let mut receiver = bus_manager.get_or_create_state_bus(view_id, key).await;
1508    let published = receiver.borrow_and_update().published;
1509    let snapshot = snapshot().await;
1510    (receiver, published, snapshot)
1511}
1512
1513async fn subscribe_list_then_snapshot<F, Fut, T>(
1514    bus_manager: &BusManager,
1515    view_id: &str,
1516    snapshot: F,
1517) -> (broadcast::Receiver<Arc<BusMessage>>, T)
1518where
1519    F: FnOnce() -> Fut,
1520    Fut: Future<Output = T>,
1521{
1522    let receiver = bus_manager.get_or_create_list_bus(view_id).await;
1523    let snapshot = snapshot().await;
1524    (receiver, snapshot)
1525}
1526
1527async fn attach_client_to_bus(
1528    context: &SubscriptionContext,
1529    mut subscription: Subscription,
1530    cancel_token: CancellationToken,
1531) -> Result<()> {
1532    let view_spec = context
1533        .view_index
1534        .get_view(&subscription.query.view)
1535        .cloned()
1536        .ok_or_else(|| anyhow::anyhow!("unknown view: {}", subscription.query.view))?;
1537
1538    if view_spec.mode == Mode::State && !view_spec.is_derived() && subscription.query.key.is_none()
1539    {
1540        return Err(anyhow::anyhow!("state subscriptions require query.key"));
1541    }
1542    if view_spec.is_derived() && subscription.query.take.is_none() {
1543        subscription.query.take = view_spec
1544            .pipeline
1545            .as_ref()
1546            .and_then(|pipeline| pipeline.limit);
1547    }
1548
1549    // A retained tape takes precedence for append views: it is the only
1550    // delivery that can honour a cursor. Without one, fall through to the
1551    // previous latest-state behaviour.
1552    let journal = context
1553        .journal
1554        .clone()
1555        .filter(|journal| journal.is_enabled() && view_spec.mode == Mode::Append);
1556    if let Some(journal) = journal {
1557        return attach_journal_subscription(
1558            context,
1559            subscription,
1560            view_spec,
1561            journal,
1562            cancel_token,
1563        )
1564        .await;
1565    }
1566
1567    if view_spec.mode == Mode::State && !view_spec.is_derived() {
1568        attach_state_subscription(context, subscription, view_spec, cancel_token).await
1569    } else {
1570        attach_collection_subscription(context, subscription, view_spec, cancel_token).await
1571    }
1572}
1573
1574async fn attach_state_subscription(
1575    context: &SubscriptionContext,
1576    subscription: Subscription,
1577    view_spec: ViewSpec,
1578    cancel_token: CancellationToken,
1579) -> Result<()> {
1580    let view_id = subscription.query.view.clone();
1581    let key = subscription.query.key.clone().unwrap_or_default();
1582    let query = subscription.query.clone();
1583    let cache = context.entity_cache.clone();
1584    let view_spec_for_snapshot = view_spec.clone();
1585    let (mut receiver, mut seen, initial) =
1586        subscribe_state_then_snapshot(&context.bus_manager, &view_id, &key, move || async move {
1587            load_query_entities(&cache, None, &view_spec_for_snapshot, &query, false).await
1588        })
1589        .await;
1590
1591    let mut snapshot_entities = initial;
1592    if let Some(limit) = subscription.query.snapshot_limit {
1593        snapshot_entities.truncate(limit);
1594    }
1595    enforce_snapshot_limit(context, snapshot_entities.len())?;
1596    send_subscribed_frame(context, &subscription, &view_spec)?;
1597    // The client holds the entity only if the snapshot actually carried it.
1598    // A disabled snapshot, or one `snapshotLimit` cut to nothing, leaves it
1599    // without a copy, so its first change must be a full `upsert`: a patch
1600    // would have nothing to merge into.
1601    let delivered = subscription.snapshot.enabled && !snapshot_entities.is_empty();
1602    // The entity as the client last received it whole, which a catch-up is
1603    // measured against.
1604    let synced = if delivered {
1605        snapshot_entities
1606            .first()
1607            .map(|(_, entity)| Synced::Entity(entity.clone()))
1608    } else {
1609        None
1610    };
1611    if subscription.snapshot.enabled {
1612        send_snapshot_batches(
1613            context,
1614            &subscription,
1615            &snapshot_entities,
1616            &view_spec,
1617            SnapshotPurpose::Initial,
1618            &context.entity_cache.snapshot_config(),
1619        )
1620        .await?;
1621    }
1622
1623    let task_context = context.clone();
1624    let subscription_id = subscription.subscription_id.clone();
1625    let query = subscription.query.clone();
1626    let view_spec_task = view_spec.clone();
1627    let span_view = view_id.clone();
1628    let span_key = key.clone();
1629    tokio::spawn(
1630        async move {
1631            // Whether the client holds this key, not whether it exists.
1632            let mut member = delivered;
1633            // The entity as the client last received it whole (its snapshot
1634            // row, an upsert or the last catch-up; see `Synced`). A catch-up
1635            // sends what changed since, and every field a patch forwarded in
1636            // between set (`touched`, see `mark_fields`): the entity may have
1637            // set it back, and a client may have dropped the patch.
1638            let mut synced = synced;
1639            let mut touched = Value::Null;
1640            // Whether the client lacks fields of frames the bus overwrote
1641            // before this task read them. The cached entity is sent in their
1642            // place, but while the cache lacks the key there is none, so the
1643            // debt waits for the next frame that finds it cached.
1644            let mut behind = false;
1645            // Whether, while behind, a frame the client missed replaced the
1646            // entity (a delete, or a whole upsert): its copy may then hold
1647            // fields the entity no longer has, which merging keeps, so the
1648            // catch-up must replace it.
1649            let mut replaced = false;
1650            loop {
1651                tokio::select! {
1652                    _ = cancel_token.cancelled() => break,
1653                    changed = receiver.changed() => {
1654                        if changed.is_err() {
1655                            break;
1656                        }
1657                        let (payload, published, replaced_at) = {
1658                            let update = receiver.borrow_and_update();
1659                            (update.payload.clone(), update.published, update.replaced)
1660                        };
1661                        // The bus keeps only the latest frame. If more than one
1662                        // was published since the last read, the earlier ones
1663                        // were overwritten and the latest patch alone would
1664                        // drop their fields, so send the cached entity instead.
1665                        behind |= published > seen + 1;
1666                        replaced |= behind && replaced_at > seen;
1667                        seen = published;
1668                        let metadata = source_frame_metadata(&payload);
1669                        if metadata.op == "delete" {
1670                            if !task_context.entity_cache
1671                                .deletion_is_current(&query.view, &key)
1672                                .await
1673                            {
1674                                behind = true;
1675                                replaced = true;
1676                                continue;
1677                            }
1678                            if member && send_membership_frame(
1679                                &task_context,
1680                                &subscription_id,
1681                                &view_spec_task,
1682                                "delete",
1683                                &key,
1684                                Value::Null,
1685                                metadata.seq,
1686                            ).is_err() {
1687                                break;
1688                            }
1689                            member = false;
1690                            behind = false;
1691                            replaced = false;
1692                            synced = None;
1693                            touched = Value::Null;
1694                            continue;
1695                        }
1696
1697                        // A key missing from the cache was evicted there or
1698                        // never held whole, not deleted (deletes arrive
1699                        // above): the cache refused this patch and the whole
1700                        // entity is on its way (see `EntityResync`). The
1701                        // client's copy is still good, so a holder gets the
1702                        // patch and stays a holder; `remove` is only for an
1703                        // entity that stopped matching. If it is behind, it
1704                        // stays behind: the resend that brings the entity
1705                        // back carries the seq of the latest change, which it
1706                        // already has, so only the catch-up below reaches it.
1707                        let Some(cached) = task_context
1708                            .entity_cache
1709                            .get_shared(&view_spec_task.id, &key)
1710                            .await
1711                        else {
1712                            // A copy of a replaced entity takes no patch: it
1713                            // waits for the whole entity to replace it.
1714                            if member && !replaced {
1715                                record_forwarded(&payload, &metadata.op, &mut synced, &mut touched);
1716                                if send_scoped_source_payload(
1717                                    &task_context,
1718                                    &subscription_id,
1719                                    &query.view,
1720                                    payload,
1721                                ).is_err() {
1722                                    break;
1723                                }
1724                            }
1725                            continue;
1726                        };
1727                        let selected =
1728                            select_query_entities(vec![(key.clone(), cached)], &query, true, false);
1729                        let is_member = !selected.is_empty();
1730                        let caught_up = !std::mem::take(&mut behind);
1731                        let replace = std::mem::take(&mut replaced);
1732                        let result = match (member, selected.into_iter().next()) {
1733                            (true, Some(_)) if caught_up => {
1734                                record_forwarded(&payload, &metadata.op, &mut synced, &mut touched);
1735                                send_scoped_source_payload(
1736                                    &task_context,
1737                                    &subscription_id,
1738                                    &query.view,
1739                                    payload,
1740                                )
1741                            }
1742                            // The catch-up: a patch of the cached entity's
1743                            // fields that changed since the client last received
1744                            // it whole or that a forwarded patch set, which
1745                            // merges into what the client holds.
1746                            // If a missed frame replaced the entity, it is the
1747                            // whole entity as an upsert instead, which replaces
1748                            // the client's copy: merging would keep the fields of
1749                            // the entity that was deleted or replaced. Either
1750                            // carries no seq, because it is newer than anything
1751                            // this subscriber was sent and seqs are not ordered
1752                            // within a slot: account updates and instructions
1753                            // number themselves differently, so the latest
1754                            // patch's seq can sort below one already delivered.
1755                            (true, Some((entity_key, data))) => {
1756                                let wire_format = &view_spec_task.wire_format;
1757                                let base = synced
1758                                    .as_ref()
1759                                    .filter(|_| !replace)
1760                                    .and_then(|synced| synced.wire(wire_format));
1761                                let result = match base {
1762                                    Some(base) => match changed_fields(
1763                                        &base,
1764                                        &wire_value(&data, wire_format),
1765                                        &touched,
1766                                    ) {
1767                                        Some(body) => send_membership_frame(
1768                                            &task_context,
1769                                            &subscription_id,
1770                                            &view_spec_task,
1771                                            "patch",
1772                                            &entity_key,
1773                                            body,
1774                                            None,
1775                                        ),
1776                                        None => Ok(()),
1777                                    },
1778                                    // Nothing to measure against: the whole
1779                                    // entity, straight from the cache.
1780                                    None => send_entity_frame(
1781                                        &task_context,
1782                                        &subscription_id,
1783                                        &view_spec_task,
1784                                        if replace { "upsert" } else { "patch" },
1785                                        &entity_key,
1786                                        &data,
1787                                        None,
1788                                    ),
1789                                };
1790                                synced = Some(Synced::Entity(data));
1791                                touched = Value::Null;
1792                                result
1793                            }
1794                            (false, Some((entity_key, data))) => {
1795                                let result = send_entity_frame(
1796                                    &task_context,
1797                                    &subscription_id,
1798                                    &view_spec_task,
1799                                    "upsert",
1800                                    &entity_key,
1801                                    &data,
1802                                    metadata.seq.as_deref(),
1803                                );
1804                                synced = Some(Synced::Entity(data));
1805                                touched = Value::Null;
1806                                result
1807                            }
1808                            (true, None) => {
1809                                synced = None;
1810                                touched = Value::Null;
1811                                send_membership_frame(
1812                                &task_context,
1813                                &subscription_id,
1814                                &view_spec_task,
1815                                "remove",
1816                                &key,
1817                                Value::Null,
1818                                metadata.seq,
1819                                )
1820                            }
1821                            (false, None) => Ok(()),
1822                        };
1823                        if result.is_err() {
1824                            break;
1825                        }
1826                        member = is_member;
1827                    }
1828                }
1829            }
1830        }
1831        .instrument(info_span!("ws.subscribe.state", client_id = %context.client_id, view = %span_view, key = %span_key)),
1832    );
1833    Ok(())
1834}
1835
1836async fn attach_collection_subscription(
1837    context: &SubscriptionContext,
1838    subscription: Subscription,
1839    view_spec: ViewSpec,
1840    cancel_token: CancellationToken,
1841) -> Result<()> {
1842    let view_id = subscription.query.view.clone();
1843    let source_view_id = view_spec
1844        .source_view
1845        .clone()
1846        .unwrap_or_else(|| view_id.clone());
1847    let (mut receiver, initial_membership) = subscribe_collection_then_snapshot(
1848        context,
1849        &source_view_id,
1850        &view_spec,
1851        &subscription.query,
1852    )
1853    .await;
1854
1855    let mut snapshot_entities = initial_membership.clone();
1856    if let Some(limit) = subscription.query.snapshot_limit {
1857        snapshot_entities.truncate(limit);
1858    }
1859    enforce_snapshot_limit(context, snapshot_entities.len())?;
1860    send_subscribed_frame(context, &subscription, &view_spec)?;
1861    // Window members the client was never sent: everything past a
1862    // `snapshotLimit` cut, or the whole window when the snapshot is disabled.
1863    // They stay in the window (take/skip positions still count them) but the
1864    // client holds no copy, so none of them may ever be sent a patch.
1865    let delivered_rows = if subscription.snapshot.enabled {
1866        snapshot_entities.len()
1867    } else {
1868        0
1869    };
1870    let mut undelivered: HashSet<String> = initial_membership
1871        .iter()
1872        .skip(delivered_rows)
1873        .map(|(key, _)| key.clone())
1874        .collect();
1875    if subscription.snapshot.enabled {
1876        send_snapshot_batches(
1877            context,
1878            &subscription,
1879            &snapshot_entities,
1880            &view_spec,
1881            SnapshotPurpose::Initial,
1882            &context.entity_cache.snapshot_config(),
1883        )
1884        .await?;
1885    }
1886
1887    let task_context = context.clone();
1888    let task_subscription = subscription.clone();
1889    let subscription_id = subscription.subscription_id.clone();
1890    let query = subscription.query.clone();
1891    let view_spec_task = view_spec.clone();
1892    let span_view = view_id.clone();
1893    tokio::spawn(
1894        async move {
1895            let mut current = initial_membership;
1896            // Append views are event tapes: collapsing two records would lose
1897            // observable history. Coalescing is only valid for latest-state
1898            // list membership, where a full final entity preserves meaning.
1899            let coalesce_ms = (view_spec_task.mode == Mode::List)
1900                .then(|| {
1901                    view_spec_task
1902                        .delivery
1903                        .coalesce_ms
1904                        .or(task_context.delivery.collection_coalesce_ms)
1905                        .filter(|milliseconds| *milliseconds > 0)
1906                })
1907                .flatten();
1908            let mut flush_interval = coalesce_ms.map(|milliseconds| {
1909                let period = Duration::from_millis(milliseconds);
1910                let mut interval = tokio::time::interval_at(
1911                    tokio::time::Instant::now() + period,
1912                    period,
1913                );
1914                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
1915                interval
1916            });
1917            let mut pending = HashMap::<String, Arc<BusMessage>>::new();
1918            let mut pending_updates = 0_u64;
1919
1920            loop {
1921                tokio::select! {
1922                    _ = cancel_token.cancelled() => break,
1923                    _ = async {
1924                        flush_interval
1925                            .as_mut()
1926                            .expect("coalescing interval is guarded")
1927                            .tick()
1928                            .await;
1929                    }, if flush_interval.is_some() => {
1930                        if pending.is_empty() {
1931                            continue;
1932                        }
1933                        let sorted_caches = view_spec_task
1934                            .is_derived()
1935                            .then(|| task_context.view_index.sorted_caches());
1936                        let next = load_query_entities(
1937                            &task_context.entity_cache,
1938                            sorted_caches,
1939                            &view_spec_task,
1940                            &query,
1941                            false,
1942                        ).await;
1943                        if emit_coalesced_collection_delta(
1944                            &task_context,
1945                            &subscription_id,
1946                            &view_spec_task,
1947                            &current,
1948                            &next,
1949                            &pending,
1950                            &mut undelivered,
1951                        ).is_err() {
1952                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
1953                            break;
1954                        }
1955                        task_context
1956                            .metrics
1957                            .collection_coalesced(&view_id, pending_updates);
1958                        current = next;
1959                        pending.clear();
1960                        pending_updates = 0;
1961                    }
1962                    received = receiver.recv() => {
1963                        let envelope = match received {
1964                            Ok(envelope) => envelope,
1965                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
1966                                task_context.metrics.subscription_lagged(&view_id, skipped);
1967                                warn!(
1968                                    "Subscription {} lagged by {} updates",
1969                                    subscription_id, skipped
1970                                );
1971                                if view_spec_task.mode == Mode::Append {
1972                                    let _ = send_control_frame(
1973                                        &task_context,
1974                                        &SocketIssueMessage::append_subscription_lagged(
1975                                            subscription_id.clone(),
1976                                            skipped,
1977                                        ),
1978                                        &view_id,
1979                                    );
1980                                    task_context.metrics.delivery_stopped(&view_id, "append-lagged-without-replay");
1981                                    break;
1982                                }
1983                                if !task_subscription.snapshot.enabled {
1984                                    let _ = send_control_frame(
1985                                        &task_context,
1986                                        &SocketIssueMessage::subscription_lagged(
1987                                            subscription_id.clone(),
1988                                            skipped,
1989                                        ),
1990                                        &view_id,
1991                                    );
1992                                    task_context.metrics.delivery_stopped(&view_id, "lagged-without-snapshot");
1993                                    break;
1994                                }
1995                                info!(
1996                                    "Subscription {} is recovering from an authoritative snapshot",
1997                                    subscription_id
1998                                );
1999                                match recover_collection_subscription(
2000                                    &task_context,
2001                                    &task_subscription,
2002                                    &view_spec_task,
2003                                    &source_view_id,
2004                                ).await {
2005                                    Ok((next_receiver, recovered)) => {
2006                                        receiver = next_receiver;
2007                                        current = recovered;
2008                                        // Recovery replaces the whole live
2009                                        // membership, untruncated.
2010                                        undelivered.clear();
2011                                        pending.clear();
2012                                        pending_updates = 0;
2013                                        task_context.metrics.subscription_resnapshot(&view_id);
2014                                        continue;
2015                                    }
2016                                    Err(error) => {
2017                                        warn!(
2018                                            "Subscription {} failed to recover from lag: {error:#}",
2019                                            subscription_id
2020                                        );
2021                                        task_context.metrics.delivery_stopped(&view_id, "resnapshot-failed");
2022                                        break;
2023                                    }
2024                                }
2025                            }
2026                            Err(broadcast::error::RecvError::Closed) => break,
2027                        };
2028
2029                        apply_collection_source_event(
2030                            &task_context,
2031                            &source_view_id,
2032                            &view_spec_task,
2033                            &query,
2034                            &envelope,
2035                        ).await;
2036
2037                        if flush_interval.is_some() {
2038                            pending_updates = pending_updates.saturating_add(1);
2039                            pending.insert(envelope.key.clone(), envelope);
2040                            continue;
2041                        }
2042
2043                        let metadata = source_frame_metadata(&envelope.payload);
2044                        let sorted_caches = view_spec_task
2045                            .is_derived()
2046                            .then(|| task_context.view_index.sorted_caches());
2047                        let next = load_query_entities(
2048                            &task_context.entity_cache,
2049                            sorted_caches,
2050                            &view_spec_task,
2051                            &query,
2052                            false,
2053                        ).await;
2054                        if emit_collection_delta(
2055                            &task_context,
2056                            &subscription_id,
2057                            &view_spec_task,
2058                            &current,
2059                            &next,
2060                            &envelope,
2061                            &metadata,
2062                            &mut undelivered,
2063                        ).is_err() {
2064                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
2065                            break;
2066                        }
2067                        current = next;
2068                    }
2069                }
2070            }
2071        }
2072        .instrument(info_span!("ws.subscribe.collection", client_id = %context.client_id, view = %span_view)),
2073    );
2074    Ok(())
2075}
2076
2077async fn subscribe_collection_then_snapshot(
2078    context: &SubscriptionContext,
2079    source_view_id: &str,
2080    view_spec: &ViewSpec,
2081    query: &SubscriptionQuery,
2082) -> (
2083    broadcast::Receiver<Arc<BusMessage>>,
2084    Vec<(String, SharedEntity)>,
2085) {
2086    let cache = context.entity_cache.clone();
2087    let sorted_caches = view_spec
2088        .is_derived()
2089        .then(|| context.view_index.sorted_caches());
2090    let view_spec = view_spec.clone();
2091    let query = query.clone();
2092    subscribe_list_then_snapshot(&context.bus_manager, source_view_id, move || async move {
2093        load_query_entities(&cache, sorted_caches, &view_spec, &query, false).await
2094    })
2095    .await
2096}
2097
2098async fn recover_collection_subscription(
2099    context: &SubscriptionContext,
2100    subscription: &Subscription,
2101    view_spec: &ViewSpec,
2102    source_view_id: &str,
2103) -> Result<(
2104    broadcast::Receiver<Arc<BusMessage>>,
2105    Vec<(String, SharedEntity)>,
2106)> {
2107    let (receiver, membership) =
2108        subscribe_collection_then_snapshot(context, source_view_id, view_spec, &subscription.query)
2109            .await;
2110    // `snapshotLimit` caps only the initial transfer. Recovery has to replace
2111    // the complete live membership; truncating it would make the replacement
2112    // authoritative while immediately omitting members the server still
2113    // considers current.
2114    enforce_snapshot_limit(context, membership.len())?;
2115    send_snapshot_batches(
2116        context,
2117        subscription,
2118        &membership,
2119        view_spec,
2120        SnapshotPurpose::Recovery,
2121        &context.entity_cache.snapshot_config(),
2122    )
2123    .await?;
2124    Ok((receiver, membership))
2125}
2126
2127async fn apply_collection_source_event(
2128    context: &SubscriptionContext,
2129    source_view_id: &str,
2130    view_spec: &ViewSpec,
2131    query: &SubscriptionQuery,
2132    envelope: &BusMessage,
2133) {
2134    let metadata = source_frame_metadata(&envelope.payload);
2135    if metadata.op != "delete" {
2136        return;
2137    }
2138    // The projector changed the shared cache before publishing this frame.
2139    // A slow subscription must inspect that authoritative lifetime state, not
2140    // replay the delete against `_seq` and potentially erase a recreation.
2141    if !context
2142        .entity_cache
2143        .deletion_is_current(source_view_id, &envelope.key)
2144        .await
2145    {
2146        return;
2147    }
2148    if view_spec.is_derived() {
2149        let caches = context.view_index.sorted_caches();
2150        let mut guard = caches.write().await;
2151        if let Some(cache) = guard.get_mut(&query.view) {
2152            cache.remove_key(&envelope.key);
2153        }
2154    }
2155}
2156
2157/// A subscription refused for a reason the client needs spelled out.
2158///
2159/// Attach paths that return this get the registration released by the
2160/// connection loop, the same as any other failure, while the client still
2161/// receives the specific error rather than a generic `subscription-rejected`.
2162/// Sending the frame and returning `Ok` instead would leave a registered
2163/// subscription with nothing attached: it would hold a slot against the
2164/// client's limit and make the advertised "resubscribe" remediation fail with
2165/// `duplicate-subscription-id`.
2166#[derive(Debug)]
2167pub(crate) struct RejectedSubscription(pub SocketIssueMessage);
2168
2169impl RejectedSubscription {
2170    fn into_error(issue: SocketIssueMessage) -> anyhow::Error {
2171        anyhow::Error::new(Self(issue))
2172    }
2173}
2174
2175impl std::fmt::Display for RejectedSubscription {
2176    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2177        formatter.write_str(&self.0.message)
2178    }
2179}
2180
2181impl std::error::Error for RejectedSubscription {}
2182
2183/// Deliver an append view as an event tape: replay the retained records after
2184/// the cursor, then forward live frames.
2185///
2186/// This deliberately does not recompute membership from the entity cache the
2187/// way [`attach_collection_subscription`] does. The cache folds each patch
2188/// into the resident entity, so a membership diff cannot express "these three
2189/// events happened"; the retained records can.
2190///
2191/// For the same reason records are never promoted to a full `upsert` the
2192/// first time this subscription sees their key, the way a window entry is.
2193/// The cache only holds the latest state, not the state as of a record's
2194/// offset, so substituting it would hand a consumer state from after its
2195/// cursor. Records go out exactly as retained, `patch` included: a tape is
2196/// events, and the consumer that resumes from a cursor holds whatever came
2197/// before it. Clients therefore apply frames that carry an `offset` as events
2198/// and do not discard them as partial entities (see "Partial entities" in
2199/// `docs/websocket-v2-protocol.md`).
2200async fn attach_journal_subscription(
2201    context: &SubscriptionContext,
2202    subscription: Subscription,
2203    view_spec: ViewSpec,
2204    journal: Arc<crate::journal::EventJournal>,
2205    cancel_token: CancellationToken,
2206) -> Result<()> {
2207    let view_id = view_spec.id.clone();
2208    let subscription_id = subscription.subscription_id.clone();
2209
2210    // A tape has no membership window, so `take`/`skip` cannot mean what they
2211    // mean on a list view. Refuse them rather than accept and ignore them.
2212    if subscription.query.take.is_some()
2213        || subscription.query.skip.is_some()
2214        || subscription.query.snapshot_limit.is_some()
2215    {
2216        return Err(RejectedSubscription::into_error(SocketIssueMessage::protocol(
2217            Some(subscription_id),
2218            "invalid-subscription",
2219            format!(
2220                "take, skip and snapshotLimit are window options and do not apply to the replayable view {view_id}"
2221            ),
2222        )));
2223    }
2224
2225    // Reject a malformed cursor rather than silently replaying from the start,
2226    // which would look like success and duplicate everything.
2227    let cursor = match subscription.query.after.as_deref() {
2228        Some(raw) => match crate::journal::Cursor::parse(raw) {
2229            Some(cursor) => Some(cursor),
2230            None => {
2231                return Err(RejectedSubscription::into_error(
2232                    SocketIssueMessage::protocol(
2233                        Some(subscription_id),
2234                        "invalid-cursor",
2235                        format!(
2236                            "`after` must be an {{epoch}}:{{offset}} replay cursor for view {view_id}, got {raw:?}"
2237                        ),
2238                    ),
2239                ));
2240            }
2241        },
2242        None => None,
2243    };
2244
2245    // Subscribe before reading the journal so anything published during the
2246    // replay is still delivered; the offset filter below drops the overlap.
2247    let mut receiver = context.bus_manager.get_or_create_list_bus(&view_id).await;
2248
2249    let replayed = match journal.replay_after(&view_id, cursor.as_ref()).await {
2250        Ok(records) => records,
2251        Err(error) => {
2252            return Err(RejectedSubscription::into_error(
2253                SocketIssueMessage::replay_refused(Some(subscription_id), &error),
2254            ));
2255        }
2256    };
2257
2258    let frame = SubscribedFrame::new(
2259        subscription.subscription_id.clone(),
2260        subscription.query.clone(),
2261        view_spec.mode,
2262        extract_sort_config(&view_spec),
2263    )
2264    .with_replay_window(journal.window(&view_id).await);
2265    send_control_frame(context, &frame, &view_id)?;
2266
2267    // Everything past the acknowledgement runs on its own task. The replay can
2268    // be long and applies real backpressure, and `attach_client_to_bus` is
2269    // awaited directly on the connection's inbound loop — doing it there would
2270    // block unsubscribe, auth refresh and pong for the whole replay.
2271    let task_context = context.clone();
2272    let task_subscription_id = subscription.subscription_id.clone();
2273    let task_query = subscription.query.clone();
2274    let task_epoch = journal.epoch().await;
2275    let span_view = view_id.clone();
2276    tokio::spawn(
2277        async move {
2278            let mut last_sent = cursor.map(|cursor| cursor.offset);
2279            // Frames that published while the replay was still running. The
2280            // bus is a bounded broadcast, so it has to be drained as we go or
2281            // a busy view laps us before the replay finishes.
2282            let mut pending: VecDeque<Arc<BusMessage>> = VecDeque::new();
2283            let mut lagged: Option<u64> = None;
2284
2285            for record in replayed {
2286                if cancel_token.is_cancelled() {
2287                    return;
2288                }
2289                drain_available(&mut receiver, &mut pending, &mut lagged);
2290                if journal_record_matches(&task_query, &record)
2291                    && send_scoped_source_payload_async(
2292                        &task_context,
2293                        &task_subscription_id,
2294                        &span_view,
2295                        record.payload,
2296                    )
2297                    .await
2298                    .is_err()
2299                {
2300                    return;
2301                }
2302                // Advance past filtered records too: they were considered.
2303                last_sent = Some(record.offset);
2304            }
2305
2306            // Flush what arrived during the replay before going live, so the
2307            // handover keeps offset order.
2308            while let Some(envelope) = pending.pop_front() {
2309                if !forward_live_frame(
2310                    &task_context,
2311                    &task_subscription_id,
2312                    &span_view,
2313                    &task_query,
2314                    &envelope,
2315                    &mut last_sent,
2316                )
2317                .await
2318                {
2319                    return;
2320                }
2321            }
2322
2323            if let Some(skipped) = lagged {
2324                report_replay_gap(
2325                    &task_context,
2326                    &task_subscription_id,
2327                    &span_view,
2328                    &task_epoch,
2329                    skipped,
2330                    last_sent,
2331                );
2332                return;
2333            }
2334
2335            loop {
2336                tokio::select! {
2337                    _ = cancel_token.cancelled() => break,
2338                    received = receiver.recv() => {
2339                        let envelope = match received {
2340                            Ok(envelope) => envelope,
2341                            // A lagged tape is a gap. Report it with the last
2342                            // offset delivered *before* the gap and stop
2343                            // delivering on this subscription.
2344                            //
2345                            // Continuing would hand the consumer frames from
2346                            // after the gap, advancing its checkpoint past
2347                            // the skipped records so they could never be
2348                            // replayed. Stopping is not a silent stall: the
2349                            // consumer has an explicit error and a cursor
2350                            // that recovers exactly what it missed.
2351                            //
2352                            // The registration is deliberately left alone.
2353                            // Its lifecycle belongs to the connection loop,
2354                            // which holds the only handle to
2355                            // `active_subscriptions`; releasing half of it
2356                            // here would desynchronise unsubscribe, the
2357                            // duplicate-ID gate and close-time usage.
2358                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
2359                                report_replay_gap(
2360                                    &task_context,
2361                                    &task_subscription_id,
2362                                    &span_view,
2363                                    &task_epoch,
2364                                    skipped,
2365                                    last_sent,
2366                                );
2367                                break;
2368                            }
2369                            Err(broadcast::error::RecvError::Closed) => break,
2370                        };
2371
2372                        if !forward_live_frame(
2373                            &task_context,
2374                            &task_subscription_id,
2375                            &span_view,
2376                            &task_query,
2377                            &envelope,
2378                            &mut last_sent,
2379                        )
2380                        .await
2381                        {
2382                            break;
2383                        }
2384                    }
2385                }
2386            }
2387        }
2388        .instrument(info_span!(
2389            "ws.subscribe.replay",
2390            client_id = %context.client_id,
2391            view = %view_id
2392        )),
2393    );
2394    Ok(())
2395}
2396
2397/// Take whatever the bus already has without waiting, so a long replay cannot
2398/// be lapped by a busy view.
2399fn drain_available(
2400    receiver: &mut broadcast::Receiver<Arc<BusMessage>>,
2401    pending: &mut VecDeque<Arc<BusMessage>>,
2402    lagged: &mut Option<u64>,
2403) {
2404    // Once a gap is known, everything still on the bus is on the far side of
2405    // it. Buffering it would put post-gap frames in front of the lag report,
2406    // advancing `last_sent` past the hole and making `recoverFrom` point
2407    // after the very records it is supposed to recover.
2408    if lagged.is_some() {
2409        return;
2410    }
2411    // Bounded so a view publishing faster than the client drains cannot turn
2412    // the buffer into an unbounded queue. Filling it is not itself a gap:
2413    // nothing has been skipped at that instant, the buffer simply stopped
2414    // accepting. Stop buffering and let the bus report the loss, with the
2415    // count it actually measures, when delivery reaches it.
2416    const MAX_PENDING: usize = 8_192;
2417    loop {
2418        if pending.len() >= MAX_PENDING {
2419            return;
2420        }
2421        match receiver.try_recv() {
2422            Ok(envelope) => pending.push_back(envelope),
2423            Err(broadcast::error::TryRecvError::Empty)
2424            | Err(broadcast::error::TryRecvError::Closed) => return,
2425            Err(broadcast::error::TryRecvError::Lagged(skipped)) => {
2426                *lagged = Some(skipped);
2427                return;
2428            }
2429        }
2430    }
2431}
2432
2433/// Whether a live frame was already delivered by the replay that preceded it.
2434///
2435/// The bus is subscribed before the tape is read, so a record published in
2436/// between appears on both paths; without this the consumer sees it twice.
2437/// Advances the high-water mark as a side effect.
2438fn already_delivered(offset: Option<u64>, last_sent: &mut Option<u64>) -> bool {
2439    let Some(offset) = offset else {
2440        // A frame with no offset predates the tape, so it cannot have been
2441        // replayed and must not move the mark.
2442        return false;
2443    };
2444    if last_sent.is_some_and(|last| offset <= last) {
2445        return true;
2446    }
2447    *last_sent = Some(offset);
2448    false
2449}
2450
2451/// Deliver one live frame, skipping anything the replay already sent.
2452///
2453/// Returns false when the subscription should end.
2454async fn forward_live_frame(
2455    context: &SubscriptionContext,
2456    subscription_id: &str,
2457    view_id: &str,
2458    query: &SubscriptionQuery,
2459    envelope: &Arc<BusMessage>,
2460    last_sent: &mut Option<u64>,
2461) -> bool {
2462    let metadata = source_frame_metadata(&envelope.payload);
2463    if already_delivered(metadata.offset, last_sent) {
2464        return true;
2465    }
2466    if !live_frame_matches(query, &envelope.key, &envelope.payload) {
2467        return true;
2468    }
2469    send_scoped_source_payload(context, subscription_id, view_id, envelope.payload.clone()).is_ok()
2470}
2471
2472fn report_replay_gap(
2473    context: &SubscriptionContext,
2474    subscription_id: &str,
2475    view_id: &str,
2476    epoch: &crate::journal::JournalEpoch,
2477    skipped: u64,
2478    last_sent: Option<u64>,
2479) {
2480    warn!(
2481        "Replay subscription {} lagged past {} records; stopping with a recovery cursor",
2482        subscription_id, skipped
2483    );
2484    let recover_from = last_sent.map(|offset| crate::journal::Cursor {
2485        epoch: epoch.clone(),
2486        offset,
2487    });
2488    let _ = send_control_frame(
2489        context,
2490        &SocketIssueMessage::replay_lagged(
2491            Some(subscription_id.to_string()),
2492            skipped,
2493            recover_from,
2494        ),
2495        view_id,
2496    );
2497}
2498
2499/// Apply the subscription's `key`, `partition` and `filters` to a retained
2500/// record. A replay must honour the same predicates a live subscription does.
2501fn journal_record_matches(
2502    query: &SubscriptionQuery,
2503    record: &crate::journal::JournalRecord,
2504) -> bool {
2505    live_frame_matches(query, &record.key, &record.payload)
2506}
2507
2508fn live_frame_matches(query: &SubscriptionQuery, key: &str, payload: &[u8]) -> bool {
2509    if !query.matches_key(key) {
2510        return false;
2511    }
2512    if query.partition.is_none() && query.filters.is_empty() {
2513        return true;
2514    }
2515    let Some(data) = source_frame_data(payload) else {
2516        return false;
2517    };
2518    let data = &data;
2519    if let Some(partition) = &query.partition {
2520        if value_at_dot_path(data, "_partition").as_deref()
2521            != Some(&Value::String(partition.clone()))
2522        {
2523            return false;
2524        }
2525    }
2526    query
2527        .filters
2528        .iter()
2529        .all(|(path, expected)| value_at_dot_path(data, path).as_deref() == Some(expected))
2530}
2531
2532/// Awaiting variant of [`send_scoped_source_payload`], for replays that can
2533/// exceed the client's send queue.
2534async fn send_scoped_source_payload_async(
2535    context: &SubscriptionContext,
2536    subscription_id: &str,
2537    view_id: &str,
2538    payload: Arc<Bytes>,
2539) -> Result<()> {
2540    let json = SourceFields::scoped(&payload, subscription_id)?;
2541    let compressed = maybe_compress(&json);
2542    let bytes = compressed.as_bytes().len();
2543    context
2544        .client_manager
2545        .send_compressed_async(context.client_id, compressed)
2546        .await
2547        .map_err(|error| anyhow::anyhow!("failed to send replayed frame: {error}"))?;
2548    context.metrics.message_sent();
2549    emit_update_sent_for_client(
2550        &context.usage_emitter,
2551        &context.client_manager,
2552        context.client_id,
2553        view_id,
2554        bytes,
2555    );
2556    Ok(())
2557}
2558
2559#[derive(Default)]
2560struct SourceFrameMetadata {
2561    op: String,
2562    seq: Option<String>,
2563    offset: Option<u64>,
2564}
2565
2566/// A source frame's `op`, `seq` and `offset`, read without building its
2567/// `data`.
2568fn source_frame_metadata(payload: &[u8]) -> SourceFrameMetadata {
2569    SourceFields::parse(payload)
2570        .map(|fields| SourceFrameMetadata {
2571            op: fields
2572                .value("op")
2573                .as_ref()
2574                .and_then(Value::as_str)
2575                .unwrap_or_default()
2576                .to_string(),
2577            seq: fields
2578                .value("seq")
2579                .as_ref()
2580                .and_then(Value::as_str)
2581                .map(str::to_string),
2582            offset: fields.value("offset").as_ref().and_then(Value::as_u64),
2583        })
2584        .unwrap_or_default()
2585}
2586
2587/// The `data` of a source frame, as a client receives it.
2588fn source_frame_data(payload: &[u8]) -> Option<Value> {
2589    SourceFields::parse(payload)?.value("data")
2590}
2591
2592/// What a state subscriber last received whole, which a catch-up is
2593/// measured against. Both kinds are shared, with the entity cache or the bus,
2594/// so a subscriber keeps no copy of its own. A subscriber does keep the
2595/// version it was sent alive while the cache moves on; subscribers sent the
2596/// same version share it.
2597enum Synced {
2598    /// The cached entity, as sent in a snapshot row, an upsert or a catch-up.
2599    Entity(SharedEntity),
2600    /// A forwarded `upsert` source frame, whose `data` the client holds.
2601    Frame(Arc<Bytes>),
2602}
2603
2604impl Synced {
2605    /// What the client holds, as it received it: in wire format.
2606    fn wire(&self, wire_format: &WireFormat) -> Option<Value> {
2607        match self {
2608            Synced::Entity(entity) => Some(wire_value(entity, wire_format)),
2609            Synced::Frame(frame) => source_frame_data(frame),
2610        }
2611    }
2612}
2613
2614/// Tracks what a client holds after a frame forwarded to it as published. An
2615/// upsert hands it the entity whole, which a catch-up is then measured
2616/// against; a patch sets fields (see [`mark_fields`]).
2617fn record_forwarded(
2618    payload: &Arc<Bytes>,
2619    op: &str,
2620    synced: &mut Option<Synced>,
2621    touched: &mut Value,
2622) {
2623    if op == "upsert" {
2624        *synced = Some(Synced::Frame(payload.clone()));
2625        *touched = Value::Null;
2626    } else if let Some(data) = source_frame_data(payload) {
2627        mark_fields(touched, &data);
2628    }
2629}
2630
2631/// Marks in `touched` the fields `patch` sets: objects field by field, any
2632/// other value whole (`true`). `null` marks nothing.
2633///
2634/// A client holds whatever a forwarded patch set, which the entity may have
2635/// set back before the next catch-up (or the client may have dropped the
2636/// patch), so a catch-up sends these fields even when they match the entity
2637/// the client last received whole.
2638fn mark_fields(touched: &mut Value, patch: &Value) {
2639    let Value::Object(fields) = patch else {
2640        *touched = Value::Bool(true);
2641        return;
2642    };
2643    if touched.is_null() {
2644        *touched = Value::Object(serde_json::Map::new());
2645    }
2646    let Value::Object(marked) = touched else {
2647        // Already marked whole.
2648        return;
2649    };
2650    for (field, value) in fields {
2651        mark_fields(marked.entry(field.clone()).or_insert(Value::Null), value);
2652    }
2653}
2654
2655/// The fields of `current` that differ from `base` or that `touched` marks
2656/// (see [`mark_fields`]), as a patch that turns what the client holds into
2657/// `current` when merged: objects compare field by field, any other value is
2658/// sent whole, and a field `current` no longer has is sent as `null`. `None`
2659/// when nothing needs sending.
2660fn changed_fields(base: &Value, current: &Value, touched: &Value) -> Option<Value> {
2661    match (base, current, touched) {
2662        (_, _, Value::Bool(true)) => Some(current.clone()),
2663        (Value::Object(base), Value::Object(current), _) => {
2664            let marked = touched.as_object();
2665            let mut patch = serde_json::Map::new();
2666            for (field, value) in current {
2667                let touched = marked
2668                    .and_then(|marked| marked.get(field))
2669                    .unwrap_or(&Value::Null);
2670                match base.get(field) {
2671                    Some(old) => {
2672                        if let Some(changed) = changed_fields(old, value, touched) {
2673                            patch.insert(field.clone(), changed);
2674                        }
2675                    }
2676                    None => {
2677                        patch.insert(field.clone(), value.clone());
2678                    }
2679                }
2680            }
2681            for field in base
2682                .keys()
2683                .chain(marked.into_iter().flat_map(|marked| marked.keys()))
2684            {
2685                if !current.contains_key(field) {
2686                    patch.insert(field.clone(), Value::Null);
2687                }
2688            }
2689            (!patch.is_empty()).then_some(Value::Object(patch))
2690        }
2691        _ => (base != current || !touched.is_null()).then(|| current.clone()),
2692    }
2693}
2694
2695fn send_scoped_source_payload(
2696    context: &SubscriptionContext,
2697    subscription_id: &str,
2698    view_id: &str,
2699    payload: Arc<Bytes>,
2700) -> Result<()> {
2701    let encoded = Arc::new(Bytes::from(SourceFields::scoped(
2702        &payload,
2703        subscription_id,
2704    )?));
2705    let bytes = encoded.len();
2706    context
2707        .client_manager
2708        .send_to_client(context.client_id, encoded)
2709        .map_err(|error| anyhow::anyhow!("failed to send live frame: {error}"))?;
2710    context.metrics.message_sent();
2711    emit_update_sent_for_client(
2712        &context.usage_emitter,
2713        &context.client_manager,
2714        context.client_id,
2715        view_id,
2716        bytes,
2717    );
2718    Ok(())
2719}
2720
2721fn send_membership_frame(
2722    context: &SubscriptionContext,
2723    subscription_id: &str,
2724    view_spec: &ViewSpec,
2725    op: &str,
2726    key: &str,
2727    mut data: Value,
2728    seq: Option<String>,
2729) -> Result<()> {
2730    apply_wire_format(&mut data, &view_spec.wire_format);
2731    let frame = Frame::scoped(
2732        subscription_id,
2733        view_spec.mode,
2734        &view_spec.id,
2735        op,
2736        key,
2737        data,
2738        seq,
2739    );
2740    send_membership_json(context, view_spec, serde_json::to_vec(&frame)?)
2741}
2742
2743/// A scoped frame whose data is a cached entity, serialized straight from
2744/// it: the same JSON as [`Frame::scoped`] with a formatted copy of the entity.
2745#[derive(Serialize)]
2746#[serde(rename_all = "camelCase")]
2747struct ScopedEntityFrame<'a> {
2748    protocol_version: u8,
2749    subscription_id: &'a str,
2750    mode: Mode,
2751    #[serde(rename = "entity")]
2752    export: &'a str,
2753    op: &'a str,
2754    key: &'a str,
2755    data: WireEntity<'a>,
2756    #[serde(skip_serializing_if = "Option::is_none")]
2757    seq: Option<&'a str>,
2758}
2759
2760/// [`send_membership_frame`] for a cached entity, without copying it.
2761fn send_entity_frame(
2762    context: &SubscriptionContext,
2763    subscription_id: &str,
2764    view_spec: &ViewSpec,
2765    op: &str,
2766    key: &str,
2767    entity: &SharedEntity,
2768    seq: Option<&str>,
2769) -> Result<()> {
2770    let frame = ScopedEntityFrame {
2771        protocol_version: PROTOCOL_VERSION,
2772        subscription_id,
2773        mode: view_spec.mode,
2774        export: &view_spec.id,
2775        op,
2776        key,
2777        data: WireEntity {
2778            entity,
2779            wire_format: &view_spec.wire_format,
2780        },
2781        seq,
2782    };
2783    send_membership_json(context, view_spec, serde_json::to_vec(&frame)?)
2784}
2785
2786fn send_membership_json(
2787    context: &SubscriptionContext,
2788    view_spec: &ViewSpec,
2789    json: Vec<u8>,
2790) -> Result<()> {
2791    let encoded = Arc::new(Bytes::from(json));
2792    let bytes = encoded.len();
2793    context
2794        .client_manager
2795        .send_to_client(context.client_id, encoded)
2796        .map_err(|error| anyhow::anyhow!("failed to send membership frame: {error}"))?;
2797    context.metrics.message_sent();
2798    emit_update_sent_for_client(
2799        &context.usage_emitter,
2800        &context.client_manager,
2801        context.client_id,
2802        &view_spec.id,
2803        bytes,
2804    );
2805    Ok(())
2806}
2807
2808/// Send what one source mutation owes a subscriber, and keep `undelivered`
2809/// (window members the client holds no copy of) in step with what was sent.
2810#[allow(clippy::too_many_arguments)]
2811fn emit_collection_delta(
2812    context: &SubscriptionContext,
2813    subscription_id: &str,
2814    view_spec: &ViewSpec,
2815    current: &[(String, SharedEntity)],
2816    next: &[(String, SharedEntity)],
2817    envelope: &BusMessage,
2818    metadata: &SourceFrameMetadata,
2819    undelivered: &mut HashSet<String>,
2820) -> Result<()> {
2821    let current_keys: HashSet<&str> = current.iter().map(|(key, _)| key.as_str()).collect();
2822    let next_set: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2823
2824    // In window order, as the frames go out.
2825    for key in current
2826        .iter()
2827        .map(|(key, _)| key.as_str())
2828        .filter(|key| !next_set.contains(key))
2829    {
2830        // A key the client was never sent has nothing to remove.
2831        if undelivered.remove(key) {
2832            continue;
2833        }
2834        let op = if metadata.op == "delete" && key == envelope.key {
2835            "delete"
2836        } else {
2837            "remove"
2838        };
2839        send_membership_frame(
2840            context,
2841            subscription_id,
2842            view_spec,
2843            op,
2844            key,
2845            Value::Null,
2846            metadata.seq.clone(),
2847        )?;
2848    }
2849
2850    for (key, data) in next.iter() {
2851        let in_window = current_keys.contains(key.as_str());
2852        match member_action(
2853            in_window,
2854            in_window && !undelivered.contains(key),
2855            key == &envelope.key,
2856            view_spec.is_derived(),
2857            &metadata.op,
2858        ) {
2859            MemberAction::Skip => {}
2860            MemberAction::ForwardPatch => send_scoped_source_payload(
2861                context,
2862                subscription_id,
2863                &view_spec.id,
2864                envelope.payload.clone(),
2865            )?,
2866            MemberAction::Upsert => {
2867                let seq = metadata
2868                    .seq
2869                    .as_deref()
2870                    .or_else(|| data.field("_seq").and_then(Value::as_str));
2871                send_entity_frame(
2872                    context,
2873                    subscription_id,
2874                    view_spec,
2875                    "upsert",
2876                    key,
2877                    data,
2878                    seq,
2879                )?;
2880                undelivered.remove(key);
2881            }
2882        }
2883    }
2884    Ok(())
2885}
2886
2887/// Emit the net effect of every source mutation seen during one coalescing
2888/// interval. Full entities are used for changed members because intermediate
2889/// sparse patches were intentionally discarded and can no longer be merged
2890/// safely by a client.
2891fn emit_coalesced_collection_delta(
2892    context: &SubscriptionContext,
2893    subscription_id: &str,
2894    view_spec: &ViewSpec,
2895    current: &[(String, SharedEntity)],
2896    next: &[(String, SharedEntity)],
2897    pending: &HashMap<String, Arc<BusMessage>>,
2898    undelivered: &mut HashSet<String>,
2899) -> Result<()> {
2900    let changes = plan_coalesced_collection_delta(current, next, pending, undelivered);
2901    // Every key that left the window or was just upserted is settled: the
2902    // client either never needs it or now holds all of it.
2903    let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2904    undelivered.retain(|key| next_keys.contains(key.as_str()));
2905    for change in changes {
2906        if change.op == "upsert" {
2907            undelivered.remove(&change.key);
2908        }
2909        match &change.data {
2910            Some(entity) => send_entity_frame(
2911                context,
2912                subscription_id,
2913                view_spec,
2914                change.op,
2915                &change.key,
2916                entity,
2917                change.seq.as_deref(),
2918            )?,
2919            None => send_membership_frame(
2920                context,
2921                subscription_id,
2922                view_spec,
2923                change.op,
2924                &change.key,
2925                Value::Null,
2926                change.seq,
2927            )?,
2928        }
2929    }
2930    Ok(())
2931}
2932
2933#[derive(Debug, PartialEq)]
2934struct CollectionChange {
2935    op: &'static str,
2936    key: String,
2937    /// The entity an `upsert` sends; `None` for a `remove` or `delete`.
2938    data: Option<SharedEntity>,
2939    seq: Option<String>,
2940}
2941
2942/// Plan one coalescing interval. `undelivered` holds window members the
2943/// client was never sent; they get no `remove`/`delete` when they leave, and
2944/// like every changed key they are sent whole when they change.
2945fn plan_coalesced_collection_delta(
2946    current: &[(String, SharedEntity)],
2947    next: &[(String, SharedEntity)],
2948    pending: &HashMap<String, Arc<BusMessage>>,
2949    undelivered: &HashSet<String>,
2950) -> Vec<CollectionChange> {
2951    let current_by_key: HashMap<&str, &SharedEntity> = current
2952        .iter()
2953        .map(|(key, data)| (key.as_str(), data))
2954        .collect();
2955    let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2956    let latest_seq = pending
2957        .values()
2958        .filter_map(|envelope| source_frame_metadata(&envelope.payload).seq)
2959        .max_by(|left, right| cmp_seq(left, right));
2960    let mut changes = Vec::new();
2961
2962    for (key, _) in current
2963        .iter()
2964        .filter(|(key, _)| !next_keys.contains(key.as_str()) && !undelivered.contains(key.as_str()))
2965    {
2966        let metadata = pending
2967            .get(key)
2968            .map(|envelope| source_frame_metadata(&envelope.payload));
2969        let op = if metadata
2970            .as_ref()
2971            .is_some_and(|metadata| metadata.op == "delete")
2972        {
2973            "delete"
2974        } else {
2975            "remove"
2976        };
2977        let seq = metadata
2978            .and_then(|metadata| metadata.seq)
2979            .or_else(|| latest_seq.clone());
2980        changes.push(CollectionChange {
2981            op,
2982            key: key.clone(),
2983            data: None,
2984            seq,
2985        });
2986    }
2987
2988    for (key, data) in next {
2989        let changed = current_by_key
2990            .get(key.as_str())
2991            .is_none_or(|previous| *previous != data);
2992        if !changed && !pending.contains_key(key) {
2993            continue;
2994        }
2995        let seq = data
2996            .field("_seq")
2997            .and_then(Value::as_str)
2998            .map(str::to_string)
2999            .or_else(|| {
3000                pending
3001                    .get(key)
3002                    .and_then(|envelope| source_frame_metadata(&envelope.payload).seq)
3003            })
3004            .or_else(|| latest_seq.clone());
3005        changes.push(CollectionChange {
3006            op: "upsert",
3007            key: key.clone(),
3008            data: Some(data.clone()),
3009            seq,
3010        });
3011    }
3012    changes
3013}
3014
3015/// What one in-window key owes a subscriber after a source mutation.
3016#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3017enum MemberAction {
3018    /// Send nothing: the subscriber already holds this entity and its data did
3019    /// not change.
3020    Skip,
3021    /// Forward the source patch verbatim.
3022    ForwardPatch,
3023    /// Send the whole entity.
3024    Upsert,
3025}
3026
3027/// Decide what to send for one key in the query window.
3028///
3029/// `in_window` says the key was in the window before this mutation; `held`
3030/// says the client also has a copy of it. The two differ after a snapshot
3031/// that `snapshotLimit` truncated or that was disabled: those keys count for
3032/// the window's positions, but the client was never sent them.
3033///
3034/// Only the mutated key's data changed, so a key that was already in the
3035/// window has at most moved index — and index is not the server's to
3036/// communicate. `subscribed` announces the window's sort (see
3037/// [`extract_sort_config`], which is `_seq` descending for a plain list) and
3038/// every SDK re-sorts locally from it, so resending an unchanged entity to
3039/// convey its new position is pure waste. This matters because on a
3040/// `_seq`-ordered list the mutated entity jumps to the front on *every*
3041/// mutation: keying the decision on position change meant rebroadcasting the
3042/// whole window each time, and the mutated entity itself never kept its
3043/// position long enough to have its patch forwarded.
3044fn member_action(
3045    in_window: bool,
3046    held: bool,
3047    is_mutated_key: bool,
3048    is_derived: bool,
3049    op: &str,
3050) -> MemberAction {
3051    if !is_mutated_key {
3052        // A key entering the window has no local state to merge into, so it
3053        // needs the whole entity. One already in the window is unchanged:
3054        // if the client holds it there is nothing new, and if it doesn't the
3055        // snapshot limit deliberately left it out, so it waits for its own
3056        // next change.
3057        return if in_window {
3058            MemberAction::Skip
3059        } else {
3060            MemberAction::Upsert
3061        };
3062    }
3063    // The mutated entity rides its own patch through untouched, but only when
3064    // the subscriber already holds a copy; a patch for a key the client never
3065    // received is not an entity. Derived views still send whole entities: the
3066    // patch on the bus is scoped to the source view, not this one (see
3067    // A4-150).
3068    if held && !is_derived && op != "delete" {
3069        MemberAction::ForwardPatch
3070    } else {
3071        MemberAction::Upsert
3072    }
3073}
3074
3075/// A copy of `entity` with the view's wire format applied, for code that
3076/// needs it as a value (a catch-up diff) rather than serialized.
3077fn wire_value(entity: &SharedEntity, wire_format: &WireFormat) -> Value {
3078    let mut value = entity.to_value();
3079    apply_wire_format(&mut value, wire_format);
3080    value
3081}
3082
3083async fn load_query_entities(
3084    entity_cache: &EntityCache,
3085    sorted_caches: Option<
3086        Arc<tokio::sync::RwLock<HashMap<String, crate::sorted_cache::SortedViewCache>>>,
3087    >,
3088    view_spec: &ViewSpec,
3089    query: &SubscriptionQuery,
3090    apply_snapshot_limit: bool,
3091) -> Vec<(String, SharedEntity)> {
3092    // Rows share their fields with the caches; nothing here copies an entity.
3093    let mut query = query.clone();
3094    if let Some(limit) = view_spec
3095        .pipeline
3096        .as_ref()
3097        .and_then(|pipeline| pipeline.limit)
3098    {
3099        query.take = Some(query.take.unwrap_or(limit).min(limit));
3100    }
3101    if let Some(sorted_caches) = sorted_caches {
3102        // A sorted cache is read in order, and only as far as the window
3103        // reaches: rows past it are neither visited nor shared.
3104        let caches = sorted_caches.read().await;
3105        if let Some(cache) = caches.get(&view_spec.id) {
3106            return select_window(
3107                cache
3108                    .iter_ordered()
3109                    .filter(|(key, data)| query_matches_entity(&query, key, *data))
3110                    .map(|(key, data)| (key.to_string(), data.clone())),
3111                &query,
3112                apply_snapshot_limit,
3113            );
3114        }
3115    }
3116    let (entities, preordered) = if view_spec.is_derived() {
3117        // Empty and filter-only pipelines have no sorted cache. Evaluate the
3118        // source rows, retaining the pipeline predicate before query selection.
3119        let mut entities = entity_cache
3120            .get_all_shared(view_spec.source_view.as_deref().unwrap_or(&view_spec.id))
3121            .await;
3122        if let Some(filter) = view_spec
3123            .pipeline
3124            .as_ref()
3125            .and_then(|pipeline| pipeline.filter.as_ref())
3126        {
3127            entities.retain(|(_, data)| filter.matches(data));
3128        }
3129        (entities, false)
3130    } else if view_spec.mode == Mode::State {
3131        let entity = match query.key.as_deref() {
3132            Some(key) => entity_cache
3133                .get_shared(&view_spec.id, key)
3134                .await
3135                .map(|data| vec![(key.to_string(), data)])
3136                .unwrap_or_default(),
3137            None => vec![],
3138        };
3139        (entity, true)
3140    } else {
3141        (entity_cache.get_all_shared(&view_spec.id).await, false)
3142    };
3143    select_query_entities(entities, &query, preordered, apply_snapshot_limit)
3144}
3145
3146fn select_query_entities<E: EntityFields>(
3147    mut entities: Vec<(String, E)>,
3148    query: &SubscriptionQuery,
3149    preordered: bool,
3150    apply_snapshot_limit: bool,
3151) -> Vec<(String, E)> {
3152    entities.retain(|(key, data)| query_matches_entity(query, key, data));
3153    if !preordered {
3154        entities.sort_by(|left, right| {
3155            let left_seq = left.1.field("_seq").and_then(Value::as_str).unwrap_or("");
3156            let right_seq = right.1.field("_seq").and_then(Value::as_str).unwrap_or("");
3157            let order = if query.after.is_some() {
3158                cmp_seq(left_seq, right_seq)
3159            } else {
3160                cmp_seq(right_seq, left_seq)
3161            };
3162            order.then_with(|| left.0.cmp(&right.0))
3163        });
3164    }
3165
3166    select_window(entities.into_iter(), query, apply_snapshot_limit)
3167}
3168
3169/// The `[skip, skip + take)` window of rows already filtered and in order,
3170/// cut to the snapshot limit when it applies. Takes no row past the window.
3171fn select_window<E>(
3172    rows: impl Iterator<Item = (String, E)>,
3173    query: &SubscriptionQuery,
3174    apply_snapshot_limit: bool,
3175) -> Vec<(String, E)> {
3176    let skip = query.skip.unwrap_or(0);
3177    let mut take = query.take.unwrap_or(usize::MAX);
3178    if apply_snapshot_limit {
3179        if let Some(limit) = query.snapshot_limit {
3180            take = take.min(limit);
3181        }
3182    }
3183    rows.skip(skip).take(take).collect()
3184}
3185
3186fn query_matches_entity<E: EntityFields + ?Sized>(
3187    query: &SubscriptionQuery,
3188    key: &str,
3189    data: &E,
3190) -> bool {
3191    if !query.matches_key(key) {
3192        return false;
3193    }
3194    if let Some(partition) = &query.partition {
3195        if value_at_dot_path(data, "_partition").as_deref()
3196            != Some(&Value::String(partition.clone()))
3197        {
3198            return false;
3199        }
3200    }
3201    if let Some(after) = &query.after {
3202        let Some(seq) = data.field("_seq").and_then(Value::as_str) else {
3203            return false;
3204        };
3205        if cmp_seq(seq, after) != std::cmp::Ordering::Greater {
3206            return false;
3207        }
3208    }
3209    query
3210        .filters
3211        .iter()
3212        .all(|(path, expected)| value_at_dot_path(data, path).as_deref() == Some(expected))
3213}
3214
3215fn value_at_dot_path<'a, E: EntityFields + ?Sized>(
3216    entity: &'a E,
3217    path: &str,
3218) -> Option<std::borrow::Cow<'a, Value>> {
3219    lookup(entity, path.split('.'))
3220}
3221
3222#[cfg(test)]
3223mod tests {
3224    use super::*;
3225    use crate::cache::{EntityCacheConfig, PatchOrigin};
3226    use crate::view::{Delivery, Filters, Projection};
3227    use serde_json::json;
3228    use tokio::sync::oneshot;
3229
3230    #[derive(Default)]
3231    struct DelayedUsageEmitter {
3232        started: tokio::sync::Notify,
3233        release: tokio::sync::Notify,
3234        events: tokio::sync::Mutex<Vec<WebSocketUsageEvent>>,
3235    }
3236
3237    #[async_trait::async_trait]
3238    impl WebSocketUsageEmitter for DelayedUsageEmitter {
3239        async fn emit(&self, event: WebSocketUsageEvent) {
3240            self.started.notify_one();
3241            self.release.notified().await;
3242            self.events.lock().await.push(event);
3243        }
3244    }
3245
3246    #[tokio::test]
3247    async fn usage_emitter_handle_waits_for_spawned_emissions() {
3248        let emitter = Arc::new(DelayedUsageEmitter::default());
3249        let handle = UsageEmitterHandle::new(emitter.clone());
3250        handle.emit(WebSocketUsageEvent::ConnectionEstablished {
3251            client_id: "client-1".to_string(),
3252            remote_addr: "127.0.0.1:1234".to_string(),
3253            deployment_id: Some("1".to_string()),
3254            identity: UsageIdentity::default(),
3255        });
3256        emitter.started.notified().await;
3257
3258        let closing_handle = handle.clone();
3259        let closing = tokio::spawn(async move {
3260            closing_handle.close_and_wait().await;
3261        });
3262        tokio::task::yield_now().await;
3263        assert!(!closing.is_finished());
3264
3265        emitter.release.notify_one();
3266        closing.await.expect("usage task wait should finish");
3267        assert_eq!(emitter.events.lock().await.len(), 1);
3268    }
3269
3270    #[test]
3271    fn a_catch_up_carries_only_the_fields_that_changed() {
3272        let base = json!({
3273            "id": 7,
3274            "state": {"deployed": [1, 2, 3], "total": 6},
3275            "results": {"square": null},
3276            "gone": true,
3277        });
3278        let current = json!({
3279            "id": 7,
3280            "state": {"deployed": [1, 2, 4], "total": 6},
3281            "results": {"square": 3},
3282            "fresh": "x",
3283        });
3284        assert_eq!(
3285            changed_fields(&base, &current, &Value::Null),
3286            Some(json!({
3287                // An array that changed is sent whole.
3288                "state": {"deployed": [1, 2, 4]},
3289                "results": {"square": 3},
3290                "fresh": "x",
3291                // A field the entity no longer has is cleared.
3292                "gone": null,
3293            }))
3294        );
3295        assert_eq!(changed_fields(&current, &current, &Value::Null), None);
3296        // Anything that is not an object on both sides is compared whole.
3297        assert_eq!(
3298            changed_fields(&json!(null), &current, &Value::Null),
3299            Some(current.clone())
3300        );
3301    }
3302
3303    #[test]
3304    fn a_catch_up_resends_the_fields_forwarded_patches_set() {
3305        let base = json!({"a": 1, "state": {"x": 1, "y": 1}, "list": [1]});
3306        let mut touched = Value::Null;
3307        mark_fields(&mut touched, &json!({"a": 2, "state": {"x": 2}}));
3308        mark_fields(&mut touched, &json!({"list": [2], "extra": true}));
3309        assert_eq!(
3310            touched,
3311            json!({"a": true, "state": {"x": true}, "list": true, "extra": true})
3312        );
3313        // The entity set `a` and `state.x` back and never kept `extra`: the
3314        // client still holds what the forwarded patches set, so a catch-up
3315        // compared with `base` alone would leave it stale.
3316        let current = json!({"a": 1, "state": {"x": 1, "y": 1}, "list": [1, 2]});
3317        assert_eq!(
3318            changed_fields(&base, &current, &Value::Null),
3319            Some(json!({"list": [1, 2]}))
3320        );
3321        assert_eq!(
3322            changed_fields(&base, &current, &touched),
3323            Some(json!({"a": 1, "state": {"x": 1}, "list": [1, 2], "extra": null}))
3324        );
3325        // A patch whose data is not an object marks everything.
3326        mark_fields(&mut touched, &json!(null));
3327        assert_eq!(
3328            changed_fields(&base, &current, &touched),
3329            Some(current.clone())
3330        );
3331    }
3332
3333    fn list_spec() -> ViewSpec {
3334        ViewSpec {
3335            id: "Thing/list".to_string(),
3336            export: "Thing".to_string(),
3337            mode: Mode::List,
3338            wire_format: Default::default(),
3339            projection: Projection::all(),
3340            filters: Filters::all(),
3341            delivery: Delivery::default(),
3342            pipeline: None,
3343            source_view: None,
3344        }
3345    }
3346
3347    /// Plan the whole window the way `emit_collection_delta` does, so a test
3348    /// can assert on what a subscriber is actually sent.
3349    fn plan_window(
3350        current_keys: &[&str],
3351        next_keys: &[&str],
3352        envelope_key: &str,
3353        is_derived: bool,
3354        op: &str,
3355    ) -> Vec<(String, MemberAction)> {
3356        plan_window_with_undelivered(current_keys, &[], next_keys, envelope_key, is_derived, op)
3357    }
3358
3359    /// [`plan_window`] for a subscriber that was never sent `undelivered`.
3360    fn plan_window_with_undelivered(
3361        current_keys: &[&str],
3362        undelivered: &[&str],
3363        next_keys: &[&str],
3364        envelope_key: &str,
3365        is_derived: bool,
3366        op: &str,
3367    ) -> Vec<(String, MemberAction)> {
3368        next_keys
3369            .iter()
3370            .map(|key| {
3371                let in_window = current_keys.contains(key);
3372                let held = in_window && !undelivered.contains(key);
3373                (
3374                    (*key).to_string(),
3375                    member_action(in_window, held, *key == envelope_key, is_derived, op),
3376                )
3377            })
3378            .collect()
3379    }
3380
3381    #[test]
3382    fn a_window_member_the_client_never_received_is_sent_whole() {
3383        // snapshotLimit cut the window [4, 3, 2, 1] to [4, 3]. "1" still
3384        // counts for positions, but a patch for it would reach a client with
3385        // nothing to merge into.
3386        let plan = plan_window_with_undelivered(
3387            &["4", "3", "2", "1"],
3388            &["2", "1"],
3389            &["1", "4", "3", "2"],
3390            "1",
3391            false,
3392            "patch",
3393        );
3394        assert_eq!(
3395            plan,
3396            vec![
3397                ("1".to_string(), MemberAction::Upsert),
3398                ("4".to_string(), MemberAction::Skip),
3399                ("3".to_string(), MemberAction::Skip),
3400                // Unchanged and never sent: it waits for its own change
3401                // rather than undoing the snapshot limit.
3402                ("2".to_string(), MemberAction::Skip),
3403            ]
3404        );
3405    }
3406
3407    #[test]
3408    fn a_held_member_still_rides_its_patch_after_a_truncated_snapshot() {
3409        let plan = plan_window_with_undelivered(
3410            &["4", "3", "2", "1"],
3411            &["2", "1"],
3412            &["4", "3", "2", "1"],
3413            "4",
3414            false,
3415            "patch",
3416        );
3417        assert_eq!(plan[0], ("4".to_string(), MemberAction::ForwardPatch));
3418    }
3419
3420    #[test]
3421    fn coalescing_sends_no_remove_for_a_key_the_client_never_received() {
3422        let current = vec![
3423            ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3424            (
3425                "unsent".to_string(),
3426                json!({"count": 1, "_seq": "10:000001"}),
3427            ),
3428            ("gone".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3429        ];
3430        let next = vec![
3431            ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3432            (
3433                "unsent".to_string(),
3434                json!({"count": 2, "_seq": "10:000003"}),
3435            ),
3436        ];
3437        let pending = HashMap::from([
3438            (
3439                "unsent".to_string(),
3440                list_message("unsent", "patch", "10:000003"),
3441            ),
3442            (
3443                "gone".to_string(),
3444                list_message("gone", "delete", "10:000002"),
3445            ),
3446        ]);
3447        let undelivered = HashSet::from(["unsent".to_string(), "gone".to_string()]);
3448
3449        assert_eq!(
3450            plan_coalesced_collection_delta(
3451                &shared(&current),
3452                &shared(&next),
3453                &pending,
3454                &undelivered
3455            ),
3456            vec![CollectionChange {
3457                op: "upsert",
3458                key: "unsent".to_string(),
3459                data: Some(SharedEntity::new(json!({"count": 2, "_seq": "10:000003"}))),
3460                seq: Some("10:000003".to_string()),
3461            }]
3462        );
3463    }
3464
3465    #[test]
3466    fn a_reordered_window_sends_only_the_mutated_entity() {
3467        // A `_seq`-descending list: mutating "1" moves it to the front and
3468        // shifts every other key down one. Only "1" changed, so only "1" is
3469        // sent — and it rides its own patch, not a full entity.
3470        let current = ["4", "3", "2", "1"];
3471        let next = ["1", "4", "3", "2"];
3472        let plan = plan_window(&current, &next, "1", false, "patch");
3473
3474        assert_eq!(
3475            plan,
3476            vec![
3477                ("1".to_string(), MemberAction::ForwardPatch),
3478                ("4".to_string(), MemberAction::Skip),
3479                ("3".to_string(), MemberAction::Skip),
3480                ("2".to_string(), MemberAction::Skip),
3481            ]
3482        );
3483    }
3484
3485    #[test]
3486    fn a_key_entering_the_window_gets_the_whole_entity() {
3487        // "5" has no local state for the subscriber to merge a patch into.
3488        let plan = plan_window(&["4", "3"], &["5", "4", "3"], "5", false, "patch");
3489        assert_eq!(
3490            plan,
3491            vec![
3492                ("5".to_string(), MemberAction::Upsert),
3493                ("4".to_string(), MemberAction::Skip),
3494                ("3".to_string(), MemberAction::Skip),
3495            ]
3496        );
3497    }
3498
3499    #[test]
3500    fn derived_views_still_send_whole_entities() {
3501        // The patch on the bus is scoped to the source view, so a derived
3502        // subscription cannot forward it verbatim (A4-150).
3503        let plan = plan_window(&["1", "2"], &["1", "2"], "1", true, "patch");
3504        assert_eq!(
3505            plan,
3506            vec![
3507                ("1".to_string(), MemberAction::Upsert),
3508                ("2".to_string(), MemberAction::Skip),
3509            ]
3510        );
3511    }
3512
3513    #[test]
3514    fn a_delete_envelope_never_forwards_a_patch() {
3515        // A surviving key on a delete envelope carries no mergeable patch.
3516        let plan = plan_window(&["1", "2"], &["1", "2"], "1", false, "delete");
3517        assert_eq!(
3518            plan,
3519            vec![
3520                ("1".to_string(), MemberAction::Upsert),
3521                ("2".to_string(), MemberAction::Skip),
3522            ]
3523        );
3524    }
3525
3526    #[test]
3527    fn an_unchanged_window_sends_one_frame_not_a_broadcast() {
3528        // The regression this guards: 500 members used to mean 500 full
3529        // entities on the wire for a single mutation.
3530        let keys: Vec<String> = (0..500).map(|index| index.to_string()).collect();
3531        let refs: Vec<&str> = keys.iter().map(String::as_str).collect();
3532        let plan = plan_window(&refs, &refs, "250", false, "patch");
3533
3534        let sent = plan
3535            .iter()
3536            .filter(|(_, action)| *action != MemberAction::Skip)
3537            .count();
3538        assert_eq!(sent, 1);
3539        assert_eq!(plan[250].1, MemberAction::ForwardPatch);
3540    }
3541
3542    fn shared(rows: &[(String, Value)]) -> Vec<(String, SharedEntity)> {
3543        rows.iter()
3544            .map(|(key, data)| (key.clone(), SharedEntity::new(data.clone())))
3545            .collect()
3546    }
3547
3548    fn list_message(key: &str, op: &str, seq: &str) -> Arc<BusMessage> {
3549        Arc::new(BusMessage {
3550            key: key.to_string(),
3551            entity: "Thing/list".to_string(),
3552            payload: Arc::new(Bytes::from(
3553                serde_json::to_vec(&json!({
3554                    "entity": "Thing/list",
3555                    "op": op,
3556                    "key": key,
3557                    "seq": seq,
3558                    "data": {},
3559                }))
3560                .unwrap(),
3561            )),
3562        })
3563    }
3564
3565    #[test]
3566    fn coalescing_emits_only_the_final_full_state_per_changed_key() {
3567        let current = vec![
3568            ("a".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3569            ("b".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3570            (
3571                "stable".to_string(),
3572                json!({"count": 1, "_seq": "10:000001"}),
3573            ),
3574        ];
3575        let next = vec![
3576            ("a".to_string(), json!({"count": 3, "_seq": "10:000004"})),
3577            ("c".to_string(), json!({"count": 1, "_seq": "10:000003"})),
3578            (
3579                "stable".to_string(),
3580                json!({"count": 1, "_seq": "10:000001"}),
3581            ),
3582        ];
3583        let pending = HashMap::from([
3584            ("a".to_string(), list_message("a", "patch", "10:000004")),
3585            ("b".to_string(), list_message("b", "delete", "10:000002")),
3586            ("c".to_string(), list_message("c", "patch", "10:000003")),
3587        ]);
3588
3589        assert_eq!(
3590            plan_coalesced_collection_delta(
3591                &shared(&current),
3592                &shared(&next),
3593                &pending,
3594                &HashSet::new()
3595            ),
3596            vec![
3597                CollectionChange {
3598                    op: "delete",
3599                    key: "b".to_string(),
3600                    data: None,
3601                    seq: Some("10:000002".to_string()),
3602                },
3603                CollectionChange {
3604                    op: "upsert",
3605                    key: "a".to_string(),
3606                    data: Some(SharedEntity::new(json!({"count": 3, "_seq": "10:000004"}))),
3607                    seq: Some("10:000004".to_string()),
3608                },
3609                CollectionChange {
3610                    op: "upsert",
3611                    key: "c".to_string(),
3612                    data: Some(SharedEntity::new(json!({"count": 1, "_seq": "10:000003"}))),
3613                    seq: Some("10:000003".to_string()),
3614                },
3615            ]
3616        );
3617    }
3618
3619    #[test]
3620    fn a_lagged_snapshotless_subscription_gets_a_fatal_retryable_error() {
3621        let issue = SocketIssueMessage::subscription_lagged("balances".to_string(), 42);
3622        assert_eq!(issue.code, "subscription-lagged");
3623        assert!(issue.retryable);
3624        assert!(issue.fatal);
3625        assert!(issue.message.contains("42"));
3626        assert!(issue
3627            .suggested_action
3628            .unwrap()
3629            .contains("snapshots enabled"));
3630
3631        let append = SocketIssueMessage::append_subscription_lagged("trades".to_string(), 9);
3632        assert!(append.fatal);
3633        assert!(append.suggested_action.unwrap().contains("retained replay"));
3634    }
3635
3636    #[test]
3637    fn snapshot_batches_share_identity_and_completion() {
3638        let entities = ["one", "two", "three"]
3639            .map(|key| (key.to_string(), SharedEntity::new(json!({"key": key}))));
3640        let wire_format = WireFormat::default();
3641        let batches = create_snapshot_batches(
3642            &entities,
3643            &wire_format,
3644            SnapshotMetadata {
3645                subscription_id: "sub-1",
3646                snapshot_id: "snapshot-1",
3647                authoritative: true,
3648                mode: Mode::List,
3649                view_id: "Thing/list",
3650                key: None,
3651            },
3652            &SnapshotBatchConfig {
3653                initial_batch_size: 2,
3654                subsequent_batch_size: 1,
3655            },
3656        );
3657        assert_eq!(batches.len(), 2);
3658        assert!(batches.iter().all(|batch| batch.subscription_id == "sub-1"));
3659        assert!(batches
3660            .iter()
3661            .all(|batch| batch.snapshot_id == "snapshot-1"));
3662        assert!(!batches[0].complete);
3663        assert!(batches[1].complete);
3664        assert!(batches.iter().all(|batch| batch.authoritative));
3665    }
3666
3667    #[test]
3668    fn empty_incremental_snapshot_is_explicitly_non_authoritative() {
3669        let wire_format = WireFormat::default();
3670        let batches = create_snapshot_batches(
3671            &[],
3672            &wire_format,
3673            SnapshotMetadata {
3674                subscription_id: "sub-1",
3675                snapshot_id: "snapshot-1",
3676                authoritative: false,
3677                mode: Mode::State,
3678                view_id: "Thing/state",
3679                key: Some("missing"),
3680            },
3681            &SnapshotBatchConfig {
3682                initial_batch_size: 1,
3683                subsequent_batch_size: 1,
3684            },
3685        );
3686        assert_eq!(batches.len(), 1);
3687        assert!(!batches[0].authoritative);
3688        assert!(batches[0].complete);
3689        assert_eq!(batches[0].key, Some("missing"));
3690    }
3691
3692    fn wide_ints() -> WireFormat {
3693        WireFormat {
3694            wide_int_paths: vec![
3695                vec!["amount".to_string()],
3696                vec!["trades".to_string(), "price".to_string()],
3697            ],
3698        }
3699    }
3700
3701    fn sample_rows() -> Vec<(String, SharedEntity)> {
3702        (0..7)
3703            .map(|index| {
3704                (
3705                    index.to_string(),
3706                    SharedEntity::new(json!({
3707                        "_seq": format!("10:{index:012}"),
3708                        "_version": format!("e:{index}"),
3709                        "amount": u64::MAX - index,
3710                        "label": "Ï€ \"quoted\" \n",
3711                        "trades": [{"price": index, "side": "buy"}, {"price": -1}],
3712                    })),
3713                )
3714            })
3715            .collect()
3716    }
3717
3718    /// Snapshot frames serialized from shared rows are byte for byte the
3719    /// frames built from formatted copies of the rows.
3720    #[test]
3721    fn snapshot_batches_serialize_as_frames_of_formatted_copies() {
3722        use crate::websocket::frame::{SnapshotEntity, SnapshotFrame};
3723        let rows = sample_rows();
3724        for (rows, key, wire_format, (initial, subsequent)) in [
3725            (&rows[..], None, wide_ints(), (2, 3)),
3726            (&rows[..], Some("3"), WireFormat::default(), (50, 100)),
3727            (&rows[..1], Some("0"), wide_ints(), (1, 1)),
3728            (&rows[..0], None, wide_ints(), (2, 3)),
3729        ] {
3730            let metadata = SnapshotMetadata {
3731                subscription_id: "sub-1",
3732                snapshot_id: "snapshot-1",
3733                authoritative: true,
3734                mode: Mode::List,
3735                view_id: "Thing/list",
3736                key,
3737            };
3738            let config = SnapshotBatchConfig {
3739                initial_batch_size: initial,
3740                subsequent_batch_size: subsequent,
3741            };
3742            let batches = create_snapshot_batches(rows, &wire_format, metadata, &config);
3743            let mut offset = 0;
3744            for batch in &batches {
3745                let copies = rows[offset..offset + batch.data.rows.len()]
3746                    .iter()
3747                    .map(|(key, entity)| SnapshotEntity {
3748                        key: key.clone(),
3749                        data: wire_value(entity, &wire_format),
3750                    })
3751                    .collect();
3752                offset += batch.data.rows.len();
3753                let frame = SnapshotFrame {
3754                    protocol_version: PROTOCOL_VERSION,
3755                    subscription_id: "sub-1".to_string(),
3756                    snapshot_id: "snapshot-1".to_string(),
3757                    authoritative: true,
3758                    mode: Mode::List,
3759                    export: "Thing/list".to_string(),
3760                    op: "snapshot",
3761                    key: key.map(str::to_string),
3762                    data: copies,
3763                    complete: batch.complete,
3764                };
3765                assert_eq!(
3766                    serde_json::to_string(batch).unwrap(),
3767                    serde_json::to_string(&frame).unwrap()
3768                );
3769            }
3770            assert_eq!(offset, rows.len());
3771        }
3772    }
3773
3774    /// A membership frame serialized from a shared entity is byte for byte
3775    /// the frame built from a formatted copy of it.
3776    #[test]
3777    fn entity_frames_serialize_as_frames_of_formatted_copies() {
3778        for (_, entity) in sample_rows() {
3779            for (wire_format, seq) in [
3780                (wide_ints(), Some("10:000000000001")),
3781                (WireFormat::default(), None),
3782            ] {
3783                let frame = ScopedEntityFrame {
3784                    protocol_version: PROTOCOL_VERSION,
3785                    subscription_id: "sub-1",
3786                    mode: Mode::List,
3787                    export: "Thing/latest",
3788                    op: "upsert",
3789                    key: "7",
3790                    data: WireEntity {
3791                        entity: &entity,
3792                        wire_format: &wire_format,
3793                    },
3794                    seq,
3795                };
3796                let copy = Frame::scoped(
3797                    "sub-1",
3798                    Mode::List,
3799                    "Thing/latest",
3800                    "upsert",
3801                    "7",
3802                    wire_value(&entity, &wire_format),
3803                    seq.map(str::to_string),
3804                );
3805                assert_eq!(
3806                    serde_json::to_string(&frame).unwrap(),
3807                    serde_json::to_string(&copy).unwrap()
3808                );
3809            }
3810        }
3811    }
3812
3813    #[test]
3814    fn lag_recovery_is_authoritative_for_an_after_query() {
3815        let subscription = Subscription {
3816            protocol_version: PROTOCOL_VERSION,
3817            subscription_id: "sub-1".to_string(),
3818            query: SubscriptionQuery {
3819                view: "Thing/list".to_string(),
3820                after: Some("40:000000000010".to_string()),
3821                ..Default::default()
3822            },
3823            snapshot: Default::default(),
3824        };
3825
3826        assert!(!SnapshotPurpose::Initial.authoritative(&subscription));
3827        assert!(SnapshotPurpose::Recovery.authoritative(&subscription));
3828    }
3829
3830    #[test]
3831    fn dot_path_filters_are_exact_and_type_sensitive() {
3832        let mut query = SubscriptionQuery {
3833            view: "Thing/list".to_string(),
3834            ..Default::default()
3835        };
3836        query
3837            .filters
3838            .insert("state.status".to_string(), json!("open"));
3839        query.filters.insert("metrics.count".to_string(), json!(2));
3840        assert!(query_matches_entity(
3841            &query,
3842            "one",
3843            &json!({"state": {"status": "open"}, "metrics": {"count": 2}}),
3844        ));
3845        assert!(!query_matches_entity(
3846            &query,
3847            "one",
3848            &json!({"state": {"status": "open"}, "metrics": {"count": "2"}}),
3849        ));
3850    }
3851
3852    #[test]
3853    fn take_and_skip_define_independent_deterministic_windows() {
3854        let entities: Vec<_> = (1..=6)
3855            .map(|id| {
3856                (
3857                    id.to_string(),
3858                    json!({"id": id, "_seq": format!("10:{id:012}")}),
3859                )
3860            })
3861            .collect();
3862        let first = SubscriptionQuery {
3863            view: "Thing/list".to_string(),
3864            take: Some(2),
3865            skip: Some(0),
3866            ..Default::default()
3867        };
3868        let second = SubscriptionQuery {
3869            skip: Some(2),
3870            ..first.clone()
3871        };
3872        let first_keys: Vec<_> = select_query_entities(entities.clone(), &first, false, false)
3873            .into_iter()
3874            .map(|(key, _)| key)
3875            .collect();
3876        let second_keys: Vec<_> = select_query_entities(entities, &second, false, false)
3877            .into_iter()
3878            .map(|(key, _)| key)
3879            .collect();
3880        assert_eq!(first_keys, ["6", "5"]);
3881        assert_eq!(second_keys, ["4", "3"]);
3882    }
3883
3884    #[tokio::test]
3885    async fn state_receiver_is_installed_before_snapshot_awaits() {
3886        let bus = BusManager::new();
3887        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3888        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3889        let bus_for_task = bus.clone();
3890        let task = tokio::spawn(async move {
3891            subscribe_state_then_snapshot(&bus_for_task, "Thing/state", "one", || async move {
3892                snapshot_started_tx.send(()).unwrap();
3893                release_snapshot_rx.await.unwrap();
3894            })
3895            .await
3896            .0
3897        });
3898        snapshot_started_rx.await.unwrap();
3899        bus.publish_state(
3900            "Thing/state",
3901            "one",
3902            Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3903        )
3904        .await;
3905        release_snapshot_tx.send(()).unwrap();
3906        let mut receiver = task.await.unwrap();
3907        receiver.changed().await.unwrap();
3908        assert!(!receiver.borrow().payload.is_empty());
3909    }
3910
3911    async fn assert_list_receiver_precedes_snapshot(view: &'static str) {
3912        let bus = BusManager::new();
3913        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3914        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3915        let bus_for_task = bus.clone();
3916        let task = tokio::spawn(async move {
3917            subscribe_list_then_snapshot(&bus_for_task, view, || async move {
3918                snapshot_started_tx.send(()).unwrap();
3919                release_snapshot_rx.await.unwrap();
3920            })
3921            .await
3922            .0
3923        });
3924        snapshot_started_rx.await.unwrap();
3925        bus.publish_list(
3926            view,
3927            Arc::new(BusMessage {
3928                key: "one".to_string(),
3929                entity: view.to_string(),
3930                payload: Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3931            }),
3932        )
3933        .await;
3934        release_snapshot_tx.send(()).unwrap();
3935        let mut receiver = task.await.unwrap();
3936        assert_eq!(receiver.recv().await.unwrap().key, "one");
3937    }
3938
3939    #[tokio::test]
3940    async fn list_receiver_is_installed_before_snapshot() {
3941        assert_list_receiver_precedes_snapshot("Thing/list").await;
3942    }
3943
3944    #[tokio::test]
3945    async fn append_receiver_is_installed_before_snapshot() {
3946        assert_list_receiver_precedes_snapshot("Thing/append").await;
3947    }
3948
3949    #[tokio::test]
3950    async fn derived_source_receiver_is_installed_before_snapshot() {
3951        assert_list_receiver_precedes_snapshot("Thing/list-source").await;
3952    }
3953
3954    /// A sorted derived view is read only as far as the window reaches, and
3955    /// gives exactly the rows filtering and windowing all of it gives.
3956    #[tokio::test]
3957    async fn sorted_derived_reads_take_the_window_of_the_whole_order() {
3958        use crate::materialized_view::{SortConfig, SortOrder, ViewPipeline};
3959        use crate::sorted_cache::{SortOrder as CacheOrder, SortedViewCache};
3960        let mut sorted =
3961            SortedViewCache::new("Thing/top".into(), vec!["score".into()], CacheOrder::Desc);
3962        for id in 0..40u64 {
3963            sorted.upsert(
3964                id.to_string(),
3965                json!({
3966                    "_seq": format!("{}:0", 100 + id * 7 % 40),
3967                    "score": (id * 13 % 17).to_string(),
3968                    "owner": if id % 3 == 0 { "alice" } else { "bob" },
3969                    "_partition": if id % 2 == 0 { "even" } else { "odd" },
3970                }),
3971            );
3972        }
3973        let whole = sorted.ordered_entities();
3974        let mut spec = list_spec();
3975        spec.id = "Thing/top".into();
3976        spec.source_view = Some("Thing/list".into());
3977        let caches = Arc::new(tokio::sync::RwLock::new(HashMap::from([(
3978            spec.id.clone(),
3979            sorted,
3980        )])));
3981        let cache = EntityCache::new();
3982        for limit in [None, Some(5)] {
3983            spec.pipeline = Some(ViewPipeline {
3984                filter: None,
3985                sort: Some(SortConfig {
3986                    field_path: vec!["score".into()],
3987                    order: SortOrder::Desc,
3988                }),
3989                limit,
3990            });
3991            for skip in [None, Some(0), Some(3), Some(50)] {
3992                for take in [None, Some(1), Some(4), Some(100)] {
3993                    for (owner, after, partition) in [
3994                        (None, None, None),
3995                        (Some("alice"), None, None),
3996                        (None, Some("120:0"), Some("even")),
3997                    ] {
3998                        let mut query = SubscriptionQuery {
3999                            view: spec.id.clone(),
4000                            skip,
4001                            take,
4002                            after: after.map(str::to_string),
4003                            partition: partition.map(str::to_string),
4004                            snapshot_limit: Some(2),
4005                            ..Default::default()
4006                        };
4007                        if let Some(owner) = owner {
4008                            query.filters.insert("owner".into(), json!(owner));
4009                        }
4010                        for apply_snapshot_limit in [false, true] {
4011                            let rows = load_query_entities(
4012                                &cache,
4013                                Some(caches.clone()),
4014                                &spec,
4015                                &query,
4016                                apply_snapshot_limit,
4017                            )
4018                            .await;
4019                            let mut expected_query = query.clone();
4020                            if let Some(limit) = limit {
4021                                expected_query.take = Some(take.unwrap_or(limit).min(limit));
4022                            }
4023                            let expected = select_query_entities(
4024                                whole.clone(),
4025                                &expected_query,
4026                                true,
4027                                apply_snapshot_limit,
4028                            );
4029                            assert_eq!(rows, expected, "{query:?} {limit:?}");
4030                        }
4031                    }
4032                }
4033            }
4034        }
4035    }
4036
4037    #[tokio::test]
4038    async fn unsorted_derived_reads_preserve_filters_order_and_snapshot_window() {
4039        use crate::materialized_view::{CompareOp, FilterConfig, ViewPipeline};
4040        let cache = EntityCache::new();
4041        for id in 1..=6 {
4042            cache
4043                .upsert(
4044                    "Thing/list",
4045                    &id.to_string(),
4046                    json!({
4047                        "_seq": format!("{}:0", id * 10), "active": id != 6,
4048                        "owner": "alice", "id": id
4049                    }),
4050                )
4051                .await;
4052        }
4053        let mut spec = list_spec();
4054        spec.id = "Thing/custom".into();
4055        spec.source_view = Some("Thing/list".into());
4056        spec.pipeline = Some(ViewPipeline::default());
4057        let caches = Arc::new(tokio::sync::RwLock::new(HashMap::new()));
4058        let mut query = SubscriptionQuery {
4059            view: spec.id.clone(),
4060            ..Default::default()
4061        };
4062        let rows = load_query_entities(&cache, Some(caches.clone()), &spec, &query, false).await;
4063        assert_eq!(
4064            rows.iter().map(|(key, _)| key.as_str()).collect::<Vec<_>>(),
4065            ["6", "5", "4", "3", "2", "1"]
4066        );
4067        spec.pipeline.as_mut().unwrap().filter = Some(FilterConfig {
4068            field_path: vec!["active".into()],
4069            op: CompareOp::Eq,
4070            value: json!(true),
4071        });
4072        query.filters.insert("owner".into(), json!("alice"));
4073        query.skip = Some(1);
4074        query.take = Some(3);
4075        query.snapshot_limit = Some(1);
4076        let rows = load_query_entities(&cache, Some(caches.clone()), &spec, &query, false).await;
4077        assert_eq!(
4078            rows.iter().map(|(key, _)| key.as_str()).collect::<Vec<_>>(),
4079            ["4", "3", "2"]
4080        );
4081        assert_eq!(
4082            load_query_entities(&cache, Some(caches.clone()), &spec, &query, true)
4083                .await
4084                .len(),
4085            1
4086        );
4087        query.skip = None;
4088        query.after = Some("20:0".into());
4089        let rows = load_query_entities(&cache, Some(caches.clone()), &spec, &query, false).await;
4090        assert_eq!(
4091            rows.iter().map(|(key, _)| key.as_str()).collect::<Vec<_>>(),
4092            ["3", "4", "5"]
4093        );
4094        query.filters.insert("owner".into(), json!("bob"));
4095        assert!(
4096            load_query_entities(&cache, Some(caches), &spec, &query, false)
4097                .await
4098                .is_empty()
4099        );
4100    }
4101
4102    #[tokio::test]
4103    async fn snapshot_limit_does_not_change_live_take_skip_membership() {
4104        let cache = EntityCache::with_config(EntityCacheConfig {
4105            max_entities_per_view: 10,
4106            ..Default::default()
4107        });
4108        for id in 1..=4 {
4109            cache
4110                .upsert(
4111                    "Thing/list",
4112                    &id.to_string(),
4113                    json!({"_seq": format!("10:{id:012}")}),
4114                )
4115                .await;
4116        }
4117        let query = SubscriptionQuery {
4118            view: "Thing/list".to_string(),
4119            take: Some(3),
4120            skip: Some(1),
4121            snapshot_limit: Some(1),
4122            ..Default::default()
4123        };
4124        let live = load_query_entities(&cache, None, &list_spec(), &query, false).await;
4125        let snapshot = load_query_entities(&cache, None, &list_spec(), &query, true).await;
4126        assert_eq!(live.len(), 3);
4127        assert_eq!(snapshot.len(), 1);
4128    }
4129
4130    #[test]
4131    fn fixture_manifest_covers_required_conformance_cases() {
4132        let manifest: Value = serde_json::from_str(include_str!(
4133            "../../../../tests/fixtures/websocket-v2/manifest.json"
4134        ))
4135        .unwrap();
4136        let names: HashSet<_> = manifest["fixtures"]
4137            .as_array()
4138            .unwrap()
4139            .iter()
4140            .filter_map(Value::as_str)
4141            .collect();
4142        for required in [
4143            "keyed-state.json",
4144            "list-windows.json",
4145            "filters.json",
4146            "multi-batch-authoritative.json",
4147            "empty-snapshot.json",
4148            "remove.json",
4149            "delete.json",
4150            "incremental-snapshot.json",
4151            "reconnect-replacement.json",
4152            "errors.json",
4153            "whole-entities.json",
4154        ] {
4155            assert!(names.contains(required), "missing fixture {required}");
4156        }
4157
4158        for document in [
4159            include_str!("../../../../tests/fixtures/websocket-v2/keyed-state.json"),
4160            include_str!("../../../../tests/fixtures/websocket-v2/list-windows.json"),
4161            include_str!("../../../../tests/fixtures/websocket-v2/filters.json"),
4162            include_str!("../../../../tests/fixtures/websocket-v2/multi-batch-authoritative.json"),
4163            include_str!("../../../../tests/fixtures/websocket-v2/empty-snapshot.json"),
4164            include_str!("../../../../tests/fixtures/websocket-v2/remove.json"),
4165            include_str!("../../../../tests/fixtures/websocket-v2/delete.json"),
4166            include_str!("../../../../tests/fixtures/websocket-v2/incremental-snapshot.json"),
4167            include_str!("../../../../tests/fixtures/websocket-v2/reconnect-replacement.json"),
4168            include_str!("../../../../tests/fixtures/websocket-v2/errors.json"),
4169            include_str!("../../../../tests/fixtures/websocket-v2/whole-entities.json"),
4170        ] {
4171            let fixture: Value = serde_json::from_str(document).unwrap();
4172            assert!(fixture["name"].is_string());
4173        }
4174    }
4175
4176    fn append_frame(key: &str, data: Value) -> Arc<Bytes> {
4177        let frame = json!({
4178            "entity": "Trade/append",
4179            "op": "patch",
4180            "key": key,
4181            "offset": 7,
4182            "data": data,
4183        });
4184        Arc::new(Bytes::from(serde_json::to_vec(&frame).unwrap()))
4185    }
4186
4187    /// A replayable subscription must honour the same predicates a live
4188    /// collection subscription does; otherwise a filtered consumer receives
4189    /// events outside the query it asked for.
4190    #[test]
4191    fn replay_delivery_applies_key_partition_and_filters() {
4192        let matching = append_frame("pool1", json!({"_partition": "us", "side": "buy"}));
4193        let other_partition = append_frame("pool1", json!({"_partition": "eu", "side": "buy"}));
4194        let other_side = append_frame("pool1", json!({"_partition": "us", "side": "sell"}));
4195
4196        let unfiltered = SubscriptionQuery {
4197            view: "Trade/append".to_string(),
4198            ..Default::default()
4199        };
4200        assert!(live_frame_matches(&unfiltered, "pool1", &matching));
4201        assert!(live_frame_matches(&unfiltered, "pool9", &matching));
4202
4203        let keyed = SubscriptionQuery {
4204            view: "Trade/append".to_string(),
4205            key: Some("pool1".to_string()),
4206            ..Default::default()
4207        };
4208        assert!(live_frame_matches(&keyed, "pool1", &matching));
4209        assert!(!live_frame_matches(&keyed, "pool2", &matching));
4210
4211        let partitioned = SubscriptionQuery {
4212            view: "Trade/append".to_string(),
4213            partition: Some("us".to_string()),
4214            ..Default::default()
4215        };
4216        assert!(live_frame_matches(&partitioned, "pool1", &matching));
4217        assert!(!live_frame_matches(&partitioned, "pool1", &other_partition));
4218
4219        let filtered = SubscriptionQuery {
4220            view: "Trade/append".to_string(),
4221            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
4222            ..Default::default()
4223        };
4224        assert!(live_frame_matches(&filtered, "pool1", &matching));
4225        assert!(!live_frame_matches(&filtered, "pool1", &other_side));
4226    }
4227
4228    #[test]
4229    fn a_frame_without_decodable_data_does_not_satisfy_a_filter() {
4230        let filtered = SubscriptionQuery {
4231            view: "Trade/append".to_string(),
4232            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
4233            ..Default::default()
4234        };
4235        let garbage = Arc::new(Bytes::from_static(b"not json"));
4236        assert!(!live_frame_matches(&filtered, "pool1", &garbage));
4237    }
4238
4239    /// The recovery cursor must be the last offset delivered *before* the
4240    /// gap. Reporting the newest offset seen would step the consumer over
4241    /// the skipped records permanently.
4242    #[test]
4243    fn replay_lagged_recovers_from_before_the_gap() {
4244        let epoch = crate::journal::JournalEpoch::new();
4245        let issue = SocketIssueMessage::replay_lagged(
4246            Some("trades".to_string()),
4247            37,
4248            Some(crate::journal::Cursor {
4249                epoch: epoch.clone(),
4250                offset: 4180,
4251            }),
4252        );
4253        assert_eq!(issue.code, "replay-lagged");
4254        assert_eq!(issue.recover_from, Some(format!("{epoch}:4180")));
4255        assert!(
4256            issue.suggested_action.unwrap().contains("4180"),
4257            "the consumer is told exactly which cursor recovers the gap"
4258        );
4259
4260        // Nothing delivered yet: there is no pre-gap offset, so the whole
4261        // retained window is the recovery.
4262        let from_scratch = SocketIssueMessage::replay_lagged(Some("trades".to_string()), 9, None);
4263        assert_eq!(from_scratch.recover_from, None);
4264        assert!(from_scratch
4265            .suggested_action
4266            .unwrap()
4267            .contains("without `after`"));
4268    }
4269
4270    fn bus_message(key: &str) -> Arc<BusMessage> {
4271        Arc::new(BusMessage {
4272            key: key.to_string(),
4273            entity: "Trade/append".to_string(),
4274            payload: Arc::new(Bytes::from_static(b"{}")),
4275        })
4276    }
4277
4278    /// A long replay must keep the bus drained. The broadcast buffer is
4279    /// bounded, so a busy view would otherwise lap the replay and the first
4280    /// live `recv` would return `Lagged` — telling the client to resubscribe,
4281    /// starting another long replay, which laps again.
4282    #[tokio::test]
4283    async fn draining_during_a_replay_keeps_a_busy_view_from_lapping_it() {
4284        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(16);
4285        let mut pending = VecDeque::new();
4286        let mut lagged = None;
4287
4288        // Publish more than the channel holds, draining as a replay would
4289        // between sends.
4290        for index in 0..48 {
4291            sender.send(bus_message(&format!("k{index}"))).unwrap();
4292            drain_available(&mut receiver, &mut pending, &mut lagged);
4293        }
4294
4295        assert_eq!(lagged, None, "draining as we go means nothing is dropped");
4296        assert_eq!(pending.len(), 48, "every published frame is buffered");
4297    }
4298
4299    #[tokio::test]
4300    async fn a_replay_that_never_drains_is_reported_as_a_gap() {
4301        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
4302        for index in 0..32 {
4303            sender.send(bus_message(&format!("k{index}"))).unwrap();
4304        }
4305
4306        let mut pending = VecDeque::new();
4307        let mut lagged = None;
4308        drain_available(&mut receiver, &mut pending, &mut lagged);
4309
4310        assert!(
4311            lagged.is_some(),
4312            "overflowing the bus is a gap, not silent truncation"
4313        );
4314    }
4315
4316    /// Everything still on the bus after a gap is on the far side of it.
4317    /// Buffering it would put those frames in front of the lag report and
4318    /// advance the recovery cursor past the records it is meant to recover.
4319    #[tokio::test]
4320    async fn nothing_after_a_gap_is_buffered_ahead_of_the_report() {
4321        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
4322        let mut pending = VecDeque::new();
4323        let mut lagged = None;
4324
4325        // Delivered and buffered normally.
4326        sender.send(bus_message("before")).unwrap();
4327        drain_available(&mut receiver, &mut pending, &mut lagged);
4328        assert_eq!(pending.len(), 1);
4329
4330        // Overflow the bus: everything published from here is past the gap.
4331        for index in 0..32 {
4332            sender.send(bus_message(&format!("lost{index}"))).unwrap();
4333        }
4334        drain_available(&mut receiver, &mut pending, &mut lagged);
4335        assert!(lagged.is_some());
4336
4337        // The bus still holds what survived the overflow, all of it past the
4338        // gap. A later iteration of the replay loop drains again.
4339        let buffered_at_gap = pending.len();
4340        drain_available(&mut receiver, &mut pending, &mut lagged);
4341        assert_eq!(
4342            pending.len(),
4343            buffered_at_gap,
4344            "post-gap frames must not join the pre-gap flush"
4345        );
4346        assert_eq!(pending.front().unwrap().key, "before");
4347    }
4348    /// The bus is subscribed before the tape is read, so a record published
4349    /// in that window arrives on both paths.
4350    #[test]
4351    fn the_seam_between_replay_and_live_neither_repeats_nor_skips() {
4352        let mut last_sent = Some(4211);
4353
4354        assert!(
4355            already_delivered(Some(4211), &mut last_sent),
4356            "the record the replay ended on must not be sent twice"
4357        );
4358        assert!(already_delivered(Some(4100), &mut last_sent));
4359        assert_eq!(last_sent, Some(4211), "a duplicate never moves the mark");
4360
4361        assert!(
4362            !already_delivered(Some(4212), &mut last_sent),
4363            "the next record is new"
4364        );
4365        assert_eq!(last_sent, Some(4212));
4366
4367        // A frame with no offset comes from a view with no tape; it cannot
4368        // have been replayed, and must not disturb the mark.
4369        assert!(!already_delivered(None, &mut last_sent));
4370        assert_eq!(last_sent, Some(4212));
4371    }
4372
4373    /// A subscription with no cursor has delivered nothing, so the first live
4374    /// frame is not a duplicate.
4375    #[test]
4376    fn a_fresh_subscription_delivers_its_first_live_frame() {
4377        let mut last_sent = None;
4378        assert!(!already_delivered(Some(0), &mut last_sent));
4379        assert_eq!(last_sent, Some(0));
4380    }
4381
4382    /// End-to-end over a real socket: the pieces above are unit-tested
4383    /// individually, but the thing a consumer actually does — reconnect with
4384    /// a stored cursor and keep reading — only exists once a subscription is
4385    /// attached to a connection.
4386    mod over_a_socket {
4387        use super::*;
4388        use crate::journal::{EventJournal, JournalConfig};
4389        use crate::projector::Projector;
4390        use crate::{MutationBatch, SlotContext};
4391        use arete_interpreter::Mutation;
4392        use futures_util::{SinkExt, StreamExt};
4393        use std::time::Duration;
4394        use tokio::net::{TcpListener, TcpStream};
4395        use tokio::sync::mpsc;
4396        use tokio_tungstenite::tungstenite::Message;
4397        use tokio_tungstenite::{client_async, WebSocketStream};
4398
4399        const RETAINED: u64 = 600;
4400
4401        fn append_index() -> ViewIndex {
4402            let mut index = ViewIndex::new();
4403            index.add_spec(ViewSpec {
4404                id: "Trade/append".to_string(),
4405                export: "Trade".to_string(),
4406                mode: Mode::Append,
4407                wire_format: Default::default(),
4408                projection: Projection::all(),
4409                filters: Filters::all(),
4410                delivery: Delivery::default(),
4411                pipeline: None,
4412                source_view: None,
4413            });
4414            index
4415        }
4416
4417        fn trade(index: u64) -> MutationBatch {
4418            MutationBatch::with_slot_context(
4419                vec![Mutation {
4420                    export: "Trade".to_string(),
4421                    key: json!(format!("pool{}", index % 4)),
4422                    patch: json!({"trade": index}),
4423                    append: vec![],
4424                }]
4425                .into_iter()
4426                .collect(),
4427                SlotContext::new(100 + index / 3, index % 3),
4428            )
4429        }
4430
4431        struct Harness {
4432            addr: SocketAddr,
4433            journal: Arc<EventJournal>,
4434            tx: mpsc::Sender<MutationBatch>,
4435        }
4436
4437        impl Harness {
4438            async fn start() -> Self {
4439                let view_index = Arc::new(append_index());
4440                let entity_cache = EntityCache::new();
4441                let bus_manager = BusManager::new();
4442                let journal = Arc::new(EventJournal::new(JournalConfig {
4443                    enabled: true,
4444                    max_bytes_per_view: u64::MAX,
4445                    max_records_per_view: 10_000,
4446                    max_age: Duration::from_secs(3_600),
4447                }));
4448
4449                let (tx, rx) = mpsc::channel::<MutationBatch>(256);
4450                tokio::spawn(
4451                    Projector::new(
4452                        view_index.clone(),
4453                        bus_manager.clone(),
4454                        entity_cache.clone(),
4455                        rx,
4456                        #[cfg(feature = "otel")]
4457                        None,
4458                    )
4459                    .with_journal(journal.clone())
4460                    .run(),
4461                );
4462
4463                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4464                let addr = listener.local_addr().unwrap();
4465                let server = WebSocketServer::new(
4466                    addr,
4467                    bus_manager,
4468                    entity_cache,
4469                    view_index,
4470                    #[cfg(feature = "otel")]
4471                    None,
4472                )
4473                .with_journal(journal.clone());
4474                let (acceptor, _cleanup) = server.into_acceptor();
4475                tokio::spawn(async move { acceptor.serve_listener(listener).await });
4476
4477                Self { addr, journal, tx }
4478            }
4479
4480            async fn publish(&self, range: std::ops::Range<u64>) {
4481                for index in range {
4482                    self.tx.send(trade(index)).await.unwrap();
4483                }
4484                let (ack, wait) = oneshot::channel();
4485                self.tx
4486                    .send(MutationBatch::flush_marker(ack))
4487                    .await
4488                    .unwrap();
4489                wait.await.unwrap();
4490            }
4491
4492            async fn connect(&self) -> WebSocketStream<TcpStream> {
4493                let stream = TcpStream::connect(self.addr).await.unwrap();
4494                client_async(format!("ws://{}/", self.addr), stream)
4495                    .await
4496                    .unwrap()
4497                    .0
4498            }
4499        }
4500
4501        async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
4502            loop {
4503                let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
4504                    .await
4505                    .expect("the server answers within the timeout")
4506                    .expect("the stream stays open")
4507                    .expect("a readable frame");
4508                // Control and data frames arrive as binary; issue frames as
4509                // text. Both are JSON.
4510                let bytes = match &message {
4511                    Message::Text(text) => text.as_bytes(),
4512                    Message::Binary(bytes) => bytes.as_ref(),
4513                    _ => continue,
4514                };
4515                return serde_json::from_slice(bytes).expect("frames are JSON");
4516            }
4517        }
4518
4519        /// Collect `count` event frames, ignoring anything else on the wire.
4520        async fn collect_trades(socket: &mut WebSocketStream<TcpStream>, count: usize) -> Vec<u64> {
4521            let mut offsets = Vec::with_capacity(count);
4522            while offsets.len() < count {
4523                let frame = next_frame(socket).await;
4524                assert_ne!(
4525                    frame["type"], "error",
4526                    "no error frame should interrupt delivery: {frame}"
4527                );
4528                if let Some(offset) = frame["offset"].as_u64() {
4529                    offsets.push(offset);
4530                }
4531            }
4532            offsets
4533        }
4534
4535        /// The headline claim: reconnecting with a stored cursor delivers every
4536        /// event published since it, in order, and then continues live without
4537        /// a duplicate or a hole at the seam.
4538        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
4539        async fn a_reconnect_replays_from_a_cursor_and_continues_live() {
4540            let harness = Harness::start().await;
4541            harness.publish(0..RETAINED).await;
4542
4543            let cursor = harness.journal.window("Trade/append").await;
4544            let stored = format!("{}:{}", cursor.epoch, 99);
4545
4546            let mut socket = harness.connect().await;
4547            socket
4548                .send(Message::Text(
4549                    json!({
4550                        "type": "subscribe",
4551                        "protocolVersion": 2,
4552                        "subscriptionId": "trades",
4553                        "query": {"view": "Trade/append", "after": stored},
4554                    })
4555                    .to_string()
4556                    .into(),
4557                ))
4558                .await
4559                .unwrap();
4560
4561            let ack = next_frame(&mut socket).await;
4562            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
4563            assert_eq!(ack["replayWindow"]["next"], json!(RETAINED));
4564
4565            // Well over the 500 a single read or buffer would cover.
4566            let replayed = collect_trades(&mut socket, (RETAINED - 100) as usize).await;
4567            assert_eq!(
4568                replayed,
4569                (100..RETAINED).collect::<Vec<_>>(),
4570                "every event after the cursor, in order, exactly once"
4571            );
4572
4573            // Published only now, so these can only arrive over the live path.
4574            harness.publish(RETAINED..RETAINED + 40).await;
4575            let live = collect_trades(&mut socket, 40).await;
4576            assert_eq!(
4577                live,
4578                (RETAINED..RETAINED + 40).collect::<Vec<_>>(),
4579                "the live stream resumes exactly where the replay stopped"
4580            );
4581
4582            socket.close(None).await.ok();
4583        }
4584
4585        /// Events published *during* the replay must still arrive. The replay
4586        /// and the live subscription are separate reads of the same tape, and
4587        /// the seam between them is where a naive implementation drops or
4588        /// repeats.
4589        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
4590        async fn events_published_during_a_replay_are_not_lost() {
4591            let harness = Harness::start().await;
4592            harness.publish(0..RETAINED).await;
4593
4594            let epoch = harness.journal.window("Trade/append").await.epoch;
4595            let mut socket = harness.connect().await;
4596            socket
4597                .send(Message::Text(
4598                    json!({
4599                        "type": "subscribe",
4600                        "protocolVersion": 2,
4601                        "subscriptionId": "trades",
4602                        "query": {"view": "Trade/append", "after": format!("{epoch}:0")},
4603                    })
4604                    .to_string()
4605                    .into(),
4606                ))
4607                .await
4608                .unwrap();
4609            let ack = next_frame(&mut socket).await;
4610            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
4611
4612            // Keep publishing while the replay is still draining.
4613            harness.publish(RETAINED..RETAINED + 200).await;
4614
4615            let total = (RETAINED + 200 - 1) as usize;
4616            let delivered = collect_trades(&mut socket, total).await;
4617            assert_eq!(
4618                delivered,
4619                (1..RETAINED + 200).collect::<Vec<_>>(),
4620                "replay and live output join without a gap or a repeat"
4621            );
4622
4623            socket.close(None).await.ok();
4624        }
4625
4626        /// A cursor from another tape lifetime is refused rather than served
4627        /// as a continuation, and the refusal releases the subscription id.
4628        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4629        async fn a_stale_epoch_is_refused_and_frees_the_subscription_id() {
4630            let harness = Harness::start().await;
4631            harness.publish(0..50).await;
4632
4633            let mut socket = harness.connect().await;
4634            socket
4635                .send(Message::Text(
4636                    json!({
4637                        "type": "subscribe",
4638                        "protocolVersion": 2,
4639                        "subscriptionId": "trades",
4640                        "query": {
4641                            "view": "Trade/append",
4642                            "after": format!("{}:10", crate::journal::JournalEpoch::new()),
4643                        },
4644                    })
4645                    .to_string()
4646                    .into(),
4647                ))
4648                .await
4649                .unwrap();
4650
4651            let error = next_frame(&mut socket).await;
4652            assert_eq!(error["type"], "error", "unexpected frame: {error}");
4653            assert_eq!(error["code"], "cursor-epoch-changed");
4654
4655            // The documented recovery is to resubscribe without a cursor. That
4656            // only works if the refused attempt released the id.
4657            socket
4658                .send(Message::Text(
4659                    json!({
4660                        "type": "subscribe",
4661                        "protocolVersion": 2,
4662                        "subscriptionId": "trades",
4663                        "query": {"view": "Trade/append"},
4664                    })
4665                    .to_string()
4666                    .into(),
4667                ))
4668                .await
4669                .unwrap();
4670            let ack = next_frame(&mut socket).await;
4671            assert_eq!(
4672                ack["op"], "subscribed",
4673                "the refused id must be reusable: {ack}"
4674            );
4675
4676            assert_eq!(collect_trades(&mut socket, 50).await.len(), 50);
4677            socket.close(None).await.ok();
4678        }
4679
4680        const ROUND: &str = "Round/state";
4681
4682        /// A state view with no projector: the test writes the cache and the
4683        /// bus itself, so it controls exactly when the subscriber can run.
4684        #[derive(Clone)]
4685        struct StateServer {
4686            addr: SocketAddr,
4687            bus_manager: BusManager,
4688            entity_cache: EntityCache,
4689        }
4690
4691        impl StateServer {
4692            async fn start() -> Self {
4693                let mut index = ViewIndex::new();
4694                index.add_spec(ViewSpec {
4695                    id: ROUND.to_string(),
4696                    export: "Round".to_string(),
4697                    mode: Mode::State,
4698                    wire_format: Default::default(),
4699                    projection: Projection::all(),
4700                    filters: Filters::all(),
4701                    delivery: Delivery::default(),
4702                    pipeline: None,
4703                    source_view: None,
4704                });
4705                let bus_manager = BusManager::new();
4706                let entity_cache = EntityCache::new();
4707                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4708                let addr = listener.local_addr().unwrap();
4709                let server = WebSocketServer::new(
4710                    addr,
4711                    bus_manager.clone(),
4712                    entity_cache.clone(),
4713                    Arc::new(index),
4714                    #[cfg(feature = "otel")]
4715                    None,
4716                );
4717                let (acceptor, _cleanup) = server.into_acceptor();
4718                tokio::spawn(async move { acceptor.serve_listener(listener).await });
4719                Self {
4720                    addr,
4721                    bus_manager,
4722                    entity_cache,
4723                }
4724            }
4725
4726            /// Publish a patch the way the projector does: stamp its `_seq`,
4727            /// write the cache, then the key's bus.
4728            async fn publish(&self, key: &str, mut patch: Value, seq: &str) {
4729                patch["_seq"] = Value::String(seq.to_string());
4730                self.entity_cache
4731                    .upsert_with_append(ROUND, key, patch.clone(), &[], PatchOrigin::Unknown)
4732                    .await;
4733                self.publish_frame(key, "patch", patch, seq).await;
4734            }
4735
4736            /// A complete new lifetime after an explicit source deletion.
4737            async fn create(&self, key: &str, mut entity: Value, seq: &str) {
4738                entity["_seq"] = Value::String(seq.to_string());
4739                self.entity_cache
4740                    .upsert_with_append(ROUND, key, entity.clone(), &[], PatchOrigin::Creation)
4741                    .await;
4742                self.publish_frame(key, "upsert", entity, seq).await;
4743            }
4744
4745            /// Publish a frame to the key's bus without touching the cache.
4746            async fn publish_frame(&self, key: &str, op: &str, data: Value, seq: &str) {
4747                let frame = json!({
4748                    "mode": "state",
4749                    "entity": ROUND,
4750                    "op": op,
4751                    "key": key,
4752                    "data": data,
4753                    "seq": seq,
4754                });
4755                self.bus_manager
4756                    .publish_state(ROUND, key, Arc::new(Bytes::from(frame.to_string())))
4757                    .await;
4758            }
4759
4760            /// Subscribe to one key and read through its snapshot.
4761            async fn subscribe(&self, key: &str) -> WebSocketStream<TcpStream> {
4762                let stream = TcpStream::connect(self.addr).await.unwrap();
4763                let mut socket = client_async(format!("ws://{}/", self.addr), stream)
4764                    .await
4765                    .unwrap()
4766                    .0;
4767                socket
4768                    .send(Message::Text(
4769                        json!({
4770                            "type": "subscribe",
4771                            "protocolVersion": 2,
4772                            "subscriptionId": "round",
4773                            "query": {"view": ROUND, "key": key},
4774                        })
4775                        .to_string()
4776                        .into(),
4777                    ))
4778                    .await
4779                    .unwrap();
4780                let ack = next_frame(&mut socket).await;
4781                assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
4782                loop {
4783                    let frame = next_frame(&mut socket).await;
4784                    assert_eq!(frame["op"], "snapshot", "unexpected frame: {frame}");
4785                    if frame["complete"] == true {
4786                        return socket;
4787                    }
4788                }
4789            }
4790        }
4791
4792        /// Two patches to one key, published before the subscriber reads the
4793        /// first, must both reach it. The bus holds only the latest frame, so
4794        /// forwarding it alone drops the fields of the one it replaced.
4795        // One worker, so tasks run one at a time. The handshake needs the
4796        // multi-threaded runtime.
4797        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4798        async fn a_state_subscriber_keeps_the_fields_of_an_overwritten_patch() {
4799            let server = StateServer::start().await;
4800            server
4801                .publish("7", json!({"id": 7}), "100:000000000001")
4802                .await;
4803            let mut socket = server.subscribe("7").await;
4804
4805            // A lone patch is forwarded as it was published. This one is an
4806            // account update, numbered by its write version.
4807            server
4808                .publish(
4809                    "7",
4810                    json!({"results": {"slot_hash": "abc"}}),
4811                    "140:003836292257",
4812                )
4813                .await;
4814            let frame = next_frame(&mut socket).await;
4815            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4816            assert_eq!(frame["seq"], "140:003836292257");
4817
4818            // Two instruction patches from later in the same slot, numbered by
4819            // transaction index. Both writes happen in one task that never
4820            // yields to the subscriber, so the second replaces the first on
4821            // the bus before the subscriber wakes.
4822            let writer = server.clone();
4823            tokio::spawn(async move {
4824                writer
4825                    .publish("7", json!({"entropy": {"seed": "def"}}), "140:000000001200")
4826                    .await;
4827                writer
4828                    .publish("7", json!({"total": 1}), "140:000000001200")
4829                    .await;
4830            })
4831            .await
4832            .unwrap();
4833            let frame = next_frame(&mut socket).await;
4834            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4835            // Everything that changed since the snapshot, the forwarded patch
4836            // included; `id` has not.
4837            assert_eq!(
4838                frame["data"],
4839                json!({
4840                    "results": {"slot_hash": "abc"},
4841                    "entropy": {"seed": "def"},
4842                    "total": 1,
4843                    "_seq": "140:000000001200",
4844                })
4845            );
4846            // The latest seq sorts below the one already delivered, and a
4847            // client would drop the frame as stale.
4848            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4849
4850            // Having caught up, the subscriber forwards patches again.
4851            server
4852                .publish("7", json!({"total": 2}), "141:000000000002")
4853                .await;
4854            let frame = next_frame(&mut socket).await;
4855            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4856            assert_eq!(
4857                frame["data"],
4858                json!({"total": 2, "_seq": "141:000000000002"})
4859            );
4860
4861            // The next catch-up is measured against the last one: only what
4862            // changed since it, not the fields that one already carried.
4863            let writer = server.clone();
4864            tokio::spawn(async move {
4865                writer
4866                    .publish("7", json!({"total": 3}), "142:000000000001")
4867                    .await;
4868                writer
4869                    .publish("7", json!({"flag": true}), "142:000000000001")
4870                    .await;
4871            })
4872            .await
4873            .unwrap();
4874            let frame = next_frame(&mut socket).await;
4875            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4876            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4877            assert_eq!(
4878                frame["data"],
4879                json!({"total": 3, "flag": true, "_seq": "142:000000000001"})
4880            );
4881
4882            socket.close(None).await.ok();
4883        }
4884
4885        /// A whole entity forwarded as it was published replaces the client's
4886        /// copy, so a later catch-up is measured against it.
4887        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4888        async fn a_forwarded_whole_entity_is_what_a_catch_up_is_measured_against() {
4889            let server = StateServer::start().await;
4890            server
4891                .publish("7", json!({"id": 7, "a": 1}), "100:000000000001")
4892                .await;
4893            let mut socket = server.subscribe("7").await;
4894
4895            let whole = json!({"id": 7, "a": 1, "b": 2, "_seq": "101:000000000001"});
4896            server
4897                .entity_cache
4898                .store_whole(ROUND, "7", whole.clone())
4899                .await;
4900            server
4901                .publish_frame("7", "upsert", whole.clone(), "101:000000000001")
4902                .await;
4903            let frame = next_frame(&mut socket).await;
4904            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4905            assert_eq!(frame["data"], whole);
4906
4907            let writer = server.clone();
4908            tokio::spawn(async move {
4909                writer
4910                    .publish("7", json!({"c": 3}), "102:000000000001")
4911                    .await;
4912                writer
4913                    .publish("7", json!({"d": 4}), "102:000000000002")
4914                    .await;
4915            })
4916            .await
4917            .unwrap();
4918            let frame = next_frame(&mut socket).await;
4919            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4920            assert_eq!(
4921                frame["data"],
4922                json!({"c": 3, "d": 4, "_seq": "102:000000000002"}),
4923                "`b` came with the whole entity, so it is not sent again"
4924            );
4925            socket.close(None).await.ok();
4926        }
4927
4928        /// Holders of the cached entity's fields beyond the cache and the
4929        /// copy this reads.
4930        async fn holders(server: &StateServer, key: &str) -> usize {
4931            let entity = server.entity_cache.get_shared(ROUND, key).await.unwrap();
4932            Arc::strong_count(entity.fields()) - 2
4933        }
4934
4935        /// Wait until `held` reports `expected`: a subscription's task takes
4936        /// over what it holds just after its snapshot reaches the client.
4937        async fn settle(mut held: impl AsyncFnMut() -> usize, expected: usize) {
4938            for _ in 0..200 {
4939                if held().await == expected {
4940                    return;
4941                }
4942                tokio::time::sleep(Duration::from_millis(5)).await;
4943            }
4944            assert_eq!(held().await, expected);
4945        }
4946
4947        /// What a state subscriber keeps to measure catch-ups against is the
4948        /// cached entity it was sent, or the frame the bus forwarded: shared,
4949        /// never a copy of its own.
4950        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4951        async fn state_subscribers_share_what_they_were_sent() {
4952            let server = StateServer::start().await;
4953            server
4954                .create(
4955                    "7",
4956                    json!({"id": 7, "detail": {"a": 1}}),
4957                    "100:000000000001",
4958                )
4959                .await;
4960            let mut sockets = Vec::new();
4961            for _ in 0..3 {
4962                sockets.push(server.subscribe("7").await);
4963            }
4964            settle(|| holders(&server, "7"), 3).await;
4965
4966            // A whole entity forwarded from the bus: each subscriber keeps the
4967            // frame the bus holds, not the entity it carries.
4968            let sent = server.entity_cache.get_shared(ROUND, "7").await.unwrap();
4969            let whole = json!({"id": 7, "detail": {"b": 2}, "_seq": "101:000000000001"});
4970            server
4971                .entity_cache
4972                .store_whole(ROUND, "7", whole.clone())
4973                .await;
4974            let frame = Arc::new(Bytes::from(
4975                json!({
4976                    "mode": "state",
4977                    "entity": ROUND,
4978                    "op": "upsert",
4979                    "key": "7",
4980                    "data": whole,
4981                    "seq": "101:000000000001",
4982                })
4983                .to_string(),
4984            ));
4985            server
4986                .bus_manager
4987                .publish_state(ROUND, "7", frame.clone())
4988                .await;
4989            for socket in &mut sockets {
4990                assert_eq!(next_frame(socket).await["op"], "upsert");
4991            }
4992            // This test's handle, the bus's, and one per subscriber.
4993            settle(async || Arc::strong_count(&frame), 2 + sockets.len()).await;
4994            // Only this test still holds the entity they were sent before.
4995            settle(async || Arc::strong_count(sent.fields()), 1).await;
4996            for mut socket in sockets {
4997                socket.close(None).await.ok();
4998            }
4999        }
5000
5001        /// A catch-up resends a field a forwarded patch set even when the
5002        /// entity has set it back to the value the client last received
5003        /// whole: the client still holds the forwarded value.
5004        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5005        async fn a_catch_up_restores_a_field_a_forwarded_patch_set() {
5006            let server = StateServer::start().await;
5007            server
5008                .publish("7", json!({"id": 7, "a": 1}), "100:000000000001")
5009                .await;
5010            let mut socket = server.subscribe("7").await;
5011
5012            server
5013                .publish("7", json!({"a": 2}), "101:000000000001")
5014                .await;
5015            let frame = next_frame(&mut socket).await;
5016            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5017            assert_eq!(frame["data"], json!({"a": 2, "_seq": "101:000000000001"}));
5018
5019            let writer = server.clone();
5020            tokio::spawn(async move {
5021                writer
5022                    .publish("7", json!({"a": 1}), "102:000000000001")
5023                    .await;
5024                writer
5025                    .publish("7", json!({"b": 1}), "102:000000000002")
5026                    .await;
5027            })
5028            .await
5029            .unwrap();
5030            let frame = next_frame(&mut socket).await;
5031            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5032            assert_eq!(
5033                frame["data"],
5034                json!({"a": 1, "b": 1, "_seq": "102:000000000002"}),
5035                "`a` matches the snapshot but the client holds 2"
5036            );
5037            socket.close(None).await.ok();
5038        }
5039
5040        /// A key deleted and created again before the subscriber reads
5041        /// either frame: the bus keeps only the new entity's. Merged into the
5042        /// client's copy it would leave the deleted entity's fields in place,
5043        /// so the client gets the new entity whole, replacing its copy.
5044        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5045        async fn a_key_deleted_and_recreated_between_reads_replaces_the_clients_copy() {
5046            let server = StateServer::start().await;
5047            server
5048                .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
5049                .await;
5050            let mut socket = server.subscribe("7").await;
5051
5052            let writer = server.clone();
5053            tokio::spawn(async move {
5054                writer.entity_cache.remove(ROUND, "7").await;
5055                writer
5056                    .publish_frame("7", "delete", Value::Null, "101:000000000001")
5057                    .await;
5058                writer
5059                    .publish("7", json!({"id": 7, "fresh": true}), "101:000000000002")
5060                    .await;
5061            })
5062            .await
5063            .unwrap();
5064            let frame = next_frame(&mut socket).await;
5065            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5066            assert_eq!(
5067                frame["data"],
5068                json!({"id": 7, "fresh": true, "_seq": "101:000000000002"})
5069            );
5070            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
5071
5072            // Deleted, with a patch overwritten before it: the delete is final.
5073            let writer = server.clone();
5074            tokio::spawn(async move {
5075                writer
5076                    .publish("7", json!({"n": 1}), "102:000000000001")
5077                    .await;
5078                writer.entity_cache.remove(ROUND, "7").await;
5079                writer
5080                    .publish_frame("7", "delete", Value::Null, "102:000000000002")
5081                    .await;
5082            })
5083            .await
5084            .unwrap();
5085            let frame = next_frame(&mut socket).await;
5086            assert_eq!(frame["op"], "delete", "unexpected frame: {frame}");
5087
5088            // Created again: the client holds nothing, so it arrives whole.
5089            server
5090                .create("7", json!({"id": 7, "third": true}), "103:000000000001")
5091                .await;
5092            let frame = next_frame(&mut socket).await;
5093            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5094            assert_eq!(frame["data"]["third"], true);
5095            assert!(frame["data"].get("fresh").is_none(), "{frame}");
5096            socket.close(None).await.ok();
5097        }
5098
5099        /// A missed delete, then a patch the cache refused (it lacks the new
5100        /// entity): the client's copy is of the deleted entity, so the patch
5101        /// is not merged into it. The new entity replaces it once it is whole.
5102        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5103        async fn a_copy_of_a_deleted_entity_takes_no_patch_until_it_is_replaced() {
5104            let server = StateServer::start().await;
5105            server
5106                .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
5107                .await;
5108            let mut socket = server.subscribe("7").await;
5109
5110            let writer = server.clone();
5111            tokio::spawn(async move {
5112                writer.entity_cache.remove(ROUND, "7").await;
5113                writer
5114                    .publish_frame("7", "delete", Value::Null, "101:000000000001")
5115                    .await;
5116                writer
5117                    .publish_frame("7", "patch", json!({"n": 1}), "101:000000000002")
5118                    .await;
5119            })
5120            .await
5121            .unwrap();
5122            // Let the subscriber read the patch, and withhold it, before the
5123            // whole entity arrives. Nothing is sent, so there is no frame to
5124            // wait for.
5125            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
5126
5127            let whole = json!({"id": 7, "n": 1, "_seq": "101:000000000002"});
5128            server
5129                .entity_cache
5130                .store_whole(ROUND, "7", whole.clone())
5131                .await;
5132            server
5133                .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
5134                .await;
5135            let frame = next_frame(&mut socket).await;
5136            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5137            assert_eq!(frame["data"], whole);
5138            socket.close(None).await.ok();
5139        }
5140
5141        /// The source replaced the entity whole (an `upsert`, dropping a field)
5142        /// and then patched it, both before the subscriber read either. The
5143        /// patch alone, or the entity merged in, would keep the dropped field.
5144        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5145        async fn a_whole_entity_among_overwritten_frames_replaces_the_clients_copy() {
5146            let server = StateServer::start().await;
5147            server
5148                .publish("7", json!({"id": 7, "dropped": 1}), "100:000000000001")
5149                .await;
5150            let mut socket = server.subscribe("7").await;
5151
5152            let writer = server.clone();
5153            tokio::spawn(async move {
5154                let whole = json!({"id": 7, "_seq": "101:000000000001"});
5155                writer
5156                    .entity_cache
5157                    .store_whole(ROUND, "7", whole.clone())
5158                    .await;
5159                writer
5160                    .publish_frame("7", "upsert", whole, "101:000000000001")
5161                    .await;
5162                writer
5163                    .publish("7", json!({"n": 1}), "101:000000000002")
5164                    .await;
5165            })
5166            .await
5167            .unwrap();
5168            let frame = next_frame(&mut socket).await;
5169            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5170            assert_eq!(
5171                frame["data"],
5172                json!({"id": 7, "n": 1, "_seq": "101:000000000002"})
5173            );
5174            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
5175            socket.close(None).await.ok();
5176        }
5177
5178        /// Patches overwritten on the bus while the cache lacks the key (it
5179        /// evicted it, and refused them) leave nothing whole to send in their
5180        /// place: the holder gets the latest patch, as a holder of an evicted
5181        /// key does, and stays behind. The resend that brings the entity back
5182        /// carries the seq of the key's latest change, which the holder
5183        /// already has and would drop as stale, so it gets the cached entity
5184        /// without a seq instead: as an upsert, since the resend is the whole
5185        /// entity and replaces what the holder has.
5186        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5187        async fn a_holder_behind_on_an_evicted_key_catches_up_when_it_returns() {
5188            let server = StateServer::start().await;
5189            server
5190                .publish("7", json!({"id": 7}), "100:000000000001")
5191                .await;
5192            let mut socket = server.subscribe("7").await;
5193            server.entity_cache.remove(ROUND, "7").await;
5194
5195            let writer = server.clone();
5196            tokio::spawn(async move {
5197                writer
5198                    .publish_frame("7", "patch", json!({"a": 1}), "101:000000000001")
5199                    .await;
5200                writer
5201                    .publish_frame("7", "patch", json!({"b": 2}), "101:000000000002")
5202                    .await;
5203            })
5204            .await
5205            .unwrap();
5206            let frame = next_frame(&mut socket).await;
5207            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5208            assert_eq!(frame["data"], json!({"b": 2}));
5209            assert_eq!(frame["seq"], "101:000000000002");
5210
5211            // The resend: the whole entity, at the latest change's position.
5212            let whole = json!({"id": 7, "a": 1, "b": 2, "_seq": "101:000000000002"});
5213            server
5214                .entity_cache
5215                .store_whole(ROUND, "7", whole.clone())
5216                .await;
5217            server
5218                .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
5219                .await;
5220            let frame = next_frame(&mut socket).await;
5221            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5222            assert_eq!(frame["data"], whole);
5223            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
5224
5225            // Caught up: patches are forwarded again.
5226            server
5227                .publish("7", json!({"c": 3}), "102:000000000001")
5228                .await;
5229            let frame = next_frame(&mut socket).await;
5230            assert_eq!(frame["data"], json!({"c": 3, "_seq": "102:000000000001"}));
5231            assert_eq!(frame["seq"], "102:000000000001");
5232            socket.close(None).await.ok();
5233        }
5234    }
5235
5236    /// A key only counts as held by the client once a frame actually carried
5237    /// it. A truncated or disabled snapshot leaves keys the client never saw,
5238    /// and their first change must arrive whole.
5239    mod partial_entities_over_a_socket {
5240        use super::*;
5241        use crate::projector::Projector;
5242        use crate::{MutationBatch, SlotContext};
5243        use arete_interpreter::{AccountPosition, Mutation};
5244        use futures_util::{SinkExt, StreamExt};
5245        use std::time::Duration;
5246        use tokio::net::{TcpListener, TcpStream};
5247        use tokio::sync::mpsc;
5248        use tokio_tungstenite::tungstenite::Message;
5249        use tokio_tungstenite::{client_async, WebSocketStream};
5250
5251        fn thing_index() -> ViewIndex {
5252            let mut index = ViewIndex::new();
5253            index.add_spec(list_spec());
5254            index.add_spec(ViewSpec {
5255                id: "Thing/state".to_string(),
5256                mode: Mode::State,
5257                ..list_spec()
5258            });
5259            index
5260        }
5261
5262        struct Harness {
5263            addr: SocketAddr,
5264            tx: mpsc::Sender<MutationBatch>,
5265            slot: std::sync::atomic::AtomicU64,
5266        }
5267
5268        impl Harness {
5269            async fn start(delivery: WebSocketDeliveryConfig) -> Self {
5270                Self::start_with_cache(delivery, EntityCache::new()).await
5271            }
5272
5273            async fn start_with_cache(
5274                delivery: WebSocketDeliveryConfig,
5275                entity_cache: EntityCache,
5276            ) -> Self {
5277                Self::start_with_index(delivery, entity_cache, thing_index()).await
5278            }
5279
5280            async fn start_with_index(
5281                delivery: WebSocketDeliveryConfig,
5282                entity_cache: EntityCache,
5283                index: ViewIndex,
5284            ) -> Self {
5285                let view_index = Arc::new(index);
5286                let bus_manager = BusManager::new();
5287                let (tx, rx) = mpsc::channel::<MutationBatch>(64);
5288                // Unconstrained, so Tokio's cooperative budget never makes the
5289                // projector yield partway through a batch: a test on one
5290                // worker then knows no subscriber runs until the whole batch
5291                // is published (see
5292                // `snapshot_rows_patches_and_catch_ups_carry_the_version`).
5293                tokio::spawn(tokio::task::unconstrained(
5294                    Projector::new(
5295                        view_index.clone(),
5296                        bus_manager.clone(),
5297                        entity_cache.clone(),
5298                        rx,
5299                        #[cfg(feature = "otel")]
5300                        None,
5301                    )
5302                    .run(),
5303                ));
5304                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5305                let addr = listener.local_addr().unwrap();
5306                let server = WebSocketServer::new(
5307                    addr,
5308                    bus_manager,
5309                    entity_cache,
5310                    view_index,
5311                    #[cfg(feature = "otel")]
5312                    None,
5313                )
5314                .with_delivery_config(delivery);
5315                let (acceptor, _cleanup) = server.into_acceptor();
5316                tokio::spawn(async move { acceptor.serve_listener(listener).await });
5317                Self {
5318                    addr,
5319                    tx,
5320                    slot: std::sync::atomic::AtomicU64::new(100),
5321                }
5322            }
5323
5324            /// Apply one source patch and wait until the projector published it.
5325            async fn patch(&self, key: &str, patch: Value) {
5326                self.mutate(Mutation {
5327                    export: "Thing".to_string(),
5328                    key: json!(key),
5329                    patch,
5330                    append: vec![],
5331                })
5332                .await;
5333            }
5334
5335            async fn mutate(&self, mutation: Mutation) {
5336                let slot = self.slot.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
5337                self.mutate_with_context(mutation, SlotContext::new(slot, 0))
5338                    .await;
5339            }
5340
5341            async fn mutate_with_context(&self, mutation: Mutation, context: SlotContext) {
5342                self.tx
5343                    .send(MutationBatch::with_slot_context(
5344                        vec![mutation].into_iter().collect(),
5345                        context,
5346                    ))
5347                    .await
5348                    .unwrap();
5349                let (ack, wait) = oneshot::channel();
5350                self.tx
5351                    .send(MutationBatch::flush_marker(ack))
5352                    .await
5353                    .unwrap();
5354                wait.await.unwrap();
5355            }
5356
5357            async fn subscribe(&self, query: Value, snapshot: bool) -> WebSocketStream<TcpStream> {
5358                let stream = TcpStream::connect(self.addr).await.unwrap();
5359                let mut socket = client_async(format!("ws://{}/", self.addr), stream)
5360                    .await
5361                    .unwrap()
5362                    .0;
5363                socket
5364                    .send(Message::Text(
5365                        json!({
5366                            "type": "subscribe",
5367                            "protocolVersion": 2,
5368                            "subscriptionId": "things",
5369                            "query": query,
5370                            "snapshot": {"enabled": snapshot},
5371                        })
5372                        .to_string()
5373                        .into(),
5374                    ))
5375                    .await
5376                    .unwrap();
5377                let ack = next_frame(&mut socket).await;
5378                assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
5379                assert_eq!(
5380                    ack["wholeEntities"],
5381                    json!(true),
5382                    "the ack advertises the whole-entity guarantee: {ack}"
5383                );
5384                socket
5385            }
5386        }
5387
5388        async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
5389            loop {
5390                let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
5391                    .await
5392                    .expect("the server answers within the timeout")
5393                    .expect("the stream stays open")
5394                    .expect("a readable frame");
5395                // Every frame here is far below the compression threshold.
5396                let bytes = match &message {
5397                    Message::Text(text) => text.as_bytes(),
5398                    Message::Binary(bytes) => bytes.as_ref(),
5399                    _ => continue,
5400                };
5401                return serde_json::from_slice(bytes).expect("frames are JSON");
5402            }
5403        }
5404
5405        /// The snapshot rows, in order, once the snapshot completes.
5406        async fn snapshot_keys(socket: &mut WebSocketStream<TcpStream>) -> Vec<String> {
5407            let mut keys = Vec::new();
5408            loop {
5409                let frame = next_frame(socket).await;
5410                assert_eq!(frame["op"], "snapshot", "expected a snapshot: {frame}");
5411                keys.extend(
5412                    frame["data"]
5413                        .as_array()
5414                        .unwrap()
5415                        .iter()
5416                        .map(|row| row["key"].as_str().unwrap().to_string()),
5417                );
5418                if frame["complete"] == json!(true) {
5419                    return keys;
5420                }
5421            }
5422        }
5423
5424        async fn seed(harness: &Harness) {
5425            for key in ["a", "b", "c"] {
5426                harness
5427                    .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
5428                    .await;
5429            }
5430        }
5431
5432        fn derived_index() -> ViewIndex {
5433            use crate::materialized_view::{
5434                CompareOp, FilterConfig, SortConfig, SortOrder, ViewPipeline,
5435            };
5436            let mut index = thing_index();
5437            for name in ["empty", "filtered", "sorted"] {
5438                let mut pipeline = ViewPipeline::default();
5439                if name != "empty" {
5440                    pipeline.filter = Some(FilterConfig {
5441                        field_path: vec!["owner".into()],
5442                        op: CompareOp::Eq,
5443                        value: json!("alice"),
5444                    });
5445                    pipeline.limit = Some(2);
5446                }
5447                if name == "sorted" {
5448                    pipeline.sort = Some(SortConfig {
5449                        field_path: vec!["amount".into()],
5450                        order: SortOrder::Desc,
5451                    });
5452                }
5453                index.add_spec(ViewSpec {
5454                    id: format!("Thing/{name}"),
5455                    source_view: Some("Thing/list".into()),
5456                    pipeline: Some(pipeline),
5457                    ..list_spec()
5458                });
5459            }
5460            index
5461        }
5462
5463        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5464        async fn unsorted_derived_idle_bootstrap_and_reconnect_are_authoritative() {
5465            let harness = Harness::start_with_index(
5466                WebSocketDeliveryConfig::default(),
5467                EntityCache::new(),
5468                derived_index(),
5469            )
5470            .await;
5471            harness
5472                .patch("a", json!({"owner":"alice", "amount":1}))
5473                .await;
5474            harness.patch("b", json!({"owner":"bob", "amount":2})).await;
5475            harness
5476                .patch("c", json!({"owner":"alice", "amount":3}))
5477                .await;
5478            harness
5479                .patch("d", json!({"owner":"alice", "amount":4}))
5480                .await;
5481            for _ in 0..2 {
5482                for (view, expected) in [
5483                    ("empty", vec!["d", "c", "b", "a"]),
5484                    ("filtered", vec!["d", "c"]),
5485                    ("sorted", vec!["d", "c"]),
5486                ] {
5487                    let mut socket = harness
5488                        .subscribe(json!({"view": format!("Thing/{view}"), "take":100}), true)
5489                        .await;
5490                    assert_eq!(snapshot_keys(&mut socket).await, expected);
5491                    socket.close(None).await.unwrap();
5492                }
5493            }
5494        }
5495
5496        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5497        async fn source_delete_and_recreation_clear_all_views_across_reconnect() {
5498            let harness = Harness::start_with_index(
5499                WebSocketDeliveryConfig::default(),
5500                EntityCache::new(),
5501                derived_index(),
5502            )
5503            .await;
5504            harness
5505                .patch("a", json!({"owner":"alice", "amount":1, "oldField":true}))
5506                .await;
5507            let mut list = harness
5508                .subscribe(json!({"view":"Thing/filtered"}), true)
5509                .await;
5510            assert_eq!(snapshot_keys(&mut list).await, ["a"]);
5511            let mut state = harness
5512                .subscribe(json!({"view":"Thing/state", "key":"a"}), true)
5513                .await;
5514            assert_eq!(snapshot_keys(&mut state).await, ["a"]);
5515            // Leaving a predicate is view eviction, not chain deletion.
5516            harness.patch("a", json!({"owner":"bob"})).await;
5517            assert_eq!(next_frame(&mut list).await["op"], "remove");
5518            harness.patch("a", json!({"owner":"alice"})).await;
5519            assert_eq!(next_frame(&mut list).await["op"], "upsert");
5520            let mut deletion = Mutation::delete("Thing", json!("a"));
5521            deletion.mark_account_position(AccountPosition::new(100, 1001));
5522            harness.mutate(deletion).await;
5523            assert_eq!(next_frame(&mut list).await["op"], "delete");
5524            // State patches may precede its deletion on the socket.
5525            loop {
5526                if next_frame(&mut state).await["op"] == "delete" {
5527                    break;
5528                }
5529            }
5530            for view in [
5531                "Thing/list",
5532                "Thing/empty",
5533                "Thing/filtered",
5534                "Thing/sorted",
5535                "Thing/state",
5536            ] {
5537                let mut socket = harness
5538                    .subscribe(json!({"view":view, "key":"a"}), true)
5539                    .await;
5540                assert!(snapshot_keys(&mut socket).await.is_empty(), "{view}");
5541                socket.close(None).await.unwrap();
5542            }
5543            // A whole resend or sparse change cannot recreate a deleted row.
5544            harness
5545                .patch("a", json!({"owner":"alice", "oldField":true}))
5546                .await;
5547            let mut stale = Mutation {
5548                export: "Thing".into(),
5549                key: json!("a"),
5550                patch: json!({"owner":"alice", "oldField":true}),
5551                append: vec![],
5552            };
5553            stale.mark_whole_entity();
5554            harness.mutate(stale).await;
5555            let mut socket = harness
5556                .subscribe(json!({"view":"Thing/filtered"}), true)
5557                .await;
5558            assert!(snapshot_keys(&mut socket).await.is_empty());
5559            socket.close(None).await.unwrap();
5560            let mut recreated = Mutation {
5561                export: "Thing".into(),
5562                key: json!("a"),
5563                patch: json!({"owner":"alice", "amount":9}),
5564                append: vec![],
5565            };
5566            recreated.mark_created();
5567            recreated.mark_account_position(AccountPosition::new(100, 1002));
5568            harness.mutate(recreated).await;
5569            for view in [
5570                "Thing/list",
5571                "Thing/empty",
5572                "Thing/filtered",
5573                "Thing/sorted",
5574                "Thing/state",
5575            ] {
5576                let mut socket = harness
5577                    .subscribe(json!({"view":view, "key":"a"}), true)
5578                    .await;
5579                let frame = next_frame(&mut socket).await;
5580                assert_eq!(frame["data"][0]["key"], "a");
5581                assert_eq!(frame["data"][0]["data"]["amount"], 9);
5582                assert!(
5583                    frame["data"][0]["data"].get("oldField").is_none(),
5584                    "{frame}"
5585                );
5586                socket.close(None).await.unwrap();
5587            }
5588        }
5589
5590        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5591        async fn resolver_offsets_do_not_suppress_later_instruction_updates() {
5592            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
5593            let mut created = Mutation {
5594                export: "Thing".into(),
5595                key: json!("a"),
5596                patch: json!({"value":"created"}),
5597                append: vec![],
5598            };
5599            created.mark_created();
5600            created.mark_account_position(AccountPosition::new(100, 9));
5601            harness
5602                .mutate_with_context(created, SlotContext::account(100, 9))
5603                .await;
5604
5605            harness
5606                .mutate_with_context(
5607                    Mutation {
5608                        export: "Thing".into(),
5609                        key: json!("a"),
5610                        patch: json!({"value":"resolver", "resolved":true}),
5611                        append: vec![],
5612                    },
5613                    SlotContext::resolver(100, 1_u64 << 63),
5614                )
5615                .await;
5616            harness
5617                .mutate_with_context(
5618                    Mutation {
5619                        export: "Thing".into(),
5620                        key: json!("a"),
5621                        patch: json!({"value":"instruction", "instruction":true}),
5622                        append: vec![],
5623                    },
5624                    SlotContext::instruction(100, 900),
5625                )
5626                .await;
5627
5628            let mut socket = harness
5629                .subscribe(json!({"view":"Thing/state", "key":"a"}), true)
5630                .await;
5631            let frame = next_frame(&mut socket).await;
5632            assert_eq!(frame["op"], "snapshot", "{frame}");
5633            let row = &frame["data"][0]["data"];
5634            assert_eq!(row["value"], "instruction", "{frame}");
5635            assert_eq!(row["resolved"], true, "{frame}");
5636            assert_eq!(row["instruction"], true, "{frame}");
5637
5638            // Recency still applies inside the instruction domain even after
5639            // activity from another source domain.
5640            harness
5641                .mutate_with_context(
5642                    Mutation {
5643                        export: "Thing".into(),
5644                        key: json!("a"),
5645                        patch: json!({"value":"stale", "stale":true}),
5646                        append: vec![],
5647                    },
5648                    SlotContext::instruction(100, 899),
5649                )
5650                .await;
5651            let mut reconnected = harness
5652                .subscribe(json!({"view":"Thing/state", "key":"a"}), true)
5653                .await;
5654            let frame = next_frame(&mut reconnected).await;
5655            let row = &frame["data"][0]["data"];
5656            assert_eq!(row["value"], "instruction", "{frame}");
5657            assert!(row.get("stale").is_none(), "{frame}");
5658        }
5659
5660        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5661        async fn projected_delete_and_recreation_between_reads_replace_state() {
5662            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
5663            harness.patch("a", json!({"old":true})).await;
5664            let mut socket = harness
5665                .subscribe(json!({"view":"Thing/state", "key":"a"}), true)
5666                .await;
5667            assert_eq!(snapshot_keys(&mut socket).await, ["a"]);
5668            let mut creation = Mutation {
5669                export: "Thing".into(),
5670                key: json!("a"),
5671                patch: json!({"fresh":true}),
5672                append: vec![],
5673            };
5674            creation.mark_created();
5675            // Both ready batches run before the watch receiver on this worker.
5676            for mutation in [Mutation::delete("Thing", json!("a")), creation] {
5677                let slot = harness
5678                    .slot
5679                    .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
5680                harness
5681                    .tx
5682                    .try_send(MutationBatch::with_slot_context(
5683                        vec![mutation].into_iter().collect(),
5684                        SlotContext::new(slot, 0),
5685                    ))
5686                    .unwrap();
5687            }
5688            let (ack, wait) = oneshot::channel();
5689            harness
5690                .tx
5691                .try_send(MutationBatch::flush_marker(ack))
5692                .unwrap();
5693            wait.await.unwrap();
5694            // Usually both batches land before the socket reads and arrive as
5695            // one upsert. Under load the socket can read between them, so the
5696            // delete may arrive on its own first. Either way the recreated
5697            // entity must not carry the deleted one's fields.
5698            let mut frame = next_frame(&mut socket).await;
5699            if frame["op"] == "delete" {
5700                assert_eq!(frame["key"], "a", "{frame}");
5701                frame = next_frame(&mut socket).await;
5702            }
5703            assert_eq!(frame["op"], "upsert", "{frame}");
5704            assert_eq!(frame["data"]["fresh"], true);
5705            assert!(frame["data"].get("old").is_none(), "{frame}");
5706        }
5707
5708        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5709        async fn a_change_to_a_key_cut_by_snapshot_limit_is_an_upsert() {
5710            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
5711            seed(&harness).await;
5712
5713            let mut socket = harness
5714                .subscribe(json!({"view": "Thing/list", "snapshotLimit": 1}), true)
5715                .await;
5716            // Newest first: only "c" was sent.
5717            assert_eq!(snapshot_keys(&mut socket).await, ["c"]);
5718
5719            harness.patch("a", json!({"count": 2})).await;
5720            let frame = next_frame(&mut socket).await;
5721            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
5722            assert_eq!(frame["op"], "upsert", "a key never sent arrives whole");
5723            assert_eq!(frame["data"]["name"], "thing-a");
5724            assert_eq!(frame["data"]["count"], 2);
5725
5726            // Now the client holds "a", so its next change is a patch again.
5727            harness.patch("a", json!({"count": 3})).await;
5728            let frame = next_frame(&mut socket).await;
5729            assert_eq!(frame["key"], "a");
5730            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5731            assert!(frame["data"].get("name").is_none());
5732
5733            // A key the snapshot did carry still rides its patch.
5734            harness.patch("c", json!({"count": 2})).await;
5735            let frame = next_frame(&mut socket).await;
5736            assert_eq!(frame["key"], "c");
5737            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5738            socket.close(None).await.ok();
5739        }
5740
5741        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5742        async fn a_list_without_a_snapshot_upserts_each_key_on_its_first_change() {
5743            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
5744            seed(&harness).await;
5745
5746            let mut socket = harness
5747                .subscribe(json!({"view": "Thing/list"}), false)
5748                .await;
5749            harness.patch("b", json!({"count": 2})).await;
5750            let frame = next_frame(&mut socket).await;
5751            assert_eq!(frame["key"], "b");
5752            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5753            assert_eq!(frame["data"]["name"], "thing-b");
5754
5755            harness.patch("b", json!({"count": 3})).await;
5756            let frame = next_frame(&mut socket).await;
5757            assert_eq!(frame["key"], "b");
5758            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5759            socket.close(None).await.ok();
5760        }
5761
5762        /// Eviction from the server's bounded cache is not a change to the
5763        /// entity: a client holding it keeps getting its patches, never a
5764        /// `remove`, and the whole entity replaces its copy when it returns.
5765        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5766        async fn a_state_holder_keeps_its_patches_when_the_cache_evicts_the_key() {
5767            let harness = Harness::start_with_cache(
5768                WebSocketDeliveryConfig::default(),
5769                EntityCache::with_config(crate::cache::EntityCacheConfig {
5770                    max_entities_per_view: 2,
5771                    ..Default::default()
5772                }),
5773            )
5774            .await;
5775            harness
5776                .patch("a", json!({"name": "thing-a", "count": 1}))
5777                .await;
5778            let mut socket = harness
5779                .subscribe(json!({"view": "Thing/state", "key": "a"}), true)
5780                .await;
5781            assert_eq!(snapshot_keys(&mut socket).await, ["a"]);
5782
5783            // "b" and "c" push "a" out of the cache.
5784            for key in ["b", "c"] {
5785                harness
5786                    .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
5787                    .await;
5788            }
5789            harness.patch("a", json!({"count": 2})).await;
5790            let frame = next_frame(&mut socket).await;
5791            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
5792            assert_eq!(frame["op"], "patch", "the holder merges it: {frame}");
5793            assert_eq!(frame["data"]["count"], 2);
5794
5795            // The whole entity comes back from the VM.
5796            let mut whole = Mutation {
5797                export: "Thing".to_string(),
5798                key: json!("a"),
5799                patch: json!({"name": "thing-a", "count": 2}),
5800                append: vec![],
5801            };
5802            whole.mark_whole_entity();
5803            harness.mutate(whole).await;
5804            let frame = next_frame(&mut socket).await;
5805            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5806            assert_eq!(frame["data"]["name"], "thing-a");
5807            assert_eq!(frame["data"]["count"], 2);
5808            assert!(frame["data"]
5809                .get(arete_interpreter::WHOLE_ENTITY_MARKER)
5810                .is_none());
5811
5812            harness.patch("a", json!({"count": 3})).await;
5813            let frame = next_frame(&mut socket).await;
5814            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5815            socket.close(None).await.ok();
5816        }
5817
5818        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5819        async fn a_state_view_without_a_snapshot_upserts_on_the_first_change() {
5820            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
5821            seed(&harness).await;
5822
5823            let mut socket = harness
5824                .subscribe(json!({"view": "Thing/state", "key": "a"}), false)
5825                .await;
5826            harness.patch("a", json!({"count": 2})).await;
5827            let frame = next_frame(&mut socket).await;
5828            assert_eq!(frame["key"], "a");
5829            assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
5830            assert_eq!(frame["data"]["name"], "thing-a");
5831            assert_eq!(frame["data"]["count"], 2);
5832
5833            harness.patch("a", json!({"count": 3})).await;
5834            let frame = next_frame(&mut socket).await;
5835            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
5836            socket.close(None).await.ok();
5837        }
5838
5839        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5840        async fn a_coalesced_list_sends_no_remove_for_a_key_it_never_sent() {
5841            let harness = Harness::start(WebSocketDeliveryConfig {
5842                collection_coalesce_ms: Some(10),
5843                ..Default::default()
5844            })
5845            .await;
5846            seed(&harness).await;
5847
5848            // A one-row window holding "c", which the disabled snapshot never
5849            // sent.
5850            let mut socket = harness
5851                .subscribe(json!({"view": "Thing/list", "take": 1}), false)
5852                .await;
5853            // "a" jumps to the front: it enters the window whole and "c",
5854            // never sent, leaves it without a `remove`.
5855            harness.patch("a", json!({"count": 2})).await;
5856            let frame = next_frame(&mut socket).await;
5857            assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
5858            assert_eq!(frame["op"], "upsert");
5859            assert_eq!(frame["data"]["name"], "thing-a");
5860
5861            // "c" re-enters whole; "a", held, leaves with a `remove`.
5862            harness.patch("c", json!({"count": 2})).await;
5863            let mut frames = [next_frame(&mut socket).await, next_frame(&mut socket).await];
5864            frames.sort_by_key(|frame| frame["key"].as_str().unwrap().to_string());
5865            assert_eq!(frames[0]["key"], "a");
5866            assert_eq!(frames[0]["op"], "remove");
5867            assert_eq!(frames[1]["key"], "c");
5868            assert_eq!(frames[1]["op"], "upsert");
5869            assert_eq!(frames[1]["data"]["name"], "thing-c");
5870            socket.close(None).await.ok();
5871        }
5872
5873        /// `(epoch, counter)` from an entity's `_version`.
5874        fn version_of(data: &Value) -> (String, u64) {
5875            let version = data["_version"]
5876                .as_str()
5877                .unwrap_or_else(|| panic!("no _version in {data}"));
5878            let (epoch, counter) = version.split_once(':').expect("epoch:counter");
5879            (epoch.to_string(), counter.parse().expect("decimal counter"))
5880        }
5881
5882        /// Every entity a subscriber receives carries the projector's
5883        /// `_version`: a snapshot row, a forwarded patch, and a state
5884        /// subscription's catch-up, which sends the cached entity after the
5885        /// bus overwrote a patch.
5886        // One worker, and a projector the cooperative budget cannot interrupt:
5887        // within a batch it awaits only locks no other task holds, so it
5888        // publishes both changes before the subscriber can run, and the
5889        // second overwrites the first on the state bus.
5890        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
5891        async fn snapshot_rows_patches_and_catch_ups_carry_the_version() {
5892            let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
5893            harness.patch("7", json!({"name": "thing-7"})).await;
5894
5895            let mut list = harness.subscribe(json!({"view": "Thing/list"}), true).await;
5896            let snapshot = next_frame(&mut list).await;
5897            assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
5898            let (epoch, _) = version_of(&snapshot["data"][0]["data"]);
5899
5900            let mut state = harness
5901                .subscribe(json!({"view": "Thing/state", "key": "7"}), true)
5902                .await;
5903            let snapshot = next_frame(&mut state).await;
5904            assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
5905            let (state_epoch, seeded) = version_of(&snapshot["data"][0]["data"]);
5906            assert_eq!(state_epoch, epoch, "one projector, one epoch");
5907
5908            harness.patch("7", json!({"count": 1})).await;
5909            let patch = next_frame(&mut list).await;
5910            assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
5911            assert_eq!(version_of(&patch["data"]).0, epoch);
5912            let patch = next_frame(&mut state).await;
5913            assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
5914            let (_, patched) = version_of(&patch["data"]);
5915            assert!(patched > seeded, "{patched} after {seeded}");
5916
5917            // Two changes to the key in one batch.
5918            let slot = harness
5919                .slot
5920                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
5921            let mutation = |patch: Value| Mutation {
5922                export: "Thing".to_string(),
5923                key: json!("7"),
5924                patch,
5925                append: vec![],
5926            };
5927            harness
5928                .tx
5929                .send(MutationBatch::with_slot_context(
5930                    vec![
5931                        mutation(json!({"count": 2})),
5932                        mutation(json!({"flag": true})),
5933                    ]
5934                    .into_iter()
5935                    .collect(),
5936                    SlotContext::new(slot, 0),
5937                ))
5938                .await
5939                .unwrap();
5940
5941            let catch_up = next_frame(&mut state).await;
5942            assert_eq!(catch_up["op"], "patch", "unexpected frame: {catch_up}");
5943            assert!(catch_up.get("seq").is_none(), "a catch-up: {catch_up}");
5944            assert_eq!(catch_up["data"]["count"], 2);
5945            assert_eq!(catch_up["data"]["flag"], true);
5946            let (_, caught_up) = version_of(&catch_up["data"]);
5947            assert!(caught_up > patched, "{caught_up} after {patched}");
5948
5949            list.close(None).await.ok();
5950            state.close(None).await.ok();
5951        }
5952    }
5953
5954    /// Session tokens that expire while their socket is open.
5955    mod session_expiry {
5956        use super::*;
5957        use crate::websocket::auth::SignedSessionAuthPlugin;
5958        use arete_auth::{KeyClass, SessionClaims, SigningKey, TokenSigner, TokenVerifier};
5959        use futures_util::{SinkExt, StreamExt};
5960        use std::time::Duration;
5961        use tokio::net::{TcpListener, TcpStream};
5962        use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
5963        use tokio_tungstenite::tungstenite::protocol::CloseFrame;
5964        use tokio_tungstenite::tungstenite::Message;
5965        use tokio_tungstenite::{client_async, WebSocketStream};
5966
5967        #[derive(Default)]
5968        struct RecordingUsageEmitter {
5969            events: tokio::sync::Mutex<Vec<WebSocketUsageEvent>>,
5970        }
5971
5972        #[async_trait::async_trait]
5973        impl WebSocketUsageEmitter for RecordingUsageEmitter {
5974            async fn emit(&self, event: WebSocketUsageEvent) {
5975                self.events.lock().await.push(event);
5976            }
5977        }
5978
5979        struct Server {
5980            addr: SocketAddr,
5981            signer: TokenSigner,
5982            usage: Arc<RecordingUsageEmitter>,
5983            metrics: Arc<DeliveryProbe>,
5984        }
5985
5986        impl Server {
5987            async fn start() -> Self {
5988                let signing_key = SigningKey::generate();
5989                let verifier =
5990                    TokenVerifier::new(signing_key.verifying_key(), "test-issuer", "test-audience");
5991                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5992                let addr = listener.local_addr().unwrap();
5993                let usage = Arc::new(RecordingUsageEmitter::default());
5994                let server = WebSocketServer::new(
5995                    addr,
5996                    BusManager::new(),
5997                    EntityCache::new(),
5998                    Arc::new(ViewIndex::new()),
5999                    #[cfg(feature = "otel")]
6000                    None,
6001                )
6002                .with_auth_plugin(Arc::new(SignedSessionAuthPlugin::new(verifier)))
6003                .with_usage_emitter(usage.clone());
6004                let metrics = Arc::new(DeliveryProbe::default());
6005                let (acceptor, _cleanup) = server.into_acceptor();
6006                let acceptor = acceptor.with_delivery_probe(metrics.clone());
6007                tokio::spawn(async move { acceptor.serve_listener(listener).await });
6008                Self {
6009                    addr,
6010                    signer: TokenSigner::new(signing_key, "test-issuer"),
6011                    usage,
6012                    metrics,
6013                }
6014            }
6015
6016            fn token(&self, ttl_seconds: u64) -> String {
6017                let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
6018                    .with_scope("read")
6019                    .with_key_class(KeyClass::Secret)
6020                    .with_ttl(ttl_seconds)
6021                    .build();
6022                self.signer.sign(claims).unwrap()
6023            }
6024
6025            fn account_token(&self, ttl_seconds: u64, account: &str) -> String {
6026                let claims = SessionClaims::builder("test-issuer", account, "test-audience")
6027                    .with_scope("read")
6028                    .with_key_class(KeyClass::Secret)
6029                    .with_metering_key(account)
6030                    .with_plan("agent_trial")
6031                    .with_actor_key(account)
6032                    .with_account_key(account)
6033                    .with_consumer_key(format!("consumer:{account}"))
6034                    .with_policy_version(3)
6035                    .with_account_limits(Default::default())
6036                    .with_ttl(ttl_seconds)
6037                    .build();
6038                self.signer.sign(claims).unwrap()
6039            }
6040
6041            async fn connect(&self, token: &str) -> WebSocketStream<TcpStream> {
6042                let stream = TcpStream::connect(self.addr).await.unwrap();
6043                client_async(format!("ws://{}/?hs_token={token}", self.addr), stream)
6044                    .await
6045                    .unwrap()
6046                    .0
6047            }
6048        }
6049
6050        async fn send_json(socket: &mut WebSocketStream<TcpStream>, message: Value) {
6051            socket
6052                .send(Message::Text(message.to_string().into()))
6053                .await
6054                .unwrap();
6055        }
6056
6057        /// The close frame, if the server closes the socket within `wait`.
6058        async fn close_within(
6059            socket: &mut WebSocketStream<TcpStream>,
6060            wait: Duration,
6061        ) -> Option<Option<CloseFrame>> {
6062            tokio::time::timeout(wait, async {
6063                while let Some(Ok(message)) = socket.next().await {
6064                    if let Message::Close(frame) = message {
6065                        return Some(frame);
6066                    }
6067                }
6068                Some(None)
6069            })
6070            .await
6071            .ok()
6072            .flatten()
6073        }
6074
6075        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6076        async fn an_expired_session_is_closed_with_the_reason() {
6077            let server = Server::start().await;
6078            let mut socket = server.connect(&server.token(2)).await;
6079
6080            tokio::time::sleep(Duration::from_secs(3)).await;
6081            send_json(&mut socket, json!({"type": "ping"})).await;
6082
6083            let frame = close_within(&mut socket, Duration::from_secs(5))
6084                .await
6085                .expect("the server closes the expired session")
6086                .expect("the close frame carries a reason");
6087            assert_eq!(frame.code, CloseCode::Policy);
6088            assert_eq!(
6089                frame.reason.as_str(),
6090                "token-expired: Authentication token expired"
6091            );
6092        }
6093
6094        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6095        async fn a_session_refreshed_in_band_outlives_its_first_token() {
6096            let server = Server::start().await;
6097            let mut socket = server.connect(&server.token(2)).await;
6098
6099            send_json(
6100                &mut socket,
6101                json!({"type": "refresh_auth", "token": server.token(3_600)}),
6102            )
6103            .await;
6104            let reply = tokio::time::timeout(Duration::from_secs(5), async {
6105                while let Some(Ok(message)) = socket.next().await {
6106                    if let Message::Text(text) = message {
6107                        return serde_json::from_str::<Value>(text.as_str()).ok();
6108                    }
6109                }
6110                None
6111            })
6112            .await
6113            .expect("the server answers the refresh")
6114            .expect("the answer is JSON");
6115            assert_eq!(reply["success"], true, "refresh accepted: {reply}");
6116
6117            tokio::time::sleep(Duration::from_secs(3)).await;
6118            send_json(&mut socket, json!({"type": "ping"})).await;
6119            assert!(
6120                close_within(&mut socket, Duration::from_millis(1_500))
6121                    .await
6122                    .is_none(),
6123                "the socket stays open on the refreshed token"
6124            );
6125        }
6126
6127        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6128        async fn usage_after_auth_refresh_uses_the_refreshed_account() {
6129            let server = Server::start().await;
6130            let mut socket = server
6131                .connect(&server.account_token(3_600, "account:1"))
6132                .await;
6133
6134            send_json(
6135                &mut socket,
6136                json!({
6137                    "type": "refresh_auth",
6138                    "token": server.account_token(3_600, "account:2")
6139                }),
6140            )
6141            .await;
6142            let reply = tokio::time::timeout(Duration::from_secs(5), async {
6143                while let Some(Ok(message)) = socket.next().await {
6144                    if let Message::Text(text) = message {
6145                        return serde_json::from_str::<Value>(text.as_str()).ok();
6146                    }
6147                }
6148                None
6149            })
6150            .await
6151            .expect("the server answers the refresh")
6152            .expect("the answer is JSON");
6153            assert_eq!(reply["success"], true, "refresh accepted: {reply}");
6154            socket.close(None).await.unwrap();
6155
6156            let account = tokio::time::timeout(Duration::from_secs(5), async {
6157                loop {
6158                    let events = server.usage.events.lock().await;
6159                    if let Some(account) = events.iter().find_map(|event| match event {
6160                        WebSocketUsageEvent::ConnectionClosed { identity, .. } => {
6161                            identity.account_key.clone()
6162                        }
6163                        _ => None,
6164                    }) {
6165                        return account;
6166                    }
6167                    drop(events);
6168                    tokio::task::yield_now().await;
6169                }
6170            })
6171            .await
6172            .expect("connection close usage event");
6173            assert_eq!(account, "account:2");
6174            let active = server
6175                .metrics
6176                .active_connections
6177                .lock()
6178                .expect("delivery probe lock poisoned");
6179            assert_eq!(active.get("account:1"), Some(&0));
6180            assert_eq!(active.get("account:2"), None);
6181        }
6182    }
6183}