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 if subscription.snapshot.enabled {
1395 send_snapshot_batches(
1396 context,
1397 &subscription,
1398 &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1399 view_spec.mode,
1400 SnapshotPurpose::Initial,
1401 &context.entity_cache.snapshot_config(),
1402 )
1403 .await?;
1404 }
1405
1406 let task_context = context.clone();
1407 let subscription_id = subscription.subscription_id.clone();
1408 let query = subscription.query.clone();
1409 let view_spec_task = view_spec.clone();
1410 let span_view = view_id.clone();
1411 let span_key = key.clone();
1412 tokio::spawn(
1413 async move {
1414 let mut member = delivered;
1416 let mut behind = false;
1421 let mut replaced = false;
1426 loop {
1427 tokio::select! {
1428 _ = cancel_token.cancelled() => break,
1429 changed = receiver.changed() => {
1430 if changed.is_err() {
1431 break;
1432 }
1433 let (payload, published, replaced_at) = {
1434 let update = receiver.borrow_and_update();
1435 (update.payload.clone(), update.published, update.replaced)
1436 };
1437 behind |= published > seen + 1;
1442 replaced |= behind && replaced_at > seen;
1443 seen = published;
1444 let metadata = source_frame_metadata(&payload);
1445 if metadata.op == "delete" {
1446 task_context.entity_cache.remove(&query.view, &key).await;
1447 if member && send_membership_frame(
1448 &task_context,
1449 &subscription_id,
1450 &view_spec_task,
1451 "delete",
1452 &key,
1453 Value::Null,
1454 metadata.seq,
1455 ).is_err() {
1456 break;
1457 }
1458 member = false;
1459 behind = false;
1460 replaced = false;
1461 continue;
1462 }
1463
1464 let Some(cached) = task_context
1475 .entity_cache
1476 .get(&view_spec_task.id, &key)
1477 .await
1478 else {
1479 if member && !replaced && send_scoped_source_payload(
1482 &task_context,
1483 &subscription_id,
1484 &query.view,
1485 payload,
1486 ).is_err() {
1487 break;
1488 }
1489 continue;
1490 };
1491 let selected =
1492 select_query_entities(vec![(key.clone(), cached)], &query, true, false);
1493 let is_member = !selected.is_empty();
1494 let caught_up = !std::mem::take(&mut behind);
1495 let replace = std::mem::take(&mut replaced);
1496 let result = match (member, selected.into_iter().next()) {
1497 (true, Some(_)) if caught_up => send_scoped_source_payload(
1498 &task_context,
1499 &subscription_id,
1500 &query.view,
1501 payload,
1502 ),
1503 (true, Some((entity_key, data))) => send_membership_frame(
1516 &task_context,
1517 &subscription_id,
1518 &view_spec_task,
1519 if replace { "upsert" } else { "patch" },
1520 &entity_key,
1521 data,
1522 None,
1523 ),
1524 (false, Some((entity_key, data))) => send_membership_frame(
1525 &task_context,
1526 &subscription_id,
1527 &view_spec_task,
1528 "upsert",
1529 &entity_key,
1530 data,
1531 metadata.seq,
1532 ),
1533 (true, None) => send_membership_frame(
1534 &task_context,
1535 &subscription_id,
1536 &view_spec_task,
1537 "remove",
1538 &key,
1539 Value::Null,
1540 metadata.seq,
1541 ),
1542 (false, None) => Ok(()),
1543 };
1544 if result.is_err() {
1545 break;
1546 }
1547 member = is_member;
1548 }
1549 }
1550 }
1551 }
1552 .instrument(info_span!("ws.subscribe.state", client_id = %context.client_id, view = %span_view, key = %span_key)),
1553 );
1554 Ok(())
1555}
1556
1557async fn attach_collection_subscription(
1558 context: &SubscriptionContext,
1559 subscription: Subscription,
1560 view_spec: ViewSpec,
1561 cancel_token: CancellationToken,
1562) -> Result<()> {
1563 let view_id = subscription.query.view.clone();
1564 let source_view_id = view_spec
1565 .source_view
1566 .clone()
1567 .unwrap_or_else(|| view_id.clone());
1568 let (mut receiver, initial_membership) = subscribe_collection_then_snapshot(
1569 context,
1570 &source_view_id,
1571 &view_spec,
1572 &subscription.query,
1573 )
1574 .await;
1575
1576 let mut snapshot_entities = initial_membership.clone();
1577 if let Some(limit) = subscription.query.snapshot_limit {
1578 snapshot_entities.truncate(limit);
1579 }
1580 enforce_snapshot_limit(context, snapshot_entities.len())?;
1581 send_subscribed_frame(context, &subscription, &view_spec)?;
1582 let delivered_rows = if subscription.snapshot.enabled {
1587 snapshot_entities.len()
1588 } else {
1589 0
1590 };
1591 let mut undelivered: HashSet<String> = initial_membership
1592 .iter()
1593 .skip(delivered_rows)
1594 .map(|(key, _)| key.clone())
1595 .collect();
1596 if subscription.snapshot.enabled {
1597 send_snapshot_batches(
1598 context,
1599 &subscription,
1600 &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1601 view_spec.mode,
1602 SnapshotPurpose::Initial,
1603 &context.entity_cache.snapshot_config(),
1604 )
1605 .await?;
1606 }
1607
1608 let task_context = context.clone();
1609 let task_subscription = subscription.clone();
1610 let subscription_id = subscription.subscription_id.clone();
1611 let query = subscription.query.clone();
1612 let view_spec_task = view_spec.clone();
1613 let span_view = view_id.clone();
1614 tokio::spawn(
1615 async move {
1616 let mut current = initial_membership;
1617 let coalesce_ms = (view_spec_task.mode == Mode::List)
1621 .then(|| {
1622 view_spec_task
1623 .delivery
1624 .coalesce_ms
1625 .or(task_context.delivery.collection_coalesce_ms)
1626 .filter(|milliseconds| *milliseconds > 0)
1627 })
1628 .flatten();
1629 let mut flush_interval = coalesce_ms.map(|milliseconds| {
1630 let period = Duration::from_millis(milliseconds);
1631 let mut interval = tokio::time::interval_at(
1632 tokio::time::Instant::now() + period,
1633 period,
1634 );
1635 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
1636 interval
1637 });
1638 let mut pending = HashMap::<String, Arc<BusMessage>>::new();
1639 let mut pending_updates = 0_u64;
1640
1641 loop {
1642 tokio::select! {
1643 _ = cancel_token.cancelled() => break,
1644 _ = async {
1645 flush_interval
1646 .as_mut()
1647 .expect("coalescing interval is guarded")
1648 .tick()
1649 .await;
1650 }, if flush_interval.is_some() => {
1651 if pending.is_empty() {
1652 continue;
1653 }
1654 let sorted_caches = view_spec_task
1655 .is_derived()
1656 .then(|| task_context.view_index.sorted_caches());
1657 let next = load_query_entities(
1658 &task_context.entity_cache,
1659 sorted_caches,
1660 &view_spec_task,
1661 &query,
1662 false,
1663 ).await;
1664 if emit_coalesced_collection_delta(
1665 &task_context,
1666 &subscription_id,
1667 &view_spec_task,
1668 ¤t,
1669 &next,
1670 &pending,
1671 &mut undelivered,
1672 ).is_err() {
1673 task_context.metrics.delivery_stopped(&view_id, "send-failed");
1674 break;
1675 }
1676 task_context
1677 .metrics
1678 .collection_coalesced(&view_id, pending_updates);
1679 current = next;
1680 pending.clear();
1681 pending_updates = 0;
1682 }
1683 received = receiver.recv() => {
1684 let envelope = match received {
1685 Ok(envelope) => envelope,
1686 Err(broadcast::error::RecvError::Lagged(skipped)) => {
1687 task_context.metrics.subscription_lagged(&view_id, skipped);
1688 warn!(
1689 "Subscription {} lagged by {} updates",
1690 subscription_id, skipped
1691 );
1692 if view_spec_task.mode == Mode::Append {
1693 let _ = send_control_frame(
1694 &task_context,
1695 &SocketIssueMessage::append_subscription_lagged(
1696 subscription_id.clone(),
1697 skipped,
1698 ),
1699 &view_id,
1700 );
1701 task_context.metrics.delivery_stopped(&view_id, "append-lagged-without-replay");
1702 break;
1703 }
1704 if !task_subscription.snapshot.enabled {
1705 let _ = send_control_frame(
1706 &task_context,
1707 &SocketIssueMessage::subscription_lagged(
1708 subscription_id.clone(),
1709 skipped,
1710 ),
1711 &view_id,
1712 );
1713 task_context.metrics.delivery_stopped(&view_id, "lagged-without-snapshot");
1714 break;
1715 }
1716 info!(
1717 "Subscription {} is recovering from an authoritative snapshot",
1718 subscription_id
1719 );
1720 match recover_collection_subscription(
1721 &task_context,
1722 &task_subscription,
1723 &view_spec_task,
1724 &source_view_id,
1725 ).await {
1726 Ok((next_receiver, recovered)) => {
1727 receiver = next_receiver;
1728 current = recovered;
1729 undelivered.clear();
1732 pending.clear();
1733 pending_updates = 0;
1734 task_context.metrics.subscription_resnapshot(&view_id);
1735 continue;
1736 }
1737 Err(error) => {
1738 warn!(
1739 "Subscription {} failed to recover from lag: {error:#}",
1740 subscription_id
1741 );
1742 task_context.metrics.delivery_stopped(&view_id, "resnapshot-failed");
1743 break;
1744 }
1745 }
1746 }
1747 Err(broadcast::error::RecvError::Closed) => break,
1748 };
1749
1750 apply_collection_source_event(
1751 &task_context,
1752 &source_view_id,
1753 &view_spec_task,
1754 &query,
1755 &envelope,
1756 ).await;
1757
1758 if flush_interval.is_some() {
1759 pending_updates = pending_updates.saturating_add(1);
1760 pending.insert(envelope.key.clone(), envelope);
1761 continue;
1762 }
1763
1764 let metadata = source_frame_metadata(&envelope.payload);
1765 let sorted_caches = view_spec_task
1766 .is_derived()
1767 .then(|| task_context.view_index.sorted_caches());
1768 let next = load_query_entities(
1769 &task_context.entity_cache,
1770 sorted_caches,
1771 &view_spec_task,
1772 &query,
1773 false,
1774 ).await;
1775 if emit_collection_delta(
1776 &task_context,
1777 &subscription_id,
1778 &view_spec_task,
1779 ¤t,
1780 &next,
1781 &envelope,
1782 &metadata,
1783 &mut undelivered,
1784 ).is_err() {
1785 task_context.metrics.delivery_stopped(&view_id, "send-failed");
1786 break;
1787 }
1788 current = next;
1789 }
1790 }
1791 }
1792 }
1793 .instrument(info_span!("ws.subscribe.collection", client_id = %context.client_id, view = %span_view)),
1794 );
1795 Ok(())
1796}
1797
1798async fn subscribe_collection_then_snapshot(
1799 context: &SubscriptionContext,
1800 source_view_id: &str,
1801 view_spec: &ViewSpec,
1802 query: &SubscriptionQuery,
1803) -> (broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>) {
1804 let cache = context.entity_cache.clone();
1805 let sorted_caches = view_spec
1806 .is_derived()
1807 .then(|| context.view_index.sorted_caches());
1808 let view_spec = view_spec.clone();
1809 let query = query.clone();
1810 subscribe_list_then_snapshot(&context.bus_manager, source_view_id, move || async move {
1811 load_query_entities(&cache, sorted_caches, &view_spec, &query, false).await
1812 })
1813 .await
1814}
1815
1816async fn recover_collection_subscription(
1817 context: &SubscriptionContext,
1818 subscription: &Subscription,
1819 view_spec: &ViewSpec,
1820 source_view_id: &str,
1821) -> Result<(broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>)> {
1822 let (receiver, membership) =
1823 subscribe_collection_then_snapshot(context, source_view_id, view_spec, &subscription.query)
1824 .await;
1825 enforce_snapshot_limit(context, membership.len())?;
1830 send_snapshot_batches(
1831 context,
1832 subscription,
1833 &to_wire_snapshot_entities(membership.clone(), view_spec),
1834 view_spec.mode,
1835 SnapshotPurpose::Recovery,
1836 &context.entity_cache.snapshot_config(),
1837 )
1838 .await?;
1839 Ok((receiver, membership))
1840}
1841
1842async fn apply_collection_source_event(
1843 context: &SubscriptionContext,
1844 source_view_id: &str,
1845 view_spec: &ViewSpec,
1846 query: &SubscriptionQuery,
1847 envelope: &BusMessage,
1848) {
1849 let metadata = source_frame_metadata(&envelope.payload);
1850 if metadata.op != "delete" {
1851 return;
1852 }
1853 let current = context
1857 .entity_cache
1858 .get(source_view_id, &envelope.key)
1859 .await;
1860 if source_delete_is_stale(current.as_ref(), metadata.seq.as_deref()) {
1861 return;
1862 }
1863 context
1864 .entity_cache
1865 .remove(source_view_id, &envelope.key)
1866 .await;
1867 if view_spec.is_derived() {
1868 let caches = context.view_index.sorted_caches();
1869 let mut guard = caches.write().await;
1870 if let Some(cache) = guard.get_mut(&query.view) {
1871 cache.remove(&envelope.key);
1872 }
1873 }
1874}
1875
1876fn source_delete_is_stale(current: Option<&Value>, delete_seq: Option<&str>) -> bool {
1877 match (current, delete_seq) {
1878 (Some(current), Some(delete_seq)) => current
1879 .get("_seq")
1880 .and_then(Value::as_str)
1881 .is_some_and(|current_seq| {
1882 cmp_seq(current_seq, delete_seq) == std::cmp::Ordering::Greater
1883 }),
1884 _ => false,
1885 }
1886}
1887
1888#[derive(Debug)]
1898pub(crate) struct RejectedSubscription(pub SocketIssueMessage);
1899
1900impl RejectedSubscription {
1901 fn into_error(issue: SocketIssueMessage) -> anyhow::Error {
1902 anyhow::Error::new(Self(issue))
1903 }
1904}
1905
1906impl std::fmt::Display for RejectedSubscription {
1907 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1908 formatter.write_str(&self.0.message)
1909 }
1910}
1911
1912impl std::error::Error for RejectedSubscription {}
1913
1914async fn attach_journal_subscription(
1932 context: &SubscriptionContext,
1933 subscription: Subscription,
1934 view_spec: ViewSpec,
1935 journal: Arc<crate::journal::EventJournal>,
1936 cancel_token: CancellationToken,
1937) -> Result<()> {
1938 let view_id = view_spec.id.clone();
1939 let subscription_id = subscription.subscription_id.clone();
1940
1941 if subscription.query.take.is_some()
1944 || subscription.query.skip.is_some()
1945 || subscription.query.snapshot_limit.is_some()
1946 {
1947 return Err(RejectedSubscription::into_error(SocketIssueMessage::protocol(
1948 Some(subscription_id),
1949 "invalid-subscription",
1950 format!(
1951 "take, skip and snapshotLimit are window options and do not apply to the replayable view {view_id}"
1952 ),
1953 )));
1954 }
1955
1956 let cursor = match subscription.query.after.as_deref() {
1959 Some(raw) => match crate::journal::Cursor::parse(raw) {
1960 Some(cursor) => Some(cursor),
1961 None => {
1962 return Err(RejectedSubscription::into_error(
1963 SocketIssueMessage::protocol(
1964 Some(subscription_id),
1965 "invalid-cursor",
1966 format!(
1967 "`after` must be an {{epoch}}:{{offset}} replay cursor for view {view_id}, got {raw:?}"
1968 ),
1969 ),
1970 ));
1971 }
1972 },
1973 None => None,
1974 };
1975
1976 let mut receiver = context.bus_manager.get_or_create_list_bus(&view_id).await;
1979
1980 let replayed = match journal.replay_after(&view_id, cursor.as_ref()).await {
1981 Ok(records) => records,
1982 Err(error) => {
1983 return Err(RejectedSubscription::into_error(
1984 SocketIssueMessage::replay_refused(Some(subscription_id), &error),
1985 ));
1986 }
1987 };
1988
1989 let frame = SubscribedFrame::new(
1990 subscription.subscription_id.clone(),
1991 subscription.query.clone(),
1992 view_spec.mode,
1993 extract_sort_config(&view_spec),
1994 )
1995 .with_replay_window(journal.window(&view_id).await);
1996 send_control_frame(context, &frame, &view_id)?;
1997
1998 let task_context = context.clone();
2003 let task_subscription_id = subscription.subscription_id.clone();
2004 let task_query = subscription.query.clone();
2005 let task_epoch = journal.epoch().await;
2006 let span_view = view_id.clone();
2007 tokio::spawn(
2008 async move {
2009 let mut last_sent = cursor.map(|cursor| cursor.offset);
2010 let mut pending: VecDeque<Arc<BusMessage>> = VecDeque::new();
2014 let mut lagged: Option<u64> = None;
2015
2016 for record in replayed {
2017 if cancel_token.is_cancelled() {
2018 return;
2019 }
2020 drain_available(&mut receiver, &mut pending, &mut lagged);
2021 if journal_record_matches(&task_query, &record)
2022 && send_scoped_source_payload_async(
2023 &task_context,
2024 &task_subscription_id,
2025 &span_view,
2026 record.payload,
2027 )
2028 .await
2029 .is_err()
2030 {
2031 return;
2032 }
2033 last_sent = Some(record.offset);
2035 }
2036
2037 while let Some(envelope) = pending.pop_front() {
2040 if !forward_live_frame(
2041 &task_context,
2042 &task_subscription_id,
2043 &span_view,
2044 &task_query,
2045 &envelope,
2046 &mut last_sent,
2047 )
2048 .await
2049 {
2050 return;
2051 }
2052 }
2053
2054 if let Some(skipped) = lagged {
2055 report_replay_gap(
2056 &task_context,
2057 &task_subscription_id,
2058 &span_view,
2059 &task_epoch,
2060 skipped,
2061 last_sent,
2062 );
2063 return;
2064 }
2065
2066 loop {
2067 tokio::select! {
2068 _ = cancel_token.cancelled() => break,
2069 received = receiver.recv() => {
2070 let envelope = match received {
2071 Ok(envelope) => envelope,
2072 Err(broadcast::error::RecvError::Lagged(skipped)) => {
2090 report_replay_gap(
2091 &task_context,
2092 &task_subscription_id,
2093 &span_view,
2094 &task_epoch,
2095 skipped,
2096 last_sent,
2097 );
2098 break;
2099 }
2100 Err(broadcast::error::RecvError::Closed) => break,
2101 };
2102
2103 if !forward_live_frame(
2104 &task_context,
2105 &task_subscription_id,
2106 &span_view,
2107 &task_query,
2108 &envelope,
2109 &mut last_sent,
2110 )
2111 .await
2112 {
2113 break;
2114 }
2115 }
2116 }
2117 }
2118 }
2119 .instrument(info_span!(
2120 "ws.subscribe.replay",
2121 client_id = %context.client_id,
2122 view = %view_id
2123 )),
2124 );
2125 Ok(())
2126}
2127
2128fn drain_available(
2131 receiver: &mut broadcast::Receiver<Arc<BusMessage>>,
2132 pending: &mut VecDeque<Arc<BusMessage>>,
2133 lagged: &mut Option<u64>,
2134) {
2135 if lagged.is_some() {
2140 return;
2141 }
2142 const MAX_PENDING: usize = 8_192;
2148 loop {
2149 if pending.len() >= MAX_PENDING {
2150 return;
2151 }
2152 match receiver.try_recv() {
2153 Ok(envelope) => pending.push_back(envelope),
2154 Err(broadcast::error::TryRecvError::Empty)
2155 | Err(broadcast::error::TryRecvError::Closed) => return,
2156 Err(broadcast::error::TryRecvError::Lagged(skipped)) => {
2157 *lagged = Some(skipped);
2158 return;
2159 }
2160 }
2161 }
2162}
2163
2164fn already_delivered(offset: Option<u64>, last_sent: &mut Option<u64>) -> bool {
2170 let Some(offset) = offset else {
2171 return false;
2174 };
2175 if last_sent.is_some_and(|last| offset <= last) {
2176 return true;
2177 }
2178 *last_sent = Some(offset);
2179 false
2180}
2181
2182async fn forward_live_frame(
2186 context: &SubscriptionContext,
2187 subscription_id: &str,
2188 view_id: &str,
2189 query: &SubscriptionQuery,
2190 envelope: &Arc<BusMessage>,
2191 last_sent: &mut Option<u64>,
2192) -> bool {
2193 let metadata = source_frame_metadata(&envelope.payload);
2194 if already_delivered(metadata.offset, last_sent) {
2195 return true;
2196 }
2197 if !live_frame_matches(query, &envelope.key, &envelope.payload) {
2198 return true;
2199 }
2200 send_scoped_source_payload(context, subscription_id, view_id, envelope.payload.clone()).is_ok()
2201}
2202
2203fn report_replay_gap(
2204 context: &SubscriptionContext,
2205 subscription_id: &str,
2206 view_id: &str,
2207 epoch: &crate::journal::JournalEpoch,
2208 skipped: u64,
2209 last_sent: Option<u64>,
2210) {
2211 warn!(
2212 "Replay subscription {} lagged past {} records; stopping with a recovery cursor",
2213 subscription_id, skipped
2214 );
2215 let recover_from = last_sent.map(|offset| crate::journal::Cursor {
2216 epoch: epoch.clone(),
2217 offset,
2218 });
2219 let _ = send_control_frame(
2220 context,
2221 &SocketIssueMessage::replay_lagged(
2222 Some(subscription_id.to_string()),
2223 skipped,
2224 recover_from,
2225 ),
2226 view_id,
2227 );
2228}
2229
2230fn journal_record_matches(
2233 query: &SubscriptionQuery,
2234 record: &crate::journal::JournalRecord,
2235) -> bool {
2236 live_frame_matches(query, &record.key, &record.payload)
2237}
2238
2239fn live_frame_matches(query: &SubscriptionQuery, key: &str, payload: &[u8]) -> bool {
2240 if !query.matches_key(key) {
2241 return false;
2242 }
2243 if query.partition.is_none() && query.filters.is_empty() {
2244 return true;
2245 }
2246 let Ok(frame) = serde_json::from_slice::<Value>(payload) else {
2247 return false;
2248 };
2249 let Some(data) = frame.get("data") else {
2250 return false;
2251 };
2252 if let Some(partition) = &query.partition {
2253 if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2254 return false;
2255 }
2256 }
2257 query
2258 .filters
2259 .iter()
2260 .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2261}
2262
2263async fn send_scoped_source_payload_async(
2266 context: &SubscriptionContext,
2267 subscription_id: &str,
2268 view_id: &str,
2269 payload: Arc<Bytes>,
2270) -> Result<()> {
2271 let mut value: Value = serde_json::from_slice(&payload)?;
2272 let object = value
2273 .as_object_mut()
2274 .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2275 object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2276 object.insert(
2277 "subscriptionId".to_string(),
2278 Value::String(subscription_id.to_string()),
2279 );
2280 let json = serde_json::to_vec(&value)?;
2281 let compressed = maybe_compress(&json);
2282 let bytes = compressed.as_bytes().len();
2283 context
2284 .client_manager
2285 .send_compressed_async(context.client_id, compressed)
2286 .await
2287 .map_err(|error| anyhow::anyhow!("failed to send replayed frame: {error}"))?;
2288 context.metrics.message_sent();
2289 emit_update_sent_for_client(
2290 &context.usage_emitter,
2291 &context.client_manager,
2292 context.client_id,
2293 view_id,
2294 bytes,
2295 );
2296 Ok(())
2297}
2298
2299#[derive(Default)]
2300struct SourceFrameMetadata {
2301 op: String,
2302 seq: Option<String>,
2303 offset: Option<u64>,
2304}
2305
2306fn source_frame_metadata(payload: &[u8]) -> SourceFrameMetadata {
2307 serde_json::from_slice::<Value>(payload)
2308 .ok()
2309 .map(|value| SourceFrameMetadata {
2310 op: value
2311 .get("op")
2312 .and_then(Value::as_str)
2313 .unwrap_or_default()
2314 .to_string(),
2315 seq: value.get("seq").and_then(Value::as_str).map(str::to_string),
2316 offset: value.get("offset").and_then(Value::as_u64),
2317 })
2318 .unwrap_or_default()
2319}
2320
2321fn send_scoped_source_payload(
2322 context: &SubscriptionContext,
2323 subscription_id: &str,
2324 view_id: &str,
2325 payload: Arc<Bytes>,
2326) -> Result<()> {
2327 let mut value: Value = serde_json::from_slice(&payload)?;
2328 let object = value
2329 .as_object_mut()
2330 .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2331 object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2332 object.insert(
2333 "subscriptionId".to_string(),
2334 Value::String(subscription_id.to_string()),
2335 );
2336 let encoded = Arc::new(Bytes::from(serde_json::to_vec(&value)?));
2337 let bytes = encoded.len();
2338 context
2339 .client_manager
2340 .send_to_client(context.client_id, encoded)
2341 .map_err(|error| anyhow::anyhow!("failed to send live frame: {error}"))?;
2342 context.metrics.message_sent();
2343 emit_update_sent_for_client(
2344 &context.usage_emitter,
2345 &context.client_manager,
2346 context.client_id,
2347 view_id,
2348 bytes,
2349 );
2350 Ok(())
2351}
2352
2353fn send_membership_frame(
2354 context: &SubscriptionContext,
2355 subscription_id: &str,
2356 view_spec: &ViewSpec,
2357 op: &str,
2358 key: &str,
2359 mut data: Value,
2360 seq: Option<String>,
2361) -> Result<()> {
2362 apply_wire_format(&mut data, &view_spec.wire_format);
2363 let frame = Frame::scoped(
2364 subscription_id,
2365 view_spec.mode,
2366 &view_spec.id,
2367 op,
2368 key,
2369 data,
2370 seq,
2371 );
2372 let encoded = Arc::new(Bytes::from(serde_json::to_vec(&frame)?));
2373 let bytes = encoded.len();
2374 context
2375 .client_manager
2376 .send_to_client(context.client_id, encoded)
2377 .map_err(|error| anyhow::anyhow!("failed to send membership frame: {error}"))?;
2378 context.metrics.message_sent();
2379 emit_update_sent_for_client(
2380 &context.usage_emitter,
2381 &context.client_manager,
2382 context.client_id,
2383 &view_spec.id,
2384 bytes,
2385 );
2386 Ok(())
2387}
2388
2389#[allow(clippy::too_many_arguments)]
2392fn emit_collection_delta(
2393 context: &SubscriptionContext,
2394 subscription_id: &str,
2395 view_spec: &ViewSpec,
2396 current: &[(String, Value)],
2397 next: &[(String, Value)],
2398 envelope: &BusMessage,
2399 metadata: &SourceFrameMetadata,
2400 undelivered: &mut HashSet<String>,
2401) -> Result<()> {
2402 let current_keys: Vec<&str> = current.iter().map(|(key, _)| key.as_str()).collect();
2403 let next_keys: Vec<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2404 let next_set: HashSet<&str> = next_keys.iter().copied().collect();
2405
2406 for key in current_keys
2407 .iter()
2408 .copied()
2409 .filter(|key| !next_set.contains(key))
2410 {
2411 if undelivered.remove(key) {
2413 continue;
2414 }
2415 let op = if metadata.op == "delete" && key == envelope.key {
2416 "delete"
2417 } else {
2418 "remove"
2419 };
2420 send_membership_frame(
2421 context,
2422 subscription_id,
2423 view_spec,
2424 op,
2425 key,
2426 Value::Null,
2427 metadata.seq.clone(),
2428 )?;
2429 }
2430
2431 for (key, data) in next.iter() {
2432 let in_window = current_keys.iter().any(|candidate| *candidate == key);
2433 match member_action(
2434 in_window,
2435 in_window && !undelivered.contains(key),
2436 key == &envelope.key,
2437 view_spec.is_derived(),
2438 &metadata.op,
2439 ) {
2440 MemberAction::Skip => {}
2441 MemberAction::ForwardPatch => send_scoped_source_payload(
2442 context,
2443 subscription_id,
2444 &view_spec.id,
2445 envelope.payload.clone(),
2446 )?,
2447 MemberAction::Upsert => {
2448 let seq = metadata
2449 .seq
2450 .clone()
2451 .or_else(|| data.get("_seq").and_then(Value::as_str).map(str::to_string));
2452 send_membership_frame(
2453 context,
2454 subscription_id,
2455 view_spec,
2456 "upsert",
2457 key,
2458 data.clone(),
2459 seq,
2460 )?;
2461 undelivered.remove(key);
2462 }
2463 }
2464 }
2465 Ok(())
2466}
2467
2468fn emit_coalesced_collection_delta(
2473 context: &SubscriptionContext,
2474 subscription_id: &str,
2475 view_spec: &ViewSpec,
2476 current: &[(String, Value)],
2477 next: &[(String, Value)],
2478 pending: &HashMap<String, Arc<BusMessage>>,
2479 undelivered: &mut HashSet<String>,
2480) -> Result<()> {
2481 let changes = plan_coalesced_collection_delta(current, next, pending, undelivered);
2482 let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2485 undelivered.retain(|key| next_keys.contains(key.as_str()));
2486 for change in changes {
2487 if change.op == "upsert" {
2488 undelivered.remove(&change.key);
2489 }
2490 send_membership_frame(
2491 context,
2492 subscription_id,
2493 view_spec,
2494 change.op,
2495 &change.key,
2496 change.data,
2497 change.seq,
2498 )?;
2499 }
2500 Ok(())
2501}
2502
2503#[derive(Debug, PartialEq)]
2504struct CollectionChange {
2505 op: &'static str,
2506 key: String,
2507 data: Value,
2508 seq: Option<String>,
2509}
2510
2511fn plan_coalesced_collection_delta(
2515 current: &[(String, Value)],
2516 next: &[(String, Value)],
2517 pending: &HashMap<String, Arc<BusMessage>>,
2518 undelivered: &HashSet<String>,
2519) -> Vec<CollectionChange> {
2520 let current_by_key: HashMap<&str, &Value> = current
2521 .iter()
2522 .map(|(key, data)| (key.as_str(), data))
2523 .collect();
2524 let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2525 let latest_seq = pending
2526 .values()
2527 .filter_map(|envelope| source_frame_metadata(&envelope.payload).seq)
2528 .max_by(|left, right| cmp_seq(left, right));
2529 let mut changes = Vec::new();
2530
2531 for (key, _) in current
2532 .iter()
2533 .filter(|(key, _)| !next_keys.contains(key.as_str()) && !undelivered.contains(key.as_str()))
2534 {
2535 let metadata = pending
2536 .get(key)
2537 .map(|envelope| source_frame_metadata(&envelope.payload));
2538 let op = if metadata
2539 .as_ref()
2540 .is_some_and(|metadata| metadata.op == "delete")
2541 {
2542 "delete"
2543 } else {
2544 "remove"
2545 };
2546 let seq = metadata
2547 .and_then(|metadata| metadata.seq)
2548 .or_else(|| latest_seq.clone());
2549 changes.push(CollectionChange {
2550 op,
2551 key: key.clone(),
2552 data: Value::Null,
2553 seq,
2554 });
2555 }
2556
2557 for (key, data) in next {
2558 let changed = current_by_key
2559 .get(key.as_str())
2560 .is_none_or(|previous| *previous != data);
2561 if !changed && !pending.contains_key(key) {
2562 continue;
2563 }
2564 let seq = data
2565 .get("_seq")
2566 .and_then(Value::as_str)
2567 .map(str::to_string)
2568 .or_else(|| {
2569 pending
2570 .get(key)
2571 .and_then(|envelope| source_frame_metadata(&envelope.payload).seq)
2572 })
2573 .or_else(|| latest_seq.clone());
2574 changes.push(CollectionChange {
2575 op: "upsert",
2576 key: key.clone(),
2577 data: data.clone(),
2578 seq,
2579 });
2580 }
2581 changes
2582}
2583
2584#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2586enum MemberAction {
2587 Skip,
2590 ForwardPatch,
2592 Upsert,
2594}
2595
2596fn member_action(
2614 in_window: bool,
2615 held: bool,
2616 is_mutated_key: bool,
2617 is_derived: bool,
2618 op: &str,
2619) -> MemberAction {
2620 if !is_mutated_key {
2621 return if in_window {
2627 MemberAction::Skip
2628 } else {
2629 MemberAction::Upsert
2630 };
2631 }
2632 if held && !is_derived && op != "delete" {
2638 MemberAction::ForwardPatch
2639 } else {
2640 MemberAction::Upsert
2641 }
2642}
2643
2644fn to_wire_snapshot_entities(
2645 entities: Vec<(String, Value)>,
2646 view_spec: &ViewSpec,
2647) -> Vec<SnapshotEntity> {
2648 entities
2649 .into_iter()
2650 .map(|(key, mut data)| {
2651 apply_wire_format(&mut data, &view_spec.wire_format);
2652 SnapshotEntity { key, data }
2653 })
2654 .collect()
2655}
2656
2657async fn load_query_entities(
2658 entity_cache: &EntityCache,
2659 sorted_caches: Option<
2660 Arc<tokio::sync::RwLock<HashMap<String, crate::sorted_cache::SortedViewCache>>>,
2661 >,
2662 view_spec: &ViewSpec,
2663 query: &SubscriptionQuery,
2664 apply_snapshot_limit: bool,
2665) -> Vec<(String, Value)> {
2666 let (entities, preordered) = if let Some(sorted_caches) = sorted_caches {
2667 let mut caches = sorted_caches.write().await;
2668 let entities = caches
2669 .get_mut(&view_spec.id)
2670 .map(|cache| cache.get_all_ordered())
2671 .unwrap_or_default();
2672 (entities, true)
2673 } else if view_spec.mode == Mode::State {
2674 let entity = match query.key.as_deref() {
2675 Some(key) => entity_cache
2676 .get(&view_spec.id, key)
2677 .await
2678 .map(|data| vec![(key.to_string(), data)])
2679 .unwrap_or_default(),
2680 None => vec![],
2681 };
2682 (entity, true)
2683 } else {
2684 (entity_cache.get_all(&view_spec.id).await, false)
2685 };
2686 select_query_entities(entities, query, preordered, apply_snapshot_limit)
2687}
2688
2689fn select_query_entities(
2690 mut entities: Vec<(String, Value)>,
2691 query: &SubscriptionQuery,
2692 preordered: bool,
2693 apply_snapshot_limit: bool,
2694) -> Vec<(String, Value)> {
2695 entities.retain(|(key, data)| query_matches_entity(query, key, data));
2696 if !preordered {
2697 entities.sort_by(|left, right| {
2698 let left_seq = left.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2699 let right_seq = right.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2700 let order = if query.after.is_some() {
2701 cmp_seq(left_seq, right_seq)
2702 } else {
2703 cmp_seq(right_seq, left_seq)
2704 };
2705 order.then_with(|| left.0.cmp(&right.0))
2706 });
2707 }
2708
2709 let skip = query.skip.unwrap_or(0);
2710 let take = query.take.unwrap_or(usize::MAX);
2711 let mut selected: Vec<_> = entities.into_iter().skip(skip).take(take).collect();
2712 if apply_snapshot_limit {
2713 if let Some(limit) = query.snapshot_limit {
2714 selected.truncate(limit);
2715 }
2716 }
2717 selected
2718}
2719
2720fn query_matches_entity(query: &SubscriptionQuery, key: &str, data: &Value) -> bool {
2721 if !query.matches_key(key) {
2722 return false;
2723 }
2724 if let Some(partition) = &query.partition {
2725 if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2726 return false;
2727 }
2728 }
2729 if let Some(after) = &query.after {
2730 let Some(seq) = data.get("_seq").and_then(Value::as_str) else {
2731 return false;
2732 };
2733 if cmp_seq(seq, after) != std::cmp::Ordering::Greater {
2734 return false;
2735 }
2736 }
2737 query
2738 .filters
2739 .iter()
2740 .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2741}
2742
2743fn value_at_dot_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
2744 path.split('.')
2745 .try_fold(value, |current, segment| current.get(segment))
2746}
2747
2748#[cfg(test)]
2749mod tests {
2750 use super::*;
2751 use crate::cache::{EntityCacheConfig, PatchOrigin};
2752 use crate::view::{Delivery, Filters, Projection};
2753 use serde_json::json;
2754 use tokio::sync::oneshot;
2755
2756 fn list_spec() -> ViewSpec {
2757 ViewSpec {
2758 id: "Thing/list".to_string(),
2759 export: "Thing".to_string(),
2760 mode: Mode::List,
2761 wire_format: Default::default(),
2762 projection: Projection::all(),
2763 filters: Filters::all(),
2764 delivery: Delivery::default(),
2765 pipeline: None,
2766 source_view: None,
2767 }
2768 }
2769
2770 fn plan_window(
2773 current_keys: &[&str],
2774 next_keys: &[&str],
2775 envelope_key: &str,
2776 is_derived: bool,
2777 op: &str,
2778 ) -> Vec<(String, MemberAction)> {
2779 plan_window_with_undelivered(current_keys, &[], next_keys, envelope_key, is_derived, op)
2780 }
2781
2782 fn plan_window_with_undelivered(
2784 current_keys: &[&str],
2785 undelivered: &[&str],
2786 next_keys: &[&str],
2787 envelope_key: &str,
2788 is_derived: bool,
2789 op: &str,
2790 ) -> Vec<(String, MemberAction)> {
2791 next_keys
2792 .iter()
2793 .map(|key| {
2794 let in_window = current_keys.contains(key);
2795 let held = in_window && !undelivered.contains(key);
2796 (
2797 (*key).to_string(),
2798 member_action(in_window, held, *key == envelope_key, is_derived, op),
2799 )
2800 })
2801 .collect()
2802 }
2803
2804 #[test]
2805 fn a_window_member_the_client_never_received_is_sent_whole() {
2806 let plan = plan_window_with_undelivered(
2810 &["4", "3", "2", "1"],
2811 &["2", "1"],
2812 &["1", "4", "3", "2"],
2813 "1",
2814 false,
2815 "patch",
2816 );
2817 assert_eq!(
2818 plan,
2819 vec![
2820 ("1".to_string(), MemberAction::Upsert),
2821 ("4".to_string(), MemberAction::Skip),
2822 ("3".to_string(), MemberAction::Skip),
2823 ("2".to_string(), MemberAction::Skip),
2826 ]
2827 );
2828 }
2829
2830 #[test]
2831 fn a_held_member_still_rides_its_patch_after_a_truncated_snapshot() {
2832 let plan = plan_window_with_undelivered(
2833 &["4", "3", "2", "1"],
2834 &["2", "1"],
2835 &["4", "3", "2", "1"],
2836 "4",
2837 false,
2838 "patch",
2839 );
2840 assert_eq!(plan[0], ("4".to_string(), MemberAction::ForwardPatch));
2841 }
2842
2843 #[test]
2844 fn coalescing_sends_no_remove_for_a_key_the_client_never_received() {
2845 let current = vec![
2846 ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2847 (
2848 "unsent".to_string(),
2849 json!({"count": 1, "_seq": "10:000001"}),
2850 ),
2851 ("gone".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2852 ];
2853 let next = vec![
2854 ("held".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2855 (
2856 "unsent".to_string(),
2857 json!({"count": 2, "_seq": "10:000003"}),
2858 ),
2859 ];
2860 let pending = HashMap::from([
2861 (
2862 "unsent".to_string(),
2863 list_message("unsent", "patch", "10:000003"),
2864 ),
2865 (
2866 "gone".to_string(),
2867 list_message("gone", "delete", "10:000002"),
2868 ),
2869 ]);
2870 let undelivered = HashSet::from(["unsent".to_string(), "gone".to_string()]);
2871
2872 assert_eq!(
2873 plan_coalesced_collection_delta(¤t, &next, &pending, &undelivered),
2874 vec![CollectionChange {
2875 op: "upsert",
2876 key: "unsent".to_string(),
2877 data: json!({"count": 2, "_seq": "10:000003"}),
2878 seq: Some("10:000003".to_string()),
2879 }]
2880 );
2881 }
2882
2883 #[test]
2884 fn a_reordered_window_sends_only_the_mutated_entity() {
2885 let current = ["4", "3", "2", "1"];
2889 let next = ["1", "4", "3", "2"];
2890 let plan = plan_window(¤t, &next, "1", false, "patch");
2891
2892 assert_eq!(
2893 plan,
2894 vec![
2895 ("1".to_string(), MemberAction::ForwardPatch),
2896 ("4".to_string(), MemberAction::Skip),
2897 ("3".to_string(), MemberAction::Skip),
2898 ("2".to_string(), MemberAction::Skip),
2899 ]
2900 );
2901 }
2902
2903 #[test]
2904 fn a_key_entering_the_window_gets_the_whole_entity() {
2905 let plan = plan_window(&["4", "3"], &["5", "4", "3"], "5", false, "patch");
2907 assert_eq!(
2908 plan,
2909 vec![
2910 ("5".to_string(), MemberAction::Upsert),
2911 ("4".to_string(), MemberAction::Skip),
2912 ("3".to_string(), MemberAction::Skip),
2913 ]
2914 );
2915 }
2916
2917 #[test]
2918 fn derived_views_still_send_whole_entities() {
2919 let plan = plan_window(&["1", "2"], &["1", "2"], "1", true, "patch");
2922 assert_eq!(
2923 plan,
2924 vec![
2925 ("1".to_string(), MemberAction::Upsert),
2926 ("2".to_string(), MemberAction::Skip),
2927 ]
2928 );
2929 }
2930
2931 #[test]
2932 fn a_delete_envelope_never_forwards_a_patch() {
2933 let plan = plan_window(&["1", "2"], &["1", "2"], "1", false, "delete");
2935 assert_eq!(
2936 plan,
2937 vec![
2938 ("1".to_string(), MemberAction::Upsert),
2939 ("2".to_string(), MemberAction::Skip),
2940 ]
2941 );
2942 }
2943
2944 #[test]
2945 fn an_unchanged_window_sends_one_frame_not_a_broadcast() {
2946 let keys: Vec<String> = (0..500).map(|index| index.to_string()).collect();
2949 let refs: Vec<&str> = keys.iter().map(String::as_str).collect();
2950 let plan = plan_window(&refs, &refs, "250", false, "patch");
2951
2952 let sent = plan
2953 .iter()
2954 .filter(|(_, action)| *action != MemberAction::Skip)
2955 .count();
2956 assert_eq!(sent, 1);
2957 assert_eq!(plan[250].1, MemberAction::ForwardPatch);
2958 }
2959
2960 fn list_message(key: &str, op: &str, seq: &str) -> Arc<BusMessage> {
2961 Arc::new(BusMessage {
2962 key: key.to_string(),
2963 entity: "Thing/list".to_string(),
2964 payload: Arc::new(Bytes::from(
2965 serde_json::to_vec(&json!({
2966 "entity": "Thing/list",
2967 "op": op,
2968 "key": key,
2969 "seq": seq,
2970 "data": {},
2971 }))
2972 .unwrap(),
2973 )),
2974 })
2975 }
2976
2977 #[test]
2978 fn coalescing_emits_only_the_final_full_state_per_changed_key() {
2979 let current = vec![
2980 ("a".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2981 ("b".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2982 (
2983 "stable".to_string(),
2984 json!({"count": 1, "_seq": "10:000001"}),
2985 ),
2986 ];
2987 let next = vec![
2988 ("a".to_string(), json!({"count": 3, "_seq": "10:000004"})),
2989 ("c".to_string(), json!({"count": 1, "_seq": "10:000003"})),
2990 (
2991 "stable".to_string(),
2992 json!({"count": 1, "_seq": "10:000001"}),
2993 ),
2994 ];
2995 let pending = HashMap::from([
2996 ("a".to_string(), list_message("a", "patch", "10:000004")),
2997 ("b".to_string(), list_message("b", "delete", "10:000002")),
2998 ("c".to_string(), list_message("c", "patch", "10:000003")),
2999 ]);
3000
3001 assert_eq!(
3002 plan_coalesced_collection_delta(¤t, &next, &pending, &HashSet::new()),
3003 vec![
3004 CollectionChange {
3005 op: "delete",
3006 key: "b".to_string(),
3007 data: Value::Null,
3008 seq: Some("10:000002".to_string()),
3009 },
3010 CollectionChange {
3011 op: "upsert",
3012 key: "a".to_string(),
3013 data: json!({"count": 3, "_seq": "10:000004"}),
3014 seq: Some("10:000004".to_string()),
3015 },
3016 CollectionChange {
3017 op: "upsert",
3018 key: "c".to_string(),
3019 data: json!({"count": 1, "_seq": "10:000003"}),
3020 seq: Some("10:000003".to_string()),
3021 },
3022 ]
3023 );
3024 }
3025
3026 #[test]
3027 fn a_delayed_delete_cannot_erase_a_newer_recreated_entity() {
3028 let recreated = json!({"balance": 2, "_seq": "10:000004"});
3029 assert!(source_delete_is_stale(Some(&recreated), Some("10:000002")));
3030 assert!(!source_delete_is_stale(Some(&recreated), Some("10:000004")));
3031 assert!(!source_delete_is_stale(Some(&recreated), None));
3032 }
3033
3034 #[test]
3035 fn a_lagged_snapshotless_subscription_gets_a_fatal_retryable_error() {
3036 let issue = SocketIssueMessage::subscription_lagged("balances".to_string(), 42);
3037 assert_eq!(issue.code, "subscription-lagged");
3038 assert!(issue.retryable);
3039 assert!(issue.fatal);
3040 assert!(issue.message.contains("42"));
3041 assert!(issue
3042 .suggested_action
3043 .unwrap()
3044 .contains("snapshots enabled"));
3045
3046 let append = SocketIssueMessage::append_subscription_lagged("trades".to_string(), 9);
3047 assert!(append.fatal);
3048 assert!(append.suggested_action.unwrap().contains("retained replay"));
3049 }
3050
3051 #[test]
3052 fn snapshot_batches_share_identity_and_completion() {
3053 let entities = ["one", "two", "three"].map(|key| SnapshotEntity {
3054 key: key.to_string(),
3055 data: json!({"key": key}),
3056 });
3057 let batches = create_snapshot_batches(
3058 &entities,
3059 SnapshotMetadata {
3060 subscription_id: "sub-1",
3061 snapshot_id: "snapshot-1",
3062 authoritative: true,
3063 mode: Mode::List,
3064 view_id: "Thing/list",
3065 key: None,
3066 },
3067 &SnapshotBatchConfig {
3068 initial_batch_size: 2,
3069 subsequent_batch_size: 1,
3070 },
3071 );
3072 assert_eq!(batches.len(), 2);
3073 assert!(batches.iter().all(|batch| batch.subscription_id == "sub-1"));
3074 assert!(batches
3075 .iter()
3076 .all(|batch| batch.snapshot_id == "snapshot-1"));
3077 assert!(!batches[0].complete);
3078 assert!(batches[1].complete);
3079 assert!(batches.iter().all(|batch| batch.authoritative));
3080 }
3081
3082 #[test]
3083 fn empty_incremental_snapshot_is_explicitly_non_authoritative() {
3084 let batches = create_snapshot_batches(
3085 &[],
3086 SnapshotMetadata {
3087 subscription_id: "sub-1",
3088 snapshot_id: "snapshot-1",
3089 authoritative: false,
3090 mode: Mode::State,
3091 view_id: "Thing/state",
3092 key: Some("missing"),
3093 },
3094 &SnapshotBatchConfig {
3095 initial_batch_size: 1,
3096 subsequent_batch_size: 1,
3097 },
3098 );
3099 assert_eq!(batches.len(), 1);
3100 assert!(!batches[0].authoritative);
3101 assert!(batches[0].complete);
3102 assert_eq!(batches[0].key.as_deref(), Some("missing"));
3103 }
3104
3105 #[test]
3106 fn lag_recovery_is_authoritative_for_an_after_query() {
3107 let subscription = Subscription {
3108 protocol_version: PROTOCOL_VERSION,
3109 subscription_id: "sub-1".to_string(),
3110 query: SubscriptionQuery {
3111 view: "Thing/list".to_string(),
3112 after: Some("40:000000000010".to_string()),
3113 ..Default::default()
3114 },
3115 snapshot: Default::default(),
3116 };
3117
3118 assert!(!SnapshotPurpose::Initial.authoritative(&subscription));
3119 assert!(SnapshotPurpose::Recovery.authoritative(&subscription));
3120 }
3121
3122 #[test]
3123 fn dot_path_filters_are_exact_and_type_sensitive() {
3124 let mut query = SubscriptionQuery {
3125 view: "Thing/list".to_string(),
3126 ..Default::default()
3127 };
3128 query
3129 .filters
3130 .insert("state.status".to_string(), json!("open"));
3131 query.filters.insert("metrics.count".to_string(), json!(2));
3132 assert!(query_matches_entity(
3133 &query,
3134 "one",
3135 &json!({"state": {"status": "open"}, "metrics": {"count": 2}}),
3136 ));
3137 assert!(!query_matches_entity(
3138 &query,
3139 "one",
3140 &json!({"state": {"status": "open"}, "metrics": {"count": "2"}}),
3141 ));
3142 }
3143
3144 #[test]
3145 fn take_and_skip_define_independent_deterministic_windows() {
3146 let entities: Vec<_> = (1..=6)
3147 .map(|id| {
3148 (
3149 id.to_string(),
3150 json!({"id": id, "_seq": format!("10:{id:012}")}),
3151 )
3152 })
3153 .collect();
3154 let first = SubscriptionQuery {
3155 view: "Thing/list".to_string(),
3156 take: Some(2),
3157 skip: Some(0),
3158 ..Default::default()
3159 };
3160 let second = SubscriptionQuery {
3161 skip: Some(2),
3162 ..first.clone()
3163 };
3164 let first_keys: Vec<_> = select_query_entities(entities.clone(), &first, false, false)
3165 .into_iter()
3166 .map(|(key, _)| key)
3167 .collect();
3168 let second_keys: Vec<_> = select_query_entities(entities, &second, false, false)
3169 .into_iter()
3170 .map(|(key, _)| key)
3171 .collect();
3172 assert_eq!(first_keys, ["6", "5"]);
3173 assert_eq!(second_keys, ["4", "3"]);
3174 }
3175
3176 #[tokio::test]
3177 async fn state_receiver_is_installed_before_snapshot_awaits() {
3178 let bus = BusManager::new();
3179 let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3180 let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3181 let bus_for_task = bus.clone();
3182 let task = tokio::spawn(async move {
3183 subscribe_state_then_snapshot(&bus_for_task, "Thing/state", "one", || async move {
3184 snapshot_started_tx.send(()).unwrap();
3185 release_snapshot_rx.await.unwrap();
3186 })
3187 .await
3188 .0
3189 });
3190 snapshot_started_rx.await.unwrap();
3191 bus.publish_state(
3192 "Thing/state",
3193 "one",
3194 Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3195 )
3196 .await;
3197 release_snapshot_tx.send(()).unwrap();
3198 let mut receiver = task.await.unwrap();
3199 receiver.changed().await.unwrap();
3200 assert!(!receiver.borrow().payload.is_empty());
3201 }
3202
3203 async fn assert_list_receiver_precedes_snapshot(view: &'static str) {
3204 let bus = BusManager::new();
3205 let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3206 let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3207 let bus_for_task = bus.clone();
3208 let task = tokio::spawn(async move {
3209 subscribe_list_then_snapshot(&bus_for_task, view, || async move {
3210 snapshot_started_tx.send(()).unwrap();
3211 release_snapshot_rx.await.unwrap();
3212 })
3213 .await
3214 .0
3215 });
3216 snapshot_started_rx.await.unwrap();
3217 bus.publish_list(
3218 view,
3219 Arc::new(BusMessage {
3220 key: "one".to_string(),
3221 entity: view.to_string(),
3222 payload: Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3223 }),
3224 )
3225 .await;
3226 release_snapshot_tx.send(()).unwrap();
3227 let mut receiver = task.await.unwrap();
3228 assert_eq!(receiver.recv().await.unwrap().key, "one");
3229 }
3230
3231 #[tokio::test]
3232 async fn list_receiver_is_installed_before_snapshot() {
3233 assert_list_receiver_precedes_snapshot("Thing/list").await;
3234 }
3235
3236 #[tokio::test]
3237 async fn append_receiver_is_installed_before_snapshot() {
3238 assert_list_receiver_precedes_snapshot("Thing/append").await;
3239 }
3240
3241 #[tokio::test]
3242 async fn derived_source_receiver_is_installed_before_snapshot() {
3243 assert_list_receiver_precedes_snapshot("Thing/list-source").await;
3244 }
3245
3246 #[tokio::test]
3247 async fn snapshot_limit_does_not_change_live_take_skip_membership() {
3248 let cache = EntityCache::with_config(EntityCacheConfig {
3249 max_entities_per_view: 10,
3250 ..Default::default()
3251 });
3252 for id in 1..=4 {
3253 cache
3254 .upsert(
3255 "Thing/list",
3256 &id.to_string(),
3257 json!({"_seq": format!("10:{id:012}")}),
3258 )
3259 .await;
3260 }
3261 let query = SubscriptionQuery {
3262 view: "Thing/list".to_string(),
3263 take: Some(3),
3264 skip: Some(1),
3265 snapshot_limit: Some(1),
3266 ..Default::default()
3267 };
3268 let live = load_query_entities(&cache, None, &list_spec(), &query, false).await;
3269 let snapshot = load_query_entities(&cache, None, &list_spec(), &query, true).await;
3270 assert_eq!(live.len(), 3);
3271 assert_eq!(snapshot.len(), 1);
3272 }
3273
3274 #[test]
3275 fn fixture_manifest_covers_required_conformance_cases() {
3276 let manifest: Value = serde_json::from_str(include_str!(
3277 "../../../../tests/fixtures/websocket-v2/manifest.json"
3278 ))
3279 .unwrap();
3280 let names: HashSet<_> = manifest["fixtures"]
3281 .as_array()
3282 .unwrap()
3283 .iter()
3284 .filter_map(Value::as_str)
3285 .collect();
3286 for required in [
3287 "keyed-state.json",
3288 "list-windows.json",
3289 "filters.json",
3290 "multi-batch-authoritative.json",
3291 "empty-snapshot.json",
3292 "remove.json",
3293 "delete.json",
3294 "incremental-snapshot.json",
3295 "reconnect-replacement.json",
3296 "errors.json",
3297 "whole-entities.json",
3298 ] {
3299 assert!(names.contains(required), "missing fixture {required}");
3300 }
3301
3302 for document in [
3303 include_str!("../../../../tests/fixtures/websocket-v2/keyed-state.json"),
3304 include_str!("../../../../tests/fixtures/websocket-v2/list-windows.json"),
3305 include_str!("../../../../tests/fixtures/websocket-v2/filters.json"),
3306 include_str!("../../../../tests/fixtures/websocket-v2/multi-batch-authoritative.json"),
3307 include_str!("../../../../tests/fixtures/websocket-v2/empty-snapshot.json"),
3308 include_str!("../../../../tests/fixtures/websocket-v2/remove.json"),
3309 include_str!("../../../../tests/fixtures/websocket-v2/delete.json"),
3310 include_str!("../../../../tests/fixtures/websocket-v2/incremental-snapshot.json"),
3311 include_str!("../../../../tests/fixtures/websocket-v2/reconnect-replacement.json"),
3312 include_str!("../../../../tests/fixtures/websocket-v2/errors.json"),
3313 include_str!("../../../../tests/fixtures/websocket-v2/whole-entities.json"),
3314 ] {
3315 let fixture: Value = serde_json::from_str(document).unwrap();
3316 assert!(fixture["name"].is_string());
3317 }
3318 }
3319
3320 fn append_frame(key: &str, data: Value) -> Arc<Bytes> {
3321 let frame = json!({
3322 "entity": "Trade/append",
3323 "op": "patch",
3324 "key": key,
3325 "offset": 7,
3326 "data": data,
3327 });
3328 Arc::new(Bytes::from(serde_json::to_vec(&frame).unwrap()))
3329 }
3330
3331 #[test]
3335 fn replay_delivery_applies_key_partition_and_filters() {
3336 let matching = append_frame("pool1", json!({"_partition": "us", "side": "buy"}));
3337 let other_partition = append_frame("pool1", json!({"_partition": "eu", "side": "buy"}));
3338 let other_side = append_frame("pool1", json!({"_partition": "us", "side": "sell"}));
3339
3340 let unfiltered = SubscriptionQuery {
3341 view: "Trade/append".to_string(),
3342 ..Default::default()
3343 };
3344 assert!(live_frame_matches(&unfiltered, "pool1", &matching));
3345 assert!(live_frame_matches(&unfiltered, "pool9", &matching));
3346
3347 let keyed = SubscriptionQuery {
3348 view: "Trade/append".to_string(),
3349 key: Some("pool1".to_string()),
3350 ..Default::default()
3351 };
3352 assert!(live_frame_matches(&keyed, "pool1", &matching));
3353 assert!(!live_frame_matches(&keyed, "pool2", &matching));
3354
3355 let partitioned = SubscriptionQuery {
3356 view: "Trade/append".to_string(),
3357 partition: Some("us".to_string()),
3358 ..Default::default()
3359 };
3360 assert!(live_frame_matches(&partitioned, "pool1", &matching));
3361 assert!(!live_frame_matches(&partitioned, "pool1", &other_partition));
3362
3363 let filtered = SubscriptionQuery {
3364 view: "Trade/append".to_string(),
3365 filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3366 ..Default::default()
3367 };
3368 assert!(live_frame_matches(&filtered, "pool1", &matching));
3369 assert!(!live_frame_matches(&filtered, "pool1", &other_side));
3370 }
3371
3372 #[test]
3373 fn a_frame_without_decodable_data_does_not_satisfy_a_filter() {
3374 let filtered = SubscriptionQuery {
3375 view: "Trade/append".to_string(),
3376 filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3377 ..Default::default()
3378 };
3379 let garbage = Arc::new(Bytes::from_static(b"not json"));
3380 assert!(!live_frame_matches(&filtered, "pool1", &garbage));
3381 }
3382
3383 #[test]
3387 fn replay_lagged_recovers_from_before_the_gap() {
3388 let epoch = crate::journal::JournalEpoch::new();
3389 let issue = SocketIssueMessage::replay_lagged(
3390 Some("trades".to_string()),
3391 37,
3392 Some(crate::journal::Cursor {
3393 epoch: epoch.clone(),
3394 offset: 4180,
3395 }),
3396 );
3397 assert_eq!(issue.code, "replay-lagged");
3398 assert_eq!(issue.recover_from, Some(format!("{epoch}:4180")));
3399 assert!(
3400 issue.suggested_action.unwrap().contains("4180"),
3401 "the consumer is told exactly which cursor recovers the gap"
3402 );
3403
3404 let from_scratch = SocketIssueMessage::replay_lagged(Some("trades".to_string()), 9, None);
3407 assert_eq!(from_scratch.recover_from, None);
3408 assert!(from_scratch
3409 .suggested_action
3410 .unwrap()
3411 .contains("without `after`"));
3412 }
3413
3414 fn bus_message(key: &str) -> Arc<BusMessage> {
3415 Arc::new(BusMessage {
3416 key: key.to_string(),
3417 entity: "Trade/append".to_string(),
3418 payload: Arc::new(Bytes::from_static(b"{}")),
3419 })
3420 }
3421
3422 #[tokio::test]
3427 async fn draining_during_a_replay_keeps_a_busy_view_from_lapping_it() {
3428 let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(16);
3429 let mut pending = VecDeque::new();
3430 let mut lagged = None;
3431
3432 for index in 0..48 {
3435 sender.send(bus_message(&format!("k{index}"))).unwrap();
3436 drain_available(&mut receiver, &mut pending, &mut lagged);
3437 }
3438
3439 assert_eq!(lagged, None, "draining as we go means nothing is dropped");
3440 assert_eq!(pending.len(), 48, "every published frame is buffered");
3441 }
3442
3443 #[tokio::test]
3444 async fn a_replay_that_never_drains_is_reported_as_a_gap() {
3445 let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3446 for index in 0..32 {
3447 sender.send(bus_message(&format!("k{index}"))).unwrap();
3448 }
3449
3450 let mut pending = VecDeque::new();
3451 let mut lagged = None;
3452 drain_available(&mut receiver, &mut pending, &mut lagged);
3453
3454 assert!(
3455 lagged.is_some(),
3456 "overflowing the bus is a gap, not silent truncation"
3457 );
3458 }
3459
3460 #[tokio::test]
3464 async fn nothing_after_a_gap_is_buffered_ahead_of_the_report() {
3465 let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3466 let mut pending = VecDeque::new();
3467 let mut lagged = None;
3468
3469 sender.send(bus_message("before")).unwrap();
3471 drain_available(&mut receiver, &mut pending, &mut lagged);
3472 assert_eq!(pending.len(), 1);
3473
3474 for index in 0..32 {
3476 sender.send(bus_message(&format!("lost{index}"))).unwrap();
3477 }
3478 drain_available(&mut receiver, &mut pending, &mut lagged);
3479 assert!(lagged.is_some());
3480
3481 let buffered_at_gap = pending.len();
3484 drain_available(&mut receiver, &mut pending, &mut lagged);
3485 assert_eq!(
3486 pending.len(),
3487 buffered_at_gap,
3488 "post-gap frames must not join the pre-gap flush"
3489 );
3490 assert_eq!(pending.front().unwrap().key, "before");
3491 }
3492 #[test]
3495 fn the_seam_between_replay_and_live_neither_repeats_nor_skips() {
3496 let mut last_sent = Some(4211);
3497
3498 assert!(
3499 already_delivered(Some(4211), &mut last_sent),
3500 "the record the replay ended on must not be sent twice"
3501 );
3502 assert!(already_delivered(Some(4100), &mut last_sent));
3503 assert_eq!(last_sent, Some(4211), "a duplicate never moves the mark");
3504
3505 assert!(
3506 !already_delivered(Some(4212), &mut last_sent),
3507 "the next record is new"
3508 );
3509 assert_eq!(last_sent, Some(4212));
3510
3511 assert!(!already_delivered(None, &mut last_sent));
3514 assert_eq!(last_sent, Some(4212));
3515 }
3516
3517 #[test]
3520 fn a_fresh_subscription_delivers_its_first_live_frame() {
3521 let mut last_sent = None;
3522 assert!(!already_delivered(Some(0), &mut last_sent));
3523 assert_eq!(last_sent, Some(0));
3524 }
3525
3526 mod over_a_socket {
3531 use super::*;
3532 use crate::journal::{EventJournal, JournalConfig};
3533 use crate::projector::Projector;
3534 use crate::{MutationBatch, SlotContext};
3535 use arete_interpreter::Mutation;
3536 use futures_util::{SinkExt, StreamExt};
3537 use std::time::Duration;
3538 use tokio::net::{TcpListener, TcpStream};
3539 use tokio::sync::mpsc;
3540 use tokio_tungstenite::tungstenite::Message;
3541 use tokio_tungstenite::{client_async, WebSocketStream};
3542
3543 const RETAINED: u64 = 600;
3544
3545 fn append_index() -> ViewIndex {
3546 let mut index = ViewIndex::new();
3547 index.add_spec(ViewSpec {
3548 id: "Trade/append".to_string(),
3549 export: "Trade".to_string(),
3550 mode: Mode::Append,
3551 wire_format: Default::default(),
3552 projection: Projection::all(),
3553 filters: Filters::all(),
3554 delivery: Delivery::default(),
3555 pipeline: None,
3556 source_view: None,
3557 });
3558 index
3559 }
3560
3561 fn trade(index: u64) -> MutationBatch {
3562 MutationBatch::with_slot_context(
3563 vec![Mutation {
3564 export: "Trade".to_string(),
3565 key: json!(format!("pool{}", index % 4)),
3566 patch: json!({"trade": index}),
3567 append: vec![],
3568 }]
3569 .into_iter()
3570 .collect(),
3571 SlotContext::new(100 + index / 3, index % 3),
3572 )
3573 }
3574
3575 struct Harness {
3576 addr: SocketAddr,
3577 journal: Arc<EventJournal>,
3578 tx: mpsc::Sender<MutationBatch>,
3579 }
3580
3581 impl Harness {
3582 async fn start() -> Self {
3583 let view_index = Arc::new(append_index());
3584 let entity_cache = EntityCache::new();
3585 let bus_manager = BusManager::new();
3586 let journal = Arc::new(EventJournal::new(JournalConfig {
3587 enabled: true,
3588 max_bytes_per_view: u64::MAX,
3589 max_records_per_view: 10_000,
3590 max_age: Duration::from_secs(3_600),
3591 }));
3592
3593 let (tx, rx) = mpsc::channel::<MutationBatch>(256);
3594 tokio::spawn(
3595 Projector::new(
3596 view_index.clone(),
3597 bus_manager.clone(),
3598 entity_cache.clone(),
3599 rx,
3600 #[cfg(feature = "otel")]
3601 None,
3602 )
3603 .with_journal(journal.clone())
3604 .run(),
3605 );
3606
3607 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3608 let addr = listener.local_addr().unwrap();
3609 let server = WebSocketServer::new(
3610 addr,
3611 bus_manager,
3612 entity_cache,
3613 view_index,
3614 #[cfg(feature = "otel")]
3615 None,
3616 )
3617 .with_journal(journal.clone());
3618 let (acceptor, _cleanup) = server.into_acceptor();
3619 tokio::spawn(async move { acceptor.serve_listener(listener).await });
3620
3621 Self { addr, journal, tx }
3622 }
3623
3624 async fn publish(&self, range: std::ops::Range<u64>) {
3625 for index in range {
3626 self.tx.send(trade(index)).await.unwrap();
3627 }
3628 let (ack, wait) = oneshot::channel();
3629 self.tx
3630 .send(MutationBatch::flush_marker(ack))
3631 .await
3632 .unwrap();
3633 wait.await.unwrap();
3634 }
3635
3636 async fn connect(&self) -> WebSocketStream<TcpStream> {
3637 let stream = TcpStream::connect(self.addr).await.unwrap();
3638 client_async(format!("ws://{}/", self.addr), stream)
3639 .await
3640 .unwrap()
3641 .0
3642 }
3643 }
3644
3645 async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
3646 loop {
3647 let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
3648 .await
3649 .expect("the server answers within the timeout")
3650 .expect("the stream stays open")
3651 .expect("a readable frame");
3652 let bytes = match &message {
3655 Message::Text(text) => text.as_bytes(),
3656 Message::Binary(bytes) => bytes.as_ref(),
3657 _ => continue,
3658 };
3659 return serde_json::from_slice(bytes).expect("frames are JSON");
3660 }
3661 }
3662
3663 async fn collect_trades(socket: &mut WebSocketStream<TcpStream>, count: usize) -> Vec<u64> {
3665 let mut offsets = Vec::with_capacity(count);
3666 while offsets.len() < count {
3667 let frame = next_frame(socket).await;
3668 assert_ne!(
3669 frame["type"], "error",
3670 "no error frame should interrupt delivery: {frame}"
3671 );
3672 if let Some(offset) = frame["offset"].as_u64() {
3673 offsets.push(offset);
3674 }
3675 }
3676 offsets
3677 }
3678
3679 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3683 async fn a_reconnect_replays_from_a_cursor_and_continues_live() {
3684 let harness = Harness::start().await;
3685 harness.publish(0..RETAINED).await;
3686
3687 let cursor = harness.journal.window("Trade/append").await;
3688 let stored = format!("{}:{}", cursor.epoch, 99);
3689
3690 let mut socket = harness.connect().await;
3691 socket
3692 .send(Message::Text(
3693 json!({
3694 "type": "subscribe",
3695 "protocolVersion": 2,
3696 "subscriptionId": "trades",
3697 "query": {"view": "Trade/append", "after": stored},
3698 })
3699 .to_string()
3700 .into(),
3701 ))
3702 .await
3703 .unwrap();
3704
3705 let ack = next_frame(&mut socket).await;
3706 assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3707 assert_eq!(ack["replayWindow"]["next"], json!(RETAINED));
3708
3709 let replayed = collect_trades(&mut socket, (RETAINED - 100) as usize).await;
3711 assert_eq!(
3712 replayed,
3713 (100..RETAINED).collect::<Vec<_>>(),
3714 "every event after the cursor, in order, exactly once"
3715 );
3716
3717 harness.publish(RETAINED..RETAINED + 40).await;
3719 let live = collect_trades(&mut socket, 40).await;
3720 assert_eq!(
3721 live,
3722 (RETAINED..RETAINED + 40).collect::<Vec<_>>(),
3723 "the live stream resumes exactly where the replay stopped"
3724 );
3725
3726 socket.close(None).await.ok();
3727 }
3728
3729 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3734 async fn events_published_during_a_replay_are_not_lost() {
3735 let harness = Harness::start().await;
3736 harness.publish(0..RETAINED).await;
3737
3738 let epoch = harness.journal.window("Trade/append").await.epoch;
3739 let mut socket = harness.connect().await;
3740 socket
3741 .send(Message::Text(
3742 json!({
3743 "type": "subscribe",
3744 "protocolVersion": 2,
3745 "subscriptionId": "trades",
3746 "query": {"view": "Trade/append", "after": format!("{epoch}:0")},
3747 })
3748 .to_string()
3749 .into(),
3750 ))
3751 .await
3752 .unwrap();
3753 let ack = next_frame(&mut socket).await;
3754 assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3755
3756 harness.publish(RETAINED..RETAINED + 200).await;
3758
3759 let total = (RETAINED + 200 - 1) as usize;
3760 let delivered = collect_trades(&mut socket, total).await;
3761 assert_eq!(
3762 delivered,
3763 (1..RETAINED + 200).collect::<Vec<_>>(),
3764 "replay and live output join without a gap or a repeat"
3765 );
3766
3767 socket.close(None).await.ok();
3768 }
3769
3770 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3773 async fn a_stale_epoch_is_refused_and_frees_the_subscription_id() {
3774 let harness = Harness::start().await;
3775 harness.publish(0..50).await;
3776
3777 let mut socket = harness.connect().await;
3778 socket
3779 .send(Message::Text(
3780 json!({
3781 "type": "subscribe",
3782 "protocolVersion": 2,
3783 "subscriptionId": "trades",
3784 "query": {
3785 "view": "Trade/append",
3786 "after": format!("{}:10", crate::journal::JournalEpoch::new()),
3787 },
3788 })
3789 .to_string()
3790 .into(),
3791 ))
3792 .await
3793 .unwrap();
3794
3795 let error = next_frame(&mut socket).await;
3796 assert_eq!(error["type"], "error", "unexpected frame: {error}");
3797 assert_eq!(error["code"], "cursor-epoch-changed");
3798
3799 socket
3802 .send(Message::Text(
3803 json!({
3804 "type": "subscribe",
3805 "protocolVersion": 2,
3806 "subscriptionId": "trades",
3807 "query": {"view": "Trade/append"},
3808 })
3809 .to_string()
3810 .into(),
3811 ))
3812 .await
3813 .unwrap();
3814 let ack = next_frame(&mut socket).await;
3815 assert_eq!(
3816 ack["op"], "subscribed",
3817 "the refused id must be reusable: {ack}"
3818 );
3819
3820 assert_eq!(collect_trades(&mut socket, 50).await.len(), 50);
3821 socket.close(None).await.ok();
3822 }
3823
3824 const ROUND: &str = "Round/state";
3825
3826 #[derive(Clone)]
3829 struct StateServer {
3830 addr: SocketAddr,
3831 bus_manager: BusManager,
3832 entity_cache: EntityCache,
3833 }
3834
3835 impl StateServer {
3836 async fn start() -> Self {
3837 let mut index = ViewIndex::new();
3838 index.add_spec(ViewSpec {
3839 id: ROUND.to_string(),
3840 export: "Round".to_string(),
3841 mode: Mode::State,
3842 wire_format: Default::default(),
3843 projection: Projection::all(),
3844 filters: Filters::all(),
3845 delivery: Delivery::default(),
3846 pipeline: None,
3847 source_view: None,
3848 });
3849 let bus_manager = BusManager::new();
3850 let entity_cache = EntityCache::new();
3851 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3852 let addr = listener.local_addr().unwrap();
3853 let server = WebSocketServer::new(
3854 addr,
3855 bus_manager.clone(),
3856 entity_cache.clone(),
3857 Arc::new(index),
3858 #[cfg(feature = "otel")]
3859 None,
3860 );
3861 let (acceptor, _cleanup) = server.into_acceptor();
3862 tokio::spawn(async move { acceptor.serve_listener(listener).await });
3863 Self {
3864 addr,
3865 bus_manager,
3866 entity_cache,
3867 }
3868 }
3869
3870 async fn publish(&self, key: &str, mut patch: Value, seq: &str) {
3873 patch["_seq"] = Value::String(seq.to_string());
3874 self.entity_cache
3875 .upsert_with_append(ROUND, key, patch.clone(), &[], PatchOrigin::Unknown)
3876 .await;
3877 self.publish_frame(key, "patch", patch, seq).await;
3878 }
3879
3880 async fn publish_frame(&self, key: &str, op: &str, data: Value, seq: &str) {
3882 let frame = json!({
3883 "mode": "state",
3884 "entity": ROUND,
3885 "op": op,
3886 "key": key,
3887 "data": data,
3888 "seq": seq,
3889 });
3890 self.bus_manager
3891 .publish_state(ROUND, key, Arc::new(Bytes::from(frame.to_string())))
3892 .await;
3893 }
3894
3895 async fn subscribe(&self, key: &str) -> WebSocketStream<TcpStream> {
3897 let stream = TcpStream::connect(self.addr).await.unwrap();
3898 let mut socket = client_async(format!("ws://{}/", self.addr), stream)
3899 .await
3900 .unwrap()
3901 .0;
3902 socket
3903 .send(Message::Text(
3904 json!({
3905 "type": "subscribe",
3906 "protocolVersion": 2,
3907 "subscriptionId": "round",
3908 "query": {"view": ROUND, "key": key},
3909 })
3910 .to_string()
3911 .into(),
3912 ))
3913 .await
3914 .unwrap();
3915 let ack = next_frame(&mut socket).await;
3916 assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3917 loop {
3918 let frame = next_frame(&mut socket).await;
3919 assert_eq!(frame["op"], "snapshot", "unexpected frame: {frame}");
3920 if frame["complete"] == true {
3921 return socket;
3922 }
3923 }
3924 }
3925 }
3926
3927 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
3933 async fn a_state_subscriber_keeps_the_fields_of_an_overwritten_patch() {
3934 let server = StateServer::start().await;
3935 server
3936 .publish("7", json!({"id": 7}), "100:000000000001")
3937 .await;
3938 let mut socket = server.subscribe("7").await;
3939
3940 server
3943 .publish(
3944 "7",
3945 json!({"results": {"slot_hash": "abc"}}),
3946 "140:003836292257",
3947 )
3948 .await;
3949 let frame = next_frame(&mut socket).await;
3950 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3951 assert_eq!(frame["seq"], "140:003836292257");
3952
3953 let writer = server.clone();
3958 tokio::spawn(async move {
3959 writer
3960 .publish("7", json!({"entropy": {"seed": "def"}}), "140:000000001200")
3961 .await;
3962 writer
3963 .publish("7", json!({"total": 1}), "140:000000001200")
3964 .await;
3965 })
3966 .await
3967 .unwrap();
3968 let frame = next_frame(&mut socket).await;
3969 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3970 assert_eq!(
3971 frame["data"],
3972 json!({
3973 "id": 7,
3974 "results": {"slot_hash": "abc"},
3975 "entropy": {"seed": "def"},
3976 "total": 1,
3977 "_seq": "140:000000001200",
3978 })
3979 );
3980 assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
3983
3984 server
3986 .publish("7", json!({"total": 2}), "141:000000000002")
3987 .await;
3988 let frame = next_frame(&mut socket).await;
3989 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3990 assert_eq!(
3991 frame["data"],
3992 json!({"total": 2, "_seq": "141:000000000002"})
3993 );
3994
3995 socket.close(None).await.ok();
3996 }
3997
3998 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4003 async fn a_key_deleted_and_recreated_between_reads_replaces_the_clients_copy() {
4004 let server = StateServer::start().await;
4005 server
4006 .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
4007 .await;
4008 let mut socket = server.subscribe("7").await;
4009
4010 let writer = server.clone();
4011 tokio::spawn(async move {
4012 writer.entity_cache.remove(ROUND, "7").await;
4013 writer
4014 .publish_frame("7", "delete", Value::Null, "101:000000000001")
4015 .await;
4016 writer
4017 .publish("7", json!({"id": 7, "fresh": true}), "101:000000000002")
4018 .await;
4019 })
4020 .await
4021 .unwrap();
4022 let frame = next_frame(&mut socket).await;
4023 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4024 assert_eq!(
4025 frame["data"],
4026 json!({"id": 7, "fresh": true, "_seq": "101:000000000002"})
4027 );
4028 assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4029
4030 let writer = server.clone();
4032 tokio::spawn(async move {
4033 writer
4034 .publish("7", json!({"n": 1}), "102:000000000001")
4035 .await;
4036 writer.entity_cache.remove(ROUND, "7").await;
4037 writer
4038 .publish_frame("7", "delete", Value::Null, "102:000000000002")
4039 .await;
4040 })
4041 .await
4042 .unwrap();
4043 let frame = next_frame(&mut socket).await;
4044 assert_eq!(frame["op"], "delete", "unexpected frame: {frame}");
4045
4046 server
4048 .publish("7", json!({"id": 7, "third": true}), "103:000000000001")
4049 .await;
4050 let frame = next_frame(&mut socket).await;
4051 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4052 assert_eq!(frame["data"]["third"], true);
4053 assert!(frame["data"].get("fresh").is_none(), "{frame}");
4054 socket.close(None).await.ok();
4055 }
4056
4057 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4061 async fn a_copy_of_a_deleted_entity_takes_no_patch_until_it_is_replaced() {
4062 let server = StateServer::start().await;
4063 server
4064 .publish("7", json!({"id": 7, "old": true}), "100:000000000001")
4065 .await;
4066 let mut socket = server.subscribe("7").await;
4067
4068 let writer = server.clone();
4069 tokio::spawn(async move {
4070 writer.entity_cache.remove(ROUND, "7").await;
4071 writer
4072 .publish_frame("7", "delete", Value::Null, "101:000000000001")
4073 .await;
4074 writer
4075 .publish_frame("7", "patch", json!({"n": 1}), "101:000000000002")
4076 .await;
4077 })
4078 .await
4079 .unwrap();
4080 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
4084
4085 let whole = json!({"id": 7, "n": 1, "_seq": "101:000000000002"});
4086 server
4087 .entity_cache
4088 .store_whole(ROUND, "7", whole.clone())
4089 .await;
4090 server
4091 .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
4092 .await;
4093 let frame = next_frame(&mut socket).await;
4094 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4095 assert_eq!(frame["data"], whole);
4096 socket.close(None).await.ok();
4097 }
4098
4099 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4103 async fn a_whole_entity_among_overwritten_frames_replaces_the_clients_copy() {
4104 let server = StateServer::start().await;
4105 server
4106 .publish("7", json!({"id": 7, "dropped": 1}), "100:000000000001")
4107 .await;
4108 let mut socket = server.subscribe("7").await;
4109
4110 let writer = server.clone();
4111 tokio::spawn(async move {
4112 let whole = json!({"id": 7, "_seq": "101:000000000001"});
4113 writer
4114 .entity_cache
4115 .store_whole(ROUND, "7", whole.clone())
4116 .await;
4117 writer
4118 .publish_frame("7", "upsert", whole, "101:000000000001")
4119 .await;
4120 writer
4121 .publish("7", json!({"n": 1}), "101:000000000002")
4122 .await;
4123 })
4124 .await
4125 .unwrap();
4126 let frame = next_frame(&mut socket).await;
4127 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4128 assert_eq!(
4129 frame["data"],
4130 json!({"id": 7, "n": 1, "_seq": "101:000000000002"})
4131 );
4132 assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4133 socket.close(None).await.ok();
4134 }
4135
4136 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4145 async fn a_holder_behind_on_an_evicted_key_catches_up_when_it_returns() {
4146 let server = StateServer::start().await;
4147 server
4148 .publish("7", json!({"id": 7}), "100:000000000001")
4149 .await;
4150 let mut socket = server.subscribe("7").await;
4151 server.entity_cache.remove(ROUND, "7").await;
4152
4153 let writer = server.clone();
4154 tokio::spawn(async move {
4155 writer
4156 .publish_frame("7", "patch", json!({"a": 1}), "101:000000000001")
4157 .await;
4158 writer
4159 .publish_frame("7", "patch", json!({"b": 2}), "101:000000000002")
4160 .await;
4161 })
4162 .await
4163 .unwrap();
4164 let frame = next_frame(&mut socket).await;
4165 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4166 assert_eq!(frame["data"], json!({"b": 2}));
4167 assert_eq!(frame["seq"], "101:000000000002");
4168
4169 let whole = json!({"id": 7, "a": 1, "b": 2, "_seq": "101:000000000002"});
4171 server
4172 .entity_cache
4173 .store_whole(ROUND, "7", whole.clone())
4174 .await;
4175 server
4176 .publish_frame("7", "upsert", whole.clone(), "101:000000000002")
4177 .await;
4178 let frame = next_frame(&mut socket).await;
4179 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4180 assert_eq!(frame["data"], whole);
4181 assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
4182
4183 server
4185 .publish("7", json!({"c": 3}), "102:000000000001")
4186 .await;
4187 let frame = next_frame(&mut socket).await;
4188 assert_eq!(frame["data"], json!({"c": 3, "_seq": "102:000000000001"}));
4189 assert_eq!(frame["seq"], "102:000000000001");
4190 socket.close(None).await.ok();
4191 }
4192 }
4193
4194 mod partial_entities_over_a_socket {
4198 use super::*;
4199 use crate::projector::Projector;
4200 use crate::{MutationBatch, SlotContext};
4201 use arete_interpreter::Mutation;
4202 use futures_util::{SinkExt, StreamExt};
4203 use std::time::Duration;
4204 use tokio::net::{TcpListener, TcpStream};
4205 use tokio::sync::mpsc;
4206 use tokio_tungstenite::tungstenite::Message;
4207 use tokio_tungstenite::{client_async, WebSocketStream};
4208
4209 fn thing_index() -> ViewIndex {
4210 let mut index = ViewIndex::new();
4211 index.add_spec(list_spec());
4212 index.add_spec(ViewSpec {
4213 id: "Thing/state".to_string(),
4214 mode: Mode::State,
4215 ..list_spec()
4216 });
4217 index
4218 }
4219
4220 struct Harness {
4221 addr: SocketAddr,
4222 tx: mpsc::Sender<MutationBatch>,
4223 slot: std::sync::atomic::AtomicU64,
4224 }
4225
4226 impl Harness {
4227 async fn start(delivery: WebSocketDeliveryConfig) -> Self {
4228 Self::start_with_cache(delivery, EntityCache::new()).await
4229 }
4230
4231 async fn start_with_cache(
4232 delivery: WebSocketDeliveryConfig,
4233 entity_cache: EntityCache,
4234 ) -> Self {
4235 let view_index = Arc::new(thing_index());
4236 let bus_manager = BusManager::new();
4237 let (tx, rx) = mpsc::channel::<MutationBatch>(64);
4238 tokio::spawn(tokio::task::unconstrained(
4244 Projector::new(
4245 view_index.clone(),
4246 bus_manager.clone(),
4247 entity_cache.clone(),
4248 rx,
4249 #[cfg(feature = "otel")]
4250 None,
4251 )
4252 .run(),
4253 ));
4254 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4255 let addr = listener.local_addr().unwrap();
4256 let server = WebSocketServer::new(
4257 addr,
4258 bus_manager,
4259 entity_cache,
4260 view_index,
4261 #[cfg(feature = "otel")]
4262 None,
4263 )
4264 .with_delivery_config(delivery);
4265 let (acceptor, _cleanup) = server.into_acceptor();
4266 tokio::spawn(async move { acceptor.serve_listener(listener).await });
4267 Self {
4268 addr,
4269 tx,
4270 slot: std::sync::atomic::AtomicU64::new(100),
4271 }
4272 }
4273
4274 async fn patch(&self, key: &str, patch: Value) {
4276 self.mutate(Mutation {
4277 export: "Thing".to_string(),
4278 key: json!(key),
4279 patch,
4280 append: vec![],
4281 })
4282 .await;
4283 }
4284
4285 async fn mutate(&self, mutation: Mutation) {
4286 let slot = self.slot.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
4287 self.tx
4288 .send(MutationBatch::with_slot_context(
4289 vec![mutation].into_iter().collect(),
4290 SlotContext::new(slot, 0),
4291 ))
4292 .await
4293 .unwrap();
4294 let (ack, wait) = oneshot::channel();
4295 self.tx
4296 .send(MutationBatch::flush_marker(ack))
4297 .await
4298 .unwrap();
4299 wait.await.unwrap();
4300 }
4301
4302 async fn subscribe(&self, query: Value, snapshot: bool) -> WebSocketStream<TcpStream> {
4303 let stream = TcpStream::connect(self.addr).await.unwrap();
4304 let mut socket = client_async(format!("ws://{}/", self.addr), stream)
4305 .await
4306 .unwrap()
4307 .0;
4308 socket
4309 .send(Message::Text(
4310 json!({
4311 "type": "subscribe",
4312 "protocolVersion": 2,
4313 "subscriptionId": "things",
4314 "query": query,
4315 "snapshot": {"enabled": snapshot},
4316 })
4317 .to_string()
4318 .into(),
4319 ))
4320 .await
4321 .unwrap();
4322 let ack = next_frame(&mut socket).await;
4323 assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
4324 assert_eq!(
4325 ack["wholeEntities"],
4326 json!(true),
4327 "the ack advertises the whole-entity guarantee: {ack}"
4328 );
4329 socket
4330 }
4331 }
4332
4333 async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
4334 loop {
4335 let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
4336 .await
4337 .expect("the server answers within the timeout")
4338 .expect("the stream stays open")
4339 .expect("a readable frame");
4340 let bytes = match &message {
4342 Message::Text(text) => text.as_bytes(),
4343 Message::Binary(bytes) => bytes.as_ref(),
4344 _ => continue,
4345 };
4346 return serde_json::from_slice(bytes).expect("frames are JSON");
4347 }
4348 }
4349
4350 async fn snapshot_keys(socket: &mut WebSocketStream<TcpStream>) -> Vec<String> {
4352 let mut keys = Vec::new();
4353 loop {
4354 let frame = next_frame(socket).await;
4355 assert_eq!(frame["op"], "snapshot", "expected a snapshot: {frame}");
4356 keys.extend(
4357 frame["data"]
4358 .as_array()
4359 .unwrap()
4360 .iter()
4361 .map(|row| row["key"].as_str().unwrap().to_string()),
4362 );
4363 if frame["complete"] == json!(true) {
4364 return keys;
4365 }
4366 }
4367 }
4368
4369 async fn seed(harness: &Harness) {
4370 for key in ["a", "b", "c"] {
4371 harness
4372 .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
4373 .await;
4374 }
4375 }
4376
4377 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4378 async fn a_change_to_a_key_cut_by_snapshot_limit_is_an_upsert() {
4379 let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4380 seed(&harness).await;
4381
4382 let mut socket = harness
4383 .subscribe(json!({"view": "Thing/list", "snapshotLimit": 1}), true)
4384 .await;
4385 assert_eq!(snapshot_keys(&mut socket).await, ["c"]);
4387
4388 harness.patch("a", json!({"count": 2})).await;
4389 let frame = next_frame(&mut socket).await;
4390 assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4391 assert_eq!(frame["op"], "upsert", "a key never sent arrives whole");
4392 assert_eq!(frame["data"]["name"], "thing-a");
4393 assert_eq!(frame["data"]["count"], 2);
4394
4395 harness.patch("a", json!({"count": 3})).await;
4397 let frame = next_frame(&mut socket).await;
4398 assert_eq!(frame["key"], "a");
4399 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4400 assert!(frame["data"].get("name").is_none());
4401
4402 harness.patch("c", json!({"count": 2})).await;
4404 let frame = next_frame(&mut socket).await;
4405 assert_eq!(frame["key"], "c");
4406 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4407 socket.close(None).await.ok();
4408 }
4409
4410 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4411 async fn a_list_without_a_snapshot_upserts_each_key_on_its_first_change() {
4412 let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4413 seed(&harness).await;
4414
4415 let mut socket = harness
4416 .subscribe(json!({"view": "Thing/list"}), false)
4417 .await;
4418 harness.patch("b", json!({"count": 2})).await;
4419 let frame = next_frame(&mut socket).await;
4420 assert_eq!(frame["key"], "b");
4421 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4422 assert_eq!(frame["data"]["name"], "thing-b");
4423
4424 harness.patch("b", json!({"count": 3})).await;
4425 let frame = next_frame(&mut socket).await;
4426 assert_eq!(frame["key"], "b");
4427 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4428 socket.close(None).await.ok();
4429 }
4430
4431 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4435 async fn a_state_holder_keeps_its_patches_when_the_cache_evicts_the_key() {
4436 let harness = Harness::start_with_cache(
4437 WebSocketDeliveryConfig::default(),
4438 EntityCache::with_config(crate::cache::EntityCacheConfig {
4439 max_entities_per_view: 2,
4440 ..Default::default()
4441 }),
4442 )
4443 .await;
4444 harness
4445 .patch("a", json!({"name": "thing-a", "count": 1}))
4446 .await;
4447 let mut socket = harness
4448 .subscribe(json!({"view": "Thing/state", "key": "a"}), true)
4449 .await;
4450 assert_eq!(snapshot_keys(&mut socket).await, ["a"]);
4451
4452 for key in ["b", "c"] {
4454 harness
4455 .patch(key, json!({"name": format!("thing-{key}"), "count": 1}))
4456 .await;
4457 }
4458 harness.patch("a", json!({"count": 2})).await;
4459 let frame = next_frame(&mut socket).await;
4460 assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4461 assert_eq!(frame["op"], "patch", "the holder merges it: {frame}");
4462 assert_eq!(frame["data"]["count"], 2);
4463
4464 let mut whole = Mutation {
4466 export: "Thing".to_string(),
4467 key: json!("a"),
4468 patch: json!({"name": "thing-a", "count": 2}),
4469 append: vec![],
4470 };
4471 whole.mark_whole_entity();
4472 harness.mutate(whole).await;
4473 let frame = next_frame(&mut socket).await;
4474 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4475 assert_eq!(frame["data"]["name"], "thing-a");
4476 assert_eq!(frame["data"]["count"], 2);
4477 assert!(frame["data"]
4478 .get(arete_interpreter::WHOLE_ENTITY_MARKER)
4479 .is_none());
4480
4481 harness.patch("a", json!({"count": 3})).await;
4482 let frame = next_frame(&mut socket).await;
4483 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4484 socket.close(None).await.ok();
4485 }
4486
4487 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4488 async fn a_state_view_without_a_snapshot_upserts_on_the_first_change() {
4489 let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4490 seed(&harness).await;
4491
4492 let mut socket = harness
4493 .subscribe(json!({"view": "Thing/state", "key": "a"}), false)
4494 .await;
4495 harness.patch("a", json!({"count": 2})).await;
4496 let frame = next_frame(&mut socket).await;
4497 assert_eq!(frame["key"], "a");
4498 assert_eq!(frame["op"], "upsert", "unexpected frame: {frame}");
4499 assert_eq!(frame["data"]["name"], "thing-a");
4500 assert_eq!(frame["data"]["count"], 2);
4501
4502 harness.patch("a", json!({"count": 3})).await;
4503 let frame = next_frame(&mut socket).await;
4504 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
4505 socket.close(None).await.ok();
4506 }
4507
4508 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4509 async fn a_coalesced_list_sends_no_remove_for_a_key_it_never_sent() {
4510 let harness = Harness::start(WebSocketDeliveryConfig {
4511 collection_coalesce_ms: Some(10),
4512 ..Default::default()
4513 })
4514 .await;
4515 seed(&harness).await;
4516
4517 let mut socket = harness
4520 .subscribe(json!({"view": "Thing/list", "take": 1}), false)
4521 .await;
4522 harness.patch("a", json!({"count": 2})).await;
4525 let frame = next_frame(&mut socket).await;
4526 assert_eq!(frame["key"], "a", "unexpected frame: {frame}");
4527 assert_eq!(frame["op"], "upsert");
4528 assert_eq!(frame["data"]["name"], "thing-a");
4529
4530 harness.patch("c", json!({"count": 2})).await;
4532 let mut frames = [next_frame(&mut socket).await, next_frame(&mut socket).await];
4533 frames.sort_by_key(|frame| frame["key"].as_str().unwrap().to_string());
4534 assert_eq!(frames[0]["key"], "a");
4535 assert_eq!(frames[0]["op"], "remove");
4536 assert_eq!(frames[1]["key"], "c");
4537 assert_eq!(frames[1]["op"], "upsert");
4538 assert_eq!(frames[1]["data"]["name"], "thing-c");
4539 socket.close(None).await.ok();
4540 }
4541
4542 fn version_of(data: &Value) -> (String, u64) {
4544 let version = data["_version"]
4545 .as_str()
4546 .unwrap_or_else(|| panic!("no _version in {data}"));
4547 let (epoch, counter) = version.split_once(':').expect("epoch:counter");
4548 (epoch.to_string(), counter.parse().expect("decimal counter"))
4549 }
4550
4551 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
4560 async fn snapshot_rows_patches_and_catch_ups_carry_the_version() {
4561 let harness = Harness::start(WebSocketDeliveryConfig::default()).await;
4562 harness.patch("7", json!({"name": "thing-7"})).await;
4563
4564 let mut list = harness.subscribe(json!({"view": "Thing/list"}), true).await;
4565 let snapshot = next_frame(&mut list).await;
4566 assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
4567 let (epoch, _) = version_of(&snapshot["data"][0]["data"]);
4568
4569 let mut state = harness
4570 .subscribe(json!({"view": "Thing/state", "key": "7"}), true)
4571 .await;
4572 let snapshot = next_frame(&mut state).await;
4573 assert_eq!(snapshot["op"], "snapshot", "unexpected frame: {snapshot}");
4574 let (state_epoch, seeded) = version_of(&snapshot["data"][0]["data"]);
4575 assert_eq!(state_epoch, epoch, "one projector, one epoch");
4576
4577 harness.patch("7", json!({"count": 1})).await;
4578 let patch = next_frame(&mut list).await;
4579 assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
4580 assert_eq!(version_of(&patch["data"]).0, epoch);
4581 let patch = next_frame(&mut state).await;
4582 assert_eq!(patch["op"], "patch", "unexpected frame: {patch}");
4583 let (_, patched) = version_of(&patch["data"]);
4584 assert!(patched > seeded, "{patched} after {seeded}");
4585
4586 let slot = harness
4588 .slot
4589 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
4590 let mutation = |patch: Value| Mutation {
4591 export: "Thing".to_string(),
4592 key: json!("7"),
4593 patch,
4594 append: vec![],
4595 };
4596 harness
4597 .tx
4598 .send(MutationBatch::with_slot_context(
4599 vec![
4600 mutation(json!({"count": 2})),
4601 mutation(json!({"flag": true})),
4602 ]
4603 .into_iter()
4604 .collect(),
4605 SlotContext::new(slot, 0),
4606 ))
4607 .await
4608 .unwrap();
4609
4610 let catch_up = next_frame(&mut state).await;
4611 assert_eq!(catch_up["op"], "patch", "unexpected frame: {catch_up}");
4612 assert!(catch_up.get("seq").is_none(), "a catch-up: {catch_up}");
4613 assert_eq!(catch_up["data"]["count"], 2);
4614 assert_eq!(catch_up["data"]["flag"], true);
4615 let (_, caught_up) = version_of(&catch_up["data"]);
4616 assert!(caught_up > patched, "{caught_up} after {patched}");
4617
4618 list.close(None).await.ok();
4619 state.close(None).await.ok();
4620 }
4621 }
4622
4623 mod session_expiry {
4625 use super::*;
4626 use crate::websocket::auth::SignedSessionAuthPlugin;
4627 use arete_auth::{KeyClass, SessionClaims, SigningKey, TokenSigner, TokenVerifier};
4628 use futures_util::{SinkExt, StreamExt};
4629 use std::time::Duration;
4630 use tokio::net::{TcpListener, TcpStream};
4631 use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
4632 use tokio_tungstenite::tungstenite::protocol::CloseFrame;
4633 use tokio_tungstenite::tungstenite::Message;
4634 use tokio_tungstenite::{client_async, WebSocketStream};
4635
4636 struct Server {
4637 addr: SocketAddr,
4638 signer: TokenSigner,
4639 }
4640
4641 impl Server {
4642 async fn start() -> Self {
4643 let signing_key = SigningKey::generate();
4644 let verifier =
4645 TokenVerifier::new(signing_key.verifying_key(), "test-issuer", "test-audience");
4646 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4647 let addr = listener.local_addr().unwrap();
4648 let server = WebSocketServer::new(
4649 addr,
4650 BusManager::new(),
4651 EntityCache::new(),
4652 Arc::new(ViewIndex::new()),
4653 #[cfg(feature = "otel")]
4654 None,
4655 )
4656 .with_auth_plugin(Arc::new(SignedSessionAuthPlugin::new(verifier)));
4657 let (acceptor, _cleanup) = server.into_acceptor();
4658 tokio::spawn(async move { acceptor.serve_listener(listener).await });
4659 Self {
4660 addr,
4661 signer: TokenSigner::new(signing_key, "test-issuer"),
4662 }
4663 }
4664
4665 fn token(&self, ttl_seconds: u64) -> String {
4666 let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
4667 .with_scope("read")
4668 .with_key_class(KeyClass::Secret)
4669 .with_ttl(ttl_seconds)
4670 .build();
4671 self.signer.sign(claims).unwrap()
4672 }
4673
4674 async fn connect(&self, token: &str) -> WebSocketStream<TcpStream> {
4675 let stream = TcpStream::connect(self.addr).await.unwrap();
4676 client_async(format!("ws://{}/?hs_token={token}", self.addr), stream)
4677 .await
4678 .unwrap()
4679 .0
4680 }
4681 }
4682
4683 async fn send_json(socket: &mut WebSocketStream<TcpStream>, message: Value) {
4684 socket
4685 .send(Message::Text(message.to_string().into()))
4686 .await
4687 .unwrap();
4688 }
4689
4690 async fn close_within(
4692 socket: &mut WebSocketStream<TcpStream>,
4693 wait: Duration,
4694 ) -> Option<Option<CloseFrame>> {
4695 tokio::time::timeout(wait, async {
4696 while let Some(Ok(message)) = socket.next().await {
4697 if let Message::Close(frame) = message {
4698 return Some(frame);
4699 }
4700 }
4701 Some(None)
4702 })
4703 .await
4704 .ok()
4705 .flatten()
4706 }
4707
4708 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4709 async fn an_expired_session_is_closed_with_the_reason() {
4710 let server = Server::start().await;
4711 let mut socket = server.connect(&server.token(2)).await;
4712
4713 tokio::time::sleep(Duration::from_secs(3)).await;
4714 send_json(&mut socket, json!({"type": "ping"})).await;
4715
4716 let frame = close_within(&mut socket, Duration::from_secs(5))
4717 .await
4718 .expect("the server closes the expired session")
4719 .expect("the close frame carries a reason");
4720 assert_eq!(frame.code, CloseCode::Policy);
4721 assert_eq!(
4722 frame.reason.as_str(),
4723 "token-expired: Authentication token expired"
4724 );
4725 }
4726
4727 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4728 async fn a_session_refreshed_in_band_outlives_its_first_token() {
4729 let server = Server::start().await;
4730 let mut socket = server.connect(&server.token(2)).await;
4731
4732 send_json(
4733 &mut socket,
4734 json!({"type": "refresh_auth", "token": server.token(3_600)}),
4735 )
4736 .await;
4737 let reply = tokio::time::timeout(Duration::from_secs(5), async {
4738 while let Some(Ok(message)) = socket.next().await {
4739 if let Message::Text(text) = message {
4740 return serde_json::from_str::<Value>(text.as_str()).ok();
4741 }
4742 }
4743 None
4744 })
4745 .await
4746 .expect("the server answers the refresh")
4747 .expect("the answer is JSON");
4748 assert_eq!(reply["success"], true, "refresh accepted: {reply}");
4749
4750 tokio::time::sleep(Duration::from_secs(3)).await;
4751 send_json(&mut socket, json!({"type": "ping"})).await;
4752 assert!(
4753 close_within(&mut socket, Duration::from_millis(1_500))
4754 .await
4755 .is_none(),
4756 "the socket stays open on the refreshed token"
4757 );
4758 }
4759 }
4760}