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.clone();
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 if subscription.snapshot.enabled {
1390 send_snapshot_batches(
1391 context,
1392 &subscription,
1393 &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1394 view_spec.mode,
1395 SnapshotPurpose::Initial,
1396 &context.entity_cache.snapshot_config(),
1397 )
1398 .await?;
1399 }
1400
1401 let task_context = context.clone();
1402 let subscription_id = subscription.subscription_id.clone();
1403 let query = subscription.query.clone();
1404 let view_spec_task = view_spec.clone();
1405 let span_view = view_id.clone();
1406 let span_key = key.clone();
1407 tokio::spawn(
1408 async move {
1409 let mut member = !initial.is_empty();
1410 loop {
1411 tokio::select! {
1412 _ = cancel_token.cancelled() => break,
1413 changed = receiver.changed() => {
1414 if changed.is_err() {
1415 break;
1416 }
1417 let (payload, published) = {
1418 let update = receiver.borrow_and_update();
1419 (update.payload.clone(), update.published)
1420 };
1421 let missed = published > seen + 1;
1426 seen = published;
1427 let metadata = source_frame_metadata(&payload);
1428 if metadata.op == "delete" {
1429 task_context.entity_cache.remove(&query.view, &key).await;
1430 if member && send_membership_frame(
1431 &task_context,
1432 &subscription_id,
1433 &view_spec_task,
1434 "delete",
1435 &key,
1436 Value::Null,
1437 metadata.seq,
1438 ).is_err() {
1439 break;
1440 }
1441 member = false;
1442 continue;
1443 }
1444
1445 let selected = load_query_entities(
1446 &task_context.entity_cache,
1447 None,
1448 &view_spec_task,
1449 &query,
1450 false,
1451 ).await;
1452 let is_member = !selected.is_empty();
1453 let result = match (member, selected.into_iter().next()) {
1454 (true, Some(_)) if !missed => send_scoped_source_payload(
1455 &task_context,
1456 &subscription_id,
1457 &query.view,
1458 payload,
1459 ),
1460 (true, Some((entity_key, data))) => send_membership_frame(
1469 &task_context,
1470 &subscription_id,
1471 &view_spec_task,
1472 "patch",
1473 &entity_key,
1474 data,
1475 None,
1476 ),
1477 (false, Some((entity_key, data))) => send_membership_frame(
1478 &task_context,
1479 &subscription_id,
1480 &view_spec_task,
1481 "upsert",
1482 &entity_key,
1483 data,
1484 metadata.seq,
1485 ),
1486 (true, None) => send_membership_frame(
1487 &task_context,
1488 &subscription_id,
1489 &view_spec_task,
1490 "remove",
1491 &key,
1492 Value::Null,
1493 metadata.seq,
1494 ),
1495 (false, None) => Ok(()),
1496 };
1497 if result.is_err() {
1498 break;
1499 }
1500 member = is_member;
1501 }
1502 }
1503 }
1504 }
1505 .instrument(info_span!("ws.subscribe.state", client_id = %context.client_id, view = %span_view, key = %span_key)),
1506 );
1507 Ok(())
1508}
1509
1510async fn attach_collection_subscription(
1511 context: &SubscriptionContext,
1512 subscription: Subscription,
1513 view_spec: ViewSpec,
1514 cancel_token: CancellationToken,
1515) -> Result<()> {
1516 let view_id = subscription.query.view.clone();
1517 let source_view_id = view_spec
1518 .source_view
1519 .clone()
1520 .unwrap_or_else(|| view_id.clone());
1521 let (mut receiver, initial_membership) = subscribe_collection_then_snapshot(
1522 context,
1523 &source_view_id,
1524 &view_spec,
1525 &subscription.query,
1526 )
1527 .await;
1528
1529 let mut snapshot_entities = initial_membership.clone();
1530 if let Some(limit) = subscription.query.snapshot_limit {
1531 snapshot_entities.truncate(limit);
1532 }
1533 enforce_snapshot_limit(context, snapshot_entities.len())?;
1534 send_subscribed_frame(context, &subscription, &view_spec)?;
1535 if subscription.snapshot.enabled {
1536 send_snapshot_batches(
1537 context,
1538 &subscription,
1539 &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1540 view_spec.mode,
1541 SnapshotPurpose::Initial,
1542 &context.entity_cache.snapshot_config(),
1543 )
1544 .await?;
1545 }
1546
1547 let task_context = context.clone();
1548 let task_subscription = subscription.clone();
1549 let subscription_id = subscription.subscription_id.clone();
1550 let query = subscription.query.clone();
1551 let view_spec_task = view_spec.clone();
1552 let span_view = view_id.clone();
1553 tokio::spawn(
1554 async move {
1555 let mut current = initial_membership;
1556 let coalesce_ms = (view_spec_task.mode == Mode::List)
1560 .then(|| {
1561 view_spec_task
1562 .delivery
1563 .coalesce_ms
1564 .or(task_context.delivery.collection_coalesce_ms)
1565 .filter(|milliseconds| *milliseconds > 0)
1566 })
1567 .flatten();
1568 let mut flush_interval = coalesce_ms.map(|milliseconds| {
1569 let period = Duration::from_millis(milliseconds);
1570 let mut interval = tokio::time::interval_at(
1571 tokio::time::Instant::now() + period,
1572 period,
1573 );
1574 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
1575 interval
1576 });
1577 let mut pending = HashMap::<String, Arc<BusMessage>>::new();
1578 let mut pending_updates = 0_u64;
1579
1580 loop {
1581 tokio::select! {
1582 _ = cancel_token.cancelled() => break,
1583 _ = async {
1584 flush_interval
1585 .as_mut()
1586 .expect("coalescing interval is guarded")
1587 .tick()
1588 .await;
1589 }, if flush_interval.is_some() => {
1590 if pending.is_empty() {
1591 continue;
1592 }
1593 let sorted_caches = view_spec_task
1594 .is_derived()
1595 .then(|| task_context.view_index.sorted_caches());
1596 let next = load_query_entities(
1597 &task_context.entity_cache,
1598 sorted_caches,
1599 &view_spec_task,
1600 &query,
1601 false,
1602 ).await;
1603 if emit_coalesced_collection_delta(
1604 &task_context,
1605 &subscription_id,
1606 &view_spec_task,
1607 ¤t,
1608 &next,
1609 &pending,
1610 ).is_err() {
1611 task_context.metrics.delivery_stopped(&view_id, "send-failed");
1612 break;
1613 }
1614 task_context
1615 .metrics
1616 .collection_coalesced(&view_id, pending_updates);
1617 current = next;
1618 pending.clear();
1619 pending_updates = 0;
1620 }
1621 received = receiver.recv() => {
1622 let envelope = match received {
1623 Ok(envelope) => envelope,
1624 Err(broadcast::error::RecvError::Lagged(skipped)) => {
1625 task_context.metrics.subscription_lagged(&view_id, skipped);
1626 warn!(
1627 "Subscription {} lagged by {} updates",
1628 subscription_id, skipped
1629 );
1630 if view_spec_task.mode == Mode::Append {
1631 let _ = send_control_frame(
1632 &task_context,
1633 &SocketIssueMessage::append_subscription_lagged(
1634 subscription_id.clone(),
1635 skipped,
1636 ),
1637 &view_id,
1638 );
1639 task_context.metrics.delivery_stopped(&view_id, "append-lagged-without-replay");
1640 break;
1641 }
1642 if !task_subscription.snapshot.enabled {
1643 let _ = send_control_frame(
1644 &task_context,
1645 &SocketIssueMessage::subscription_lagged(
1646 subscription_id.clone(),
1647 skipped,
1648 ),
1649 &view_id,
1650 );
1651 task_context.metrics.delivery_stopped(&view_id, "lagged-without-snapshot");
1652 break;
1653 }
1654 info!(
1655 "Subscription {} is recovering from an authoritative snapshot",
1656 subscription_id
1657 );
1658 match recover_collection_subscription(
1659 &task_context,
1660 &task_subscription,
1661 &view_spec_task,
1662 &source_view_id,
1663 ).await {
1664 Ok((next_receiver, recovered)) => {
1665 receiver = next_receiver;
1666 current = recovered;
1667 pending.clear();
1668 pending_updates = 0;
1669 task_context.metrics.subscription_resnapshot(&view_id);
1670 continue;
1671 }
1672 Err(error) => {
1673 warn!(
1674 "Subscription {} failed to recover from lag: {error:#}",
1675 subscription_id
1676 );
1677 task_context.metrics.delivery_stopped(&view_id, "resnapshot-failed");
1678 break;
1679 }
1680 }
1681 }
1682 Err(broadcast::error::RecvError::Closed) => break,
1683 };
1684
1685 apply_collection_source_event(
1686 &task_context,
1687 &source_view_id,
1688 &view_spec_task,
1689 &query,
1690 &envelope,
1691 ).await;
1692
1693 if flush_interval.is_some() {
1694 pending_updates = pending_updates.saturating_add(1);
1695 pending.insert(envelope.key.clone(), envelope);
1696 continue;
1697 }
1698
1699 let metadata = source_frame_metadata(&envelope.payload);
1700 let sorted_caches = view_spec_task
1701 .is_derived()
1702 .then(|| task_context.view_index.sorted_caches());
1703 let next = load_query_entities(
1704 &task_context.entity_cache,
1705 sorted_caches,
1706 &view_spec_task,
1707 &query,
1708 false,
1709 ).await;
1710 if emit_collection_delta(
1711 &task_context,
1712 &subscription_id,
1713 &view_spec_task,
1714 ¤t,
1715 &next,
1716 &envelope,
1717 &metadata,
1718 ).is_err() {
1719 task_context.metrics.delivery_stopped(&view_id, "send-failed");
1720 break;
1721 }
1722 current = next;
1723 }
1724 }
1725 }
1726 }
1727 .instrument(info_span!("ws.subscribe.collection", client_id = %context.client_id, view = %span_view)),
1728 );
1729 Ok(())
1730}
1731
1732async fn subscribe_collection_then_snapshot(
1733 context: &SubscriptionContext,
1734 source_view_id: &str,
1735 view_spec: &ViewSpec,
1736 query: &SubscriptionQuery,
1737) -> (broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>) {
1738 let cache = context.entity_cache.clone();
1739 let sorted_caches = view_spec
1740 .is_derived()
1741 .then(|| context.view_index.sorted_caches());
1742 let view_spec = view_spec.clone();
1743 let query = query.clone();
1744 subscribe_list_then_snapshot(&context.bus_manager, source_view_id, move || async move {
1745 load_query_entities(&cache, sorted_caches, &view_spec, &query, false).await
1746 })
1747 .await
1748}
1749
1750async fn recover_collection_subscription(
1751 context: &SubscriptionContext,
1752 subscription: &Subscription,
1753 view_spec: &ViewSpec,
1754 source_view_id: &str,
1755) -> Result<(broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>)> {
1756 let (receiver, membership) =
1757 subscribe_collection_then_snapshot(context, source_view_id, view_spec, &subscription.query)
1758 .await;
1759 enforce_snapshot_limit(context, membership.len())?;
1764 send_snapshot_batches(
1765 context,
1766 subscription,
1767 &to_wire_snapshot_entities(membership.clone(), view_spec),
1768 view_spec.mode,
1769 SnapshotPurpose::Recovery,
1770 &context.entity_cache.snapshot_config(),
1771 )
1772 .await?;
1773 Ok((receiver, membership))
1774}
1775
1776async fn apply_collection_source_event(
1777 context: &SubscriptionContext,
1778 source_view_id: &str,
1779 view_spec: &ViewSpec,
1780 query: &SubscriptionQuery,
1781 envelope: &BusMessage,
1782) {
1783 let metadata = source_frame_metadata(&envelope.payload);
1784 if metadata.op != "delete" {
1785 return;
1786 }
1787 let current = context
1791 .entity_cache
1792 .get(source_view_id, &envelope.key)
1793 .await;
1794 if source_delete_is_stale(current.as_ref(), metadata.seq.as_deref()) {
1795 return;
1796 }
1797 context
1798 .entity_cache
1799 .remove(source_view_id, &envelope.key)
1800 .await;
1801 if view_spec.is_derived() {
1802 let caches = context.view_index.sorted_caches();
1803 let mut guard = caches.write().await;
1804 if let Some(cache) = guard.get_mut(&query.view) {
1805 cache.remove(&envelope.key);
1806 }
1807 }
1808}
1809
1810fn source_delete_is_stale(current: Option<&Value>, delete_seq: Option<&str>) -> bool {
1811 match (current, delete_seq) {
1812 (Some(current), Some(delete_seq)) => current
1813 .get("_seq")
1814 .and_then(Value::as_str)
1815 .is_some_and(|current_seq| {
1816 cmp_seq(current_seq, delete_seq) == std::cmp::Ordering::Greater
1817 }),
1818 _ => false,
1819 }
1820}
1821
1822#[derive(Debug)]
1832pub(crate) struct RejectedSubscription(pub SocketIssueMessage);
1833
1834impl RejectedSubscription {
1835 fn into_error(issue: SocketIssueMessage) -> anyhow::Error {
1836 anyhow::Error::new(Self(issue))
1837 }
1838}
1839
1840impl std::fmt::Display for RejectedSubscription {
1841 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1842 formatter.write_str(&self.0.message)
1843 }
1844}
1845
1846impl std::error::Error for RejectedSubscription {}
1847
1848async fn attach_journal_subscription(
1856 context: &SubscriptionContext,
1857 subscription: Subscription,
1858 view_spec: ViewSpec,
1859 journal: Arc<crate::journal::EventJournal>,
1860 cancel_token: CancellationToken,
1861) -> Result<()> {
1862 let view_id = view_spec.id.clone();
1863 let subscription_id = subscription.subscription_id.clone();
1864
1865 if subscription.query.take.is_some()
1868 || subscription.query.skip.is_some()
1869 || subscription.query.snapshot_limit.is_some()
1870 {
1871 return Err(RejectedSubscription::into_error(SocketIssueMessage::protocol(
1872 Some(subscription_id),
1873 "invalid-subscription",
1874 format!(
1875 "take, skip and snapshotLimit are window options and do not apply to the replayable view {view_id}"
1876 ),
1877 )));
1878 }
1879
1880 let cursor = match subscription.query.after.as_deref() {
1883 Some(raw) => match crate::journal::Cursor::parse(raw) {
1884 Some(cursor) => Some(cursor),
1885 None => {
1886 return Err(RejectedSubscription::into_error(
1887 SocketIssueMessage::protocol(
1888 Some(subscription_id),
1889 "invalid-cursor",
1890 format!(
1891 "`after` must be an {{epoch}}:{{offset}} replay cursor for view {view_id}, got {raw:?}"
1892 ),
1893 ),
1894 ));
1895 }
1896 },
1897 None => None,
1898 };
1899
1900 let mut receiver = context.bus_manager.get_or_create_list_bus(&view_id).await;
1903
1904 let replayed = match journal.replay_after(&view_id, cursor.as_ref()).await {
1905 Ok(records) => records,
1906 Err(error) => {
1907 return Err(RejectedSubscription::into_error(
1908 SocketIssueMessage::replay_refused(Some(subscription_id), &error),
1909 ));
1910 }
1911 };
1912
1913 let frame = SubscribedFrame::new(
1914 subscription.subscription_id.clone(),
1915 subscription.query.clone(),
1916 view_spec.mode,
1917 extract_sort_config(&view_spec),
1918 )
1919 .with_replay_window(journal.window(&view_id).await);
1920 send_control_frame(context, &frame, &view_id)?;
1921
1922 let task_context = context.clone();
1927 let task_subscription_id = subscription.subscription_id.clone();
1928 let task_query = subscription.query.clone();
1929 let task_epoch = journal.epoch().await;
1930 let span_view = view_id.clone();
1931 tokio::spawn(
1932 async move {
1933 let mut last_sent = cursor.map(|cursor| cursor.offset);
1934 let mut pending: VecDeque<Arc<BusMessage>> = VecDeque::new();
1938 let mut lagged: Option<u64> = None;
1939
1940 for record in replayed {
1941 if cancel_token.is_cancelled() {
1942 return;
1943 }
1944 drain_available(&mut receiver, &mut pending, &mut lagged);
1945 if journal_record_matches(&task_query, &record)
1946 && send_scoped_source_payload_async(
1947 &task_context,
1948 &task_subscription_id,
1949 &span_view,
1950 record.payload,
1951 )
1952 .await
1953 .is_err()
1954 {
1955 return;
1956 }
1957 last_sent = Some(record.offset);
1959 }
1960
1961 while let Some(envelope) = pending.pop_front() {
1964 if !forward_live_frame(
1965 &task_context,
1966 &task_subscription_id,
1967 &span_view,
1968 &task_query,
1969 &envelope,
1970 &mut last_sent,
1971 )
1972 .await
1973 {
1974 return;
1975 }
1976 }
1977
1978 if let Some(skipped) = lagged {
1979 report_replay_gap(
1980 &task_context,
1981 &task_subscription_id,
1982 &span_view,
1983 &task_epoch,
1984 skipped,
1985 last_sent,
1986 );
1987 return;
1988 }
1989
1990 loop {
1991 tokio::select! {
1992 _ = cancel_token.cancelled() => break,
1993 received = receiver.recv() => {
1994 let envelope = match received {
1995 Ok(envelope) => envelope,
1996 Err(broadcast::error::RecvError::Lagged(skipped)) => {
2014 report_replay_gap(
2015 &task_context,
2016 &task_subscription_id,
2017 &span_view,
2018 &task_epoch,
2019 skipped,
2020 last_sent,
2021 );
2022 break;
2023 }
2024 Err(broadcast::error::RecvError::Closed) => break,
2025 };
2026
2027 if !forward_live_frame(
2028 &task_context,
2029 &task_subscription_id,
2030 &span_view,
2031 &task_query,
2032 &envelope,
2033 &mut last_sent,
2034 )
2035 .await
2036 {
2037 break;
2038 }
2039 }
2040 }
2041 }
2042 }
2043 .instrument(info_span!(
2044 "ws.subscribe.replay",
2045 client_id = %context.client_id,
2046 view = %view_id
2047 )),
2048 );
2049 Ok(())
2050}
2051
2052fn drain_available(
2055 receiver: &mut broadcast::Receiver<Arc<BusMessage>>,
2056 pending: &mut VecDeque<Arc<BusMessage>>,
2057 lagged: &mut Option<u64>,
2058) {
2059 if lagged.is_some() {
2064 return;
2065 }
2066 const MAX_PENDING: usize = 8_192;
2072 loop {
2073 if pending.len() >= MAX_PENDING {
2074 return;
2075 }
2076 match receiver.try_recv() {
2077 Ok(envelope) => pending.push_back(envelope),
2078 Err(broadcast::error::TryRecvError::Empty)
2079 | Err(broadcast::error::TryRecvError::Closed) => return,
2080 Err(broadcast::error::TryRecvError::Lagged(skipped)) => {
2081 *lagged = Some(skipped);
2082 return;
2083 }
2084 }
2085 }
2086}
2087
2088fn already_delivered(offset: Option<u64>, last_sent: &mut Option<u64>) -> bool {
2094 let Some(offset) = offset else {
2095 return false;
2098 };
2099 if last_sent.is_some_and(|last| offset <= last) {
2100 return true;
2101 }
2102 *last_sent = Some(offset);
2103 false
2104}
2105
2106async fn forward_live_frame(
2110 context: &SubscriptionContext,
2111 subscription_id: &str,
2112 view_id: &str,
2113 query: &SubscriptionQuery,
2114 envelope: &Arc<BusMessage>,
2115 last_sent: &mut Option<u64>,
2116) -> bool {
2117 let metadata = source_frame_metadata(&envelope.payload);
2118 if already_delivered(metadata.offset, last_sent) {
2119 return true;
2120 }
2121 if !live_frame_matches(query, &envelope.key, &envelope.payload) {
2122 return true;
2123 }
2124 send_scoped_source_payload(context, subscription_id, view_id, envelope.payload.clone()).is_ok()
2125}
2126
2127fn report_replay_gap(
2128 context: &SubscriptionContext,
2129 subscription_id: &str,
2130 view_id: &str,
2131 epoch: &crate::journal::JournalEpoch,
2132 skipped: u64,
2133 last_sent: Option<u64>,
2134) {
2135 warn!(
2136 "Replay subscription {} lagged past {} records; stopping with a recovery cursor",
2137 subscription_id, skipped
2138 );
2139 let recover_from = last_sent.map(|offset| crate::journal::Cursor {
2140 epoch: epoch.clone(),
2141 offset,
2142 });
2143 let _ = send_control_frame(
2144 context,
2145 &SocketIssueMessage::replay_lagged(
2146 Some(subscription_id.to_string()),
2147 skipped,
2148 recover_from,
2149 ),
2150 view_id,
2151 );
2152}
2153
2154fn journal_record_matches(
2157 query: &SubscriptionQuery,
2158 record: &crate::journal::JournalRecord,
2159) -> bool {
2160 live_frame_matches(query, &record.key, &record.payload)
2161}
2162
2163fn live_frame_matches(query: &SubscriptionQuery, key: &str, payload: &[u8]) -> bool {
2164 if !query.matches_key(key) {
2165 return false;
2166 }
2167 if query.partition.is_none() && query.filters.is_empty() {
2168 return true;
2169 }
2170 let Ok(frame) = serde_json::from_slice::<Value>(payload) else {
2171 return false;
2172 };
2173 let Some(data) = frame.get("data") else {
2174 return false;
2175 };
2176 if let Some(partition) = &query.partition {
2177 if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2178 return false;
2179 }
2180 }
2181 query
2182 .filters
2183 .iter()
2184 .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2185}
2186
2187async fn send_scoped_source_payload_async(
2190 context: &SubscriptionContext,
2191 subscription_id: &str,
2192 view_id: &str,
2193 payload: Arc<Bytes>,
2194) -> Result<()> {
2195 let mut value: Value = serde_json::from_slice(&payload)?;
2196 let object = value
2197 .as_object_mut()
2198 .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2199 object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2200 object.insert(
2201 "subscriptionId".to_string(),
2202 Value::String(subscription_id.to_string()),
2203 );
2204 let json = serde_json::to_vec(&value)?;
2205 let compressed = maybe_compress(&json);
2206 let bytes = compressed.as_bytes().len();
2207 context
2208 .client_manager
2209 .send_compressed_async(context.client_id, compressed)
2210 .await
2211 .map_err(|error| anyhow::anyhow!("failed to send replayed frame: {error}"))?;
2212 context.metrics.message_sent();
2213 emit_update_sent_for_client(
2214 &context.usage_emitter,
2215 &context.client_manager,
2216 context.client_id,
2217 view_id,
2218 bytes,
2219 );
2220 Ok(())
2221}
2222
2223#[derive(Default)]
2224struct SourceFrameMetadata {
2225 op: String,
2226 seq: Option<String>,
2227 offset: Option<u64>,
2228}
2229
2230fn source_frame_metadata(payload: &[u8]) -> SourceFrameMetadata {
2231 serde_json::from_slice::<Value>(payload)
2232 .ok()
2233 .map(|value| SourceFrameMetadata {
2234 op: value
2235 .get("op")
2236 .and_then(Value::as_str)
2237 .unwrap_or_default()
2238 .to_string(),
2239 seq: value.get("seq").and_then(Value::as_str).map(str::to_string),
2240 offset: value.get("offset").and_then(Value::as_u64),
2241 })
2242 .unwrap_or_default()
2243}
2244
2245fn send_scoped_source_payload(
2246 context: &SubscriptionContext,
2247 subscription_id: &str,
2248 view_id: &str,
2249 payload: Arc<Bytes>,
2250) -> Result<()> {
2251 let mut value: Value = serde_json::from_slice(&payload)?;
2252 let object = value
2253 .as_object_mut()
2254 .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2255 object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2256 object.insert(
2257 "subscriptionId".to_string(),
2258 Value::String(subscription_id.to_string()),
2259 );
2260 let encoded = Arc::new(Bytes::from(serde_json::to_vec(&value)?));
2261 let bytes = encoded.len();
2262 context
2263 .client_manager
2264 .send_to_client(context.client_id, encoded)
2265 .map_err(|error| anyhow::anyhow!("failed to send live frame: {error}"))?;
2266 context.metrics.message_sent();
2267 emit_update_sent_for_client(
2268 &context.usage_emitter,
2269 &context.client_manager,
2270 context.client_id,
2271 view_id,
2272 bytes,
2273 );
2274 Ok(())
2275}
2276
2277fn send_membership_frame(
2278 context: &SubscriptionContext,
2279 subscription_id: &str,
2280 view_spec: &ViewSpec,
2281 op: &str,
2282 key: &str,
2283 mut data: Value,
2284 seq: Option<String>,
2285) -> Result<()> {
2286 apply_wire_format(&mut data, &view_spec.wire_format);
2287 let frame = Frame::scoped(
2288 subscription_id,
2289 view_spec.mode,
2290 &view_spec.id,
2291 op,
2292 key,
2293 data,
2294 seq,
2295 );
2296 let encoded = Arc::new(Bytes::from(serde_json::to_vec(&frame)?));
2297 let bytes = encoded.len();
2298 context
2299 .client_manager
2300 .send_to_client(context.client_id, encoded)
2301 .map_err(|error| anyhow::anyhow!("failed to send membership frame: {error}"))?;
2302 context.metrics.message_sent();
2303 emit_update_sent_for_client(
2304 &context.usage_emitter,
2305 &context.client_manager,
2306 context.client_id,
2307 &view_spec.id,
2308 bytes,
2309 );
2310 Ok(())
2311}
2312
2313fn emit_collection_delta(
2314 context: &SubscriptionContext,
2315 subscription_id: &str,
2316 view_spec: &ViewSpec,
2317 current: &[(String, Value)],
2318 next: &[(String, Value)],
2319 envelope: &BusMessage,
2320 metadata: &SourceFrameMetadata,
2321) -> Result<()> {
2322 let current_keys: Vec<&str> = current.iter().map(|(key, _)| key.as_str()).collect();
2323 let next_keys: Vec<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2324 let next_set: HashSet<&str> = next_keys.iter().copied().collect();
2325
2326 for key in current_keys
2327 .iter()
2328 .copied()
2329 .filter(|key| !next_set.contains(key))
2330 {
2331 let op = if metadata.op == "delete" && key == envelope.key {
2332 "delete"
2333 } else {
2334 "remove"
2335 };
2336 send_membership_frame(
2337 context,
2338 subscription_id,
2339 view_spec,
2340 op,
2341 key,
2342 Value::Null,
2343 metadata.seq.clone(),
2344 )?;
2345 }
2346
2347 for (key, data) in next.iter() {
2348 let was_member = current_keys.iter().any(|candidate| *candidate == key);
2349 match member_action(
2350 was_member,
2351 key == &envelope.key,
2352 view_spec.is_derived(),
2353 &metadata.op,
2354 ) {
2355 MemberAction::Skip => {}
2356 MemberAction::ForwardPatch => send_scoped_source_payload(
2357 context,
2358 subscription_id,
2359 &view_spec.id,
2360 envelope.payload.clone(),
2361 )?,
2362 MemberAction::Upsert => {
2363 let seq = metadata
2364 .seq
2365 .clone()
2366 .or_else(|| data.get("_seq").and_then(Value::as_str).map(str::to_string));
2367 send_membership_frame(
2368 context,
2369 subscription_id,
2370 view_spec,
2371 "upsert",
2372 key,
2373 data.clone(),
2374 seq,
2375 )?;
2376 }
2377 }
2378 }
2379 Ok(())
2380}
2381
2382fn emit_coalesced_collection_delta(
2387 context: &SubscriptionContext,
2388 subscription_id: &str,
2389 view_spec: &ViewSpec,
2390 current: &[(String, Value)],
2391 next: &[(String, Value)],
2392 pending: &HashMap<String, Arc<BusMessage>>,
2393) -> Result<()> {
2394 for change in plan_coalesced_collection_delta(current, next, pending) {
2395 send_membership_frame(
2396 context,
2397 subscription_id,
2398 view_spec,
2399 change.op,
2400 &change.key,
2401 change.data,
2402 change.seq,
2403 )?;
2404 }
2405 Ok(())
2406}
2407
2408#[derive(Debug, PartialEq)]
2409struct CollectionChange {
2410 op: &'static str,
2411 key: String,
2412 data: Value,
2413 seq: Option<String>,
2414}
2415
2416fn plan_coalesced_collection_delta(
2417 current: &[(String, Value)],
2418 next: &[(String, Value)],
2419 pending: &HashMap<String, Arc<BusMessage>>,
2420) -> Vec<CollectionChange> {
2421 let current_by_key: HashMap<&str, &Value> = current
2422 .iter()
2423 .map(|(key, data)| (key.as_str(), data))
2424 .collect();
2425 let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2426 let latest_seq = pending
2427 .values()
2428 .filter_map(|envelope| source_frame_metadata(&envelope.payload).seq)
2429 .max_by(|left, right| cmp_seq(left, right));
2430 let mut changes = Vec::new();
2431
2432 for (key, _) in current
2433 .iter()
2434 .filter(|(key, _)| !next_keys.contains(key.as_str()))
2435 {
2436 let metadata = pending
2437 .get(key)
2438 .map(|envelope| source_frame_metadata(&envelope.payload));
2439 let op = if metadata
2440 .as_ref()
2441 .is_some_and(|metadata| metadata.op == "delete")
2442 {
2443 "delete"
2444 } else {
2445 "remove"
2446 };
2447 let seq = metadata
2448 .and_then(|metadata| metadata.seq)
2449 .or_else(|| latest_seq.clone());
2450 changes.push(CollectionChange {
2451 op,
2452 key: key.clone(),
2453 data: Value::Null,
2454 seq,
2455 });
2456 }
2457
2458 for (key, data) in next {
2459 let changed = current_by_key
2460 .get(key.as_str())
2461 .is_none_or(|previous| *previous != data);
2462 if !changed && !pending.contains_key(key) {
2463 continue;
2464 }
2465 let seq = data
2466 .get("_seq")
2467 .and_then(Value::as_str)
2468 .map(str::to_string)
2469 .or_else(|| {
2470 pending
2471 .get(key)
2472 .and_then(|envelope| source_frame_metadata(&envelope.payload).seq)
2473 })
2474 .or_else(|| latest_seq.clone());
2475 changes.push(CollectionChange {
2476 op: "upsert",
2477 key: key.clone(),
2478 data: data.clone(),
2479 seq,
2480 });
2481 }
2482 changes
2483}
2484
2485#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2487enum MemberAction {
2488 Skip,
2491 ForwardPatch,
2493 Upsert,
2495}
2496
2497fn member_action(
2510 was_member: bool,
2511 is_mutated_key: bool,
2512 is_derived: bool,
2513 op: &str,
2514) -> MemberAction {
2515 if !is_mutated_key {
2516 return if was_member {
2519 MemberAction::Skip
2520 } else {
2521 MemberAction::Upsert
2522 };
2523 }
2524 if was_member && !is_derived && op != "delete" {
2529 MemberAction::ForwardPatch
2530 } else {
2531 MemberAction::Upsert
2532 }
2533}
2534
2535fn to_wire_snapshot_entities(
2536 entities: Vec<(String, Value)>,
2537 view_spec: &ViewSpec,
2538) -> Vec<SnapshotEntity> {
2539 entities
2540 .into_iter()
2541 .map(|(key, mut data)| {
2542 apply_wire_format(&mut data, &view_spec.wire_format);
2543 SnapshotEntity { key, data }
2544 })
2545 .collect()
2546}
2547
2548async fn load_query_entities(
2549 entity_cache: &EntityCache,
2550 sorted_caches: Option<
2551 Arc<tokio::sync::RwLock<HashMap<String, crate::sorted_cache::SortedViewCache>>>,
2552 >,
2553 view_spec: &ViewSpec,
2554 query: &SubscriptionQuery,
2555 apply_snapshot_limit: bool,
2556) -> Vec<(String, Value)> {
2557 let (entities, preordered) = if let Some(sorted_caches) = sorted_caches {
2558 let mut caches = sorted_caches.write().await;
2559 let entities = caches
2560 .get_mut(&view_spec.id)
2561 .map(|cache| cache.get_all_ordered())
2562 .unwrap_or_default();
2563 (entities, true)
2564 } else if view_spec.mode == Mode::State {
2565 let entity = match query.key.as_deref() {
2566 Some(key) => entity_cache
2567 .get(&view_spec.id, key)
2568 .await
2569 .map(|data| vec![(key.to_string(), data)])
2570 .unwrap_or_default(),
2571 None => vec![],
2572 };
2573 (entity, true)
2574 } else {
2575 (entity_cache.get_all(&view_spec.id).await, false)
2576 };
2577 select_query_entities(entities, query, preordered, apply_snapshot_limit)
2578}
2579
2580fn select_query_entities(
2581 mut entities: Vec<(String, Value)>,
2582 query: &SubscriptionQuery,
2583 preordered: bool,
2584 apply_snapshot_limit: bool,
2585) -> Vec<(String, Value)> {
2586 entities.retain(|(key, data)| query_matches_entity(query, key, data));
2587 if !preordered {
2588 entities.sort_by(|left, right| {
2589 let left_seq = left.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2590 let right_seq = right.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2591 let order = if query.after.is_some() {
2592 cmp_seq(left_seq, right_seq)
2593 } else {
2594 cmp_seq(right_seq, left_seq)
2595 };
2596 order.then_with(|| left.0.cmp(&right.0))
2597 });
2598 }
2599
2600 let skip = query.skip.unwrap_or(0);
2601 let take = query.take.unwrap_or(usize::MAX);
2602 let mut selected: Vec<_> = entities.into_iter().skip(skip).take(take).collect();
2603 if apply_snapshot_limit {
2604 if let Some(limit) = query.snapshot_limit {
2605 selected.truncate(limit);
2606 }
2607 }
2608 selected
2609}
2610
2611fn query_matches_entity(query: &SubscriptionQuery, key: &str, data: &Value) -> bool {
2612 if !query.matches_key(key) {
2613 return false;
2614 }
2615 if let Some(partition) = &query.partition {
2616 if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2617 return false;
2618 }
2619 }
2620 if let Some(after) = &query.after {
2621 let Some(seq) = data.get("_seq").and_then(Value::as_str) else {
2622 return false;
2623 };
2624 if cmp_seq(seq, after) != std::cmp::Ordering::Greater {
2625 return false;
2626 }
2627 }
2628 query
2629 .filters
2630 .iter()
2631 .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2632}
2633
2634fn value_at_dot_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
2635 path.split('.')
2636 .try_fold(value, |current, segment| current.get(segment))
2637}
2638
2639#[cfg(test)]
2640mod tests {
2641 use super::*;
2642 use crate::cache::EntityCacheConfig;
2643 use crate::view::{Delivery, Filters, Projection};
2644 use serde_json::json;
2645 use tokio::sync::oneshot;
2646
2647 fn list_spec() -> ViewSpec {
2648 ViewSpec {
2649 id: "Thing/list".to_string(),
2650 export: "Thing".to_string(),
2651 mode: Mode::List,
2652 wire_format: Default::default(),
2653 projection: Projection::all(),
2654 filters: Filters::all(),
2655 delivery: Delivery::default(),
2656 pipeline: None,
2657 source_view: None,
2658 }
2659 }
2660
2661 fn plan_window(
2664 current_keys: &[&str],
2665 next_keys: &[&str],
2666 envelope_key: &str,
2667 is_derived: bool,
2668 op: &str,
2669 ) -> Vec<(String, MemberAction)> {
2670 next_keys
2671 .iter()
2672 .map(|key| {
2673 let was_member = current_keys.contains(key);
2674 (
2675 (*key).to_string(),
2676 member_action(was_member, *key == envelope_key, is_derived, op),
2677 )
2678 })
2679 .collect()
2680 }
2681
2682 #[test]
2683 fn a_reordered_window_sends_only_the_mutated_entity() {
2684 let current = ["4", "3", "2", "1"];
2688 let next = ["1", "4", "3", "2"];
2689 let plan = plan_window(¤t, &next, "1", false, "patch");
2690
2691 assert_eq!(
2692 plan,
2693 vec![
2694 ("1".to_string(), MemberAction::ForwardPatch),
2695 ("4".to_string(), MemberAction::Skip),
2696 ("3".to_string(), MemberAction::Skip),
2697 ("2".to_string(), MemberAction::Skip),
2698 ]
2699 );
2700 }
2701
2702 #[test]
2703 fn a_key_entering_the_window_gets_the_whole_entity() {
2704 let plan = plan_window(&["4", "3"], &["5", "4", "3"], "5", false, "patch");
2706 assert_eq!(
2707 plan,
2708 vec![
2709 ("5".to_string(), MemberAction::Upsert),
2710 ("4".to_string(), MemberAction::Skip),
2711 ("3".to_string(), MemberAction::Skip),
2712 ]
2713 );
2714 }
2715
2716 #[test]
2717 fn derived_views_still_send_whole_entities() {
2718 let plan = plan_window(&["1", "2"], &["1", "2"], "1", true, "patch");
2721 assert_eq!(
2722 plan,
2723 vec![
2724 ("1".to_string(), MemberAction::Upsert),
2725 ("2".to_string(), MemberAction::Skip),
2726 ]
2727 );
2728 }
2729
2730 #[test]
2731 fn a_delete_envelope_never_forwards_a_patch() {
2732 let plan = plan_window(&["1", "2"], &["1", "2"], "1", false, "delete");
2734 assert_eq!(
2735 plan,
2736 vec![
2737 ("1".to_string(), MemberAction::Upsert),
2738 ("2".to_string(), MemberAction::Skip),
2739 ]
2740 );
2741 }
2742
2743 #[test]
2744 fn an_unchanged_window_sends_one_frame_not_a_broadcast() {
2745 let keys: Vec<String> = (0..500).map(|index| index.to_string()).collect();
2748 let refs: Vec<&str> = keys.iter().map(String::as_str).collect();
2749 let plan = plan_window(&refs, &refs, "250", false, "patch");
2750
2751 let sent = plan
2752 .iter()
2753 .filter(|(_, action)| *action != MemberAction::Skip)
2754 .count();
2755 assert_eq!(sent, 1);
2756 assert_eq!(plan[250].1, MemberAction::ForwardPatch);
2757 }
2758
2759 fn list_message(key: &str, op: &str, seq: &str) -> Arc<BusMessage> {
2760 Arc::new(BusMessage {
2761 key: key.to_string(),
2762 entity: "Thing/list".to_string(),
2763 payload: Arc::new(Bytes::from(
2764 serde_json::to_vec(&json!({
2765 "entity": "Thing/list",
2766 "op": op,
2767 "key": key,
2768 "seq": seq,
2769 "data": {},
2770 }))
2771 .unwrap(),
2772 )),
2773 })
2774 }
2775
2776 #[test]
2777 fn coalescing_emits_only_the_final_full_state_per_changed_key() {
2778 let current = vec![
2779 ("a".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2780 ("b".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2781 (
2782 "stable".to_string(),
2783 json!({"count": 1, "_seq": "10:000001"}),
2784 ),
2785 ];
2786 let next = vec![
2787 ("a".to_string(), json!({"count": 3, "_seq": "10:000004"})),
2788 ("c".to_string(), json!({"count": 1, "_seq": "10:000003"})),
2789 (
2790 "stable".to_string(),
2791 json!({"count": 1, "_seq": "10:000001"}),
2792 ),
2793 ];
2794 let pending = HashMap::from([
2795 ("a".to_string(), list_message("a", "patch", "10:000004")),
2796 ("b".to_string(), list_message("b", "delete", "10:000002")),
2797 ("c".to_string(), list_message("c", "patch", "10:000003")),
2798 ]);
2799
2800 assert_eq!(
2801 plan_coalesced_collection_delta(¤t, &next, &pending),
2802 vec![
2803 CollectionChange {
2804 op: "delete",
2805 key: "b".to_string(),
2806 data: Value::Null,
2807 seq: Some("10:000002".to_string()),
2808 },
2809 CollectionChange {
2810 op: "upsert",
2811 key: "a".to_string(),
2812 data: json!({"count": 3, "_seq": "10:000004"}),
2813 seq: Some("10:000004".to_string()),
2814 },
2815 CollectionChange {
2816 op: "upsert",
2817 key: "c".to_string(),
2818 data: json!({"count": 1, "_seq": "10:000003"}),
2819 seq: Some("10:000003".to_string()),
2820 },
2821 ]
2822 );
2823 }
2824
2825 #[test]
2826 fn a_delayed_delete_cannot_erase_a_newer_recreated_entity() {
2827 let recreated = json!({"balance": 2, "_seq": "10:000004"});
2828 assert!(source_delete_is_stale(Some(&recreated), Some("10:000002")));
2829 assert!(!source_delete_is_stale(Some(&recreated), Some("10:000004")));
2830 assert!(!source_delete_is_stale(Some(&recreated), None));
2831 }
2832
2833 #[test]
2834 fn a_lagged_snapshotless_subscription_gets_a_fatal_retryable_error() {
2835 let issue = SocketIssueMessage::subscription_lagged("balances".to_string(), 42);
2836 assert_eq!(issue.code, "subscription-lagged");
2837 assert!(issue.retryable);
2838 assert!(issue.fatal);
2839 assert!(issue.message.contains("42"));
2840 assert!(issue
2841 .suggested_action
2842 .unwrap()
2843 .contains("snapshots enabled"));
2844
2845 let append = SocketIssueMessage::append_subscription_lagged("trades".to_string(), 9);
2846 assert!(append.fatal);
2847 assert!(append.suggested_action.unwrap().contains("retained replay"));
2848 }
2849
2850 #[test]
2851 fn snapshot_batches_share_identity_and_completion() {
2852 let entities = ["one", "two", "three"].map(|key| SnapshotEntity {
2853 key: key.to_string(),
2854 data: json!({"key": key}),
2855 });
2856 let batches = create_snapshot_batches(
2857 &entities,
2858 SnapshotMetadata {
2859 subscription_id: "sub-1",
2860 snapshot_id: "snapshot-1",
2861 authoritative: true,
2862 mode: Mode::List,
2863 view_id: "Thing/list",
2864 key: None,
2865 },
2866 &SnapshotBatchConfig {
2867 initial_batch_size: 2,
2868 subsequent_batch_size: 1,
2869 },
2870 );
2871 assert_eq!(batches.len(), 2);
2872 assert!(batches.iter().all(|batch| batch.subscription_id == "sub-1"));
2873 assert!(batches
2874 .iter()
2875 .all(|batch| batch.snapshot_id == "snapshot-1"));
2876 assert!(!batches[0].complete);
2877 assert!(batches[1].complete);
2878 assert!(batches.iter().all(|batch| batch.authoritative));
2879 }
2880
2881 #[test]
2882 fn empty_incremental_snapshot_is_explicitly_non_authoritative() {
2883 let batches = create_snapshot_batches(
2884 &[],
2885 SnapshotMetadata {
2886 subscription_id: "sub-1",
2887 snapshot_id: "snapshot-1",
2888 authoritative: false,
2889 mode: Mode::State,
2890 view_id: "Thing/state",
2891 key: Some("missing"),
2892 },
2893 &SnapshotBatchConfig {
2894 initial_batch_size: 1,
2895 subsequent_batch_size: 1,
2896 },
2897 );
2898 assert_eq!(batches.len(), 1);
2899 assert!(!batches[0].authoritative);
2900 assert!(batches[0].complete);
2901 assert_eq!(batches[0].key.as_deref(), Some("missing"));
2902 }
2903
2904 #[test]
2905 fn lag_recovery_is_authoritative_for_an_after_query() {
2906 let subscription = Subscription {
2907 protocol_version: PROTOCOL_VERSION,
2908 subscription_id: "sub-1".to_string(),
2909 query: SubscriptionQuery {
2910 view: "Thing/list".to_string(),
2911 after: Some("40:000000000010".to_string()),
2912 ..Default::default()
2913 },
2914 snapshot: Default::default(),
2915 };
2916
2917 assert!(!SnapshotPurpose::Initial.authoritative(&subscription));
2918 assert!(SnapshotPurpose::Recovery.authoritative(&subscription));
2919 }
2920
2921 #[test]
2922 fn dot_path_filters_are_exact_and_type_sensitive() {
2923 let mut query = SubscriptionQuery {
2924 view: "Thing/list".to_string(),
2925 ..Default::default()
2926 };
2927 query
2928 .filters
2929 .insert("state.status".to_string(), json!("open"));
2930 query.filters.insert("metrics.count".to_string(), json!(2));
2931 assert!(query_matches_entity(
2932 &query,
2933 "one",
2934 &json!({"state": {"status": "open"}, "metrics": {"count": 2}}),
2935 ));
2936 assert!(!query_matches_entity(
2937 &query,
2938 "one",
2939 &json!({"state": {"status": "open"}, "metrics": {"count": "2"}}),
2940 ));
2941 }
2942
2943 #[test]
2944 fn take_and_skip_define_independent_deterministic_windows() {
2945 let entities: Vec<_> = (1..=6)
2946 .map(|id| {
2947 (
2948 id.to_string(),
2949 json!({"id": id, "_seq": format!("10:{id:012}")}),
2950 )
2951 })
2952 .collect();
2953 let first = SubscriptionQuery {
2954 view: "Thing/list".to_string(),
2955 take: Some(2),
2956 skip: Some(0),
2957 ..Default::default()
2958 };
2959 let second = SubscriptionQuery {
2960 skip: Some(2),
2961 ..first.clone()
2962 };
2963 let first_keys: Vec<_> = select_query_entities(entities.clone(), &first, false, false)
2964 .into_iter()
2965 .map(|(key, _)| key)
2966 .collect();
2967 let second_keys: Vec<_> = select_query_entities(entities, &second, false, false)
2968 .into_iter()
2969 .map(|(key, _)| key)
2970 .collect();
2971 assert_eq!(first_keys, ["6", "5"]);
2972 assert_eq!(second_keys, ["4", "3"]);
2973 }
2974
2975 #[tokio::test]
2976 async fn state_receiver_is_installed_before_snapshot_awaits() {
2977 let bus = BusManager::new();
2978 let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
2979 let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
2980 let bus_for_task = bus.clone();
2981 let task = tokio::spawn(async move {
2982 subscribe_state_then_snapshot(&bus_for_task, "Thing/state", "one", || async move {
2983 snapshot_started_tx.send(()).unwrap();
2984 release_snapshot_rx.await.unwrap();
2985 })
2986 .await
2987 .0
2988 });
2989 snapshot_started_rx.await.unwrap();
2990 bus.publish_state(
2991 "Thing/state",
2992 "one",
2993 Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
2994 )
2995 .await;
2996 release_snapshot_tx.send(()).unwrap();
2997 let mut receiver = task.await.unwrap();
2998 receiver.changed().await.unwrap();
2999 assert!(!receiver.borrow().payload.is_empty());
3000 }
3001
3002 async fn assert_list_receiver_precedes_snapshot(view: &'static str) {
3003 let bus = BusManager::new();
3004 let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3005 let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3006 let bus_for_task = bus.clone();
3007 let task = tokio::spawn(async move {
3008 subscribe_list_then_snapshot(&bus_for_task, view, || async move {
3009 snapshot_started_tx.send(()).unwrap();
3010 release_snapshot_rx.await.unwrap();
3011 })
3012 .await
3013 .0
3014 });
3015 snapshot_started_rx.await.unwrap();
3016 bus.publish_list(
3017 view,
3018 Arc::new(BusMessage {
3019 key: "one".to_string(),
3020 entity: view.to_string(),
3021 payload: Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3022 }),
3023 )
3024 .await;
3025 release_snapshot_tx.send(()).unwrap();
3026 let mut receiver = task.await.unwrap();
3027 assert_eq!(receiver.recv().await.unwrap().key, "one");
3028 }
3029
3030 #[tokio::test]
3031 async fn list_receiver_is_installed_before_snapshot() {
3032 assert_list_receiver_precedes_snapshot("Thing/list").await;
3033 }
3034
3035 #[tokio::test]
3036 async fn append_receiver_is_installed_before_snapshot() {
3037 assert_list_receiver_precedes_snapshot("Thing/append").await;
3038 }
3039
3040 #[tokio::test]
3041 async fn derived_source_receiver_is_installed_before_snapshot() {
3042 assert_list_receiver_precedes_snapshot("Thing/list-source").await;
3043 }
3044
3045 #[tokio::test]
3046 async fn snapshot_limit_does_not_change_live_take_skip_membership() {
3047 let cache = EntityCache::with_config(EntityCacheConfig {
3048 max_entities_per_view: 10,
3049 ..Default::default()
3050 });
3051 for id in 1..=4 {
3052 cache
3053 .upsert(
3054 "Thing/list",
3055 &id.to_string(),
3056 json!({"_seq": format!("10:{id:012}")}),
3057 )
3058 .await;
3059 }
3060 let query = SubscriptionQuery {
3061 view: "Thing/list".to_string(),
3062 take: Some(3),
3063 skip: Some(1),
3064 snapshot_limit: Some(1),
3065 ..Default::default()
3066 };
3067 let live = load_query_entities(&cache, None, &list_spec(), &query, false).await;
3068 let snapshot = load_query_entities(&cache, None, &list_spec(), &query, true).await;
3069 assert_eq!(live.len(), 3);
3070 assert_eq!(snapshot.len(), 1);
3071 }
3072
3073 #[test]
3074 fn fixture_manifest_covers_required_conformance_cases() {
3075 let manifest: Value = serde_json::from_str(include_str!(
3076 "../../../../tests/fixtures/websocket-v2/manifest.json"
3077 ))
3078 .unwrap();
3079 let names: HashSet<_> = manifest["fixtures"]
3080 .as_array()
3081 .unwrap()
3082 .iter()
3083 .filter_map(Value::as_str)
3084 .collect();
3085 for required in [
3086 "keyed-state.json",
3087 "list-windows.json",
3088 "filters.json",
3089 "multi-batch-authoritative.json",
3090 "empty-snapshot.json",
3091 "remove.json",
3092 "delete.json",
3093 "incremental-snapshot.json",
3094 "reconnect-replacement.json",
3095 "errors.json",
3096 ] {
3097 assert!(names.contains(required), "missing fixture {required}");
3098 }
3099
3100 for document in [
3101 include_str!("../../../../tests/fixtures/websocket-v2/keyed-state.json"),
3102 include_str!("../../../../tests/fixtures/websocket-v2/list-windows.json"),
3103 include_str!("../../../../tests/fixtures/websocket-v2/filters.json"),
3104 include_str!("../../../../tests/fixtures/websocket-v2/multi-batch-authoritative.json"),
3105 include_str!("../../../../tests/fixtures/websocket-v2/empty-snapshot.json"),
3106 include_str!("../../../../tests/fixtures/websocket-v2/remove.json"),
3107 include_str!("../../../../tests/fixtures/websocket-v2/delete.json"),
3108 include_str!("../../../../tests/fixtures/websocket-v2/incremental-snapshot.json"),
3109 include_str!("../../../../tests/fixtures/websocket-v2/reconnect-replacement.json"),
3110 include_str!("../../../../tests/fixtures/websocket-v2/errors.json"),
3111 ] {
3112 let fixture: Value = serde_json::from_str(document).unwrap();
3113 assert!(fixture["name"].is_string());
3114 }
3115 }
3116
3117 fn append_frame(key: &str, data: Value) -> Arc<Bytes> {
3118 let frame = json!({
3119 "entity": "Trade/append",
3120 "op": "patch",
3121 "key": key,
3122 "offset": 7,
3123 "data": data,
3124 });
3125 Arc::new(Bytes::from(serde_json::to_vec(&frame).unwrap()))
3126 }
3127
3128 #[test]
3132 fn replay_delivery_applies_key_partition_and_filters() {
3133 let matching = append_frame("pool1", json!({"_partition": "us", "side": "buy"}));
3134 let other_partition = append_frame("pool1", json!({"_partition": "eu", "side": "buy"}));
3135 let other_side = append_frame("pool1", json!({"_partition": "us", "side": "sell"}));
3136
3137 let unfiltered = SubscriptionQuery {
3138 view: "Trade/append".to_string(),
3139 ..Default::default()
3140 };
3141 assert!(live_frame_matches(&unfiltered, "pool1", &matching));
3142 assert!(live_frame_matches(&unfiltered, "pool9", &matching));
3143
3144 let keyed = SubscriptionQuery {
3145 view: "Trade/append".to_string(),
3146 key: Some("pool1".to_string()),
3147 ..Default::default()
3148 };
3149 assert!(live_frame_matches(&keyed, "pool1", &matching));
3150 assert!(!live_frame_matches(&keyed, "pool2", &matching));
3151
3152 let partitioned = SubscriptionQuery {
3153 view: "Trade/append".to_string(),
3154 partition: Some("us".to_string()),
3155 ..Default::default()
3156 };
3157 assert!(live_frame_matches(&partitioned, "pool1", &matching));
3158 assert!(!live_frame_matches(&partitioned, "pool1", &other_partition));
3159
3160 let filtered = SubscriptionQuery {
3161 view: "Trade/append".to_string(),
3162 filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3163 ..Default::default()
3164 };
3165 assert!(live_frame_matches(&filtered, "pool1", &matching));
3166 assert!(!live_frame_matches(&filtered, "pool1", &other_side));
3167 }
3168
3169 #[test]
3170 fn a_frame_without_decodable_data_does_not_satisfy_a_filter() {
3171 let filtered = SubscriptionQuery {
3172 view: "Trade/append".to_string(),
3173 filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3174 ..Default::default()
3175 };
3176 let garbage = Arc::new(Bytes::from_static(b"not json"));
3177 assert!(!live_frame_matches(&filtered, "pool1", &garbage));
3178 }
3179
3180 #[test]
3184 fn replay_lagged_recovers_from_before_the_gap() {
3185 let epoch = crate::journal::JournalEpoch::new();
3186 let issue = SocketIssueMessage::replay_lagged(
3187 Some("trades".to_string()),
3188 37,
3189 Some(crate::journal::Cursor {
3190 epoch: epoch.clone(),
3191 offset: 4180,
3192 }),
3193 );
3194 assert_eq!(issue.code, "replay-lagged");
3195 assert_eq!(issue.recover_from, Some(format!("{epoch}:4180")));
3196 assert!(
3197 issue.suggested_action.unwrap().contains("4180"),
3198 "the consumer is told exactly which cursor recovers the gap"
3199 );
3200
3201 let from_scratch = SocketIssueMessage::replay_lagged(Some("trades".to_string()), 9, None);
3204 assert_eq!(from_scratch.recover_from, None);
3205 assert!(from_scratch
3206 .suggested_action
3207 .unwrap()
3208 .contains("without `after`"));
3209 }
3210
3211 fn bus_message(key: &str) -> Arc<BusMessage> {
3212 Arc::new(BusMessage {
3213 key: key.to_string(),
3214 entity: "Trade/append".to_string(),
3215 payload: Arc::new(Bytes::from_static(b"{}")),
3216 })
3217 }
3218
3219 #[tokio::test]
3224 async fn draining_during_a_replay_keeps_a_busy_view_from_lapping_it() {
3225 let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(16);
3226 let mut pending = VecDeque::new();
3227 let mut lagged = None;
3228
3229 for index in 0..48 {
3232 sender.send(bus_message(&format!("k{index}"))).unwrap();
3233 drain_available(&mut receiver, &mut pending, &mut lagged);
3234 }
3235
3236 assert_eq!(lagged, None, "draining as we go means nothing is dropped");
3237 assert_eq!(pending.len(), 48, "every published frame is buffered");
3238 }
3239
3240 #[tokio::test]
3241 async fn a_replay_that_never_drains_is_reported_as_a_gap() {
3242 let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3243 for index in 0..32 {
3244 sender.send(bus_message(&format!("k{index}"))).unwrap();
3245 }
3246
3247 let mut pending = VecDeque::new();
3248 let mut lagged = None;
3249 drain_available(&mut receiver, &mut pending, &mut lagged);
3250
3251 assert!(
3252 lagged.is_some(),
3253 "overflowing the bus is a gap, not silent truncation"
3254 );
3255 }
3256
3257 #[tokio::test]
3261 async fn nothing_after_a_gap_is_buffered_ahead_of_the_report() {
3262 let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3263 let mut pending = VecDeque::new();
3264 let mut lagged = None;
3265
3266 sender.send(bus_message("before")).unwrap();
3268 drain_available(&mut receiver, &mut pending, &mut lagged);
3269 assert_eq!(pending.len(), 1);
3270
3271 for index in 0..32 {
3273 sender.send(bus_message(&format!("lost{index}"))).unwrap();
3274 }
3275 drain_available(&mut receiver, &mut pending, &mut lagged);
3276 assert!(lagged.is_some());
3277
3278 let buffered_at_gap = pending.len();
3281 drain_available(&mut receiver, &mut pending, &mut lagged);
3282 assert_eq!(
3283 pending.len(),
3284 buffered_at_gap,
3285 "post-gap frames must not join the pre-gap flush"
3286 );
3287 assert_eq!(pending.front().unwrap().key, "before");
3288 }
3289 #[test]
3292 fn the_seam_between_replay_and_live_neither_repeats_nor_skips() {
3293 let mut last_sent = Some(4211);
3294
3295 assert!(
3296 already_delivered(Some(4211), &mut last_sent),
3297 "the record the replay ended on must not be sent twice"
3298 );
3299 assert!(already_delivered(Some(4100), &mut last_sent));
3300 assert_eq!(last_sent, Some(4211), "a duplicate never moves the mark");
3301
3302 assert!(
3303 !already_delivered(Some(4212), &mut last_sent),
3304 "the next record is new"
3305 );
3306 assert_eq!(last_sent, Some(4212));
3307
3308 assert!(!already_delivered(None, &mut last_sent));
3311 assert_eq!(last_sent, Some(4212));
3312 }
3313
3314 #[test]
3317 fn a_fresh_subscription_delivers_its_first_live_frame() {
3318 let mut last_sent = None;
3319 assert!(!already_delivered(Some(0), &mut last_sent));
3320 assert_eq!(last_sent, Some(0));
3321 }
3322
3323 mod over_a_socket {
3328 use super::*;
3329 use crate::journal::{EventJournal, JournalConfig};
3330 use crate::projector::Projector;
3331 use crate::{MutationBatch, SlotContext};
3332 use arete_interpreter::Mutation;
3333 use futures_util::{SinkExt, StreamExt};
3334 use std::time::Duration;
3335 use tokio::net::{TcpListener, TcpStream};
3336 use tokio::sync::mpsc;
3337 use tokio_tungstenite::tungstenite::Message;
3338 use tokio_tungstenite::{client_async, WebSocketStream};
3339
3340 const RETAINED: u64 = 600;
3341
3342 fn append_index() -> ViewIndex {
3343 let mut index = ViewIndex::new();
3344 index.add_spec(ViewSpec {
3345 id: "Trade/append".to_string(),
3346 export: "Trade".to_string(),
3347 mode: Mode::Append,
3348 wire_format: Default::default(),
3349 projection: Projection::all(),
3350 filters: Filters::all(),
3351 delivery: Delivery::default(),
3352 pipeline: None,
3353 source_view: None,
3354 });
3355 index
3356 }
3357
3358 fn trade(index: u64) -> MutationBatch {
3359 MutationBatch::with_slot_context(
3360 vec![Mutation {
3361 export: "Trade".to_string(),
3362 key: json!(format!("pool{}", index % 4)),
3363 patch: json!({"trade": index}),
3364 append: vec![],
3365 }]
3366 .into_iter()
3367 .collect(),
3368 SlotContext::new(100 + index / 3, index % 3),
3369 )
3370 }
3371
3372 struct Harness {
3373 addr: SocketAddr,
3374 journal: Arc<EventJournal>,
3375 tx: mpsc::Sender<MutationBatch>,
3376 }
3377
3378 impl Harness {
3379 async fn start() -> Self {
3380 let view_index = Arc::new(append_index());
3381 let entity_cache = EntityCache::new();
3382 let bus_manager = BusManager::new();
3383 let journal = Arc::new(EventJournal::new(JournalConfig {
3384 enabled: true,
3385 max_bytes_per_view: u64::MAX,
3386 max_records_per_view: 10_000,
3387 max_age: Duration::from_secs(3_600),
3388 }));
3389
3390 let (tx, rx) = mpsc::channel::<MutationBatch>(256);
3391 tokio::spawn(
3392 Projector::new(
3393 view_index.clone(),
3394 bus_manager.clone(),
3395 entity_cache.clone(),
3396 rx,
3397 #[cfg(feature = "otel")]
3398 None,
3399 )
3400 .with_journal(journal.clone())
3401 .run(),
3402 );
3403
3404 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3405 let addr = listener.local_addr().unwrap();
3406 let server = WebSocketServer::new(
3407 addr,
3408 bus_manager,
3409 entity_cache,
3410 view_index,
3411 #[cfg(feature = "otel")]
3412 None,
3413 )
3414 .with_journal(journal.clone());
3415 let (acceptor, _cleanup) = server.into_acceptor();
3416 tokio::spawn(async move { acceptor.serve_listener(listener).await });
3417
3418 Self { addr, journal, tx }
3419 }
3420
3421 async fn publish(&self, range: std::ops::Range<u64>) {
3422 for index in range {
3423 self.tx.send(trade(index)).await.unwrap();
3424 }
3425 let (ack, wait) = oneshot::channel();
3426 self.tx
3427 .send(MutationBatch::flush_marker(ack))
3428 .await
3429 .unwrap();
3430 wait.await.unwrap();
3431 }
3432
3433 async fn connect(&self) -> WebSocketStream<TcpStream> {
3434 let stream = TcpStream::connect(self.addr).await.unwrap();
3435 client_async(format!("ws://{}/", self.addr), stream)
3436 .await
3437 .unwrap()
3438 .0
3439 }
3440 }
3441
3442 async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
3443 loop {
3444 let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
3445 .await
3446 .expect("the server answers within the timeout")
3447 .expect("the stream stays open")
3448 .expect("a readable frame");
3449 let bytes = match &message {
3452 Message::Text(text) => text.as_bytes(),
3453 Message::Binary(bytes) => bytes.as_ref(),
3454 _ => continue,
3455 };
3456 return serde_json::from_slice(bytes).expect("frames are JSON");
3457 }
3458 }
3459
3460 async fn collect_trades(socket: &mut WebSocketStream<TcpStream>, count: usize) -> Vec<u64> {
3462 let mut offsets = Vec::with_capacity(count);
3463 while offsets.len() < count {
3464 let frame = next_frame(socket).await;
3465 assert_ne!(
3466 frame["type"], "error",
3467 "no error frame should interrupt delivery: {frame}"
3468 );
3469 if let Some(offset) = frame["offset"].as_u64() {
3470 offsets.push(offset);
3471 }
3472 }
3473 offsets
3474 }
3475
3476 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3480 async fn a_reconnect_replays_from_a_cursor_and_continues_live() {
3481 let harness = Harness::start().await;
3482 harness.publish(0..RETAINED).await;
3483
3484 let cursor = harness.journal.window("Trade/append").await;
3485 let stored = format!("{}:{}", cursor.epoch, 99);
3486
3487 let mut socket = harness.connect().await;
3488 socket
3489 .send(Message::Text(
3490 json!({
3491 "type": "subscribe",
3492 "protocolVersion": 2,
3493 "subscriptionId": "trades",
3494 "query": {"view": "Trade/append", "after": stored},
3495 })
3496 .to_string()
3497 .into(),
3498 ))
3499 .await
3500 .unwrap();
3501
3502 let ack = next_frame(&mut socket).await;
3503 assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3504 assert_eq!(ack["replayWindow"]["next"], json!(RETAINED));
3505
3506 let replayed = collect_trades(&mut socket, (RETAINED - 100) as usize).await;
3508 assert_eq!(
3509 replayed,
3510 (100..RETAINED).collect::<Vec<_>>(),
3511 "every event after the cursor, in order, exactly once"
3512 );
3513
3514 harness.publish(RETAINED..RETAINED + 40).await;
3516 let live = collect_trades(&mut socket, 40).await;
3517 assert_eq!(
3518 live,
3519 (RETAINED..RETAINED + 40).collect::<Vec<_>>(),
3520 "the live stream resumes exactly where the replay stopped"
3521 );
3522
3523 socket.close(None).await.ok();
3524 }
3525
3526 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3531 async fn events_published_during_a_replay_are_not_lost() {
3532 let harness = Harness::start().await;
3533 harness.publish(0..RETAINED).await;
3534
3535 let epoch = harness.journal.window("Trade/append").await.epoch;
3536 let mut socket = harness.connect().await;
3537 socket
3538 .send(Message::Text(
3539 json!({
3540 "type": "subscribe",
3541 "protocolVersion": 2,
3542 "subscriptionId": "trades",
3543 "query": {"view": "Trade/append", "after": format!("{epoch}:0")},
3544 })
3545 .to_string()
3546 .into(),
3547 ))
3548 .await
3549 .unwrap();
3550 let ack = next_frame(&mut socket).await;
3551 assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3552
3553 harness.publish(RETAINED..RETAINED + 200).await;
3555
3556 let total = (RETAINED + 200 - 1) as usize;
3557 let delivered = collect_trades(&mut socket, total).await;
3558 assert_eq!(
3559 delivered,
3560 (1..RETAINED + 200).collect::<Vec<_>>(),
3561 "replay and live output join without a gap or a repeat"
3562 );
3563
3564 socket.close(None).await.ok();
3565 }
3566
3567 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3570 async fn a_stale_epoch_is_refused_and_frees_the_subscription_id() {
3571 let harness = Harness::start().await;
3572 harness.publish(0..50).await;
3573
3574 let mut socket = harness.connect().await;
3575 socket
3576 .send(Message::Text(
3577 json!({
3578 "type": "subscribe",
3579 "protocolVersion": 2,
3580 "subscriptionId": "trades",
3581 "query": {
3582 "view": "Trade/append",
3583 "after": format!("{}:10", crate::journal::JournalEpoch::new()),
3584 },
3585 })
3586 .to_string()
3587 .into(),
3588 ))
3589 .await
3590 .unwrap();
3591
3592 let error = next_frame(&mut socket).await;
3593 assert_eq!(error["type"], "error", "unexpected frame: {error}");
3594 assert_eq!(error["code"], "cursor-epoch-changed");
3595
3596 socket
3599 .send(Message::Text(
3600 json!({
3601 "type": "subscribe",
3602 "protocolVersion": 2,
3603 "subscriptionId": "trades",
3604 "query": {"view": "Trade/append"},
3605 })
3606 .to_string()
3607 .into(),
3608 ))
3609 .await
3610 .unwrap();
3611 let ack = next_frame(&mut socket).await;
3612 assert_eq!(
3613 ack["op"], "subscribed",
3614 "the refused id must be reusable: {ack}"
3615 );
3616
3617 assert_eq!(collect_trades(&mut socket, 50).await.len(), 50);
3618 socket.close(None).await.ok();
3619 }
3620
3621 const ROUND: &str = "Round/state";
3622
3623 #[derive(Clone)]
3626 struct StateServer {
3627 addr: SocketAddr,
3628 bus_manager: BusManager,
3629 entity_cache: EntityCache,
3630 }
3631
3632 impl StateServer {
3633 async fn start() -> Self {
3634 let mut index = ViewIndex::new();
3635 index.add_spec(ViewSpec {
3636 id: ROUND.to_string(),
3637 export: "Round".to_string(),
3638 mode: Mode::State,
3639 wire_format: Default::default(),
3640 projection: Projection::all(),
3641 filters: Filters::all(),
3642 delivery: Delivery::default(),
3643 pipeline: None,
3644 source_view: None,
3645 });
3646 let bus_manager = BusManager::new();
3647 let entity_cache = EntityCache::new();
3648 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3649 let addr = listener.local_addr().unwrap();
3650 let server = WebSocketServer::new(
3651 addr,
3652 bus_manager.clone(),
3653 entity_cache.clone(),
3654 Arc::new(index),
3655 #[cfg(feature = "otel")]
3656 None,
3657 );
3658 let (acceptor, _cleanup) = server.into_acceptor();
3659 tokio::spawn(async move { acceptor.serve_listener(listener).await });
3660 Self {
3661 addr,
3662 bus_manager,
3663 entity_cache,
3664 }
3665 }
3666
3667 async fn publish(&self, key: &str, mut patch: Value, seq: &str) {
3670 patch["_seq"] = Value::String(seq.to_string());
3671 self.entity_cache
3672 .upsert_with_append(ROUND, key, patch.clone(), &[])
3673 .await;
3674 let frame = json!({
3675 "mode": "state",
3676 "entity": ROUND,
3677 "op": "patch",
3678 "key": key,
3679 "data": patch,
3680 "seq": seq,
3681 });
3682 self.bus_manager
3683 .publish_state(ROUND, key, Arc::new(Bytes::from(frame.to_string())))
3684 .await;
3685 }
3686
3687 async fn subscribe(&self, key: &str) -> WebSocketStream<TcpStream> {
3689 let stream = TcpStream::connect(self.addr).await.unwrap();
3690 let mut socket = client_async(format!("ws://{}/", self.addr), stream)
3691 .await
3692 .unwrap()
3693 .0;
3694 socket
3695 .send(Message::Text(
3696 json!({
3697 "type": "subscribe",
3698 "protocolVersion": 2,
3699 "subscriptionId": "round",
3700 "query": {"view": ROUND, "key": key},
3701 })
3702 .to_string()
3703 .into(),
3704 ))
3705 .await
3706 .unwrap();
3707 let ack = next_frame(&mut socket).await;
3708 assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3709 loop {
3710 let frame = next_frame(&mut socket).await;
3711 assert_eq!(frame["op"], "snapshot", "unexpected frame: {frame}");
3712 if frame["complete"] == true {
3713 return socket;
3714 }
3715 }
3716 }
3717 }
3718
3719 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
3725 async fn a_state_subscriber_keeps_the_fields_of_an_overwritten_patch() {
3726 let server = StateServer::start().await;
3727 server
3728 .publish("7", json!({"id": 7}), "100:000000000001")
3729 .await;
3730 let mut socket = server.subscribe("7").await;
3731
3732 server
3735 .publish(
3736 "7",
3737 json!({"results": {"slot_hash": "abc"}}),
3738 "140:003836292257",
3739 )
3740 .await;
3741 let frame = next_frame(&mut socket).await;
3742 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3743 assert_eq!(frame["seq"], "140:003836292257");
3744
3745 let writer = server.clone();
3750 tokio::spawn(async move {
3751 writer
3752 .publish("7", json!({"entropy": {"seed": "def"}}), "140:000000001200")
3753 .await;
3754 writer
3755 .publish("7", json!({"total": 1}), "140:000000001200")
3756 .await;
3757 })
3758 .await
3759 .unwrap();
3760 let frame = next_frame(&mut socket).await;
3761 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3762 assert_eq!(
3763 frame["data"],
3764 json!({
3765 "id": 7,
3766 "results": {"slot_hash": "abc"},
3767 "entropy": {"seed": "def"},
3768 "total": 1,
3769 "_seq": "140:000000001200",
3770 })
3771 );
3772 assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
3775
3776 server
3778 .publish("7", json!({"total": 2}), "141:000000000002")
3779 .await;
3780 let frame = next_frame(&mut socket).await;
3781 assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3782 assert_eq!(
3783 frame["data"],
3784 json!({"total": 2, "_seq": "141:000000000002"})
3785 );
3786
3787 socket.close(None).await.ok();
3788 }
3789 }
3790
3791 mod session_expiry {
3793 use super::*;
3794 use crate::websocket::auth::SignedSessionAuthPlugin;
3795 use arete_auth::{KeyClass, SessionClaims, SigningKey, TokenSigner, TokenVerifier};
3796 use futures_util::{SinkExt, StreamExt};
3797 use std::time::Duration;
3798 use tokio::net::{TcpListener, TcpStream};
3799 use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
3800 use tokio_tungstenite::tungstenite::protocol::CloseFrame;
3801 use tokio_tungstenite::tungstenite::Message;
3802 use tokio_tungstenite::{client_async, WebSocketStream};
3803
3804 struct Server {
3805 addr: SocketAddr,
3806 signer: TokenSigner,
3807 }
3808
3809 impl Server {
3810 async fn start() -> Self {
3811 let signing_key = SigningKey::generate();
3812 let verifier =
3813 TokenVerifier::new(signing_key.verifying_key(), "test-issuer", "test-audience");
3814 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3815 let addr = listener.local_addr().unwrap();
3816 let server = WebSocketServer::new(
3817 addr,
3818 BusManager::new(),
3819 EntityCache::new(),
3820 Arc::new(ViewIndex::new()),
3821 #[cfg(feature = "otel")]
3822 None,
3823 )
3824 .with_auth_plugin(Arc::new(SignedSessionAuthPlugin::new(verifier)));
3825 let (acceptor, _cleanup) = server.into_acceptor();
3826 tokio::spawn(async move { acceptor.serve_listener(listener).await });
3827 Self {
3828 addr,
3829 signer: TokenSigner::new(signing_key, "test-issuer"),
3830 }
3831 }
3832
3833 fn token(&self, ttl_seconds: u64) -> String {
3834 let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
3835 .with_scope("read")
3836 .with_key_class(KeyClass::Secret)
3837 .with_ttl(ttl_seconds)
3838 .build();
3839 self.signer.sign(claims).unwrap()
3840 }
3841
3842 async fn connect(&self, token: &str) -> WebSocketStream<TcpStream> {
3843 let stream = TcpStream::connect(self.addr).await.unwrap();
3844 client_async(format!("ws://{}/?hs_token={token}", self.addr), stream)
3845 .await
3846 .unwrap()
3847 .0
3848 }
3849 }
3850
3851 async fn send_json(socket: &mut WebSocketStream<TcpStream>, message: Value) {
3852 socket
3853 .send(Message::Text(message.to_string().into()))
3854 .await
3855 .unwrap();
3856 }
3857
3858 async fn close_within(
3860 socket: &mut WebSocketStream<TcpStream>,
3861 wait: Duration,
3862 ) -> Option<Option<CloseFrame>> {
3863 tokio::time::timeout(wait, async {
3864 while let Some(Ok(message)) = socket.next().await {
3865 if let Message::Close(frame) = message {
3866 return Some(frame);
3867 }
3868 }
3869 Some(None)
3870 })
3871 .await
3872 .ok()
3873 .flatten()
3874 }
3875
3876 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3877 async fn an_expired_session_is_closed_with_the_reason() {
3878 let server = Server::start().await;
3879 let mut socket = server.connect(&server.token(2)).await;
3880
3881 tokio::time::sleep(Duration::from_secs(3)).await;
3882 send_json(&mut socket, json!({"type": "ping"})).await;
3883
3884 let frame = close_within(&mut socket, Duration::from_secs(5))
3885 .await
3886 .expect("the server closes the expired session")
3887 .expect("the close frame carries a reason");
3888 assert_eq!(frame.code, CloseCode::Policy);
3889 assert_eq!(
3890 frame.reason.as_str(),
3891 "token-expired: Authentication token expired"
3892 );
3893 }
3894
3895 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3896 async fn a_session_refreshed_in_band_outlives_its_first_token() {
3897 let server = Server::start().await;
3898 let mut socket = server.connect(&server.token(2)).await;
3899
3900 send_json(
3901 &mut socket,
3902 json!({"type": "refresh_auth", "token": server.token(3_600)}),
3903 )
3904 .await;
3905 let reply = tokio::time::timeout(Duration::from_secs(5), async {
3906 while let Some(Ok(message)) = socket.next().await {
3907 if let Message::Text(text) = message {
3908 return serde_json::from_str::<Value>(text.as_str()).ok();
3909 }
3910 }
3911 None
3912 })
3913 .await
3914 .expect("the server answers the refresh")
3915 .expect("the answer is JSON");
3916 assert_eq!(reply["success"], true, "refresh accepted: {reply}");
3917
3918 tokio::time::sleep(Duration::from_secs(3)).await;
3919 send_json(&mut socket, json!({"type": "ping"})).await;
3920 assert!(
3921 close_within(&mut socket, Duration::from_millis(1_500))
3922 .await
3923 .is_none(),
3924 "the socket stays open on the refreshed token"
3925 );
3926 }
3927 }
3928}