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