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