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