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