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