1use crate::bus::{BusManager, BusMessage, StateUpdate};
2use crate::cache::{cmp_seq, EntityCache, SnapshotBatchConfig};
3use crate::compression::maybe_compress;
4use crate::view::{ViewIndex, ViewSpec};
5use crate::websocket::auth::{
6 AuthContext, AuthDecision, AuthDeny, ConnectionAuthRequest, WebSocketAuthPlugin,
7};
8use crate::websocket::client_manager::{ClientManager, RateLimitConfig};
9use crate::websocket::frame::{
10 apply_wire_format, Frame, Mode, SnapshotEntity, SnapshotFrame, SortConfig, SortOrder,
11 SubscribedFrame, UnsubscribedFrame,
12};
13use crate::websocket::subscription::{
14 ClientMessage, RefreshAuthRequest, RefreshAuthResponse, SocketIssueMessage, Subscription,
15 SubscriptionQuery, Unsubscription, PROTOCOL_VERSION,
16};
17use crate::websocket::usage::{WebSocketUsageEmitter, WebSocketUsageEvent};
18use crate::WebSocketDeliveryConfig;
19use anyhow::Result;
20use bytes::Bytes;
21use futures_util::StreamExt;
22use serde::Serialize;
23use serde_json::Value;
24use std::collections::{HashMap, HashSet, VecDeque};
25use std::future::Future;
26use std::net::SocketAddr;
27use std::sync::Arc;
28use std::time::{Duration, Instant};
29use tokio::net::{TcpListener, TcpStream};
30use tokio::sync::{broadcast, watch};
31use tokio_tungstenite::{
32 accept_hdr_async,
33 tungstenite::{
34 handshake::server::{ErrorResponse as HandshakeErrorResponse, Request, Response},
35 http::{header::CONTENT_TYPE, StatusCode},
36 Error as WsError,
37 },
38};
39use tokio_util::sync::CancellationToken;
40use tokio_util::task::TaskTracker;
41use tracing::{debug, error, info, info_span, warn, Instrument};
42use uuid::Uuid;
43
44#[cfg(feature = "otel")]
45use crate::metrics::Metrics;
46
47#[derive(Clone, Default)]
48struct WsMetrics {
49 #[cfg(feature = "otel")]
50 inner: Option<Arc<Metrics>>,
51 #[cfg(test)]
52 probe: Option<Arc<DeliveryProbe>>,
53}
54
55#[cfg(test)]
59#[derive(Debug, Default)]
60pub(crate) struct DeliveryProbe {
61 pub(crate) messages_sent: std::sync::atomic::AtomicU64,
63 pub(crate) lag_events: std::sync::atomic::AtomicU64,
65 pub(crate) dropped_updates: std::sync::atomic::AtomicU64,
67 pub(crate) resnapshots: std::sync::atomic::AtomicU64,
69 pub(crate) coalesced_updates: std::sync::atomic::AtomicU64,
71 pub(crate) coalesced_flushes: std::sync::atomic::AtomicU64,
73 pub(crate) delivery_stopped: std::sync::Mutex<std::collections::BTreeMap<&'static str, u64>>,
75}
76
77#[cfg(test)]
78impl DeliveryProbe {
79 fn add(counter: &std::sync::atomic::AtomicU64, value: u64) {
80 counter.fetch_add(value, std::sync::atomic::Ordering::Relaxed);
81 }
82}
83
84impl WsMetrics {
85 #[cfg(feature = "otel")]
86 fn new(inner: Option<Arc<Metrics>>) -> Self {
87 Self {
88 inner,
89 #[cfg(test)]
90 probe: None,
91 }
92 }
93
94 #[cfg(test)]
95 fn probe(&self, record: impl FnOnce(&DeliveryProbe)) {
96 if let Some(probe) = &self.probe {
97 record(probe);
98 }
99 }
100
101 fn connection_opened(&self, metering_key: Option<&str>) {
102 #[cfg(not(feature = "otel"))]
103 let _ = metering_key;
104 #[cfg(feature = "otel")]
105 if let Some(metrics) = &self.inner {
106 if let Some(metering_key) = metering_key {
107 metrics.record_ws_connection_with_metering(metering_key);
108 } else {
109 metrics.record_ws_connection();
110 }
111 }
112 }
113
114 fn connection_closed(&self, duration_secs: f64, metering_key: Option<&str>) {
115 #[cfg(not(feature = "otel"))]
116 let _ = (duration_secs, metering_key);
117 #[cfg(feature = "otel")]
118 if let Some(metrics) = &self.inner {
119 if let Some(metering_key) = metering_key {
120 metrics.record_ws_disconnection_with_metering(duration_secs, metering_key);
121 } else {
122 metrics.record_ws_disconnection(duration_secs);
123 }
124 }
125 }
126
127 fn message_received(&self, metering_key: Option<&str>) {
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_message_received_with_metering(metering_key);
134 } else {
135 metrics.record_ws_message_received();
136 }
137 }
138 }
139
140 fn message_sent(&self) {
141 #[cfg(test)]
142 self.probe(|probe| DeliveryProbe::add(&probe.messages_sent, 1));
143 #[cfg(feature = "otel")]
144 if let Some(metrics) = &self.inner {
145 metrics.record_ws_message_sent();
146 }
147 }
148
149 fn subscription_created(&self, view: &str, metering_key: Option<&str>) {
150 #[cfg(not(feature = "otel"))]
151 let _ = (view, metering_key);
152 #[cfg(feature = "otel")]
153 if let Some(metrics) = &self.inner {
154 if let Some(metering_key) = metering_key {
155 metrics.record_subscription_created_with_metering(view, metering_key);
156 } else {
157 metrics.record_subscription_created(view);
158 }
159 }
160 }
161
162 fn subscription_removed(&self, view: &str, metering_key: Option<&str>) {
163 #[cfg(not(feature = "otel"))]
164 let _ = (view, metering_key);
165 #[cfg(feature = "otel")]
166 if let Some(metrics) = &self.inner {
167 if let Some(metering_key) = metering_key {
168 metrics.record_subscription_removed_with_metering(view, metering_key);
169 } else {
170 metrics.record_subscription_removed(view);
171 }
172 }
173 }
174
175 fn protocol_error(&self, code: &str) {
176 #[cfg(not(feature = "otel"))]
177 let _ = code;
178 #[cfg(feature = "otel")]
179 if let Some(metrics) = &self.inner {
180 metrics.record_ws_protocol_error(code);
181 }
182 }
183
184 fn subscription_lagged(&self, view: &str, skipped: u64) {
185 #[cfg(test)]
186 self.probe(|probe| {
187 DeliveryProbe::add(&probe.lag_events, 1);
188 DeliveryProbe::add(&probe.dropped_updates, skipped);
189 });
190 #[cfg(not(feature = "otel"))]
191 let _ = (view, skipped);
192 #[cfg(feature = "otel")]
193 if let Some(metrics) = &self.inner {
194 metrics.record_ws_subscription_lagged(view, skipped);
195 }
196 }
197
198 fn subscription_resnapshot(&self, view: &str) {
199 #[cfg(test)]
200 self.probe(|probe| DeliveryProbe::add(&probe.resnapshots, 1));
201 #[cfg(not(feature = "otel"))]
202 let _ = view;
203 #[cfg(feature = "otel")]
204 if let Some(metrics) = &self.inner {
205 metrics.record_ws_subscription_resnapshot(view);
206 }
207 }
208
209 fn collection_coalesced(&self, view: &str, updates: u64) {
210 #[cfg(test)]
211 self.probe(|probe| {
212 DeliveryProbe::add(&probe.coalesced_updates, updates);
213 DeliveryProbe::add(&probe.coalesced_flushes, 1);
214 });
215 #[cfg(not(feature = "otel"))]
216 let _ = (view, updates);
217 #[cfg(feature = "otel")]
218 if let Some(metrics) = &self.inner {
219 metrics.record_ws_collection_coalesced(view, updates);
220 }
221 }
222
223 fn delivery_stopped(&self, view: &str, reason: &'static str) {
224 #[cfg(test)]
225 self.probe(|probe| {
226 *probe
227 .delivery_stopped
228 .lock()
229 .expect("delivery probe lock poisoned")
230 .entry(reason)
231 .or_default() += 1;
232 });
233 #[cfg(not(feature = "otel"))]
234 let _ = (view, reason);
235 #[cfg(feature = "otel")]
236 if let Some(metrics) = &self.inner {
237 metrics.record_ws_delivery_stopped(view, reason);
238 }
239 }
240}
241
242async fn handle_refresh_auth(
243 client_id: Uuid,
244 refresh_req: &RefreshAuthRequest,
245 client_manager: &ClientManager,
246 auth_plugin: &Arc<dyn WebSocketAuthPlugin>,
247) {
248 let refresh_result: Result<AuthContext, String> = if let Some(signed_plugin) = auth_plugin
249 .as_any()
250 .downcast_ref::<crate::websocket::auth::SignedSessionAuthPlugin>()
251 {
252 signed_plugin
253 .verify_refresh_token(&refresh_req.token)
254 .await
255 .map_err(|error| error.reason)
256 } else {
257 Err("In-band auth refresh not supported with current auth plugin".to_string())
258 };
259
260 let response = match refresh_result {
261 Ok(new_context) => {
262 let expires_at = new_context.expires_at;
263 if client_manager.update_client_auth(client_id, new_context) {
264 RefreshAuthResponse {
265 success: true,
266 error: None,
267 expires_at: Some(expires_at),
268 }
269 } else {
270 RefreshAuthResponse {
271 success: false,
272 error: Some("client-not-found".to_string()),
273 expires_at: None,
274 }
275 }
276 }
277 Err(error) => {
278 let code = if error.contains("expired") {
279 "token-expired"
280 } else if error.contains("signature") {
281 "token-invalid-signature"
282 } else if error.contains("issuer") {
283 "token-invalid-issuer"
284 } else if error.contains("audience") {
285 "token-invalid-audience"
286 } else {
287 "token-invalid"
288 };
289 RefreshAuthResponse {
290 success: false,
291 error: Some(code.to_string()),
292 expires_at: None,
293 }
294 }
295 };
296
297 if let Ok(json) = serde_json::to_string(&response) {
298 let _ = client_manager.send_text_to_client(client_id, json).await;
299 }
300}
301
302async fn send_socket_issue(
303 client_id: Uuid,
304 client_manager: &ClientManager,
305 deny: &AuthDeny,
306 fatal: bool,
307 subscription_id: Option<String>,
308) {
309 let message = SocketIssueMessage::from_auth_deny(deny, fatal, subscription_id);
310 if let Ok(json) = serde_json::to_string(&message) {
311 let _ = client_manager.send_text_to_client(client_id, json).await;
312 }
313}
314
315async fn send_protocol_issue(
316 client_id: Uuid,
317 client_manager: &ClientManager,
318 metrics: &WsMetrics,
319 subscription_id: Option<String>,
320 code: &str,
321 message: impl Into<String>,
322) {
323 metrics.protocol_error(code);
324 let issue = SocketIssueMessage::protocol(subscription_id, code, message);
325 if let Ok(json) = serde_json::to_string(&issue) {
326 let _ = client_manager.send_text_to_client(client_id, json).await;
327 }
328}
329
330async fn send_prepared_issue(
332 client_id: Uuid,
333 client_manager: &ClientManager,
334 metrics: &WsMetrics,
335 issue: SocketIssueMessage,
336) {
337 metrics.protocol_error(&issue.code);
338 if let Ok(json) = serde_json::to_string(&issue) {
339 let _ = client_manager.send_text_to_client(client_id, json).await;
340 }
341}
342
343fn key_class_label(key_class: arete_auth::KeyClass) -> &'static str {
344 match key_class {
345 arete_auth::KeyClass::Secret => "secret",
346 arete_auth::KeyClass::Publishable => "publishable",
347 }
348}
349
350fn emit_usage_event(
351 usage_emitter: &Option<Arc<dyn WebSocketUsageEmitter>>,
352 event: WebSocketUsageEvent,
353) {
354 if let Some(emitter) = usage_emitter.clone() {
355 tokio::spawn(async move {
356 emitter.emit(event).await;
357 });
358 }
359}
360
361fn usage_identity(
362 auth_context: Option<&AuthContext>,
363) -> (
364 Option<String>,
365 Option<String>,
366 Option<String>,
367 Option<String>,
368) {
369 match auth_context {
370 Some(context) => (
371 Some(context.metering_key.clone()),
372 Some(context.subject.clone()),
373 Some(key_class_label(context.key_class).to_string()),
374 context.deployment_id.clone(),
375 ),
376 None => (None, None, None, None),
377 }
378}
379
380fn emit_update_sent_for_client(
381 usage_emitter: &Option<Arc<dyn WebSocketUsageEmitter>>,
382 client_manager: &ClientManager,
383 client_id: Uuid,
384 view_id: &str,
385 bytes: usize,
386) {
387 let auth_context = client_manager.get_auth_context(client_id);
388 let (metering_key, subject, _, deployment_id) = usage_identity(auth_context.as_ref());
389 emit_usage_event(
390 usage_emitter,
391 WebSocketUsageEvent::UpdateSent {
392 client_id: client_id.to_string(),
393 deployment_id,
394 metering_key,
395 subject,
396 view_id: view_id.to_string(),
397 messages: 1,
398 bytes: bytes as u64,
399 },
400 );
401}
402
403#[derive(Clone)]
404struct SubscriptionContext {
405 client_id: Uuid,
406 client_manager: ClientManager,
407 bus_manager: BusManager,
408 entity_cache: EntityCache,
409 view_index: Arc<ViewIndex>,
410 usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
411 journal: Option<Arc<crate::journal::EventJournal>>,
412 metrics: WsMetrics,
413 delivery: WebSocketDeliveryConfig,
414 shutdown: CancellationToken,
417}
418
419pub struct WebSocketServer {
420 bind_addr: SocketAddr,
421 client_manager: ClientManager,
422 bus_manager: BusManager,
423 entity_cache: EntityCache,
424 view_index: Arc<ViewIndex>,
425 max_clients: usize,
426 auth_plugin: Arc<dyn WebSocketAuthPlugin>,
427 usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
428 rate_limit_config: Option<RateLimitConfig>,
429 journal: Option<Arc<crate::journal::EventJournal>>,
430 delivery: WebSocketDeliveryConfig,
431 #[cfg(feature = "otel")]
432 metrics: Option<Arc<Metrics>>,
433}
434
435impl WebSocketServer {
436 #[cfg(feature = "otel")]
437 pub fn new(
438 bind_addr: SocketAddr,
439 bus_manager: BusManager,
440 entity_cache: EntityCache,
441 view_index: Arc<ViewIndex>,
442 metrics: Option<Arc<Metrics>>,
443 ) -> Self {
444 Self {
445 bind_addr,
446 client_manager: ClientManager::new(),
447 bus_manager,
448 entity_cache,
449 view_index,
450 max_clients: 10_000,
451 auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
452 usage_emitter: None,
453 rate_limit_config: None,
454 journal: None,
455 delivery: WebSocketDeliveryConfig::default(),
456 metrics,
457 }
458 }
459
460 #[cfg(not(feature = "otel"))]
461 pub fn new(
462 bind_addr: SocketAddr,
463 bus_manager: BusManager,
464 entity_cache: EntityCache,
465 view_index: Arc<ViewIndex>,
466 ) -> Self {
467 Self {
468 bind_addr,
469 client_manager: ClientManager::new(),
470 bus_manager,
471 entity_cache,
472 view_index,
473 max_clients: 10_000,
474 auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
475 usage_emitter: None,
476 rate_limit_config: None,
477 journal: None,
478 delivery: WebSocketDeliveryConfig::default(),
479 }
480 }
481
482 pub fn with_max_clients(mut self, max_clients: usize) -> Self {
483 self.max_clients = max_clients;
484 self
485 }
486
487 pub fn with_auth_plugin(mut self, auth_plugin: Arc<dyn WebSocketAuthPlugin>) -> Self {
488 self.auth_plugin = auth_plugin;
489 self
490 }
491
492 pub fn with_usage_emitter(mut self, usage_emitter: Arc<dyn WebSocketUsageEmitter>) -> Self {
493 self.usage_emitter = Some(usage_emitter);
494 self
495 }
496
497 pub fn with_journal(mut self, journal: Arc<crate::journal::EventJournal>) -> Self {
499 self.journal = Some(journal);
500 self
501 }
502
503 pub fn with_rate_limit_config(mut self, config: RateLimitConfig) -> Self {
504 self.rate_limit_config = Some(config);
505 self
506 }
507
508 pub fn with_delivery_config(mut self, config: WebSocketDeliveryConfig) -> Self {
509 self.delivery = config;
510 self
511 }
512
513 pub async fn start(self) -> Result<()> {
517 info!(
518 "Starting WebSocket server on {} (max_clients: {})",
519 self.bind_addr, self.max_clients
520 );
521 let listener = TcpListener::bind(&self.bind_addr).await?;
522 let (acceptor, _cleanup) = self.into_acceptor();
523 acceptor.serve_listener(listener).await
524 }
525
526 pub(crate) fn into_acceptor(self) -> (ConnectionAcceptor, tokio::task::JoinHandle<()>) {
532 let client_manager = self
533 .rate_limit_config
534 .map(ClientManager::with_config)
535 .unwrap_or(self.client_manager);
536 let cleanup = client_manager.start_cleanup_task();
537
538 #[cfg(feature = "otel")]
539 let metrics = WsMetrics::new(self.metrics.clone());
540 #[cfg(not(feature = "otel"))]
541 let metrics = WsMetrics::default();
542
543 let acceptor = ConnectionAcceptor {
544 client_manager,
545 bus_manager: self.bus_manager,
546 entity_cache: self.entity_cache,
547 view_index: self.view_index,
548 max_clients: self.max_clients,
549 auth_plugin: self.auth_plugin,
550 usage_emitter: self.usage_emitter,
551 journal: self.journal,
552 delivery: self.delivery,
553 metrics,
554 shutdown: CancellationToken::new(),
555 sessions: TaskTracker::new(),
556 };
557 (acceptor, cleanup)
558 }
559}
560
561#[derive(Clone)]
569pub(crate) struct ConnectionAcceptor {
570 client_manager: ClientManager,
571 bus_manager: BusManager,
572 entity_cache: EntityCache,
573 view_index: Arc<ViewIndex>,
574 max_clients: usize,
575 auth_plugin: Arc<dyn WebSocketAuthPlugin>,
576 usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
577 journal: Option<Arc<crate::journal::EventJournal>>,
578 delivery: WebSocketDeliveryConfig,
579 metrics: WsMetrics,
580 shutdown: CancellationToken,
581 sessions: TaskTracker,
585}
586
587impl ConnectionAcceptor {
588 #[cfg(test)]
591 pub(crate) fn with_delivery_probe(mut self, probe: Arc<DeliveryProbe>) -> Self {
592 self.metrics.probe = Some(probe);
593 self
594 }
595
596 pub(crate) fn client_count(&self) -> usize {
598 self.client_manager.client_count()
599 }
600
601 pub(crate) fn shutdown(&self) {
607 self.shutdown.cancel();
608 self.sessions.close();
609 }
610
611 pub(crate) async fn wait_for_sessions(&self) {
614 self.sessions.wait().await;
615 }
616
617 pub(crate) async fn serve(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
623 if self.client_manager.client_count() >= self.max_clients {
624 warn!(
625 "Rejecting connection from {}: max clients reached",
626 remote_addr
627 );
628 return Ok(());
629 }
630
631 let context = SubscriptionContext {
632 client_id: Uuid::nil(),
633 client_manager: self.client_manager.clone(),
634 bus_manager: self.bus_manager.clone(),
635 entity_cache: self.entity_cache.clone(),
636 view_index: self.view_index.clone(),
637 usage_emitter: self.usage_emitter.clone(),
638 journal: self.journal.clone(),
639 delivery: self.delivery.clone(),
640 metrics: self.metrics.clone(),
641 shutdown: self.shutdown.clone(),
642 };
643 handle_connection(stream, context, remote_addr, self.auth_plugin.clone()).await
644 }
645
646 pub(crate) async fn serve_listener(self, listener: TcpListener) -> Result<()> {
649 loop {
650 let accepted = tokio::select! {
651 _ = self.shutdown.cancelled() => return Ok(()),
652 accepted = listener.accept() => accepted,
653 };
654 match accepted {
655 Ok((stream, addr)) => {
656 let acceptor = self.clone();
657 self.sessions.spawn(
658 async move {
659 if let Err(error) = acceptor.serve(stream, addr).await {
660 error!("WebSocket connection error: {}", error);
661 }
662 }
663 .instrument(info_span!("ws.connection", %addr)),
664 );
665 }
666 Err(error) => error!("Failed to accept connection: {}", error),
667 }
668 }
669 }
670}
671
672#[derive(Debug, Clone)]
673struct HandshakeReject {
674 status: StatusCode,
675 body: crate::websocket::auth::ErrorResponse,
676 error_code: String,
677 retry_after_secs: Option<u64>,
678}
679
680impl HandshakeReject {
681 fn from_deny(deny: &AuthDeny) -> Self {
682 let retry_after_secs = match deny.retry_policy {
683 crate::websocket::auth::RetryPolicy::RetryAfter(duration) => Some(duration.as_secs()),
684 _ => None,
685 };
686 Self {
687 status: StatusCode::from_u16(deny.http_status).unwrap_or(StatusCode::UNAUTHORIZED),
688 body: deny.to_error_response(),
689 error_code: deny.code.to_string(),
690 retry_after_secs,
691 }
692 }
693}
694
695fn build_handshake_error_response(
696 response: &Response,
697 reject: &HandshakeReject,
698) -> HandshakeErrorResponse {
699 let mut builder = Response::builder()
700 .status(reject.status)
701 .version(response.version())
702 .header(CONTENT_TYPE, "application/json; charset=utf-8")
703 .header("X-Error-Code", &reject.error_code)
704 .header("Cache-Control", "no-store");
705 if let Some(retry_after_secs) = reject.retry_after_secs {
706 builder = builder.header("Retry-After", retry_after_secs.to_string());
707 }
708 let body = serde_json::to_string(&reject.body).unwrap_or_else(|_| {
709 format!(
710 r#"{{"error":"{}","message":"{}","code":"{}","retryable":false}}"#,
711 reject.body.error, reject.body.message, reject.body.code
712 )
713 });
714 builder
715 .body(Some(body))
716 .expect("handshake rejection response should build")
717}
718
719#[allow(clippy::result_large_err)]
720async fn accept_authorized_connection(
721 stream: TcpStream,
722 remote_addr: SocketAddr,
723 auth_plugin: Arc<dyn WebSocketAuthPlugin>,
724 client_manager: ClientManager,
725) -> Result<Option<(tokio_tungstenite::WebSocketStream<TcpStream>, AuthContext)>> {
726 use std::sync::Mutex;
727
728 let capture: Arc<Mutex<Option<Result<AuthContext, HandshakeReject>>>> =
729 Arc::new(Mutex::new(None));
730 let capture_ref = capture.clone();
731 let auth_plugin_ref = auth_plugin.clone();
732 let manager_ref = client_manager.clone();
733
734 let handshake_result = accept_hdr_async(stream, move |request: &Request, response| {
735 let request = ConnectionAuthRequest::from_http_request(remote_addr, request);
736 let result = tokio::task::block_in_place(|| {
737 tokio::runtime::Handle::current().block_on(async {
738 match auth_plugin_ref.authorize(&request).await {
739 AuthDecision::Allow(context) => manager_ref
740 .check_connection_allowed(remote_addr, &Some(context.clone()))
741 .await
742 .map(|()| context)
743 .map_err(|deny| HandshakeReject::from_deny(&deny)),
744 AuthDecision::Deny(deny) => Err(HandshakeReject::from_deny(&deny)),
745 }
746 })
747 });
748 *capture_ref.lock().expect("capture lock poisoned") = Some(result.clone());
749 match result {
750 Ok(_) => Ok(response),
751 Err(reject) => Err(build_handshake_error_response(&response, &reject)),
752 }
753 })
754 .await;
755
756 let auth_result = capture.lock().expect("capture lock poisoned").take();
757 match handshake_result {
758 Ok(stream) => match auth_result {
759 Some(Ok(context)) => Ok(Some((stream, context))),
760 Some(Err(reject)) => Err(anyhow::anyhow!(
761 "handshake unexpectedly succeeded after rejection: {}",
762 reject.body.message
763 )),
764 None => Err(anyhow::anyhow!("no auth result captured during handshake")),
765 },
766 Err(WsError::Http(_)) => Ok(None),
767 Err(error) => Err(error.into()),
768 }
769}
770
771async fn handle_connection(
772 stream: TcpStream,
773 mut context: SubscriptionContext,
774 remote_addr: SocketAddr,
775 auth_plugin: Arc<dyn WebSocketAuthPlugin>,
776) -> Result<()> {
777 let accepted = tokio::select! {
780 _ = context.shutdown.cancelled() => return Ok(()),
781 accepted = accept_authorized_connection(
782 stream,
783 remote_addr,
784 auth_plugin.clone(),
785 context.client_manager.clone(),
786 ) => accepted?,
787 };
788 let Some((ws_stream, auth_context)) = accepted else {
789 return Ok(());
790 };
791
792 let client_id = Uuid::new_v4();
793 context.client_id = client_id;
794 let connection_start = Instant::now();
795 let (metering_key, subject, key_class, deployment_id) = usage_identity(Some(&auth_context));
796 context.metrics.connection_opened(metering_key.as_deref());
797
798 let (ws_sender, mut ws_receiver) = ws_stream.split();
799 context
800 .client_manager
801 .add_client(client_id, ws_sender, Some(auth_context), remote_addr);
802 emit_usage_event(
803 &context.usage_emitter,
804 WebSocketUsageEvent::ConnectionEstablished {
805 client_id: client_id.to_string(),
806 remote_addr: remote_addr.to_string(),
807 deployment_id: deployment_id.clone(),
808 metering_key: metering_key.clone(),
809 subject: subject.clone(),
810 key_class,
811 },
812 );
813
814 let mut active_subscriptions: HashMap<String, String> = HashMap::new();
815 loop {
816 let message = tokio::select! {
817 _ = context.shutdown.cancelled() => break,
818 next = ws_receiver.next() => match next {
819 Some(message) => message,
820 None => break,
821 },
822 };
823 let message = match message {
824 Ok(message) => message,
825 Err(error) => {
826 warn!("WebSocket error for client {}: {}", client_id, error);
827 break;
828 }
829 };
830 if message.is_close() {
831 break;
832 }
833 context.client_manager.update_client_last_seen(client_id);
834 if !message.is_text() {
835 continue;
836 }
837 if let Err(deny) = context
838 .client_manager
839 .check_inbound_message_allowed(client_id)
840 {
841 send_socket_issue(client_id, &context.client_manager, &deny, true, None).await;
842 break;
843 }
844 context.metrics.message_received(metering_key.as_deref());
845
846 let text = match message.to_text() {
847 Ok(text) => text,
848 Err(_) => continue,
849 };
850 let client_message = match serde_json::from_str::<ClientMessage>(text) {
851 Ok(message) => message,
852 Err(parse_error) => {
853 let subscription_id = extract_subscription_id(text);
854 send_protocol_issue(
855 client_id,
856 &context.client_manager,
857 &context.metrics,
858 subscription_id,
859 "malformed-message",
860 format!("invalid protocol v2 message: {parse_error}"),
861 )
862 .await;
863 continue;
864 }
865 };
866
867 match client_message {
868 ClientMessage::Subscribe(subscription) => {
869 let subscription_id = subscription.subscription_id.clone();
870 if let Err(message) = subscription.validate() {
871 send_protocol_issue(
872 client_id,
873 &context.client_manager,
874 &context.metrics,
875 Some(subscription_id),
876 "invalid-subscription",
877 message,
878 )
879 .await;
880 continue;
881 }
882 if let Err(deny) = context
883 .client_manager
884 .check_subscription_allowed(client_id)
885 .await
886 {
887 send_socket_issue(
888 client_id,
889 &context.client_manager,
890 &deny,
891 false,
892 Some(subscription_id),
893 )
894 .await;
895 continue;
896 }
897
898 let cancel_token = CancellationToken::new();
899 if !context
900 .client_manager
901 .add_client_subscription(
902 client_id,
903 subscription_id.clone(),
904 cancel_token.clone(),
905 )
906 .await
907 {
908 send_protocol_issue(
909 client_id,
910 &context.client_manager,
911 &context.metrics,
912 Some(subscription_id),
913 "duplicate-subscription-id",
914 "subscriptionId is already active on this connection",
915 )
916 .await;
917 continue;
918 }
919
920 let view = subscription.query.view.clone();
921 if let Err(error) = attach_client_to_bus(&context, subscription, cancel_token).await
922 {
923 context
924 .client_manager
925 .remove_client_subscription(client_id, &subscription_id)
926 .await;
927 match error.downcast::<RejectedSubscription>() {
931 Ok(rejected) => {
932 send_prepared_issue(
933 client_id,
934 &context.client_manager,
935 &context.metrics,
936 rejected.0,
937 )
938 .await;
939 }
940 Err(error) => {
941 send_protocol_issue(
942 client_id,
943 &context.client_manager,
944 &context.metrics,
945 Some(subscription_id),
946 "subscription-rejected",
947 error.to_string(),
948 )
949 .await;
950 }
951 }
952 continue;
953 }
954
955 active_subscriptions.insert(subscription_id, view.clone());
956 context
957 .metrics
958 .subscription_created(&view, metering_key.as_deref());
959 emit_usage_event(
960 &context.usage_emitter,
961 WebSocketUsageEvent::SubscriptionCreated {
962 client_id: client_id.to_string(),
963 deployment_id: deployment_id.clone(),
964 metering_key: metering_key.clone(),
965 subject: subject.clone(),
966 view_id: view,
967 },
968 );
969 }
970 ClientMessage::Unsubscribe(unsubscription) => {
971 handle_unsubscribe(
972 &context,
973 unsubscription,
974 &mut active_subscriptions,
975 metering_key.as_deref(),
976 &deployment_id,
977 &metering_key,
978 &subject,
979 )
980 .await;
981 }
982 ClientMessage::Ping => debug!("Received ping from client {}", client_id),
983 ClientMessage::RefreshAuth(request) => {
984 handle_refresh_auth(client_id, &request, &context.client_manager, &auth_plugin)
985 .await;
986 }
987 }
988 }
989
990 context
991 .client_manager
992 .cancel_all_client_subscriptions(client_id)
993 .await;
994 context.client_manager.remove_client(client_id);
995 if let Some(rate_limiter) = context.client_manager.rate_limiter().cloned() {
996 rate_limiter.remove_client_buckets(client_id).await;
997 }
998 for view in active_subscriptions.values() {
999 context
1000 .metrics
1001 .subscription_removed(view, metering_key.as_deref());
1002 emit_usage_event(
1003 &context.usage_emitter,
1004 WebSocketUsageEvent::SubscriptionRemoved {
1005 client_id: client_id.to_string(),
1006 deployment_id: deployment_id.clone(),
1007 metering_key: metering_key.clone(),
1008 subject: subject.clone(),
1009 view_id: view.clone(),
1010 },
1011 );
1012 }
1013 let duration = connection_start.elapsed().as_secs_f64();
1014 context
1015 .metrics
1016 .connection_closed(duration, metering_key.as_deref());
1017 emit_usage_event(
1018 &context.usage_emitter,
1019 WebSocketUsageEvent::ConnectionClosed {
1020 client_id: client_id.to_string(),
1021 deployment_id,
1022 metering_key,
1023 subject,
1024 duration_secs: Some(duration),
1025 subscription_count: u32::try_from(active_subscriptions.len()).unwrap_or(u32::MAX),
1026 },
1027 );
1028 Ok(())
1029}
1030
1031#[allow(clippy::too_many_arguments)]
1032async fn handle_unsubscribe(
1033 context: &SubscriptionContext,
1034 unsubscription: Unsubscription,
1035 active_subscriptions: &mut HashMap<String, String>,
1036 metrics_metering_key: Option<&str>,
1037 deployment_id: &Option<String>,
1038 usage_metering_key: &Option<String>,
1039 subject: &Option<String>,
1040) {
1041 let subscription_id = unsubscription.subscription_id.clone();
1042 if let Err(message) = unsubscription.validate() {
1043 send_protocol_issue(
1044 context.client_id,
1045 &context.client_manager,
1046 &context.metrics,
1047 Some(subscription_id),
1048 "invalid-unsubscription",
1049 message,
1050 )
1051 .await;
1052 return;
1053 }
1054
1055 if !context
1056 .client_manager
1057 .remove_client_subscription(context.client_id, &subscription_id)
1058 .await
1059 {
1060 send_protocol_issue(
1061 context.client_id,
1062 &context.client_manager,
1063 &context.metrics,
1064 Some(subscription_id),
1065 "unknown-subscription-id",
1066 "subscriptionId is not active on this connection",
1067 )
1068 .await;
1069 return;
1070 }
1071
1072 let Some(view) = active_subscriptions.remove(&subscription_id) else {
1073 return;
1074 };
1075 let _ = send_control_frame(context, &UnsubscribedFrame::new(subscription_id), &view);
1076 context
1077 .metrics
1078 .subscription_removed(&view, metrics_metering_key);
1079 emit_usage_event(
1080 &context.usage_emitter,
1081 WebSocketUsageEvent::SubscriptionRemoved {
1082 client_id: context.client_id.to_string(),
1083 deployment_id: deployment_id.clone(),
1084 metering_key: usage_metering_key.clone(),
1085 subject: subject.clone(),
1086 view_id: view,
1087 },
1088 );
1089}
1090
1091fn extract_subscription_id(text: &str) -> Option<String> {
1092 serde_json::from_str::<Value>(text)
1093 .ok()?
1094 .get("subscriptionId")?
1095 .as_str()
1096 .map(str::to_string)
1097}
1098
1099struct SnapshotMetadata<'a> {
1100 subscription_id: &'a str,
1101 snapshot_id: &'a str,
1102 authoritative: bool,
1103 mode: Mode,
1104 view_id: &'a str,
1105 key: Option<&'a str>,
1106}
1107
1108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1109enum SnapshotPurpose {
1110 Initial,
1111 Recovery,
1112}
1113
1114impl SnapshotPurpose {
1115 fn authoritative(self, subscription: &Subscription) -> bool {
1116 match self {
1117 Self::Initial => subscription.query.after.is_none(),
1118 Self::Recovery => true,
1122 }
1123 }
1124}
1125
1126fn create_snapshot_batches(
1127 entities: &[SnapshotEntity],
1128 metadata: SnapshotMetadata<'_>,
1129 batch_config: &SnapshotBatchConfig,
1130) -> Vec<SnapshotFrame> {
1131 if entities.is_empty() {
1132 return vec![SnapshotFrame {
1133 protocol_version: PROTOCOL_VERSION,
1134 subscription_id: metadata.subscription_id.to_string(),
1135 snapshot_id: metadata.snapshot_id.to_string(),
1136 authoritative: metadata.authoritative,
1137 mode: metadata.mode,
1138 export: metadata.view_id.to_string(),
1139 op: "snapshot",
1140 key: metadata.key.map(str::to_string),
1141 data: vec![],
1142 complete: true,
1143 }];
1144 }
1145
1146 let mut batches = Vec::new();
1147 let mut offset = 0;
1148 while offset < entities.len() {
1149 let configured_size = if offset == 0 {
1150 batch_config.initial_batch_size
1151 } else {
1152 batch_config.subsequent_batch_size
1153 };
1154 let end = (offset + configured_size.max(1)).min(entities.len());
1155 batches.push(SnapshotFrame {
1156 protocol_version: PROTOCOL_VERSION,
1157 subscription_id: metadata.subscription_id.to_string(),
1158 snapshot_id: metadata.snapshot_id.to_string(),
1159 authoritative: metadata.authoritative,
1160 mode: metadata.mode,
1161 export: metadata.view_id.to_string(),
1162 op: "snapshot",
1163 key: metadata.key.map(str::to_string),
1164 data: entities[offset..end].to_vec(),
1165 complete: end == entities.len(),
1166 });
1167 offset = end;
1168 }
1169 batches
1170}
1171
1172async fn send_snapshot_batches(
1173 context: &SubscriptionContext,
1174 subscription: &Subscription,
1175 entities: &[SnapshotEntity],
1176 mode: Mode,
1177 purpose: SnapshotPurpose,
1178 batch_config: &SnapshotBatchConfig,
1179) -> Result<()> {
1180 let snapshot_id = Uuid::new_v4().to_string();
1181 let authoritative = purpose.authoritative(subscription);
1182 let frames = create_snapshot_batches(
1183 entities,
1184 SnapshotMetadata {
1185 subscription_id: &subscription.subscription_id,
1186 snapshot_id: &snapshot_id,
1187 authoritative,
1188 mode,
1189 view_id: &subscription.query.view,
1190 key: subscription.query.key.as_deref(),
1191 },
1192 batch_config,
1193 );
1194
1195 for frame in frames {
1196 let rows = frame.data.len() as u32;
1197 let json = serde_json::to_vec(&frame)?;
1198 let payload = maybe_compress(&json);
1199 let bytes = payload.as_bytes().len() as u64;
1200 context
1201 .client_manager
1202 .send_compressed_async(context.client_id, payload)
1203 .await
1204 .map_err(|error| anyhow::anyhow!("failed to send snapshot: {error}"))?;
1205 context.metrics.message_sent();
1206
1207 let auth_context = context.client_manager.get_auth_context(context.client_id);
1208 let (metering_key, subject, _, deployment_id) = usage_identity(auth_context.as_ref());
1209 emit_usage_event(
1210 &context.usage_emitter,
1211 WebSocketUsageEvent::SnapshotSent {
1212 client_id: context.client_id.to_string(),
1213 deployment_id,
1214 metering_key,
1215 subject,
1216 view_id: subscription.query.view.clone(),
1217 rows,
1218 messages: 1,
1219 bytes,
1220 },
1221 );
1222 }
1223 Ok(())
1224}
1225
1226fn extract_sort_config(view_spec: &ViewSpec) -> Option<SortConfig> {
1227 if let Some(sort) = view_spec
1228 .pipeline
1229 .as_ref()
1230 .and_then(|pipeline| pipeline.sort.as_ref())
1231 {
1232 return Some(SortConfig {
1233 field: sort.field_path.clone(),
1234 order: match sort.order {
1235 crate::materialized_view::SortOrder::Asc => SortOrder::Asc,
1236 crate::materialized_view::SortOrder::Desc => SortOrder::Desc,
1237 },
1238 });
1239 }
1240 (view_spec.mode == Mode::List).then(|| SortConfig {
1241 field: vec!["_seq".to_string()],
1242 order: SortOrder::Desc,
1243 })
1244}
1245
1246fn send_control_frame<T: Serialize>(
1247 context: &SubscriptionContext,
1248 frame: &T,
1249 view_id: &str,
1250) -> Result<()> {
1251 let json = serde_json::to_vec(frame)?;
1252 let bytes = json.len();
1253 context
1254 .client_manager
1255 .send_to_client(context.client_id, Arc::new(Bytes::from(json)))
1256 .map_err(|error| anyhow::anyhow!("failed to send control frame: {error}"))?;
1257 context.metrics.message_sent();
1258 emit_update_sent_for_client(
1259 &context.usage_emitter,
1260 &context.client_manager,
1261 context.client_id,
1262 view_id,
1263 bytes,
1264 );
1265 Ok(())
1266}
1267
1268fn send_subscribed_frame(
1269 context: &SubscriptionContext,
1270 subscription: &Subscription,
1271 view_spec: &ViewSpec,
1272) -> Result<()> {
1273 let frame = SubscribedFrame::new(
1274 subscription.subscription_id.clone(),
1275 subscription.query.clone(),
1276 view_spec.mode,
1277 extract_sort_config(view_spec),
1278 );
1279 send_control_frame(context, &frame, &subscription.query.view)
1280}
1281
1282fn enforce_snapshot_limit(context: &SubscriptionContext, rows: usize) -> Result<()> {
1283 context
1284 .client_manager
1285 .check_snapshot_allowed(context.client_id, u32::try_from(rows).unwrap_or(u32::MAX))
1286 .map_err(|deny| anyhow::anyhow!(deny.reason))
1287}
1288
1289async fn subscribe_state_then_snapshot<F, Fut, T>(
1290 bus_manager: &BusManager,
1291 view_id: &str,
1292 key: &str,
1293 snapshot: F,
1294) -> (watch::Receiver<StateUpdate>, u64, T)
1295where
1296 F: FnOnce() -> Fut,
1297 Fut: Future<Output = T>,
1298{
1299 let mut receiver = bus_manager.get_or_create_state_bus(view_id, key).await;
1300 let published = receiver.borrow_and_update().published;
1301 let snapshot = snapshot().await;
1302 (receiver, published, snapshot)
1303}
1304
1305async fn subscribe_list_then_snapshot<F, Fut, T>(
1306 bus_manager: &BusManager,
1307 view_id: &str,
1308 snapshot: F,
1309) -> (broadcast::Receiver<Arc<BusMessage>>, T)
1310where
1311 F: FnOnce() -> Fut,
1312 Fut: Future<Output = T>,
1313{
1314 let receiver = bus_manager.get_or_create_list_bus(view_id).await;
1315 let snapshot = snapshot().await;
1316 (receiver, snapshot)
1317}
1318
1319async fn attach_client_to_bus(
1320 context: &SubscriptionContext,
1321 mut subscription: Subscription,
1322 cancel_token: CancellationToken,
1323) -> Result<()> {
1324 let view_spec = context
1325 .view_index
1326 .get_view(&subscription.query.view)
1327 .cloned()
1328 .ok_or_else(|| anyhow::anyhow!("unknown view: {}", subscription.query.view))?;
1329
1330 if view_spec.mode == Mode::State && !view_spec.is_derived() && subscription.query.key.is_none()
1331 {
1332 return Err(anyhow::anyhow!("state subscriptions require query.key"));
1333 }
1334 if view_spec.is_derived() && subscription.query.take.is_none() {
1335 subscription.query.take = view_spec
1336 .pipeline
1337 .as_ref()
1338 .and_then(|pipeline| pipeline.limit);
1339 }
1340
1341 let journal = context
1345 .journal
1346 .clone()
1347 .filter(|journal| journal.is_enabled() && view_spec.mode == Mode::Append);
1348 if let Some(journal) = journal {
1349 return attach_journal_subscription(
1350 context,
1351 subscription,
1352 view_spec,
1353 journal,
1354 cancel_token,
1355 )
1356 .await;
1357 }
1358
1359 if view_spec.mode == Mode::State && !view_spec.is_derived() {
1360 attach_state_subscription(context, subscription, view_spec, cancel_token).await
1361 } else {
1362 attach_collection_subscription(context, subscription, view_spec, cancel_token).await
1363 }
1364}
1365
1366async fn attach_state_subscription(
1367 context: &SubscriptionContext,
1368 subscription: Subscription,
1369 view_spec: ViewSpec,
1370 cancel_token: CancellationToken,
1371) -> Result<()> {
1372 let view_id = subscription.query.view.clone();
1373 let key = subscription.query.key.clone().unwrap_or_default();
1374 let query = subscription.query.clone();
1375 let cache = context.entity_cache.clone();
1376 let view_spec_for_snapshot = view_spec.clone();
1377 let (mut receiver, mut seen, initial) =
1378 subscribe_state_then_snapshot(&context.bus_manager, &view_id, &key, move || async move {
1379 load_query_entities(&cache, None, &view_spec_for_snapshot, &query, false).await
1380 })
1381 .await;
1382
1383 let mut snapshot_entities = initial;
1384 if let Some(limit) = subscription.query.snapshot_limit {
1385 snapshot_entities.truncate(limit);
1386 }
1387 enforce_snapshot_limit(context, snapshot_entities.len())?;
1388 send_subscribed_frame(context, &subscription, &view_spec)?;
1389 let delivered = subscription.snapshot.enabled && !snapshot_entities.is_empty();
1394 let snapshot_rows = to_wire_snapshot_entities(snapshot_entities, &view_spec);
1395 let synced = if delivered {
1398 snapshot_rows.first().map(|row| row.data.clone())
1399 } else {
1400 None
1401 };
1402 if subscription.snapshot.enabled {
1403 send_snapshot_batches(
1404 context,
1405 &subscription,
1406 &snapshot_rows,
1407 view_spec.mode,
1408 SnapshotPurpose::Initial,
1409 &context.entity_cache.snapshot_config(),
1410 )
1411 .await?;
1412 }
1413
1414 let task_context = context.clone();
1415 let subscription_id = subscription.subscription_id.clone();
1416 let query = subscription.query.clone();
1417 let view_spec_task = view_spec.clone();
1418 let span_view = view_id.clone();
1419 let span_key = key.clone();
1420 tokio::spawn(
1421 async move {
1422 let mut member = delivered;
1424 let mut synced = synced;
1430 let mut touched = Value::Null;
1431 let mut behind = false;
1436 let mut replaced = false;
1441 loop {
1442 tokio::select! {
1443 _ = cancel_token.cancelled() => break,
1444 changed = receiver.changed() => {
1445 if changed.is_err() {
1446 break;
1447 }
1448 let (payload, published, replaced_at) = {
1449 let update = receiver.borrow_and_update();
1450 (update.payload.clone(), update.published, update.replaced)
1451 };
1452 behind |= published > seen + 1;
1457 replaced |= behind && replaced_at > seen;
1458 seen = published;
1459 let metadata = source_frame_metadata(&payload);
1460 if metadata.op == "delete" {
1461 task_context.entity_cache.remove(&query.view, &key).await;
1462 if member && send_membership_frame(
1463 &task_context,
1464 &subscription_id,
1465 &view_spec_task,
1466 "delete",
1467 &key,
1468 Value::Null,
1469 metadata.seq,
1470 ).is_err() {
1471 break;
1472 }
1473 member = false;
1474 behind = false;
1475 replaced = false;
1476 synced = None;
1477 touched = Value::Null;
1478 continue;
1479 }
1480
1481 let Some(cached) = task_context
1492 .entity_cache
1493 .get(&view_spec_task.id, &key)
1494 .await
1495 else {
1496 if member && !replaced {
1499 record_forwarded(&payload, &metadata.op, &mut synced, &mut touched);
1500 if send_scoped_source_payload(
1501 &task_context,
1502 &subscription_id,
1503 &query.view,
1504 payload,
1505 ).is_err() {
1506 break;
1507 }
1508 }
1509 continue;
1510 };
1511 let selected =
1512 select_query_entities(vec![(key.clone(), cached)], &query, true, false);
1513 let is_member = !selected.is_empty();
1514 let caught_up = !std::mem::take(&mut behind);
1515 let replace = std::mem::take(&mut replaced);
1516 let result = match (member, selected.into_iter().next()) {
1517 (true, Some(_)) if caught_up => {
1518 record_forwarded(&payload, &metadata.op, &mut synced, &mut touched);
1519 send_scoped_source_payload(
1520 &task_context,
1521 &subscription_id,
1522 &query.view,
1523 payload,
1524 )
1525 }
1526 (true, Some((entity_key, data))) => {
1540 let mut current = data;
1541 apply_wire_format(&mut current, &view_spec_task.wire_format);
1542 let body = match (&synced, replace) {
1543 (Some(base), false) => changed_fields(base, ¤t, &touched),
1544 _ => Some(current.clone()),
1545 };
1546 synced = Some(current);
1547 touched = Value::Null;
1548 match body {
1549 Some(body) => send_membership_frame(
1550 &task_context,
1551 &subscription_id,
1552 &view_spec_task,
1553 if replace { "upsert" } else { "patch" },
1554 &entity_key,
1555 body,
1556 None,
1557 ),
1558 None => Ok(()),
1559 }
1560 }
1561 (false, Some((entity_key, data))) => {
1562 let mut current = data;
1563 apply_wire_format(&mut current, &view_spec_task.wire_format);
1564 synced = Some(current.clone());
1565 touched = Value::Null;
1566 send_membership_frame(
1567 &task_context,
1568 &subscription_id,
1569 &view_spec_task,
1570 "upsert",
1571 &entity_key,
1572 current,
1573 metadata.seq,
1574 )
1575 }
1576 (true, None) => {
1577 synced = None;
1578 touched = Value::Null;
1579 send_membership_frame(
1580 &task_context,
1581 &subscription_id,
1582 &view_spec_task,
1583 "remove",
1584 &key,
1585 Value::Null,
1586 metadata.seq,
1587 )
1588 }
1589 (false, None) => Ok(()),
1590 };
1591 if result.is_err() {
1592 break;
1593 }
1594 member = is_member;
1595 }
1596 }
1597 }
1598 }
1599 .instrument(info_span!("ws.subscribe.state", client_id = %context.client_id, view = %span_view, key = %span_key)),
1600 );
1601 Ok(())
1602}
1603
1604async fn attach_collection_subscription(
1605 context: &SubscriptionContext,
1606 subscription: Subscription,
1607 view_spec: ViewSpec,
1608 cancel_token: CancellationToken,
1609) -> Result<()> {
1610 let view_id = subscription.query.view.clone();
1611 let source_view_id = view_spec
1612 .source_view
1613 .clone()
1614 .unwrap_or_else(|| view_id.clone());
1615 let (mut receiver, initial_membership) = subscribe_collection_then_snapshot(
1616 context,
1617 &source_view_id,
1618 &view_spec,
1619 &subscription.query,
1620 )
1621 .await;
1622
1623 let mut snapshot_entities = initial_membership.clone();
1624 if let Some(limit) = subscription.query.snapshot_limit {
1625 snapshot_entities.truncate(limit);
1626 }
1627 enforce_snapshot_limit(context, snapshot_entities.len())?;
1628 send_subscribed_frame(context, &subscription, &view_spec)?;
1629 let delivered_rows = if subscription.snapshot.enabled {
1634 snapshot_entities.len()
1635 } else {
1636 0
1637 };
1638 let mut undelivered: HashSet<String> = initial_membership
1639 .iter()
1640 .skip(delivered_rows)
1641 .map(|(key, _)| key.clone())
1642 .collect();
1643 if subscription.snapshot.enabled {
1644 send_snapshot_batches(
1645 context,
1646 &subscription,
1647 &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1648 view_spec.mode,
1649 SnapshotPurpose::Initial,
1650 &context.entity_cache.snapshot_config(),
1651 )
1652 .await?;
1653 }
1654
1655 let task_context = context.clone();
1656 let task_subscription = subscription.clone();
1657 let subscription_id = subscription.subscription_id.clone();
1658 let query = subscription.query.clone();
1659 let view_spec_task = view_spec.clone();
1660 let span_view = view_id.clone();
1661 tokio::spawn(
1662 async move {
1663 let mut current = initial_membership;
1664 let coalesce_ms = (view_spec_task.mode == Mode::List)
1668 .then(|| {
1669 view_spec_task
1670 .delivery
1671 .coalesce_ms
1672 .or(task_context.delivery.collection_coalesce_ms)
1673 .filter(|milliseconds| *milliseconds > 0)
1674 })
1675 .flatten();
1676 let mut flush_interval = coalesce_ms.map(|milliseconds| {
1677 let period = Duration::from_millis(milliseconds);
1678 let mut interval = tokio::time::interval_at(
1679 tokio::time::Instant::now() + period,
1680 period,
1681 );
1682 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
1683 interval
1684 });
1685 let mut pending = HashMap::<String, Arc<BusMessage>>::new();
1686 let mut pending_updates = 0_u64;
1687
1688 loop {
1689 tokio::select! {
1690 _ = cancel_token.cancelled() => break,
1691 _ = async {
1692 flush_interval
1693 .as_mut()
1694 .expect("coalescing interval is guarded")
1695 .tick()
1696 .await;
1697 }, if flush_interval.is_some() => {
1698 if pending.is_empty() {
1699 continue;
1700 }
1701 let sorted_caches = view_spec_task
1702 .is_derived()
1703 .then(|| task_context.view_index.sorted_caches());
1704 let next = load_query_entities(
1705 &task_context.entity_cache,
1706 sorted_caches,
1707 &view_spec_task,
1708 &query,
1709 false,
1710 ).await;
1711 if emit_coalesced_collection_delta(
1712 &task_context,
1713 &subscription_id,
1714 &view_spec_task,
1715 ¤t,
1716 &next,
1717 &pending,
1718 &mut undelivered,
1719 ).is_err() {
1720 task_context.metrics.delivery_stopped(&view_id, "send-failed");
1721 break;
1722 }
1723 task_context
1724 .metrics
1725 .collection_coalesced(&view_id, pending_updates);
1726 current = next;
1727 pending.clear();
1728 pending_updates = 0;
1729 }
1730 received = receiver.recv() => {
1731 let envelope = match received {
1732 Ok(envelope) => envelope,
1733 Err(broadcast::error::RecvError::Lagged(skipped)) => {
1734 task_context.metrics.subscription_lagged(&view_id, skipped);
1735 warn!(
1736 "Subscription {} lagged by {} updates",
1737 subscription_id, skipped
1738 );
1739 if view_spec_task.mode == Mode::Append {
1740 let _ = send_control_frame(
1741 &task_context,
1742 &SocketIssueMessage::append_subscription_lagged(
1743 subscription_id.clone(),
1744 skipped,
1745 ),
1746 &view_id,
1747 );
1748 task_context.metrics.delivery_stopped(&view_id, "append-lagged-without-replay");
1749 break;
1750 }
1751 if !task_subscription.snapshot.enabled {
1752 let _ = send_control_frame(
1753 &task_context,
1754 &SocketIssueMessage::subscription_lagged(
1755 subscription_id.clone(),
1756 skipped,
1757 ),
1758 &view_id,
1759 );
1760 task_context.metrics.delivery_stopped(&view_id, "lagged-without-snapshot");
1761 break;
1762 }
1763 info!(
1764 "Subscription {} is recovering from an authoritative snapshot",
1765 subscription_id
1766 );
1767 match recover_collection_subscription(
1768 &task_context,
1769 &task_subscription,
1770 &view_spec_task,
1771 &source_view_id,
1772 ).await {
1773 Ok((next_receiver, recovered)) => {
1774 receiver = next_receiver;
1775 current = recovered;
1776 undelivered.clear();
1779 pending.clear();
1780 pending_updates = 0;
1781 task_context.metrics.subscription_resnapshot(&view_id);
1782 continue;
1783 }
1784 Err(error) => {
1785 warn!(
1786 "Subscription {} failed to recover from lag: {error:#}",
1787 subscription_id
1788 );
1789 task_context.metrics.delivery_stopped(&view_id, "resnapshot-failed");
1790 break;
1791 }
1792 }
1793 }
1794 Err(broadcast::error::RecvError::Closed) => break,
1795 };
1796
1797 apply_collection_source_event(
1798 &task_context,
1799 &source_view_id,
1800 &view_spec_task,
1801 &query,
1802 &envelope,
1803 ).await;
1804
1805 if flush_interval.is_some() {
1806 pending_updates = pending_updates.saturating_add(1);
1807 pending.insert(envelope.key.clone(), envelope);
1808 continue;
1809 }
1810
1811 let metadata = source_frame_metadata(&envelope.payload);
1812 let sorted_caches = view_spec_task
1813 .is_derived()
1814 .then(|| task_context.view_index.sorted_caches());
1815 let next = load_query_entities(
1816 &task_context.entity_cache,
1817 sorted_caches,
1818 &view_spec_task,
1819 &query,
1820 false,
1821 ).await;
1822 if emit_collection_delta(
1823 &task_context,
1824 &subscription_id,
1825 &view_spec_task,
1826 ¤t,
1827 &next,
1828 &envelope,
1829 &metadata,
1830 &mut undelivered,
1831 ).is_err() {
1832 task_context.metrics.delivery_stopped(&view_id, "send-failed");
1833 break;
1834 }
1835 current = next;
1836 }
1837 }
1838 }
1839 }
1840 .instrument(info_span!("ws.subscribe.collection", client_id = %context.client_id, view = %span_view)),
1841 );
1842 Ok(())
1843}
1844
1845async fn subscribe_collection_then_snapshot(
1846 context: &SubscriptionContext,
1847 source_view_id: &str,
1848 view_spec: &ViewSpec,
1849 query: &SubscriptionQuery,
1850) -> (broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>) {
1851 let cache = context.entity_cache.clone();
1852 let sorted_caches = view_spec
1853 .is_derived()
1854 .then(|| context.view_index.sorted_caches());
1855 let view_spec = view_spec.clone();
1856 let query = query.clone();
1857 subscribe_list_then_snapshot(&context.bus_manager, source_view_id, move || async move {
1858 load_query_entities(&cache, sorted_caches, &view_spec, &query, false).await
1859 })
1860 .await
1861}
1862
1863async fn recover_collection_subscription(
1864 context: &SubscriptionContext,
1865 subscription: &Subscription,
1866 view_spec: &ViewSpec,
1867 source_view_id: &str,
1868) -> Result<(broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>)> {
1869 let (receiver, membership) =
1870 subscribe_collection_then_snapshot(context, source_view_id, view_spec, &subscription.query)
1871 .await;
1872 enforce_snapshot_limit(context, membership.len())?;
1877 send_snapshot_batches(
1878 context,
1879 subscription,
1880 &to_wire_snapshot_entities(membership.clone(), view_spec),
1881 view_spec.mode,
1882 SnapshotPurpose::Recovery,
1883 &context.entity_cache.snapshot_config(),
1884 )
1885 .await?;
1886 Ok((receiver, membership))
1887}
1888
1889async fn apply_collection_source_event(
1890 context: &SubscriptionContext,
1891 source_view_id: &str,
1892 view_spec: &ViewSpec,
1893 query: &SubscriptionQuery,
1894 envelope: &BusMessage,
1895) {
1896 let metadata = source_frame_metadata(&envelope.payload);
1897 if metadata.op != "delete" {
1898 return;
1899 }
1900 let current = context
1904 .entity_cache
1905 .get(source_view_id, &envelope.key)
1906 .await;
1907 if source_delete_is_stale(current.as_ref(), metadata.seq.as_deref()) {
1908 return;
1909 }
1910 context
1911 .entity_cache
1912 .remove(source_view_id, &envelope.key)
1913 .await;
1914 if view_spec.is_derived() {
1915 let caches = context.view_index.sorted_caches();
1916 let mut guard = caches.write().await;
1917 if let Some(cache) = guard.get_mut(&query.view) {
1918 cache.remove(&envelope.key);
1919 }
1920 }
1921}
1922
1923fn source_delete_is_stale(current: Option<&Value>, delete_seq: Option<&str>) -> bool {
1924 match (current, delete_seq) {
1925 (Some(current), Some(delete_seq)) => current
1926 .get("_seq")
1927 .and_then(Value::as_str)
1928 .is_some_and(|current_seq| {
1929 cmp_seq(current_seq, delete_seq) == std::cmp::Ordering::Greater
1930 }),
1931 _ => false,
1932 }
1933}
1934
1935#[derive(Debug)]
1945pub(crate) struct RejectedSubscription(pub SocketIssueMessage);
1946
1947impl RejectedSubscription {
1948 fn into_error(issue: SocketIssueMessage) -> anyhow::Error {
1949 anyhow::Error::new(Self(issue))
1950 }
1951}
1952
1953impl std::fmt::Display for RejectedSubscription {
1954 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1955 formatter.write_str(&self.0.message)
1956 }
1957}
1958
1959impl std::error::Error for RejectedSubscription {}
1960
1961async fn attach_journal_subscription(
1979 context: &SubscriptionContext,
1980 subscription: Subscription,
1981 view_spec: ViewSpec,
1982 journal: Arc<crate::journal::EventJournal>,
1983 cancel_token: CancellationToken,
1984) -> Result<()> {
1985 let view_id = view_spec.id.clone();
1986 let subscription_id = subscription.subscription_id.clone();
1987
1988 if subscription.query.take.is_some()
1991 || subscription.query.skip.is_some()
1992 || subscription.query.snapshot_limit.is_some()
1993 {
1994 return Err(RejectedSubscription::into_error(SocketIssueMessage::protocol(
1995 Some(subscription_id),
1996 "invalid-subscription",
1997 format!(
1998 "take, skip and snapshotLimit are window options and do not apply to the replayable view {view_id}"
1999 ),
2000 )));
2001 }
2002
2003 let cursor = match subscription.query.after.as_deref() {
2006 Some(raw) => match crate::journal::Cursor::parse(raw) {
2007 Some(cursor) => Some(cursor),
2008 None => {
2009 return Err(RejectedSubscription::into_error(
2010 SocketIssueMessage::protocol(
2011 Some(subscription_id),
2012 "invalid-cursor",
2013 format!(
2014 "`after` must be an {{epoch}}:{{offset}} replay cursor for view {view_id}, got {raw:?}"
2015 ),
2016 ),
2017 ));
2018 }
2019 },
2020 None => None,
2021 };
2022
2023 let mut receiver = context.bus_manager.get_or_create_list_bus(&view_id).await;
2026
2027 let replayed = match journal.replay_after(&view_id, cursor.as_ref()).await {
2028 Ok(records) => records,
2029 Err(error) => {
2030 return Err(RejectedSubscription::into_error(
2031 SocketIssueMessage::replay_refused(Some(subscription_id), &error),
2032 ));
2033 }
2034 };
2035
2036 let frame = SubscribedFrame::new(
2037 subscription.subscription_id.clone(),
2038 subscription.query.clone(),
2039 view_spec.mode,
2040 extract_sort_config(&view_spec),
2041 )
2042 .with_replay_window(journal.window(&view_id).await);
2043 send_control_frame(context, &frame, &view_id)?;
2044
2045 let task_context = context.clone();
2050 let task_subscription_id = subscription.subscription_id.clone();
2051 let task_query = subscription.query.clone();
2052 let task_epoch = journal.epoch().await;
2053 let span_view = view_id.clone();
2054 tokio::spawn(
2055 async move {
2056 let mut last_sent = cursor.map(|cursor| cursor.offset);
2057 let mut pending: VecDeque<Arc<BusMessage>> = VecDeque::new();
2061 let mut lagged: Option<u64> = None;
2062
2063 for record in replayed {
2064 if cancel_token.is_cancelled() {
2065 return;
2066 }
2067 drain_available(&mut receiver, &mut pending, &mut lagged);
2068 if journal_record_matches(&task_query, &record)
2069 && send_scoped_source_payload_async(
2070 &task_context,
2071 &task_subscription_id,
2072 &span_view,
2073 record.payload,
2074 )
2075 .await
2076 .is_err()
2077 {
2078 return;
2079 }
2080 last_sent = Some(record.offset);
2082 }
2083
2084 while let Some(envelope) = pending.pop_front() {
2087 if !forward_live_frame(
2088 &task_context,
2089 &task_subscription_id,
2090 &span_view,
2091 &task_query,
2092 &envelope,
2093 &mut last_sent,
2094 )
2095 .await
2096 {
2097 return;
2098 }
2099 }
2100
2101 if let Some(skipped) = lagged {
2102 report_replay_gap(
2103 &task_context,
2104 &task_subscription_id,
2105 &span_view,
2106 &task_epoch,
2107 skipped,
2108 last_sent,
2109 );
2110 return;
2111 }
2112
2113 loop {
2114 tokio::select! {
2115 _ = cancel_token.cancelled() => break,
2116 received = receiver.recv() => {
2117 let envelope = match received {
2118 Ok(envelope) => envelope,
2119 Err(broadcast::error::RecvError::Lagged(skipped)) => {
2137 report_replay_gap(
2138 &task_context,
2139 &task_subscription_id,
2140 &span_view,
2141 &task_epoch,
2142 skipped,
2143 last_sent,
2144 );
2145 break;
2146 }
2147 Err(broadcast::error::RecvError::Closed) => break,
2148 };
2149
2150 if !forward_live_frame(
2151 &task_context,
2152 &task_subscription_id,
2153 &span_view,
2154 &task_query,
2155 &envelope,
2156 &mut last_sent,
2157 )
2158 .await
2159 {
2160 break;
2161 }
2162 }
2163 }
2164 }
2165 }
2166 .instrument(info_span!(
2167 "ws.subscribe.replay",
2168 client_id = %context.client_id,
2169 view = %view_id
2170 )),
2171 );
2172 Ok(())
2173}
2174
2175fn drain_available(
2178 receiver: &mut broadcast::Receiver<Arc<BusMessage>>,
2179 pending: &mut VecDeque<Arc<BusMessage>>,
2180 lagged: &mut Option<u64>,
2181) {
2182 if lagged.is_some() {
2187 return;
2188 }
2189 const MAX_PENDING: usize = 8_192;
2195 loop {
2196 if pending.len() >= MAX_PENDING {
2197 return;
2198 }
2199 match receiver.try_recv() {
2200 Ok(envelope) => pending.push_back(envelope),
2201 Err(broadcast::error::TryRecvError::Empty)
2202 | Err(broadcast::error::TryRecvError::Closed) => return,
2203 Err(broadcast::error::TryRecvError::Lagged(skipped)) => {
2204 *lagged = Some(skipped);
2205 return;
2206 }
2207 }
2208 }
2209}
2210
2211fn already_delivered(offset: Option<u64>, last_sent: &mut Option<u64>) -> bool {
2217 let Some(offset) = offset else {
2218 return false;
2221 };
2222 if last_sent.is_some_and(|last| offset <= last) {
2223 return true;
2224 }
2225 *last_sent = Some(offset);
2226 false
2227}
2228
2229async fn forward_live_frame(
2233 context: &SubscriptionContext,
2234 subscription_id: &str,
2235 view_id: &str,
2236 query: &SubscriptionQuery,
2237 envelope: &Arc<BusMessage>,
2238 last_sent: &mut Option<u64>,
2239) -> bool {
2240 let metadata = source_frame_metadata(&envelope.payload);
2241 if already_delivered(metadata.offset, last_sent) {
2242 return true;
2243 }
2244 if !live_frame_matches(query, &envelope.key, &envelope.payload) {
2245 return true;
2246 }
2247 send_scoped_source_payload(context, subscription_id, view_id, envelope.payload.clone()).is_ok()
2248}
2249
2250fn report_replay_gap(
2251 context: &SubscriptionContext,
2252 subscription_id: &str,
2253 view_id: &str,
2254 epoch: &crate::journal::JournalEpoch,
2255 skipped: u64,
2256 last_sent: Option<u64>,
2257) {
2258 warn!(
2259 "Replay subscription {} lagged past {} records; stopping with a recovery cursor",
2260 subscription_id, skipped
2261 );
2262 let recover_from = last_sent.map(|offset| crate::journal::Cursor {
2263 epoch: epoch.clone(),
2264 offset,
2265 });
2266 let _ = send_control_frame(
2267 context,
2268 &SocketIssueMessage::replay_lagged(
2269 Some(subscription_id.to_string()),
2270 skipped,
2271 recover_from,
2272 ),
2273 view_id,
2274 );
2275}
2276
2277fn journal_record_matches(
2280 query: &SubscriptionQuery,
2281 record: &crate::journal::JournalRecord,
2282) -> bool {
2283 live_frame_matches(query, &record.key, &record.payload)
2284}
2285
2286fn live_frame_matches(query: &SubscriptionQuery, key: &str, payload: &[u8]) -> bool {
2287 if !query.matches_key(key) {
2288 return false;
2289 }
2290 if query.partition.is_none() && query.filters.is_empty() {
2291 return true;
2292 }
2293 let Ok(frame) = serde_json::from_slice::<Value>(payload) else {
2294 return false;
2295 };
2296 let Some(data) = frame.get("data") else {
2297 return false;
2298 };
2299 if let Some(partition) = &query.partition {
2300 if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2301 return false;
2302 }
2303 }
2304 query
2305 .filters
2306 .iter()
2307 .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2308}
2309
2310async fn send_scoped_source_payload_async(
2313 context: &SubscriptionContext,
2314 subscription_id: &str,
2315 view_id: &str,
2316 payload: Arc<Bytes>,
2317) -> Result<()> {
2318 let mut value: Value = serde_json::from_slice(&payload)?;
2319 let object = value
2320 .as_object_mut()
2321 .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2322 object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2323 object.insert(
2324 "subscriptionId".to_string(),
2325 Value::String(subscription_id.to_string()),
2326 );
2327 let json = serde_json::to_vec(&value)?;
2328 let compressed = maybe_compress(&json);
2329 let bytes = compressed.as_bytes().len();
2330 context
2331 .client_manager
2332 .send_compressed_async(context.client_id, compressed)
2333 .await
2334 .map_err(|error| anyhow::anyhow!("failed to send replayed frame: {error}"))?;
2335 context.metrics.message_sent();
2336 emit_update_sent_for_client(
2337 &context.usage_emitter,
2338 &context.client_manager,
2339 context.client_id,
2340 view_id,
2341 bytes,
2342 );
2343 Ok(())
2344}
2345
2346#[derive(Default)]
2347struct SourceFrameMetadata {
2348 op: String,
2349 seq: Option<String>,
2350 offset: Option<u64>,
2351}
2352
2353fn source_frame_metadata(payload: &[u8]) -> SourceFrameMetadata {
2354 serde_json::from_slice::<Value>(payload)
2355 .ok()
2356 .map(|value| SourceFrameMetadata {
2357 op: value
2358 .get("op")
2359 .and_then(Value::as_str)
2360 .unwrap_or_default()
2361 .to_string(),
2362 seq: value.get("seq").and_then(Value::as_str).map(str::to_string),
2363 offset: value.get("offset").and_then(Value::as_u64),
2364 })
2365 .unwrap_or_default()
2366}
2367
2368fn source_frame_data(payload: &[u8]) -> Option<Value> {
2370 let mut frame: Value = serde_json::from_slice(payload).ok()?;
2371 frame.get_mut("data").map(Value::take)
2372}
2373
2374fn record_forwarded(payload: &[u8], op: &str, synced: &mut Option<Value>, touched: &mut Value) {
2378 let data = source_frame_data(payload);
2379 if op == "upsert" {
2380 *synced = data;
2381 *touched = Value::Null;
2382 } else if let Some(data) = data {
2383 mark_fields(touched, &data);
2384 }
2385}
2386
2387fn mark_fields(touched: &mut Value, patch: &Value) {
2395 let Value::Object(fields) = patch else {
2396 *touched = Value::Bool(true);
2397 return;
2398 };
2399 if touched.is_null() {
2400 *touched = Value::Object(serde_json::Map::new());
2401 }
2402 let Value::Object(marked) = touched else {
2403 return;
2405 };
2406 for (field, value) in fields {
2407 mark_fields(marked.entry(field.clone()).or_insert(Value::Null), value);
2408 }
2409}
2410
2411fn changed_fields(base: &Value, current: &Value, touched: &Value) -> Option<Value> {
2417 match (base, current, touched) {
2418 (_, _, Value::Bool(true)) => Some(current.clone()),
2419 (Value::Object(base), Value::Object(current), _) => {
2420 let marked = touched.as_object();
2421 let mut patch = serde_json::Map::new();
2422 for (field, value) in current {
2423 let touched = marked
2424 .and_then(|marked| marked.get(field))
2425 .unwrap_or(&Value::Null);
2426 match base.get(field) {
2427 Some(old) => {
2428 if let Some(changed) = changed_fields(old, value, touched) {
2429 patch.insert(field.clone(), changed);
2430 }
2431 }
2432 None => {
2433 patch.insert(field.clone(), value.clone());
2434 }
2435 }
2436 }
2437 for field in base
2438 .keys()
2439 .chain(marked.into_iter().flat_map(|marked| marked.keys()))
2440 {
2441 if !current.contains_key(field) {
2442 patch.insert(field.clone(), Value::Null);
2443 }
2444 }
2445 (!patch.is_empty()).then_some(Value::Object(patch))
2446 }
2447 _ => (base != current || !touched.is_null()).then(|| current.clone()),
2448 }
2449}
2450
2451fn send_scoped_source_payload(
2452 context: &SubscriptionContext,
2453 subscription_id: &str,
2454 view_id: &str,
2455 payload: Arc<Bytes>,
2456) -> Result<()> {
2457 let mut value: Value = serde_json::from_slice(&payload)?;
2458 let object = value
2459 .as_object_mut()
2460 .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2461 object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2462 object.insert(
2463 "subscriptionId".to_string(),
2464 Value::String(subscription_id.to_string()),
2465 );
2466 let encoded = Arc::new(Bytes::from(serde_json::to_vec(&value)?));
2467 let bytes = encoded.len();
2468 context
2469 .client_manager
2470 .send_to_client(context.client_id, encoded)
2471 .map_err(|error| anyhow::anyhow!("failed to send live frame: {error}"))?;
2472 context.metrics.message_sent();
2473 emit_update_sent_for_client(
2474 &context.usage_emitter,
2475 &context.client_manager,
2476 context.client_id,
2477 view_id,
2478 bytes,
2479 );
2480 Ok(())
2481}
2482
2483fn send_membership_frame(
2484 context: &SubscriptionContext,
2485 subscription_id: &str,
2486 view_spec: &ViewSpec,
2487 op: &str,
2488 key: &str,
2489 mut data: Value,
2490 seq: Option<String>,
2491) -> Result<()> {
2492 apply_wire_format(&mut data, &view_spec.wire_format);
2493 let frame = Frame::scoped(
2494 subscription_id,
2495 view_spec.mode,
2496 &view_spec.id,
2497 op,
2498 key,
2499 data,
2500 seq,
2501 );
2502 let encoded = Arc::new(Bytes::from(serde_json::to_vec(&frame)?));
2503 let bytes = encoded.len();
2504 context
2505 .client_manager
2506 .send_to_client(context.client_id, encoded)
2507 .map_err(|error| anyhow::anyhow!("failed to send membership frame: {error}"))?;
2508 context.metrics.message_sent();
2509 emit_update_sent_for_client(
2510 &context.usage_emitter,
2511 &context.client_manager,
2512 context.client_id,
2513 &view_spec.id,
2514 bytes,
2515 );
2516 Ok(())
2517}
2518
2519#[allow(clippy::too_many_arguments)]
2522fn emit_collection_delta(
2523 context: &SubscriptionContext,
2524 subscription_id: &str,
2525 view_spec: &ViewSpec,
2526 current: &[(String, Value)],
2527 next: &[(String, Value)],
2528 envelope: &BusMessage,
2529 metadata: &SourceFrameMetadata,
2530 undelivered: &mut HashSet<String>,
2531) -> Result<()> {
2532 let current_keys: Vec<&str> = current.iter().map(|(key, _)| key.as_str()).collect();
2533 let next_keys: Vec<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2534 let next_set: HashSet<&str> = next_keys.iter().copied().collect();
2535
2536 for key in current_keys
2537 .iter()
2538 .copied()
2539 .filter(|key| !next_set.contains(key))
2540 {
2541 if undelivered.remove(key) {
2543 continue;
2544 }
2545 let op = if metadata.op == "delete" && key == envelope.key {
2546 "delete"
2547 } else {
2548 "remove"
2549 };
2550 send_membership_frame(
2551 context,
2552 subscription_id,
2553 view_spec,
2554 op,
2555 key,
2556 Value::Null,
2557 metadata.seq.clone(),
2558 )?;
2559 }
2560
2561 for (key, data) in next.iter() {
2562 let in_window = current_keys.iter().any(|candidate| *candidate == key);
2563 match member_action(
2564 in_window,
2565 in_window && !undelivered.contains(key),
2566 key == &envelope.key,
2567 view_spec.is_derived(),
2568 &metadata.op,
2569 ) {
2570 MemberAction::Skip => {}
2571 MemberAction::ForwardPatch => send_scoped_source_payload(
2572 context,
2573 subscription_id,
2574 &view_spec.id,
2575 envelope.payload.clone(),
2576 )?,
2577 MemberAction::Upsert => {
2578 let seq = metadata
2579 .seq
2580 .clone()
2581 .or_else(|| data.get("_seq").and_then(Value::as_str).map(str::to_string));
2582 send_membership_frame(
2583 context,
2584 subscription_id,
2585 view_spec,
2586 "upsert",
2587 key,
2588 data.clone(),
2589 seq,
2590 )?;
2591 undelivered.remove(key);
2592 }
2593 }
2594 }
2595 Ok(())
2596}
2597
2598fn emit_coalesced_collection_delta(
2603 context: &SubscriptionContext,
2604 subscription_id: &str,
2605 view_spec: &ViewSpec,
2606 current: &[(String, Value)],
2607 next: &[(String, Value)],
2608 pending: &HashMap<String, Arc<BusMessage>>,
2609 undelivered: &mut HashSet<String>,
2610) -> Result<()> {
2611 let changes = plan_coalesced_collection_delta(current, next, pending, undelivered);
2612 let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2615 undelivered.retain(|key| next_keys.contains(key.as_str()));
2616 for change in changes {
2617 if change.op == "upsert" {
2618 undelivered.remove(&change.key);
2619 }
2620 send_membership_frame(
2621 context,
2622 subscription_id,
2623 view_spec,
2624 change.op,
2625 &change.key,
2626 change.data,
2627 change.seq,
2628 )?;
2629 }
2630 Ok(())
2631}
2632
2633#[derive(Debug, PartialEq)]
2634struct CollectionChange {
2635 op: &'static str,
2636 key: String,
2637 data: Value,
2638 seq: Option<String>,
2639}
2640
2641fn plan_coalesced_collection_delta(
2645 current: &[(String, Value)],
2646 next: &[(String, Value)],
2647 pending: &HashMap<String, Arc<BusMessage>>,
2648 undelivered: &HashSet<String>,
2649) -> Vec<CollectionChange> {
2650 let current_by_key: HashMap<&str, &Value> = current
2651 .iter()
2652 .map(|(key, data)| (key.as_str(), data))
2653 .collect();
2654 let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2655 let latest_seq = pending
2656 .values()
2657 .filter_map(|envelope| source_frame_metadata(&envelope.payload).seq)
2658 .max_by(|left, right| cmp_seq(left, right));
2659 let mut changes = Vec::new();
2660
2661 for (key, _) in current
2662 .iter()
2663 .filter(|(key, _)| !next_keys.contains(key.as_str()) && !undelivered.contains(key.as_str()))
2664 {
2665 let metadata = pending
2666 .get(key)
2667 .map(|envelope| source_frame_metadata(&envelope.payload));
2668 let op = if metadata
2669 .as_ref()
2670 .is_some_and(|metadata| metadata.op == "delete")
2671 {
2672 "delete"
2673 } else {
2674 "remove"
2675 };
2676 let seq = metadata
2677 .and_then(|metadata| metadata.seq)
2678 .or_else(|| latest_seq.clone());
2679 changes.push(CollectionChange {
2680 op,
2681 key: key.clone(),
2682 data: Value::Null,
2683 seq,
2684 });
2685 }
2686
2687 for (key, data) in next {
2688 let changed = current_by_key
2689 .get(key.as_str())
2690 .is_none_or(|previous| *previous != data);
2691 if !changed && !pending.contains_key(key) {
2692 continue;
2693 }
2694 let seq = data
2695 .get("_seq")
2696 .and_then(Value::as_str)
2697 .map(str::to_string)
2698 .or_else(|| {
2699 pending
2700 .get(key)
2701 .and_then(|envelope| source_frame_metadata(&envelope.payload).seq)
2702 })
2703 .or_else(|| latest_seq.clone());
2704 changes.push(CollectionChange {
2705 op: "upsert",
2706 key: key.clone(),
2707 data: data.clone(),
2708 seq,
2709 });
2710 }
2711 changes
2712}
2713
2714#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2716enum MemberAction {
2717 Skip,
2720 ForwardPatch,
2722 Upsert,
2724}
2725
2726fn member_action(
2744 in_window: bool,
2745 held: bool,
2746 is_mutated_key: bool,
2747 is_derived: bool,
2748 op: &str,
2749) -> MemberAction {
2750 if !is_mutated_key {
2751 return if in_window {
2757 MemberAction::Skip
2758 } else {
2759 MemberAction::Upsert
2760 };
2761 }
2762 if held && !is_derived && op != "delete" {
2768 MemberAction::ForwardPatch
2769 } else {
2770 MemberAction::Upsert
2771 }
2772}
2773
2774fn to_wire_snapshot_entities(
2775 entities: Vec<(String, Value)>,
2776 view_spec: &ViewSpec,
2777) -> Vec<SnapshotEntity> {
2778 entities
2779 .into_iter()
2780 .map(|(key, mut data)| {
2781 apply_wire_format(&mut data, &view_spec.wire_format);
2782 SnapshotEntity { key, data }
2783 })
2784 .collect()
2785}
2786
2787async fn load_query_entities(
2788 entity_cache: &EntityCache,
2789 sorted_caches: Option<
2790 Arc<tokio::sync::RwLock<HashMap<String, crate::sorted_cache::SortedViewCache>>>,
2791 >,
2792 view_spec: &ViewSpec,
2793 query: &SubscriptionQuery,
2794 apply_snapshot_limit: bool,
2795) -> Vec<(String, Value)> {
2796 let (entities, preordered) = if let Some(sorted_caches) = sorted_caches {
2797 let mut caches = sorted_caches.write().await;
2798 let entities = caches
2799 .get_mut(&view_spec.id)
2800 .map(|cache| cache.get_all_ordered())
2801 .unwrap_or_default();
2802 (entities, true)
2803 } else if view_spec.mode == Mode::State {
2804 let entity = match query.key.as_deref() {
2805 Some(key) => entity_cache
2806 .get(&view_spec.id, key)
2807 .await
2808 .map(|data| vec![(key.to_string(), data)])
2809 .unwrap_or_default(),
2810 None => vec![],
2811 };
2812 (entity, true)
2813 } else {
2814 (entity_cache.get_all(&view_spec.id).await, false)
2815 };
2816 select_query_entities(entities, query, preordered, apply_snapshot_limit)
2817}
2818
2819fn select_query_entities(
2820 mut entities: Vec<(String, Value)>,
2821 query: &SubscriptionQuery,
2822 preordered: bool,
2823 apply_snapshot_limit: bool,
2824) -> Vec<(String, Value)> {
2825 entities.retain(|(key, data)| query_matches_entity(query, key, data));
2826 if !preordered {
2827 entities.sort_by(|left, right| {
2828 let left_seq = left.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2829 let right_seq = right.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2830 let order = if query.after.is_some() {
2831 cmp_seq(left_seq, right_seq)
2832 } else {
2833 cmp_seq(right_seq, left_seq)
2834 };
2835 order.then_with(|| left.0.cmp(&right.0))
2836 });
2837 }
2838
2839 let skip = query.skip.unwrap_or(0);
2840 let take = query.take.unwrap_or(usize::MAX);
2841 let mut selected: Vec<_> = entities.into_iter().skip(skip).take(take).collect();
2842 if apply_snapshot_limit {
2843 if let Some(limit) = query.snapshot_limit {
2844 selected.truncate(limit);
2845 }
2846 }
2847 selected
2848}
2849
2850fn query_matches_entity(query: &SubscriptionQuery, key: &str, data: &Value) -> bool {
2851 if !query.matches_key(key) {
2852 return false;
2853 }
2854 if let Some(partition) = &query.partition {
2855 if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2856 return false;
2857 }
2858 }
2859 if let Some(after) = &query.after {
2860 let Some(seq) = data.get("_seq").and_then(Value::as_str) else {
2861 return false;
2862 };
2863 if cmp_seq(seq, after) != std::cmp::Ordering::Greater {
2864 return false;
2865 }
2866 }
2867 query
2868 .filters
2869 .iter()
2870 .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2871}
2872
2873fn value_at_dot_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
2874 path.split('.')
2875 .try_fold(value, |current, segment| current.get(segment))
2876}
2877
2878#[cfg(test)]
2879mod tests {
2880 use super::*;
2881 use crate::cache::{EntityCacheConfig, PatchOrigin};
2882 use crate::view::{Delivery, Filters, Projection};
2883 use serde_json::json;
2884 use tokio::sync::oneshot;
2885
2886 #[test]
2887 fn a_catch_up_carries_only_the_fields_that_changed() {
2888 let base = json!({
2889 "id": 7,
2890 "state": {"deployed": [1, 2, 3], "total": 6},
2891 "results": {"square": null},
2892 "gone": true,
2893 });
2894 let current = json!({
2895 "id": 7,
2896 "state": {"deployed": [1, 2, 4], "total": 6},
2897 "results": {"square": 3},
2898 "fresh": "x",
2899 });
2900 assert_eq!(
2901 changed_fields(&base, ¤t, &Value::Null),
2902 Some(json!({
2903 "state": {"deployed": [1, 2, 4]},
2905 "results": {"square": 3},
2906 "fresh": "x",
2907 "gone": null,
2909 }))
2910 );
2911 assert_eq!(changed_fields(¤t, ¤t, &Value::Null), None);
2912 assert_eq!(
2914 changed_fields(&json!(null), ¤t, &Value::Null),
2915 Some(current.clone())
2916 );
2917 }
2918
2919 #[test]
2920 fn a_catch_up_resends_the_fields_forwarded_patches_set() {
2921 let base = json!({"a": 1, "state": {"x": 1, "y": 1}, "list": [1]});
2922 let mut touched = Value::Null;
2923 mark_fields(&mut touched, &json!({"a": 2, "state": {"x": 2}}));
2924 mark_fields(&mut touched, &json!({"list": [2], "extra": true}));
2925 assert_eq!(
2926 touched,
2927 json!({"a": true, "state": {"x": true}, "list": true, "extra": true})
2928 );
2929 let current = json!({"a": 1, "state": {"x": 1, "y": 1}, "list": [1, 2]});
2933 assert_eq!(
2934 changed_fields(&base, ¤t, &Value::Null),
2935 Some(json!({"list": [1, 2]}))
2936 );
2937 assert_eq!(
2938 changed_fields(&base, ¤t, &touched),
2939 Some(json!({"a": 1, "state": {"x": 1}, "list": [1, 2], "extra": null}))
2940 );
2941 mark_fields(&mut touched, &json!(null));
2943 assert_eq!(
2944 changed_fields(&base, ¤t, &touched),
2945 Some(current.clone())
2946 );
2947 }
2948
2949 fn list_spec() -> ViewSpec {
2950 ViewSpec {
2951 id: "Thing/list".to_string(),
2952 export: "Thing".to_string(),
2953 mode: Mode::List,
2954 wire_format: Default::default(),
2955 projection: Projection::all(),
2956 filters: Filters::all(),
2957 delivery: Delivery::default(),
2958 pipeline: None,
2959 source_view: None,
2960 }
2961 }
2962
2963 fn plan_window(
2966 current_keys: &[&str],
2967 next_keys: &[&str],
2968 envelope_key: &str,
2969 is_derived: bool,
2970 op: &str,
2971 ) -> Vec<(String, MemberAction)> {
2972 plan_window_with_undelivered(current_keys, &[], next_keys, envelope_key, is_derived, op)
2973 }
2974
2975 fn plan_window_with_undelivered(
2977 current_keys: &[&str],
2978 undelivered: &[&str],
2979 next_keys: &[&str],
2980 envelope_key: &str,
2981 is_derived: bool,
2982 op: &str,
2983 ) -> Vec<(String, MemberAction)> {
2984 next_keys
2985 .iter()
2986 .map(|key| {
2987 let in_window = current_keys.contains(key);
2988 let held = in_window && !undelivered.contains(key);
2989 (
2990 (*key).to_string(),
2991 member_action(in_window, held, *key == envelope_key, is_derived, op),
2992 )
2993 })
2994 .collect()
2995 }
2996
2997 #[test]
2998 fn a_window_member_the_client_never_received_is_sent_whole() {
2999 let plan = plan_window_with_undelivered(
3003 &["4", "3", "2", "1"],
3004 &["2", "1"],
3005 &["1", "4", "3", "2"],
3006 "1",
3007 false,
3008 "patch",
3009 );
3010 assert_eq!(
3011 plan,
3012 vec![
3013 ("1".to_string(), MemberAction::Upsert),
3014 ("4".to_string(), MemberAction::Skip),
3015 ("3".to_string(), MemberAction::Skip),
3016 ("2".to_string(), MemberAction::Skip),
3019 ]
3020 );
3021 }
3022
3023 #[test]
3024 fn a_held_member_still_rides_its_patch_after_a_truncated_snapshot() {
3025 let plan = plan_window_with_undelivered(
3026 &["4", "3", "2", "1"],
3027 &["2", "1"],
3028 &["4", "3", "2", "1"],
3029 "4",
3030 false,
3031 "patch",
3032 );
3033 assert_eq!(plan[0], ("4".to_string(), MemberAction::ForwardPatch));
3034 }
3035
3036 #[test]
3037 fn coalescing_sends_no_remove_for_a_key_the_client_never_received() {
3038 let current = vec![
3039 ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3040 (
3041 "unsent".to_string(),
3042 json!({"count": 1, "_seq": "10:000001"}),
3043 ),
3044 ("gone".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3045 ];
3046 let next = vec![
3047 ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3048 (
3049 "unsent".to_string(),
3050 json!({"count": 2, "_seq": "10:000003"}),
3051 ),
3052 ];
3053 let pending = HashMap::from([
3054 (
3055 "unsent".to_string(),
3056 list_message("unsent", "patch", "10:000003"),
3057 ),
3058 (
3059 "gone".to_string(),
3060 list_message("gone", "delete", "10:000002"),
3061 ),
3062 ]);
3063 let undelivered = HashSet::from(["unsent".to_string(), "gone".to_string()]);
3064
3065 assert_eq!(
3066 plan_coalesced_collection_delta(¤t, &next, &pending, &undelivered),
3067 vec![CollectionChange {
3068 op: "upsert",
3069 key: "unsent".to_string(),
3070 data: json!({"count": 2, "_seq": "10:000003"}),
3071 seq: Some("10:000003".to_string()),
3072 }]
3073 );
3074 }
3075
3076 #[test]
3077 fn a_reordered_window_sends_only_the_mutated_entity() {
3078 let current = ["4", "3", "2", "1"];
3082 let next = ["1", "4", "3", "2"];
3083 let plan = plan_window(¤t, &next, "1", false, "patch");
3084
3085 assert_eq!(
3086 plan,
3087 vec![
3088 ("1".to_string(), MemberAction::ForwardPatch),
3089 ("4".to_string(), MemberAction::Skip),
3090 ("3".to_string(), MemberAction::Skip),
3091 ("2".to_string(), MemberAction::Skip),
3092 ]
3093 );
3094 }
3095
3096 #[test]
3097 fn a_key_entering_the_window_gets_the_whole_entity() {
3098 let plan = plan_window(&["4", "3"], &["5", "4", "3"], "5", false, "patch");
3100 assert_eq!(
3101 plan,
3102 vec![
3103 ("5".to_string(), MemberAction::Upsert),
3104 ("4".to_string(), MemberAction::Skip),
3105 ("3".to_string(), MemberAction::Skip),
3106 ]
3107 );
3108 }
3109
3110 #[test]
3111 fn derived_views_still_send_whole_entities() {
3112 let plan = plan_window(&["1", "2"], &["1", "2"], "1", true, "patch");
3115 assert_eq!(
3116 plan,
3117 vec![
3118 ("1".to_string(), MemberAction::Upsert),
3119 ("2".to_string(), MemberAction::Skip),
3120 ]
3121 );
3122 }
3123
3124 #[test]
3125 fn a_delete_envelope_never_forwards_a_patch() {
3126 let plan = plan_window(&["1", "2"], &["1", "2"], "1", false, "delete");
3128 assert_eq!(
3129 plan,
3130 vec![
3131 ("1".to_string(), MemberAction::Upsert),
3132 ("2".to_string(), MemberAction::Skip),
3133 ]
3134 );
3135 }
3136
3137 #[test]
3138 fn an_unchanged_window_sends_one_frame_not_a_broadcast() {
3139 let keys: Vec<String> = (0..500).map(|index| index.to_string()).collect();
3142 let refs: Vec<&str> = keys.iter().map(String::as_str).collect();
3143 let plan = plan_window(&refs, &refs, "250", false, "patch");
3144
3145 let sent = plan
3146 .iter()
3147 .filter(|(_, action)| *action != MemberAction::Skip)
3148 .count();
3149 assert_eq!(sent, 1);
3150 assert_eq!(plan[250].1, MemberAction::ForwardPatch);
3151 }
3152
3153 fn list_message(key: &str, op: &str, seq: &str) -> Arc<BusMessage> {
3154 Arc::new(BusMessage {
3155 key: key.to_string(),
3156 entity: "Thing/list".to_string(),
3157 payload: Arc::new(Bytes::from(
3158 serde_json::to_vec(&json!({
3159 "entity": "Thing/list",
3160 "op": op,
3161 "key": key,
3162 "seq": seq,
3163 "data": {},
3164 }))
3165 .unwrap(),
3166 )),
3167 })
3168 }
3169
3170 #[test]
3171 fn coalescing_emits_only_the_final_full_state_per_changed_key() {
3172 let current = vec![
3173 ("a".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3174 ("b".to_string(), json!({"count": 1, "_seq": "10:000001"})),
3175 (
3176 "stable".to_string(),
3177 json!({"count": 1, "_seq": "10:000001"}),
3178 ),
3179 ];
3180 let next = vec![
3181 ("a".to_string(), json!({"count": 3, "_seq": "10:000004"})),
3182 ("c".to_string(), json!({"count": 1, "_seq": "10:000003"})),
3183 (
3184 "stable".to_string(),
3185 json!({"count": 1, "_seq": "10:000001"}),
3186 ),
3187 ];
3188 let pending = HashMap::from([
3189 ("a".to_string(), list_message("a", "patch", "10:000004")),
3190 ("b".to_string(), list_message("b", "delete", "10:000002")),
3191 ("c".to_string(), list_message("c", "patch", "10:000003")),
3192 ]);
3193
3194 assert_eq!(
3195 plan_coalesced_collection_delta(¤t, &next, &pending, &HashSet::new()),
3196 vec![
3197 CollectionChange {
3198 op: "delete",
3199 key: "b".to_string(),
3200 data: Value::Null,
3201 seq: Some("10:000002".to_string()),
3202 },
3203 CollectionChange {
3204 op: "upsert",
3205 key: "a".to_string(),
3206 data: json!({"count": 3, "_seq": "10:000004"}),
3207 seq: Some("10:000004".to_string()),
3208 },
3209 CollectionChange {
3210 op: "upsert",
3211 key: "c".to_string(),
3212 data: json!({"count": 1, "_seq": "10:000003"}),
3213 seq: Some("10:000003".to_string()),
3214 },
3215 ]
3216 );
3217 }
3218
3219 #[test]
3220 fn a_delayed_delete_cannot_erase_a_newer_recreated_entity() {
3221 let recreated = json!({"balance": 2, "_seq": "10:000004"});
3222 assert!(source_delete_is_stale(Some(&recreated), Some("10:000002")));
3223 assert!(!source_delete_is_stale(Some(&recreated), Some("10:000004")));
3224 assert!(!source_delete_is_stale(Some(&recreated), None));
3225 }
3226
3227 #[test]
3228 fn a_lagged_snapshotless_subscription_gets_a_fatal_retryable_error() {
3229 let issue = SocketIssueMessage::subscription_lagged("balances".to_string(), 42);
3230 assert_eq!(issue.code, "subscription-lagged");
3231 assert!(issue.retryable);
3232 assert!(issue.fatal);
3233 assert!(issue.message.contains("42"));
3234 assert!(issue
3235 .suggested_action
3236 .unwrap()
3237 .contains("snapshots enabled"));
3238
3239 let append = SocketIssueMessage::append_subscription_lagged("trades".to_string(), 9);
3240 assert!(append.fatal);
3241 assert!(append.suggested_action.unwrap().contains("retained replay"));
3242 }
3243
3244 #[test]
3245 fn snapshot_batches_share_identity_and_completion() {
3246 let entities = ["one", "two", "three"].map(|key| SnapshotEntity {
3247 key: key.to_string(),
3248 data: json!({"key": key}),
3249 });
3250 let batches = create_snapshot_batches(
3251 &entities,
3252 SnapshotMetadata {
3253 subscription_id: "sub-1",
3254 snapshot_id: "snapshot-1",
3255 authoritative: true,
3256 mode: Mode::List,
3257 view_id: "Thing/list",
3258 key: None,
3259 },
3260 &SnapshotBatchConfig {
3261 initial_batch_size: 2,
3262 subsequent_batch_size: 1,
3263 },
3264 );
3265 assert_eq!(batches.len(), 2);
3266 assert!(batches.iter().all(|batch| batch.subscription_id == "sub-1"));
3267 assert!(batches
3268 .iter()
3269 .all(|batch| batch.snapshot_id == "snapshot-1"));
3270 assert!(!batches[0].complete);
3271 assert!(batches[1].complete);
3272 assert!(batches.iter().all(|batch| batch.authoritative));
3273 }
3274
3275 #[test]
3276 fn empty_incremental_snapshot_is_explicitly_non_authoritative() {
3277 let batches = create_snapshot_batches(
3278 &[],
3279 SnapshotMetadata {
3280 subscription_id: "sub-1",
3281 snapshot_id: "snapshot-1",
3282 authoritative: false,
3283 mode: Mode::State,
3284 view_id: "Thing/state",
3285 key: Some("missing"),
3286 },
3287 &SnapshotBatchConfig {
3288 initial_batch_size: 1,
3289 subsequent_batch_size: 1,
3290 },
3291 );
3292 assert_eq!(batches.len(), 1);
3293 assert!(!batches[0].authoritative);
3294 assert!(batches[0].complete);
3295 assert_eq!(batches[0].key.as_deref(), Some("missing"));
3296 }
3297
3298 #[test]
3299 fn lag_recovery_is_authoritative_for_an_after_query() {
3300 let subscription = Subscription {
3301 protocol_version: PROTOCOL_VERSION,
3302 subscription_id: "sub-1".to_string(),
3303 query: SubscriptionQuery {
3304 view: "Thing/list".to_string(),
3305 after: Some("40:000000000010".to_string()),
3306 ..Default::default()
3307 },
3308 snapshot: Default::default(),
3309 };
3310
3311 assert!(!SnapshotPurpose::Initial.authoritative(&subscription));
3312 assert!(SnapshotPurpose::Recovery.authoritative(&subscription));
3313 }
3314
3315 #[test]
3316 fn dot_path_filters_are_exact_and_type_sensitive() {
3317 let mut query = SubscriptionQuery {
3318 view: "Thing/list".to_string(),
3319 ..Default::default()
3320 };
3321 query
3322 .filters
3323 .insert("state.status".to_string(), json!("open"));
3324 query.filters.insert("metrics.count".to_string(), json!(2));
3325 assert!(query_matches_entity(
3326 &query,
3327 "one",
3328 &json!({"state": {"status": "open"}, "metrics": {"count": 2}}),
3329 ));
3330 assert!(!query_matches_entity(
3331 &query,
3332 "one",
3333 &json!({"state": {"status": "open"}, "metrics": {"count": "2"}}),
3334 ));
3335 }
3336
3337 #[test]
3338 fn take_and_skip_define_independent_deterministic_windows() {
3339 let entities: Vec<_> = (1..=6)
3340 .map(|id| {
3341 (
3342 id.to_string(),
3343 json!({"id": id, "_seq": format!("10:{id:012}")}),
3344 )
3345 })
3346 .collect();
3347 let first = SubscriptionQuery {
3348 view: "Thing/list".to_string(),
3349 take: Some(2),
3350 skip: Some(0),
3351 ..Default::default()
3352 };
3353 let second = SubscriptionQuery {
3354 skip: Some(2),
3355 ..first.clone()
3356 };
3357 let first_keys: Vec<_> = select_query_entities(entities.clone(), &first, false, false)
3358 .into_iter()
3359 .map(|(key, _)| key)
3360 .collect();
3361 let second_keys: Vec<_> = select_query_entities(entities, &second, false, false)
3362 .into_iter()
3363 .map(|(key, _)| key)
3364 .collect();
3365 assert_eq!(first_keys, ["6", "5"]);
3366 assert_eq!(second_keys, ["4", "3"]);
3367 }
3368
3369 #[tokio::test]
3370 async fn state_receiver_is_installed_before_snapshot_awaits() {
3371 let bus = BusManager::new();
3372 let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3373 let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3374 let bus_for_task = bus.clone();
3375 let task = tokio::spawn(async move {
3376 subscribe_state_then_snapshot(&bus_for_task, "Thing/state", "one", || async move {
3377 snapshot_started_tx.send(()).unwrap();
3378 release_snapshot_rx.await.unwrap();
3379 })
3380 .await
3381 .0
3382 });
3383 snapshot_started_rx.await.unwrap();
3384 bus.publish_state(
3385 "Thing/state",
3386 "one",
3387 Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3388 )
3389 .await;
3390 release_snapshot_tx.send(()).unwrap();
3391 let mut receiver = task.await.unwrap();
3392 receiver.changed().await.unwrap();
3393 assert!(!receiver.borrow().payload.is_empty());
3394 }
3395
3396 async fn assert_list_receiver_precedes_snapshot(view: &'static str) {
3397 let bus = BusManager::new();
3398 let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3399 let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3400 let bus_for_task = bus.clone();
3401 let task = tokio::spawn(async move {
3402 subscribe_list_then_snapshot(&bus_for_task, view, || async move {
3403 snapshot_started_tx.send(()).unwrap();
3404 release_snapshot_rx.await.unwrap();
3405 })
3406 .await
3407 .0
3408 });
3409 snapshot_started_rx.await.unwrap();
3410 bus.publish_list(
3411 view,
3412 Arc::new(BusMessage {
3413 key: "one".to_string(),
3414 entity: view.to_string(),
3415 payload: Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3416 }),
3417 )
3418 .await;
3419 release_snapshot_tx.send(()).unwrap();
3420 let mut receiver = task.await.unwrap();
3421 assert_eq!(receiver.recv().await.unwrap().key, "one");
3422 }
3423
3424 #[tokio::test]
3425 async fn list_receiver_is_installed_before_snapshot() {
3426 assert_list_receiver_precedes_snapshot("Thing/list").await;
3427 }
3428
3429 #[tokio::test]
3430 async fn append_receiver_is_installed_before_snapshot() {
3431 assert_list_receiver_precedes_snapshot("Thing/append").await;
3432 }
3433
3434 #[tokio::test]
3435 async fn derived_source_receiver_is_installed_before_snapshot() {
3436 assert_list_receiver_precedes_snapshot("Thing/list-source").await;
3437 }
3438
3439 #[tokio::test]
3440 async fn snapshot_limit_does_not_change_live_take_skip_membership() {
3441 let cache = EntityCache::with_config(EntityCacheConfig {
3442 max_entities_per_view: 10,
3443 ..Default::default()
3444 });
3445 for id in 1..=4 {
3446 cache
3447 .upsert(
3448 "Thing/list",
3449 &id.to_string(),
3450 json!({"_seq": format!("10:{id:012}")}),
3451 )
3452 .await;
3453 }
3454 let query = SubscriptionQuery {
3455 view: "Thing/list".to_string(),
3456 take: Some(3),
3457 skip: Some(1),
3458 snapshot_limit: Some(1),
3459 ..Default::default()
3460 };
3461 let live = load_query_entities(&cache, None, &list_spec(), &query, false).await;
3462 let snapshot = load_query_entities(&cache, None, &list_spec(), &query, true).await;
3463 assert_eq!(live.len(), 3);
3464 assert_eq!(snapshot.len(), 1);
3465 }
3466
3467 #[test]
3468 fn fixture_manifest_covers_required_conformance_cases() {
3469 let manifest: Value = serde_json::from_str(include_str!(
3470 "../../../../tests/fixtures/websocket-v2/manifest.json"
3471 ))
3472 .unwrap();
3473 let names: HashSet<_> = manifest["fixtures"]
3474 .as_array()
3475 .unwrap()
3476 .iter()
3477 .filter_map(Value::as_str)
3478 .collect();
3479 for required in [
3480 "keyed-state.json",
3481 "list-windows.json",
3482 "filters.json",
3483 "multi-batch-authoritative.json",
3484 "empty-snapshot.json",
3485 "remove.json",
3486 "delete.json",
3487 "incremental-snapshot.json",
3488 "reconnect-replacement.json",
3489 "errors.json",
3490 "whole-entities.json",
3491 ] {
3492 assert!(names.contains(required), "missing fixture {required}");
3493 }
3494
3495 for document in [
3496 include_str!("../../../../tests/fixtures/websocket-v2/keyed-state.json"),
3497 include_str!("../../../../tests/fixtures/websocket-v2/list-windows.json"),
3498 include_str!("../../../../tests/fixtures/websocket-v2/filters.json"),
3499 include_str!("../../../../tests/fixtures/websocket-v2/multi-batch-authoritative.json"),
3500 include_str!("../../../../tests/fixtures/websocket-v2/empty-snapshot.json"),
3501 include_str!("../../../../tests/fixtures/websocket-v2/remove.json"),
3502 include_str!("../../../../tests/fixtures/websocket-v2/delete.json"),
3503 include_str!("../../../../tests/fixtures/websocket-v2/incremental-snapshot.json"),
3504 include_str!("../../../../tests/fixtures/websocket-v2/reconnect-replacement.json"),
3505 include_str!("../../../../tests/fixtures/websocket-v2/errors.json"),
3506 include_str!("../../../../tests/fixtures/websocket-v2/whole-entities.json"),
3507 ] {
3508 let fixture: Value = serde_json::from_str(document).unwrap();
3509 assert!(fixture["name"].is_string());
3510 }
3511 }
3512
3513 fn append_frame(key: &str, data: Value) -> Arc<Bytes> {
3514 let frame = json!({
3515 "entity": "Trade/append",
3516 "op": "patch",
3517 "key": key,
3518 "offset": 7,
3519 "data": data,
3520 });
3521 Arc::new(Bytes::from(serde_json::to_vec(&frame).unwrap()))
3522 }
3523
3524 #[test]
3528 fn replay_delivery_applies_key_partition_and_filters() {
3529 let matching = append_frame("pool1", json!({"_partition": "us", "side": "buy"}));
3530 let other_partition = append_frame("pool1", json!({"_partition": "eu", "side": "buy"}));
3531 let other_side = append_frame("pool1", json!({"_partition": "us", "side": "sell"}));
3532
3533 let unfiltered = SubscriptionQuery {
3534 view: "Trade/append".to_string(),
3535 ..Default::default()
3536 };
3537 assert!(live_frame_matches(&unfiltered, "pool1", &matching));
3538 assert!(live_frame_matches(&unfiltered, "pool9", &matching));
3539
3540 let keyed = SubscriptionQuery {
3541 view: "Trade/append".to_string(),
3542 key: Some("pool1".to_string()),
3543 ..Default::default()
3544 };
3545 assert!(live_frame_matches(&keyed, "pool1", &matching));
3546 assert!(!live_frame_matches(&keyed, "pool2", &matching));
3547
3548 let partitioned = SubscriptionQuery {
3549 view: "Trade/append".to_string(),
3550 partition: Some("us".to_string()),
3551 ..Default::default()
3552 };
3553 assert!(live_frame_matches(&partitioned, "pool1", &matching));
3554 assert!(!live_frame_matches(&partitioned, "pool1", &other_partition));
3555
3556 let filtered = SubscriptionQuery {
3557 view: "Trade/append".to_string(),
3558 filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3559 ..Default::default()
3560 };
3561 assert!(live_frame_matches(&filtered, "pool1", &matching));
3562 assert!(!live_frame_matches(&filtered, "pool1", &other_side));
3563 }
3564
3565 #[test]
3566 fn a_frame_without_decodable_data_does_not_satisfy_a_filter() {
3567 let filtered = SubscriptionQuery {
3568 view: "Trade/append".to_string(),
3569 filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3570 ..Default::default()
3571 };
3572 let garbage = Arc::new(Bytes::from_static(b"not json"));
3573 assert!(!live_frame_matches(&filtered, "pool1", &garbage));
3574 }
3575
3576 #[test]
3580 fn replay_lagged_recovers_from_before_the_gap() {
3581 let epoch = crate::journal::JournalEpoch::new();
3582 let issue = SocketIssueMessage::replay_lagged(
3583 Some("trades".to_string()),
3584 37,
3585 Some(crate::journal::Cursor {
3586 epoch: epoch.clone(),
3587 offset: 4180,
3588 }),
3589 );
3590 assert_eq!(issue.code, "replay-lagged");
3591 assert_eq!(issue.recover_from, Some(format!("{epoch}:4180")));
3592 assert!(
3593 issue.suggested_action.unwrap().contains("4180"),
3594 "the consumer is told exactly which cursor recovers the gap"
3595 );
3596
3597 let from_scratch = SocketIssueMessage::replay_lagged(Some("trades".to_string()), 9, None);
3600 assert_eq!(from_scratch.recover_from, None);
3601 assert!(from_scratch
3602 .suggested_action
3603 .unwrap()
3604 .contains("without `after`"));
3605 }
3606
3607 fn bus_message(key: &str) -> Arc<BusMessage> {
3608 Arc::new(BusMessage {
3609 key: key.to_string(),
3610 entity: "Trade/append".to_string(),
3611 payload: Arc::new(Bytes::from_static(b"{}")),
3612 })
3613 }
3614
3615 #[tokio::test]
3620 async fn draining_during_a_replay_keeps_a_busy_view_from_lapping_it() {
3621 let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(16);
3622 let mut pending = VecDeque::new();
3623 let mut lagged = None;
3624
3625 for index in 0..48 {
3628 sender.send(bus_message(&format!("k{index}"))).unwrap();
3629 drain_available(&mut receiver, &mut pending, &mut lagged);
3630 }
3631
3632 assert_eq!(lagged, None, "draining as we go means nothing is dropped");
3633 assert_eq!(pending.len(), 48, "every published frame is buffered");
3634 }
3635
3636 #[tokio::test]
3637 async fn a_replay_that_never_drains_is_reported_as_a_gap() {
3638 let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3639 for index in 0..32 {
3640 sender.send(bus_message(&format!("k{index}"))).unwrap();
3641 }
3642
3643 let mut pending = VecDeque::new();
3644 let mut lagged = None;
3645 drain_available(&mut receiver, &mut pending, &mut lagged);
3646
3647 assert!(
3648 lagged.is_some(),
3649 "overflowing the bus is a gap, not silent truncation"
3650 );
3651 }
3652
3653 #[tokio::test]
3657 async fn nothing_after_a_gap_is_buffered_ahead_of_the_report() {
3658 let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3659 let mut pending = VecDeque::new();
3660 let mut lagged = None;
3661
3662 sender.send(bus_message("before")).unwrap();
3664 drain_available(&mut receiver, &mut pending, &mut lagged);
3665 assert_eq!(pending.len(), 1);
3666
3667 for index in 0..32 {
3669 sender.send(bus_message(&format!("lost{index}"))).unwrap();
3670 }
3671 drain_available(&mut receiver, &mut pending, &mut lagged);
3672 assert!(lagged.is_some());
3673
3674 let buffered_at_gap = pending.len();
3677 drain_available(&mut receiver, &mut pending, &mut lagged);
3678 assert_eq!(
3679 pending.len(),
3680 buffered_at_gap,
3681 "post-gap frames must not join the pre-gap flush"
3682 );
3683 assert_eq!(pending.front().unwrap().key, "before");
3684 }
3685 #[test]
3688 fn the_seam_between_replay_and_live_neither_repeats_nor_skips() {
3689 let mut last_sent = Some(4211);
3690
3691 assert!(
3692 already_delivered(Some(4211), &mut last_sent),
3693 "the record the replay ended on must not be sent twice"
3694 );
3695 assert!(already_delivered(Some(4100), &mut last_sent));
3696 assert_eq!(last_sent, Some(4211), "a duplicate never moves the mark");
3697
3698 assert!(
3699 !already_delivered(Some(4212), &mut last_sent),
3700 "the next record is new"
3701 );
3702 assert_eq!(last_sent, Some(4212));
3703
3704 assert!(!already_delivered(None, &mut last_sent));
3707 assert_eq!(last_sent, Some(4212));
3708 }
3709
3710 #[test]
3713 fn a_fresh_subscription_delivers_its_first_live_frame() {
3714 let mut last_sent = None;
3715 assert!(!already_delivered(Some(0), &mut last_sent));
3716 assert_eq!(last_sent, Some(0));
3717 }
3718
3719 mod over_a_socket {
3724 use super::*;
3725 use crate::journal::{EventJournal, JournalConfig};
3726 use crate::projector::Projector;
3727 use crate::{MutationBatch, SlotContext};
3728 use arete_interpreter::Mutation;
3729 use futures_util::{SinkExt, StreamExt};
3730 use std::time::Duration;
3731 use tokio::net::{TcpListener, TcpStream};
3732 use tokio::sync::mpsc;
3733 use tokio_tungstenite::tungstenite::Message;
3734 use tokio_tungstenite::{client_async, WebSocketStream};
3735
3736 const RETAINED: u64 = 600;
3737
3738 fn append_index() -> ViewIndex {
3739 let mut index = ViewIndex::new();
3740 index.add_spec(ViewSpec {
3741 id: "Trade/append".to_string(),
3742 export: "Trade".to_string(),
3743 mode: Mode::Append,
3744 wire_format: Default::default(),
3745 projection: Projection::all(),
3746 filters: Filters::all(),
3747 delivery: Delivery::default(),
3748 pipeline: None,
3749 source_view: None,
3750 });
3751 index
3752 }
3753
3754 fn trade(index: u64) -> MutationBatch {
3755 MutationBatch::with_slot_context(
3756 vec![Mutation {
3757 export: "Trade".to_string(),
3758 key: json!(format!("pool{}", index % 4)),
3759 patch: json!({"trade": index}),
3760 append: vec![],
3761 }]
3762 .into_iter()
3763 .collect(),
3764 SlotContext::new(100 + index / 3, index % 3),
3765 )
3766 }
3767
3768 struct Harness {
3769 addr: SocketAddr,
3770 journal: Arc<EventJournal>,
3771 tx: mpsc::Sender<MutationBatch>,
3772 }
3773
3774 impl Harness {
3775 async fn start() -> Self {
3776 let view_index = Arc::new(append_index());
3777 let entity_cache = EntityCache::new();
3778 let bus_manager = BusManager::new();
3779 let journal = Arc::new(EventJournal::new(JournalConfig {
3780 enabled: true,
3781 max_bytes_per_view: u64::MAX,
3782 max_records_per_view: 10_000,
3783 max_age: Duration::from_secs(3_600),
3784 }));
3785
3786 let (tx, rx) = mpsc::channel::<MutationBatch>(256);
3787 tokio::spawn(
3788 Projector::new(
3789 view_index.clone(),
3790 bus_manager.clone(),
3791 entity_cache.clone(),
3792 rx,
3793 #[cfg(feature = "otel")]
3794 None,
3795 )
3796 .with_journal(journal.clone())
3797 .run(),
3798 );
3799
3800 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3801 let addr = listener.local_addr().unwrap();
3802 let server = WebSocketServer::new(
3803 addr,
3804 bus_manager,
3805 entity_cache,
3806 view_index,
3807 #[cfg(feature = "otel")]
3808 None,
3809 )
3810 .with_journal(journal.clone());
3811 let (acceptor, _cleanup) = server.into_acceptor();
3812 tokio::spawn(async move { acceptor.serve_listener(listener).await });
3813
3814 Self { addr, journal, tx }
3815 }
3816
3817 async fn publish(&self, range: std::ops::Range<u64>) {
3818 for index in range {
3819 self.tx.send(trade(index)).await.unwrap();
3820 }
3821 let (ack, wait) = oneshot::channel();
3822 self.tx
3823 .send(MutationBatch::flush_marker(ack))
3824 .await
3825 .unwrap();
3826 wait.await.unwrap();
3827 }
3828
3829 async fn connect(&self) -> WebSocketStream<TcpStream> {
3830 let stream = TcpStream::connect(self.addr).await.unwrap();
3831 client_async(format!("ws://{}/", self.addr), stream)
3832 .await
3833 .unwrap()
3834 .0
3835 }
3836 }
3837
3838 async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
3839 loop {
3840 let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
3841 .await
3842 .expect("the server answers within the timeout")
3843 .expect("the stream stays open")
3844 .expect("a readable frame");
3845 let bytes = match &message {
3848 Message::Text(text) => text.as_bytes(),
3849 Message::Binary(bytes) => bytes.as_ref(),
3850 _ => continue,
3851 };
3852 return serde_json::from_slice(bytes).expect("frames are JSON");
3853 }
3854 }
3855
3856 async fn collect_trades(socket: &mut WebSocketStream<TcpStream>, count: usize) -> Vec<u64> {
3858 let mut offsets = Vec::with_capacity(count);
3859 while offsets.len() < count {
3860 let frame = next_frame(socket).await;
3861 assert_ne!(
3862 frame["type"], "error",
3863 "no error frame should interrupt delivery: {frame}"
3864 );
3865 if let Some(offset) = frame["offset"].as_u64() {
3866 offsets.push(offset);
3867 }
3868 }
3869 offsets
3870 }
3871
3872 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3876 async fn a_reconnect_replays_from_a_cursor_and_continues_live() {
3877 let harness = Harness::start().await;
3878 harness.publish(0..RETAINED).await;
3879
3880 let cursor = harness.journal.window("Trade/append").await;
3881 let stored = format!("{}:{}", cursor.epoch, 99);
3882
3883 let mut socket = harness.connect().await;
3884 socket
3885 .send(Message::Text(
3886 json!({
3887 "type": "subscribe",
3888 "protocolVersion": 2,
3889 "subscriptionId": "trades",
3890 "query": {"view": "Trade/append", "after": stored},
3891 })
3892 .to_string()
3893 .into(),
3894 ))
3895 .await
3896 .unwrap();
3897
3898 let ack = next_frame(&mut socket).await;
3899 assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3900 assert_eq!(ack["replayWindow"]["next"], json!(RETAINED));
3901
3902 let replayed = collect_trades(&mut socket, (RETAINED - 100) as usize).await;
3904 assert_eq!(
3905 replayed,
3906 (100..RETAINED).collect::<Vec<_>>(),
3907 "every event after the cursor, in order, exactly once"
3908 );
3909
3910 harness.publish(RETAINED..RETAINED + 40).await;
3912 let live = collect_trades(&mut socket, 40).await;
3913 assert_eq!(
3914 live,
3915 (RETAINED..RETAINED + 40).collect::<Vec<_>>(),
3916 "the live stream resumes exactly where the replay stopped"
3917 );
3918
3919 socket.close(None).await.ok();
3920 }
3921
3922 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3927 async fn events_published_during_a_replay_are_not_lost() {
3928 let harness = Harness::start().await;
3929 harness.publish(0..RETAINED).await;
3930
3931 let epoch = harness.journal.window("Trade/append").await.epoch;
3932 let mut socket = harness.connect().await;
3933 socket
3934 .send(Message::Text(
3935 json!({
3936 "type": "subscribe",
3937 "protocolVersion": 2,
3938 "subscriptionId": "trades",
3939 "query": {"view": "Trade/append", "after": format!("{epoch}:0")},
3940 })
3941 .to_string()
3942 .into(),
3943 ))
3944 .await
3945 .unwrap();
3946 let ack = next_frame(&mut socket).await;
3947 assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3948
3949 harness.publish(RETAINED..RETAINED + 200).await;
3951
3952 let total = (RETAINED + 200 - 1) as usize;
3953 let delivered = collect_trades(&mut socket, total).await;
3954 assert_eq!(
3955 delivered,
3956 (1..RETAINED + 200).collect::<Vec<_>>(),
3957 "replay and live output join without a gap or a repeat"
3958 );
3959
3960 socket.close(None).await.ok();
3961 }
3962
3963 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3966 async fn a_stale_epoch_is_refused_and_frees_the_subscription_id() {
3967 let harness = Harness::start().await;
3968 harness.publish(0..50).await;
3969
3970 let mut socket = harness.connect().await;
3971 socket
3972 .send(Message::Text(
3973 json!({
3974 "type": "subscribe",
3975 "protocolVersion": 2,
3976 "subscriptionId": "trades",
3977 "query": {
3978 "view": "Trade/append",
3979 "after": format!("{}:10", crate::journal::JournalEpoch::new()),
3980 },
3981 })
3982 .to_string()
3983 .into(),
3984 ))
3985 .await
3986 .unwrap();
3987
3988 let error = next_frame(&mut socket).await;
3989 assert_eq!(error["type"], "error", "unexpected frame: {error}");
3990 assert_eq!(error["code"], "cursor-epoch-changed");
3991
3992 socket
3995 .send(Message::Text(
3996 json!({
3997 "type": "subscribe",
3998 "protocolVersion": 2,
3999 "subscriptionId": "trades",
4000 "query": {"view": "Trade/append"},
4001 })
4002 .to_string()
4003 .into(),
4004 ))
4005 .await
4006 .unwrap();
4007 let ack = next_frame(&mut socket).await;
4008 assert_eq!(
4009 ack["op"], "subscribed",
4010 "the refused id must be reusable: {ack}"
4011 );
4012
4013 assert_eq!(collect_trades(&mut socket, 50).await.len(), 50);
4014 socket.close(None).await.ok();
4015 }
4016
4017 const ROUND: &str = "Round/state";
4018
4019 #[derive(Clone)]
4022 struct StateServer {
4023 addr: SocketAddr,
4024 bus_manager: BusManager,
4025 entity_cache: EntityCache,
4026 }
4027
4028 impl StateServer {
4029 async fn start() -> Self {
4030 let mut index = ViewIndex::new();
4031 index.add_spec(ViewSpec {
4032 id: ROUND.to_string(),
4033 export: "Round".to_string(),
4034 mode: Mode::State,
4035 wire_format: Default::default(),
4036 projection: Projection::all(),
4037 filters: Filters::all(),
4038 delivery: Delivery::default(),
4039 pipeline: None,
4040 source_view: None,
4041 });
4042 let bus_manager = BusManager::new();
4043 let entity_cache = EntityCache::new();
4044 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4045 let addr = listener.local_addr().unwrap();
4046 let server = WebSocketServer::new(
4047 addr,
4048 bus_manager.clone(),
4049 entity_cache.clone(),
4050 Arc::new(index),
4051 #[cfg(feature = "otel")]
4052 None,
4053 );
4054 let (acceptor, _cleanup) = server.into_acceptor();
4055 tokio::spawn(async move { acceptor.serve_listener(listener).await });
4056 Self {
4057 addr,
4058 bus_manager,
4059 entity_cache,
4060 }
4061 }
4062
4063 async fn publish(&self, key: &str, mut patch: Value, seq: &str) {
4066 patch["_seq"] = Value::String(seq.to_string());
4067 self.entity_cache
4068 .upsert_with_append(ROUND, key, patch.clone(), &[], PatchOrigin::Unknown)
4069 .await;
4070 self.publish_frame(key, "patch", patch, seq).await;
4071 }
4072
4073 async fn publish_frame(&self, key: &str, op: &str, data: Value, seq: &str) {
4075 let frame = json!({
4076 "mode": "state",
4077 "entity": ROUND,
4078 "op": op,
4079 "key": key,
4080 "data": data,
4081 "seq": seq,
4082 });
4083 self.bus_manager
4084 .publish_state(ROUND, key, Arc::new(Bytes::from(frame.to_string())))
4085 .await;
4086 }
4087
4088 async fn subscribe(&self, key: &str) -> WebSocketStream<TcpStream> {
4090 let stream = TcpStream::connect(self.addr).await.unwrap();
4091 let mut socket = client_async(format!("ws://{}/", self.addr), stream)
4092 .await
4093 .unwrap()
4094 .0;
4095 socket
4096 .send(Message::Text(
4097 json!({
4098 "type": "subscribe",
4099 "protocolVersion": 2,
4100 "subscriptionId": "round",
4101 "query": {"view": ROUND, "key": key},
4102 })
4103 .to_string()
4104 .into(),
4105 ))
4106 .await
4107 .unwrap();
4108 let ack = next_frame(&mut socket).await;
4109 assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
4110 loop {
4111 let frame = next_frame(&mut socket).await;
4112 assert_eq!(frame["op"], "snapshot", "unexpected frame: {frame}");
4113 if frame["complete"] == true {
4114 return socket;
4115 }
4116 }
4117 }
4118 }
4119
4120 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4126 async fn a_state_subscriber_keeps_the_fields_of_an_overwritten_patch() {
4127 let server = StateServer::start().await;
4128 server
4129 .publish("7", json!({"id": 7}), "100:000000000001")
4130 .await;
4131 let mut socket = server.subscribe("7").await;
4132
4133 server
4136 .publish(
4137 "7",
4138 json!({"results": {"slot_hash": "abc"}}),
4139 "140:003836292257",
4140 )
4141 .await;
4142 let frame = next_frame(&mut socket).await;
4143 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4144 assert_eq!(frame["seq"], "140:003836292257");
4145
4146 let writer = server.clone();
4151 tokio::spawn(async move {
4152 writer
4153 .publish("7", json!({"entropy": {"seed": "def"}}), "140:000000001200")
4154 .await;
4155 writer
4156 .publish("7", json!({"total": 1}), "140:000000001200")
4157 .await;
4158 })
4159 .await
4160 .unwrap();
4161 let frame = next_frame(&mut socket).await;
4162 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4163 assert_eq!(
4166 frame["data"],
4167 json!({
4168 "results": {"slot_hash": "abc"},
4169 "entropy": {"seed": "def"},
4170 "total": 1,
4171 "_seq": "140:000000001200",
4172 })
4173 );
4174 assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4177
4178 server
4180 .publish("7", json!({"total": 2}), "141:000000000002")
4181 .await;
4182 let frame = next_frame(&mut socket).await;
4183 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4184 assert_eq!(
4185 frame["data"],
4186 json!({"total": 2, "_seq": "141:000000000002"})
4187 );
4188
4189 let writer = server.clone();
4192 tokio::spawn(async move {
4193 writer
4194 .publish("7", json!({"total": 3}), "142:000000000001")
4195 .await;
4196 writer
4197 .publish("7", json!({"flag": true}), "142:000000000001")
4198 .await;
4199 })
4200 .await
4201 .unwrap();
4202 let frame = next_frame(&mut socket).await;
4203 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4204 assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4205 assert_eq!(
4206 frame["data"],
4207 json!({"total": 3, "flag": true, "_seq": "142:000000000001"})
4208 );
4209
4210 socket.close(None).await.ok();
4211 }
4212
4213 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4216 async fn a_forwarded_whole_entity_is_what_a_catch_up_is_measured_against() {
4217 let server = StateServer::start().await;
4218 server
4219 .publish("7", json!({"id": 7, "a": 1}), "100:000000000001")
4220 .await;
4221 let mut socket = server.subscribe("7").await;
4222
4223 let whole = json!({"id": 7, "a": 1, "b": 2, "_seq": "101:000000000001"});
4224 server
4225 .entity_cache
4226 .store_whole(ROUND, "7", whole.clone())
4227 .await;
4228 server
4229 .publish_frame("7", "upsert", whole.clone(), "101:000000000001")
4230 .await;
4231 let frame = next_frame(&mut socket).await;
4232 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4233 assert_eq!(frame["data"], whole);
4234
4235 let writer = server.clone();
4236 tokio::spawn(async move {
4237 writer
4238 .publish("7", json!({"c": 3}), "102:000000000001")
4239 .await;
4240 writer
4241 .publish("7", json!({"d": 4}), "102:000000000002")
4242 .await;
4243 })
4244 .await
4245 .unwrap();
4246 let frame = next_frame(&mut socket).await;
4247 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4248 assert_eq!(
4249 frame["data"],
4250 json!({"c": 3, "d": 4, "_seq": "102:000000000002"}),
4251 "`b` came with the whole entity, so it is not sent again"
4252 );
4253 socket.close(None).await.ok();
4254 }
4255
4256 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4260 async fn a_catch_up_restores_a_field_a_forwarded_patch_set() {
4261 let server = StateServer::start().await;
4262 server
4263 .publish("7", json!({"id": 7, "a": 1}), "100:000000000001")
4264 .await;
4265 let mut socket = server.subscribe("7").await;
4266
4267 server
4268 .publish("7", json!({"a": 2}), "101:000000000001")
4269 .await;
4270 let frame = next_frame(&mut socket).await;
4271 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4272 assert_eq!(frame["data"], json!({"a": 2, "_seq": "101:000000000001"}));
4273
4274 let writer = server.clone();
4275 tokio::spawn(async move {
4276 writer
4277 .publish("7", json!({"a": 1}), "102:000000000001")
4278 .await;
4279 writer
4280 .publish("7", json!({"b": 1}), "102:000000000002")
4281 .await;
4282 })
4283 .await
4284 .unwrap();
4285 let frame = next_frame(&mut socket).await;
4286 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4287 assert_eq!(
4288 frame["data"],
4289 json!({"a": 1, "b": 1, "_seq": "102:000000000002"}),
4290 "`a` matches the snapshot but the client holds 2"
4291 );
4292 socket.close(None).await.ok();
4293 }
4294
4295 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4300 async fn a_key_deleted_and_recreated_between_reads_replaces_the_clients_copy() {
4301 let server = StateServer::start().await;
4302 server
4303 .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
4304 .await;
4305 let mut socket = server.subscribe("7").await;
4306
4307 let writer = server.clone();
4308 tokio::spawn(async move {
4309 writer.entity_cache.remove(ROUND, "7").await;
4310 writer
4311 .publish_frame("7", "delete", Value::Null, "101:000000000001")
4312 .await;
4313 writer
4314 .publish("7", json!({"id": 7, "fresh": true}), "101:000000000002")
4315 .await;
4316 })
4317 .await
4318 .unwrap();
4319 let frame = next_frame(&mut socket).await;
4320 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4321 assert_eq!(
4322 frame["data"],
4323 json!({"id": 7, "fresh": true, "_seq": "101:000000000002"})
4324 );
4325 assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4326
4327 let writer = server.clone();
4329 tokio::spawn(async move {
4330 writer
4331 .publish("7", json!({"n": 1}), "102:000000000001")
4332 .await;
4333 writer.entity_cache.remove(ROUND, "7").await;
4334 writer
4335 .publish_frame("7", "delete", Value::Null, "102:000000000002")
4336 .await;
4337 })
4338 .await
4339 .unwrap();
4340 let frame = next_frame(&mut socket).await;
4341 assert_eq!(frame["op"], "delete", "unexpected frame: {frame}");
4342
4343 server
4345 .publish("7", json!({"id": 7, "third": true}), "103:000000000001")
4346 .await;
4347 let frame = next_frame(&mut socket).await;
4348 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4349 assert_eq!(frame["data"]["third"], true);
4350 assert!(frame["data"].get("fresh").is_none(), "{frame}");
4351 socket.close(None).await.ok();
4352 }
4353
4354 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4358 async fn a_copy_of_a_deleted_entity_takes_no_patch_until_it_is_replaced() {
4359 let server = StateServer::start().await;
4360 server
4361 .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
4362 .await;
4363 let mut socket = server.subscribe("7").await;
4364
4365 let writer = server.clone();
4366 tokio::spawn(async move {
4367 writer.entity_cache.remove(ROUND, "7").await;
4368 writer
4369 .publish_frame("7", "delete", Value::Null, "101:000000000001")
4370 .await;
4371 writer
4372 .publish_frame("7", "patch", json!({"n": 1}), "101:000000000002")
4373 .await;
4374 })
4375 .await
4376 .unwrap();
4377 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
4381
4382 let whole = json!({"id": 7, "n": 1, "_seq": "101:000000000002"});
4383 server
4384 .entity_cache
4385 .store_whole(ROUND, "7", whole.clone())
4386 .await;
4387 server
4388 .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
4389 .await;
4390 let frame = next_frame(&mut socket).await;
4391 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4392 assert_eq!(frame["data"], whole);
4393 socket.close(None).await.ok();
4394 }
4395
4396 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4400 async fn a_whole_entity_among_overwritten_frames_replaces_the_clients_copy() {
4401 let server = StateServer::start().await;
4402 server
4403 .publish("7", json!({"id": 7, "dropped": 1}), "100:000000000001")
4404 .await;
4405 let mut socket = server.subscribe("7").await;
4406
4407 let writer = server.clone();
4408 tokio::spawn(async move {
4409 let whole = json!({"id": 7, "_seq": "101:000000000001"});
4410 writer
4411 .entity_cache
4412 .store_whole(ROUND, "7", whole.clone())
4413 .await;
4414 writer
4415 .publish_frame("7", "upsert", whole, "101:000000000001")
4416 .await;
4417 writer
4418 .publish("7", json!({"n": 1}), "101:000000000002")
4419 .await;
4420 })
4421 .await
4422 .unwrap();
4423 let frame = next_frame(&mut socket).await;
4424 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4425 assert_eq!(
4426 frame["data"],
4427 json!({"id": 7, "n": 1, "_seq": "101:000000000002"})
4428 );
4429 assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4430 socket.close(None).await.ok();
4431 }
4432
4433 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4442 async fn a_holder_behind_on_an_evicted_key_catches_up_when_it_returns() {
4443 let server = StateServer::start().await;
4444 server
4445 .publish("7", json!({"id": 7}), "100:000000000001")
4446 .await;
4447 let mut socket = server.subscribe("7").await;
4448 server.entity_cache.remove(ROUND, "7").await;
4449
4450 let writer = server.clone();
4451 tokio::spawn(async move {
4452 writer
4453 .publish_frame("7", "patch", json!({"a": 1}), "101:000000000001")
4454 .await;
4455 writer
4456 .publish_frame("7", "patch", json!({"b": 2}), "101:000000000002")
4457 .await;
4458 })
4459 .await
4460 .unwrap();
4461 let frame = next_frame(&mut socket).await;
4462 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4463 assert_eq!(frame["data"], json!({"b": 2}));
4464 assert_eq!(frame["seq"], "101:000000000002");
4465
4466 let whole = json!({"id": 7, "a": 1, "b": 2, "_seq": "101:000000000002"});
4468 server
4469 .entity_cache
4470 .store_whole(ROUND, "7", whole.clone())
4471 .await;
4472 server
4473 .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
4474 .await;
4475 let frame = next_frame(&mut socket).await;
4476 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4477 assert_eq!(frame["data"], whole);
4478 assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4479
4480 server
4482 .publish("7", json!({"c": 3}), "102:000000000001")
4483 .await;
4484 let frame = next_frame(&mut socket).await;
4485 assert_eq!(frame["data"], json!({"c": 3, "_seq": "102:000000000001"}));
4486 assert_eq!(frame["seq"], "102:000000000001");
4487 socket.close(None).await.ok();
4488 }
4489 }
4490
4491 mod partial_entities_over_a_socket {
4495 use super::*;
4496 use crate::projector::Projector;
4497 use crate::{MutationBatch, SlotContext};
4498 use arete_interpreter::Mutation;
4499 use futures_util::{SinkExt, StreamExt};
4500 use std::time::Duration;
4501 use tokio::net::{TcpListener, TcpStream};
4502 use tokio::sync::mpsc;
4503 use tokio_tungstenite::tungstenite::Message;
4504 use tokio_tungstenite::{client_async, WebSocketStream};
4505
4506 fn thing_index() -> ViewIndex {
4507 let mut index = ViewIndex::new();
4508 index.add_spec(list_spec());
4509 index.add_spec(ViewSpec {
4510 id: "Thing/state".to_string(),
4511 mode: Mode::State,
4512 ..list_spec()
4513 });
4514 index
4515 }
4516
4517 struct Harness {
4518 addr: SocketAddr,
4519 tx: mpsc::Sender<MutationBatch>,
4520 slot: std::sync::atomic::AtomicU64,
4521 }
4522
4523 impl Harness {
4524 async fn start(delivery: WebSocketDeliveryConfig) -> Self {
4525 Self::start_with_cache(delivery, EntityCache::new()).await
4526 }
4527
4528 async fn start_with_cache(
4529 delivery: WebSocketDeliveryConfig,
4530 entity_cache: EntityCache,
4531 ) -> Self {
4532 let view_index = Arc::new(thing_index());
4533 let bus_manager = BusManager::new();
4534 let (tx, rx) = mpsc::channel::<MutationBatch>(64);
4535 tokio::spawn(tokio::task::unconstrained(
4541 Projector::new(
4542 view_index.clone(),
4543 bus_manager.clone(),
4544 entity_cache.clone(),
4545 rx,
4546 #[cfg(feature = "otel")]
4547 None,
4548 )
4549 .run(),
4550 ));
4551 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4552 let addr = listener.local_addr().unwrap();
4553 let server = WebSocketServer::new(
4554 addr,
4555 bus_manager,
4556 entity_cache,
4557 view_index,
4558 #[cfg(feature = "otel")]
4559 None,
4560 )
4561 .with_delivery_config(delivery);
4562 let (acceptor, _cleanup) = server.into_acceptor();
4563 tokio::spawn(async move { acceptor.serve_listener(listener).await });
4564 Self {
4565 addr,
4566 tx,
4567 slot: std::sync::atomic::AtomicU64::new(100),
4568 }
4569 }
4570
4571 async fn patch(&self, key: &str, patch: Value) {
4573 self.mutate(Mutation {
4574 export: "Thing".to_string(),
4575 key: json!(key),
4576 patch,
4577 append: vec![],
4578 })
4579 .await;
4580 }
4581
4582 async fn mutate(&self, mutation: Mutation) {
4583 let slot = self.slot.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
4584 self.tx
4585 .send(MutationBatch::with_slot_context(
4586 vec![mutation].into_iter().collect(),
4587 SlotContext::new(slot, 0),
4588 ))
4589 .await
4590 .unwrap();
4591 let (ack, wait) = oneshot::channel();
4592 self.tx
4593 .send(MutationBatch::flush_marker(ack))
4594 .await
4595 .unwrap();
4596 wait.await.unwrap();
4597 }
4598
4599 async fn subscribe(&self, query: Value, snapshot: bool) -> WebSocketStream<TcpStream> {
4600 let stream = TcpStream::connect(self.addr).await.unwrap();
4601 let mut socket = client_async(format!("ws://{}/", self.addr), stream)
4602 .await
4603 .unwrap()
4604 .0;
4605 socket
4606 .send(Message::Text(
4607 json!({
4608 "type": "subscribe",
4609 "protocolVersion": 2,
4610 "subscriptionId": "things",
4611 "query": query,
4612 "snapshot": {"enabled": snapshot},
4613 })
4614 .to_string()
4615 .into(),
4616 ))
4617 .await
4618 .unwrap();
4619 let ack = next_frame(&mut socket).await;
4620 assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
4621 assert_eq!(
4622 ack["wholeEntities"],
4623 json!(true),
4624 "the ack advertises the whole-entity guarantee: {ack}"
4625 );
4626 socket
4627 }
4628 }
4629
4630 async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
4631 loop {
4632 let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
4633 .await
4634 .expect("the server answers within the timeout")
4635 .expect("the stream stays open")
4636 .expect("a readable frame");
4637 let bytes = match &message {
4639 Message::Text(text) => text.as_bytes(),
4640 Message::Binary(bytes) => bytes.as_ref(),
4641 _ => continue,
4642 };
4643 return serde_json::from_slice(bytes).expect("frames are JSON");
4644 }
4645 }
4646
4647 async fn snapshot_keys(socket: &mut WebSocketStream<TcpStream>) -> Vec<String> {
4649 let mut keys = Vec::new();
4650 loop {
4651 let frame = next_frame(socket).await;
4652 assert_eq!(frame["op"], "snapshot", "expected a snapshot: {frame}");
4653 keys.extend(
4654 frame["data"]
4655 .as_array()
4656 .unwrap()
4657 .iter()
4658 .map(|row| row["key"].as_str().unwrap().to_string()),
4659 );
4660 if frame["complete"] == json!(true) {
4661 return keys;
4662 }
4663 }
4664 }
4665
4666 async fn seed(harness: &Harness) {
4667 for key in ["a", "b", "c"] {
4668 harness
4669 .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
4670 .await;
4671 }
4672 }
4673
4674 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4675 async fn a_change_to_a_key_cut_by_snapshot_limit_is_an_upsert() {
4676 let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4677 seed(&harness).await;
4678
4679 let mut socket = harness
4680 .subscribe(json!({"view": "Thing/list", "snapshotLimit": 1}), true)
4681 .await;
4682 assert_eq!(snapshot_keys(&mut socket).await, ["c"]);
4684
4685 harness.patch("a", json!({"count": 2})).await;
4686 let frame = next_frame(&mut socket).await;
4687 assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4688 assert_eq!(frame["op"], "upsert", "a key never sent arrives whole");
4689 assert_eq!(frame["data"]["name"], "thing-a");
4690 assert_eq!(frame["data"]["count"], 2);
4691
4692 harness.patch("a", json!({"count": 3})).await;
4694 let frame = next_frame(&mut socket).await;
4695 assert_eq!(frame["key"], "a");
4696 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4697 assert!(frame["data"].get("name").is_none());
4698
4699 harness.patch("c", json!({"count": 2})).await;
4701 let frame = next_frame(&mut socket).await;
4702 assert_eq!(frame["key"], "c");
4703 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4704 socket.close(None).await.ok();
4705 }
4706
4707 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4708 async fn a_list_without_a_snapshot_upserts_each_key_on_its_first_change() {
4709 let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4710 seed(&harness).await;
4711
4712 let mut socket = harness
4713 .subscribe(json!({"view": "Thing/list"}), false)
4714 .await;
4715 harness.patch("b", json!({"count": 2})).await;
4716 let frame = next_frame(&mut socket).await;
4717 assert_eq!(frame["key"], "b");
4718 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4719 assert_eq!(frame["data"]["name"], "thing-b");
4720
4721 harness.patch("b", json!({"count": 3})).await;
4722 let frame = next_frame(&mut socket).await;
4723 assert_eq!(frame["key"], "b");
4724 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4725 socket.close(None).await.ok();
4726 }
4727
4728 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4732 async fn a_state_holder_keeps_its_patches_when_the_cache_evicts_the_key() {
4733 let harness = Harness::start_with_cache(
4734 WebSocketDeliveryConfig::default(),
4735 EntityCache::with_config(crate::cache::EntityCacheConfig {
4736 max_entities_per_view: 2,
4737 ..Default::default()
4738 }),
4739 )
4740 .await;
4741 harness
4742 .patch("a", json!({"name": "thing-a", "count": 1}))
4743 .await;
4744 let mut socket = harness
4745 .subscribe(json!({"view": "Thing/state", "key": "a"}), true)
4746 .await;
4747 assert_eq!(snapshot_keys(&mut socket).await, ["a"]);
4748
4749 for key in ["b", "c"] {
4751 harness
4752 .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
4753 .await;
4754 }
4755 harness.patch("a", json!({"count": 2})).await;
4756 let frame = next_frame(&mut socket).await;
4757 assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4758 assert_eq!(frame["op"], "patch", "the holder merges it: {frame}");
4759 assert_eq!(frame["data"]["count"], 2);
4760
4761 let mut whole = Mutation {
4763 export: "Thing".to_string(),
4764 key: json!("a"),
4765 patch: json!({"name": "thing-a", "count": 2}),
4766 append: vec![],
4767 };
4768 whole.mark_whole_entity();
4769 harness.mutate(whole).await;
4770 let frame = next_frame(&mut socket).await;
4771 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4772 assert_eq!(frame["data"]["name"], "thing-a");
4773 assert_eq!(frame["data"]["count"], 2);
4774 assert!(frame["data"]
4775 .get(arete_interpreter::WHOLE_ENTITY_MARKER)
4776 .is_none());
4777
4778 harness.patch("a", json!({"count": 3})).await;
4779 let frame = next_frame(&mut socket).await;
4780 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4781 socket.close(None).await.ok();
4782 }
4783
4784 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4785 async fn a_state_view_without_a_snapshot_upserts_on_the_first_change() {
4786 let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4787 seed(&harness).await;
4788
4789 let mut socket = harness
4790 .subscribe(json!({"view": "Thing/state", "key": "a"}), false)
4791 .await;
4792 harness.patch("a", json!({"count": 2})).await;
4793 let frame = next_frame(&mut socket).await;
4794 assert_eq!(frame["key"], "a");
4795 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4796 assert_eq!(frame["data"]["name"], "thing-a");
4797 assert_eq!(frame["data"]["count"], 2);
4798
4799 harness.patch("a", json!({"count": 3})).await;
4800 let frame = next_frame(&mut socket).await;
4801 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4802 socket.close(None).await.ok();
4803 }
4804
4805 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4806 async fn a_coalesced_list_sends_no_remove_for_a_key_it_never_sent() {
4807 let harness = Harness::start(WebSocketDeliveryConfig {
4808 collection_coalesce_ms: Some(10),
4809 ..Default::default()
4810 })
4811 .await;
4812 seed(&harness).await;
4813
4814 let mut socket = harness
4817 .subscribe(json!({"view": "Thing/list", "take": 1}), false)
4818 .await;
4819 harness.patch("a", json!({"count": 2})).await;
4822 let frame = next_frame(&mut socket).await;
4823 assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4824 assert_eq!(frame["op"], "upsert");
4825 assert_eq!(frame["data"]["name"], "thing-a");
4826
4827 harness.patch("c", json!({"count": 2})).await;
4829 let mut frames = [next_frame(&mut socket).await, next_frame(&mut socket).await];
4830 frames.sort_by_key(|frame| frame["key"].as_str().unwrap().to_string());
4831 assert_eq!(frames[0]["key"], "a");
4832 assert_eq!(frames[0]["op"], "remove");
4833 assert_eq!(frames[1]["key"], "c");
4834 assert_eq!(frames[1]["op"], "upsert");
4835 assert_eq!(frames[1]["data"]["name"], "thing-c");
4836 socket.close(None).await.ok();
4837 }
4838
4839 fn version_of(data: &Value) -> (String, u64) {
4841 let version = data["_version"]
4842 .as_str()
4843 .unwrap_or_else(|| panic!("no _version in {data}"));
4844 let (epoch, counter) = version.split_once(':').expect("epoch:counter");
4845 (epoch.to_string(), counter.parse().expect("decimal counter"))
4846 }
4847
4848 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4857 async fn snapshot_rows_patches_and_catch_ups_carry_the_version() {
4858 let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4859 harness.patch("7", json!({"name": "thing-7"})).await;
4860
4861 let mut list = harness.subscribe(json!({"view": "Thing/list"}), true).await;
4862 let snapshot = next_frame(&mut list).await;
4863 assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
4864 let (epoch, _) = version_of(&snapshot["data"][0]["data"]);
4865
4866 let mut state = harness
4867 .subscribe(json!({"view": "Thing/state", "key": "7"}), true)
4868 .await;
4869 let snapshot = next_frame(&mut state).await;
4870 assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
4871 let (state_epoch, seeded) = version_of(&snapshot["data"][0]["data"]);
4872 assert_eq!(state_epoch, epoch, "one projector, one epoch");
4873
4874 harness.patch("7", json!({"count": 1})).await;
4875 let patch = next_frame(&mut list).await;
4876 assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
4877 assert_eq!(version_of(&patch["data"]).0, epoch);
4878 let patch = next_frame(&mut state).await;
4879 assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
4880 let (_, patched) = version_of(&patch["data"]);
4881 assert!(patched > seeded, "{patched} after {seeded}");
4882
4883 let slot = harness
4885 .slot
4886 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
4887 let mutation = |patch: Value| Mutation {
4888 export: "Thing".to_string(),
4889 key: json!("7"),
4890 patch,
4891 append: vec![],
4892 };
4893 harness
4894 .tx
4895 .send(MutationBatch::with_slot_context(
4896 vec![
4897 mutation(json!({"count": 2})),
4898 mutation(json!({"flag": true})),
4899 ]
4900 .into_iter()
4901 .collect(),
4902 SlotContext::new(slot, 0),
4903 ))
4904 .await
4905 .unwrap();
4906
4907 let catch_up = next_frame(&mut state).await;
4908 assert_eq!(catch_up["op"], "patch", "unexpected frame: {catch_up}");
4909 assert!(catch_up.get("seq").is_none(), "a catch-up: {catch_up}");
4910 assert_eq!(catch_up["data"]["count"], 2);
4911 assert_eq!(catch_up["data"]["flag"], true);
4912 let (_, caught_up) = version_of(&catch_up["data"]);
4913 assert!(caught_up > patched, "{caught_up} after {patched}");
4914
4915 list.close(None).await.ok();
4916 state.close(None).await.ok();
4917 }
4918 }
4919
4920 mod session_expiry {
4922 use super::*;
4923 use crate::websocket::auth::SignedSessionAuthPlugin;
4924 use arete_auth::{KeyClass, SessionClaims, SigningKey, TokenSigner, TokenVerifier};
4925 use futures_util::{SinkExt, StreamExt};
4926 use std::time::Duration;
4927 use tokio::net::{TcpListener, TcpStream};
4928 use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
4929 use tokio_tungstenite::tungstenite::protocol::CloseFrame;
4930 use tokio_tungstenite::tungstenite::Message;
4931 use tokio_tungstenite::{client_async, WebSocketStream};
4932
4933 struct Server {
4934 addr: SocketAddr,
4935 signer: TokenSigner,
4936 }
4937
4938 impl Server {
4939 async fn start() -> Self {
4940 let signing_key = SigningKey::generate();
4941 let verifier =
4942 TokenVerifier::new(signing_key.verifying_key(), "test-issuer", "test-audience");
4943 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4944 let addr = listener.local_addr().unwrap();
4945 let server = WebSocketServer::new(
4946 addr,
4947 BusManager::new(),
4948 EntityCache::new(),
4949 Arc::new(ViewIndex::new()),
4950 #[cfg(feature = "otel")]
4951 None,
4952 )
4953 .with_auth_plugin(Arc::new(SignedSessionAuthPlugin::new(verifier)));
4954 let (acceptor, _cleanup) = server.into_acceptor();
4955 tokio::spawn(async move { acceptor.serve_listener(listener).await });
4956 Self {
4957 addr,
4958 signer: TokenSigner::new(signing_key, "test-issuer"),
4959 }
4960 }
4961
4962 fn token(&self, ttl_seconds: u64) -> String {
4963 let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
4964 .with_scope("read")
4965 .with_key_class(KeyClass::Secret)
4966 .with_ttl(ttl_seconds)
4967 .build();
4968 self.signer.sign(claims).unwrap()
4969 }
4970
4971 async fn connect(&self, token: &str) -> WebSocketStream<TcpStream> {
4972 let stream = TcpStream::connect(self.addr).await.unwrap();
4973 client_async(format!("ws://{}/?hs_token={token}", self.addr), stream)
4974 .await
4975 .unwrap()
4976 .0
4977 }
4978 }
4979
4980 async fn send_json(socket: &mut WebSocketStream<TcpStream>, message: Value) {
4981 socket
4982 .send(Message::Text(message.to_string().into()))
4983 .await
4984 .unwrap();
4985 }
4986
4987 async fn close_within(
4989 socket: &mut WebSocketStream<TcpStream>,
4990 wait: Duration,
4991 ) -> Option<Option<CloseFrame>> {
4992 tokio::time::timeout(wait, async {
4993 while let Some(Ok(message)) = socket.next().await {
4994 if let Message::Close(frame) = message {
4995 return Some(frame);
4996 }
4997 }
4998 Some(None)
4999 })
5000 .await
5001 .ok()
5002 .flatten()
5003 }
5004
5005 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5006 async fn an_expired_session_is_closed_with_the_reason() {
5007 let server = Server::start().await;
5008 let mut socket = server.connect(&server.token(2)).await;
5009
5010 tokio::time::sleep(Duration::from_secs(3)).await;
5011 send_json(&mut socket, json!({"type": "ping"})).await;
5012
5013 let frame = close_within(&mut socket, Duration::from_secs(5))
5014 .await
5015 .expect("the server closes the expired session")
5016 .expect("the close frame carries a reason");
5017 assert_eq!(frame.code, CloseCode::Policy);
5018 assert_eq!(
5019 frame.reason.as_str(),
5020 "token-expired: Authentication token expired"
5021 );
5022 }
5023
5024 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5025 async fn a_session_refreshed_in_band_outlives_its_first_token() {
5026 let server = Server::start().await;
5027 let mut socket = server.connect(&server.token(2)).await;
5028
5029 send_json(
5030 &mut socket,
5031 json!({"type": "refresh_auth", "token": server.token(3_600)}),
5032 )
5033 .await;
5034 let reply = tokio::time::timeout(Duration::from_secs(5), async {
5035 while let Some(Ok(message)) = socket.next().await {
5036 if let Message::Text(text) = message {
5037 return serde_json::from_str::<Value>(text.as_str()).ok();
5038 }
5039 }
5040 None
5041 })
5042 .await
5043 .expect("the server answers the refresh")
5044 .expect("the answer is JSON");
5045 assert_eq!(reply["success"], true, "refresh accepted: {reply}");
5046
5047 tokio::time::sleep(Duration::from_secs(3)).await;
5048 send_json(&mut socket, json!({"type": "ping"})).await;
5049 assert!(
5050 close_within(&mut socket, Duration::from_millis(1_500))
5051 .await
5052 .is_none(),
5053 "the socket stays open on the refreshed token"
5054 );
5055 }
5056 }
5057}