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