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