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 anyhow::Result;
19use bytes::Bytes;
20use futures_util::StreamExt;
21use serde::Serialize;
22use serde_json::Value;
23use std::collections::{HashMap, HashSet};
24use std::future::Future;
25use std::net::SocketAddr;
26use std::sync::Arc;
27use std::time::Instant;
28use tokio::net::{TcpListener, TcpStream};
29use tokio::sync::{broadcast, watch};
30use tokio_tungstenite::{
31 accept_hdr_async,
32 tungstenite::{
33 handshake::server::{ErrorResponse as HandshakeErrorResponse, Request, Response},
34 http::{header::CONTENT_TYPE, StatusCode},
35 Error as WsError,
36 },
37};
38use tokio_util::sync::CancellationToken;
39use tracing::{debug, error, info, info_span, warn, Instrument};
40use uuid::Uuid;
41
42#[cfg(feature = "otel")]
43use crate::metrics::Metrics;
44
45#[derive(Clone, Default)]
46struct WsMetrics {
47 #[cfg(feature = "otel")]
48 inner: Option<Arc<Metrics>>,
49}
50
51impl WsMetrics {
52 #[cfg(feature = "otel")]
53 fn new(inner: Option<Arc<Metrics>>) -> Self {
54 Self { inner }
55 }
56
57 fn connection_opened(&self, metering_key: Option<&str>) {
58 #[cfg(not(feature = "otel"))]
59 let _ = metering_key;
60 #[cfg(feature = "otel")]
61 if let Some(metrics) = &self.inner {
62 if let Some(metering_key) = metering_key {
63 metrics.record_ws_connection_with_metering(metering_key);
64 } else {
65 metrics.record_ws_connection();
66 }
67 }
68 }
69
70 fn connection_closed(&self, duration_secs: f64, metering_key: Option<&str>) {
71 #[cfg(not(feature = "otel"))]
72 let _ = (duration_secs, metering_key);
73 #[cfg(feature = "otel")]
74 if let Some(metrics) = &self.inner {
75 if let Some(metering_key) = metering_key {
76 metrics.record_ws_disconnection_with_metering(duration_secs, metering_key);
77 } else {
78 metrics.record_ws_disconnection(duration_secs);
79 }
80 }
81 }
82
83 fn message_received(&self, metering_key: Option<&str>) {
84 #[cfg(not(feature = "otel"))]
85 let _ = metering_key;
86 #[cfg(feature = "otel")]
87 if let Some(metrics) = &self.inner {
88 if let Some(metering_key) = metering_key {
89 metrics.record_ws_message_received_with_metering(metering_key);
90 } else {
91 metrics.record_ws_message_received();
92 }
93 }
94 }
95
96 fn message_sent(&self) {
97 #[cfg(feature = "otel")]
98 if let Some(metrics) = &self.inner {
99 metrics.record_ws_message_sent();
100 }
101 }
102
103 fn subscription_created(&self, view: &str, metering_key: Option<&str>) {
104 #[cfg(not(feature = "otel"))]
105 let _ = (view, metering_key);
106 #[cfg(feature = "otel")]
107 if let Some(metrics) = &self.inner {
108 if let Some(metering_key) = metering_key {
109 metrics.record_subscription_created_with_metering(view, metering_key);
110 } else {
111 metrics.record_subscription_created(view);
112 }
113 }
114 }
115
116 fn subscription_removed(&self, view: &str, metering_key: Option<&str>) {
117 #[cfg(not(feature = "otel"))]
118 let _ = (view, metering_key);
119 #[cfg(feature = "otel")]
120 if let Some(metrics) = &self.inner {
121 if let Some(metering_key) = metering_key {
122 metrics.record_subscription_removed_with_metering(view, metering_key);
123 } else {
124 metrics.record_subscription_removed(view);
125 }
126 }
127 }
128
129 fn protocol_error(&self, code: &str) {
130 #[cfg(not(feature = "otel"))]
131 let _ = code;
132 #[cfg(feature = "otel")]
133 if let Some(metrics) = &self.inner {
134 metrics.record_ws_protocol_error(code);
135 }
136 }
137}
138
139async fn handle_refresh_auth(
140 client_id: Uuid,
141 refresh_req: &RefreshAuthRequest,
142 client_manager: &ClientManager,
143 auth_plugin: &Arc<dyn WebSocketAuthPlugin>,
144) {
145 let refresh_result: Result<AuthContext, String> = if let Some(signed_plugin) = auth_plugin
146 .as_any()
147 .downcast_ref::<crate::websocket::auth::SignedSessionAuthPlugin>()
148 {
149 signed_plugin
150 .verify_refresh_token(&refresh_req.token)
151 .await
152 .map_err(|error| error.reason)
153 } else {
154 Err("In-band auth refresh not supported with current auth plugin".to_string())
155 };
156
157 let response = match refresh_result {
158 Ok(new_context) => {
159 let expires_at = new_context.expires_at;
160 if client_manager.update_client_auth(client_id, new_context) {
161 RefreshAuthResponse {
162 success: true,
163 error: None,
164 expires_at: Some(expires_at),
165 }
166 } else {
167 RefreshAuthResponse {
168 success: false,
169 error: Some("client-not-found".to_string()),
170 expires_at: None,
171 }
172 }
173 }
174 Err(error) => {
175 let code = if error.contains("expired") {
176 "token-expired"
177 } else if error.contains("signature") {
178 "token-invalid-signature"
179 } else if error.contains("issuer") {
180 "token-invalid-issuer"
181 } else if error.contains("audience") {
182 "token-invalid-audience"
183 } else {
184 "token-invalid"
185 };
186 RefreshAuthResponse {
187 success: false,
188 error: Some(code.to_string()),
189 expires_at: None,
190 }
191 }
192 };
193
194 if let Ok(json) = serde_json::to_string(&response) {
195 let _ = client_manager.send_text_to_client(client_id, json).await;
196 }
197}
198
199async fn send_socket_issue(
200 client_id: Uuid,
201 client_manager: &ClientManager,
202 deny: &AuthDeny,
203 fatal: bool,
204 subscription_id: Option<String>,
205) {
206 let message = SocketIssueMessage::from_auth_deny(deny, fatal, subscription_id);
207 if let Ok(json) = serde_json::to_string(&message) {
208 let _ = client_manager.send_text_to_client(client_id, json).await;
209 }
210}
211
212async fn send_protocol_issue(
213 client_id: Uuid,
214 client_manager: &ClientManager,
215 metrics: &WsMetrics,
216 subscription_id: Option<String>,
217 code: &str,
218 message: impl Into<String>,
219) {
220 metrics.protocol_error(code);
221 let issue = SocketIssueMessage::protocol(subscription_id, code, message);
222 if let Ok(json) = serde_json::to_string(&issue) {
223 let _ = client_manager.send_text_to_client(client_id, json).await;
224 }
225}
226
227fn key_class_label(key_class: arete_auth::KeyClass) -> &'static str {
228 match key_class {
229 arete_auth::KeyClass::Secret => "secret",
230 arete_auth::KeyClass::Publishable => "publishable",
231 }
232}
233
234fn emit_usage_event(
235 usage_emitter: &Option<Arc<dyn WebSocketUsageEmitter>>,
236 event: WebSocketUsageEvent,
237) {
238 if let Some(emitter) = usage_emitter.clone() {
239 tokio::spawn(async move {
240 emitter.emit(event).await;
241 });
242 }
243}
244
245fn usage_identity(
246 auth_context: Option<&AuthContext>,
247) -> (
248 Option<String>,
249 Option<String>,
250 Option<String>,
251 Option<String>,
252) {
253 match auth_context {
254 Some(context) => (
255 Some(context.metering_key.clone()),
256 Some(context.subject.clone()),
257 Some(key_class_label(context.key_class).to_string()),
258 context.deployment_id.clone(),
259 ),
260 None => (None, None, None, None),
261 }
262}
263
264fn emit_update_sent_for_client(
265 usage_emitter: &Option<Arc<dyn WebSocketUsageEmitter>>,
266 client_manager: &ClientManager,
267 client_id: Uuid,
268 view_id: &str,
269 bytes: usize,
270) {
271 let auth_context = client_manager.get_auth_context(client_id);
272 let (metering_key, subject, _, deployment_id) = usage_identity(auth_context.as_ref());
273 emit_usage_event(
274 usage_emitter,
275 WebSocketUsageEvent::UpdateSent {
276 client_id: client_id.to_string(),
277 deployment_id,
278 metering_key,
279 subject,
280 view_id: view_id.to_string(),
281 messages: 1,
282 bytes: bytes as u64,
283 },
284 );
285}
286
287#[derive(Clone)]
288struct SubscriptionContext {
289 client_id: Uuid,
290 client_manager: ClientManager,
291 bus_manager: BusManager,
292 entity_cache: EntityCache,
293 view_index: Arc<ViewIndex>,
294 usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
295 metrics: WsMetrics,
296}
297
298pub struct WebSocketServer {
299 bind_addr: SocketAddr,
300 client_manager: ClientManager,
301 bus_manager: BusManager,
302 entity_cache: EntityCache,
303 view_index: Arc<ViewIndex>,
304 max_clients: usize,
305 auth_plugin: Arc<dyn WebSocketAuthPlugin>,
306 usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
307 rate_limit_config: Option<RateLimitConfig>,
308 #[cfg(feature = "otel")]
309 metrics: Option<Arc<Metrics>>,
310}
311
312impl WebSocketServer {
313 #[cfg(feature = "otel")]
314 pub fn new(
315 bind_addr: SocketAddr,
316 bus_manager: BusManager,
317 entity_cache: EntityCache,
318 view_index: Arc<ViewIndex>,
319 metrics: Option<Arc<Metrics>>,
320 ) -> Self {
321 Self {
322 bind_addr,
323 client_manager: ClientManager::new(),
324 bus_manager,
325 entity_cache,
326 view_index,
327 max_clients: 10_000,
328 auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
329 usage_emitter: None,
330 rate_limit_config: None,
331 metrics,
332 }
333 }
334
335 #[cfg(not(feature = "otel"))]
336 pub fn new(
337 bind_addr: SocketAddr,
338 bus_manager: BusManager,
339 entity_cache: EntityCache,
340 view_index: Arc<ViewIndex>,
341 ) -> Self {
342 Self {
343 bind_addr,
344 client_manager: ClientManager::new(),
345 bus_manager,
346 entity_cache,
347 view_index,
348 max_clients: 10_000,
349 auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
350 usage_emitter: None,
351 rate_limit_config: None,
352 }
353 }
354
355 pub fn with_max_clients(mut self, max_clients: usize) -> Self {
356 self.max_clients = max_clients;
357 self
358 }
359
360 pub fn with_auth_plugin(mut self, auth_plugin: Arc<dyn WebSocketAuthPlugin>) -> Self {
361 self.auth_plugin = auth_plugin;
362 self
363 }
364
365 pub fn with_usage_emitter(mut self, usage_emitter: Arc<dyn WebSocketUsageEmitter>) -> Self {
366 self.usage_emitter = Some(usage_emitter);
367 self
368 }
369
370 pub fn with_rate_limit_config(mut self, config: RateLimitConfig) -> Self {
371 self.rate_limit_config = Some(config);
372 self
373 }
374
375 pub async fn start(self) -> Result<()> {
376 info!(
377 "Starting WebSocket server on {} (max_clients: {})",
378 self.bind_addr, self.max_clients
379 );
380 let listener = TcpListener::bind(&self.bind_addr).await?;
381 let client_manager = self
382 .rate_limit_config
383 .map(ClientManager::with_config)
384 .unwrap_or(self.client_manager);
385 client_manager.start_cleanup_task();
386
387 #[cfg(feature = "otel")]
388 let metrics = WsMetrics::new(self.metrics.clone());
389 #[cfg(not(feature = "otel"))]
390 let metrics = WsMetrics::default();
391
392 loop {
393 match listener.accept().await {
394 Ok((stream, addr)) => {
395 if client_manager.client_count() >= self.max_clients {
396 warn!("Rejecting connection from {}: max clients reached", addr);
397 continue;
398 }
399
400 let context = SubscriptionContext {
401 client_id: Uuid::nil(),
402 client_manager: client_manager.clone(),
403 bus_manager: self.bus_manager.clone(),
404 entity_cache: self.entity_cache.clone(),
405 view_index: self.view_index.clone(),
406 usage_emitter: self.usage_emitter.clone(),
407 metrics: metrics.clone(),
408 };
409 let auth_plugin = self.auth_plugin.clone();
410 tokio::spawn(
411 async move {
412 if let Err(error) =
413 handle_connection(stream, context, addr, auth_plugin).await
414 {
415 error!("WebSocket connection error: {}", error);
416 }
417 }
418 .instrument(info_span!("ws.connection", %addr)),
419 );
420 }
421 Err(error) => error!("Failed to accept connection: {}", error),
422 }
423 }
424 }
425}
426
427#[derive(Debug, Clone)]
428struct HandshakeReject {
429 status: StatusCode,
430 body: crate::websocket::auth::ErrorResponse,
431 error_code: String,
432 retry_after_secs: Option<u64>,
433}
434
435impl HandshakeReject {
436 fn from_deny(deny: &AuthDeny) -> Self {
437 let retry_after_secs = match deny.retry_policy {
438 crate::websocket::auth::RetryPolicy::RetryAfter(duration) => Some(duration.as_secs()),
439 _ => None,
440 };
441 Self {
442 status: StatusCode::from_u16(deny.http_status).unwrap_or(StatusCode::UNAUTHORIZED),
443 body: deny.to_error_response(),
444 error_code: deny.code.to_string(),
445 retry_after_secs,
446 }
447 }
448}
449
450fn build_handshake_error_response(
451 response: &Response,
452 reject: &HandshakeReject,
453) -> HandshakeErrorResponse {
454 let mut builder = Response::builder()
455 .status(reject.status)
456 .version(response.version())
457 .header(CONTENT_TYPE, "application/json; charset=utf-8")
458 .header("X-Error-Code", &reject.error_code)
459 .header("Cache-Control", "no-store");
460 if let Some(retry_after_secs) = reject.retry_after_secs {
461 builder = builder.header("Retry-After", retry_after_secs.to_string());
462 }
463 let body = serde_json::to_string(&reject.body).unwrap_or_else(|_| {
464 format!(
465 r#"{{"error":"{}","message":"{}","code":"{}","retryable":false}}"#,
466 reject.body.error, reject.body.message, reject.body.code
467 )
468 });
469 builder
470 .body(Some(body))
471 .expect("handshake rejection response should build")
472}
473
474#[allow(clippy::result_large_err)]
475async fn accept_authorized_connection(
476 stream: TcpStream,
477 remote_addr: SocketAddr,
478 auth_plugin: Arc<dyn WebSocketAuthPlugin>,
479 client_manager: ClientManager,
480) -> Result<Option<(tokio_tungstenite::WebSocketStream<TcpStream>, AuthContext)>> {
481 use std::sync::Mutex;
482
483 let capture: Arc<Mutex<Option<Result<AuthContext, HandshakeReject>>>> =
484 Arc::new(Mutex::new(None));
485 let capture_ref = capture.clone();
486 let auth_plugin_ref = auth_plugin.clone();
487 let manager_ref = client_manager.clone();
488
489 let handshake_result = accept_hdr_async(stream, move |request: &Request, response| {
490 let request = ConnectionAuthRequest::from_http_request(remote_addr, request);
491 let result = tokio::task::block_in_place(|| {
492 tokio::runtime::Handle::current().block_on(async {
493 match auth_plugin_ref.authorize(&request).await {
494 AuthDecision::Allow(context) => manager_ref
495 .check_connection_allowed(remote_addr, &Some(context.clone()))
496 .await
497 .map(|()| context)
498 .map_err(|deny| HandshakeReject::from_deny(&deny)),
499 AuthDecision::Deny(deny) => Err(HandshakeReject::from_deny(&deny)),
500 }
501 })
502 });
503 *capture_ref.lock().expect("capture lock poisoned") = Some(result.clone());
504 match result {
505 Ok(_) => Ok(response),
506 Err(reject) => Err(build_handshake_error_response(&response, &reject)),
507 }
508 })
509 .await;
510
511 let auth_result = capture.lock().expect("capture lock poisoned").take();
512 match handshake_result {
513 Ok(stream) => match auth_result {
514 Some(Ok(context)) => Ok(Some((stream, context))),
515 Some(Err(reject)) => Err(anyhow::anyhow!(
516 "handshake unexpectedly succeeded after rejection: {}",
517 reject.body.message
518 )),
519 None => Err(anyhow::anyhow!("no auth result captured during handshake")),
520 },
521 Err(WsError::Http(_)) => Ok(None),
522 Err(error) => Err(error.into()),
523 }
524}
525
526async fn handle_connection(
527 stream: TcpStream,
528 mut context: SubscriptionContext,
529 remote_addr: SocketAddr,
530 auth_plugin: Arc<dyn WebSocketAuthPlugin>,
531) -> Result<()> {
532 let Some((ws_stream, auth_context)) = accept_authorized_connection(
533 stream,
534 remote_addr,
535 auth_plugin.clone(),
536 context.client_manager.clone(),
537 )
538 .await?
539 else {
540 return Ok(());
541 };
542
543 let client_id = Uuid::new_v4();
544 context.client_id = client_id;
545 let connection_start = Instant::now();
546 let (metering_key, subject, key_class, deployment_id) = usage_identity(Some(&auth_context));
547 context.metrics.connection_opened(metering_key.as_deref());
548
549 let (ws_sender, mut ws_receiver) = ws_stream.split();
550 context
551 .client_manager
552 .add_client(client_id, ws_sender, Some(auth_context), remote_addr);
553 emit_usage_event(
554 &context.usage_emitter,
555 WebSocketUsageEvent::ConnectionEstablished {
556 client_id: client_id.to_string(),
557 remote_addr: remote_addr.to_string(),
558 deployment_id: deployment_id.clone(),
559 metering_key: metering_key.clone(),
560 subject: subject.clone(),
561 key_class,
562 },
563 );
564
565 let mut active_subscriptions: HashMap<String, String> = HashMap::new();
566 while let Some(message) = ws_receiver.next().await {
567 let message = match message {
568 Ok(message) => message,
569 Err(error) => {
570 warn!("WebSocket error for client {}: {}", client_id, error);
571 break;
572 }
573 };
574 if message.is_close() {
575 break;
576 }
577 context.client_manager.update_client_last_seen(client_id);
578 if !message.is_text() {
579 continue;
580 }
581 if let Err(deny) = context
582 .client_manager
583 .check_inbound_message_allowed(client_id)
584 {
585 send_socket_issue(client_id, &context.client_manager, &deny, true, None).await;
586 break;
587 }
588 context.metrics.message_received(metering_key.as_deref());
589
590 let text = match message.to_text() {
591 Ok(text) => text,
592 Err(_) => continue,
593 };
594 let client_message = match serde_json::from_str::<ClientMessage>(text) {
595 Ok(message) => message,
596 Err(parse_error) => {
597 let subscription_id = extract_subscription_id(text);
598 send_protocol_issue(
599 client_id,
600 &context.client_manager,
601 &context.metrics,
602 subscription_id,
603 "malformed-message",
604 format!("invalid protocol v2 message: {parse_error}"),
605 )
606 .await;
607 continue;
608 }
609 };
610
611 match client_message {
612 ClientMessage::Subscribe(subscription) => {
613 let subscription_id = subscription.subscription_id.clone();
614 if let Err(message) = subscription.validate() {
615 send_protocol_issue(
616 client_id,
617 &context.client_manager,
618 &context.metrics,
619 Some(subscription_id),
620 "invalid-subscription",
621 message,
622 )
623 .await;
624 continue;
625 }
626 if let Err(deny) = context
627 .client_manager
628 .check_subscription_allowed(client_id)
629 .await
630 {
631 send_socket_issue(
632 client_id,
633 &context.client_manager,
634 &deny,
635 false,
636 Some(subscription_id),
637 )
638 .await;
639 continue;
640 }
641
642 let cancel_token = CancellationToken::new();
643 if !context
644 .client_manager
645 .add_client_subscription(
646 client_id,
647 subscription_id.clone(),
648 cancel_token.clone(),
649 )
650 .await
651 {
652 send_protocol_issue(
653 client_id,
654 &context.client_manager,
655 &context.metrics,
656 Some(subscription_id),
657 "duplicate-subscription-id",
658 "subscriptionId is already active on this connection",
659 )
660 .await;
661 continue;
662 }
663
664 let view = subscription.query.view.clone();
665 if let Err(error) = attach_client_to_bus(&context, subscription, cancel_token).await
666 {
667 context
668 .client_manager
669 .remove_client_subscription(client_id, &subscription_id)
670 .await;
671 send_protocol_issue(
672 client_id,
673 &context.client_manager,
674 &context.metrics,
675 Some(subscription_id),
676 "subscription-rejected",
677 error.to_string(),
678 )
679 .await;
680 continue;
681 }
682
683 active_subscriptions.insert(subscription_id, view.clone());
684 context
685 .metrics
686 .subscription_created(&view, metering_key.as_deref());
687 emit_usage_event(
688 &context.usage_emitter,
689 WebSocketUsageEvent::SubscriptionCreated {
690 client_id: client_id.to_string(),
691 deployment_id: deployment_id.clone(),
692 metering_key: metering_key.clone(),
693 subject: subject.clone(),
694 view_id: view,
695 },
696 );
697 }
698 ClientMessage::Unsubscribe(unsubscription) => {
699 handle_unsubscribe(
700 &context,
701 unsubscription,
702 &mut active_subscriptions,
703 metering_key.as_deref(),
704 &deployment_id,
705 &metering_key,
706 &subject,
707 )
708 .await;
709 }
710 ClientMessage::Ping => debug!("Received ping from client {}", client_id),
711 ClientMessage::RefreshAuth(request) => {
712 handle_refresh_auth(client_id, &request, &context.client_manager, &auth_plugin)
713 .await;
714 }
715 }
716 }
717
718 context
719 .client_manager
720 .cancel_all_client_subscriptions(client_id)
721 .await;
722 context.client_manager.remove_client(client_id);
723 if let Some(rate_limiter) = context.client_manager.rate_limiter().cloned() {
724 rate_limiter.remove_client_buckets(client_id).await;
725 }
726 for view in active_subscriptions.values() {
727 context
728 .metrics
729 .subscription_removed(view, metering_key.as_deref());
730 emit_usage_event(
731 &context.usage_emitter,
732 WebSocketUsageEvent::SubscriptionRemoved {
733 client_id: client_id.to_string(),
734 deployment_id: deployment_id.clone(),
735 metering_key: metering_key.clone(),
736 subject: subject.clone(),
737 view_id: view.clone(),
738 },
739 );
740 }
741 let duration = connection_start.elapsed().as_secs_f64();
742 context
743 .metrics
744 .connection_closed(duration, metering_key.as_deref());
745 emit_usage_event(
746 &context.usage_emitter,
747 WebSocketUsageEvent::ConnectionClosed {
748 client_id: client_id.to_string(),
749 deployment_id,
750 metering_key,
751 subject,
752 duration_secs: Some(duration),
753 subscription_count: u32::try_from(active_subscriptions.len()).unwrap_or(u32::MAX),
754 },
755 );
756 Ok(())
757}
758
759#[allow(clippy::too_many_arguments)]
760async fn handle_unsubscribe(
761 context: &SubscriptionContext,
762 unsubscription: Unsubscription,
763 active_subscriptions: &mut HashMap<String, String>,
764 metrics_metering_key: Option<&str>,
765 deployment_id: &Option<String>,
766 usage_metering_key: &Option<String>,
767 subject: &Option<String>,
768) {
769 let subscription_id = unsubscription.subscription_id.clone();
770 if let Err(message) = unsubscription.validate() {
771 send_protocol_issue(
772 context.client_id,
773 &context.client_manager,
774 &context.metrics,
775 Some(subscription_id),
776 "invalid-unsubscription",
777 message,
778 )
779 .await;
780 return;
781 }
782
783 if !context
784 .client_manager
785 .remove_client_subscription(context.client_id, &subscription_id)
786 .await
787 {
788 send_protocol_issue(
789 context.client_id,
790 &context.client_manager,
791 &context.metrics,
792 Some(subscription_id),
793 "unknown-subscription-id",
794 "subscriptionId is not active on this connection",
795 )
796 .await;
797 return;
798 }
799
800 let Some(view) = active_subscriptions.remove(&subscription_id) else {
801 return;
802 };
803 let _ = send_control_frame(context, &UnsubscribedFrame::new(subscription_id), &view);
804 context
805 .metrics
806 .subscription_removed(&view, metrics_metering_key);
807 emit_usage_event(
808 &context.usage_emitter,
809 WebSocketUsageEvent::SubscriptionRemoved {
810 client_id: context.client_id.to_string(),
811 deployment_id: deployment_id.clone(),
812 metering_key: usage_metering_key.clone(),
813 subject: subject.clone(),
814 view_id: view,
815 },
816 );
817}
818
819fn extract_subscription_id(text: &str) -> Option<String> {
820 serde_json::from_str::<Value>(text)
821 .ok()?
822 .get("subscriptionId")?
823 .as_str()
824 .map(str::to_string)
825}
826
827struct SnapshotMetadata<'a> {
828 subscription_id: &'a str,
829 snapshot_id: &'a str,
830 authoritative: bool,
831 mode: Mode,
832 view_id: &'a str,
833 key: Option<&'a str>,
834}
835
836fn create_snapshot_batches(
837 entities: &[SnapshotEntity],
838 metadata: SnapshotMetadata<'_>,
839 batch_config: &SnapshotBatchConfig,
840) -> Vec<SnapshotFrame> {
841 if entities.is_empty() {
842 return vec![SnapshotFrame {
843 protocol_version: PROTOCOL_VERSION,
844 subscription_id: metadata.subscription_id.to_string(),
845 snapshot_id: metadata.snapshot_id.to_string(),
846 authoritative: metadata.authoritative,
847 mode: metadata.mode,
848 export: metadata.view_id.to_string(),
849 op: "snapshot",
850 key: metadata.key.map(str::to_string),
851 data: vec![],
852 complete: true,
853 }];
854 }
855
856 let mut batches = Vec::new();
857 let mut offset = 0;
858 while offset < entities.len() {
859 let configured_size = if offset == 0 {
860 batch_config.initial_batch_size
861 } else {
862 batch_config.subsequent_batch_size
863 };
864 let end = (offset + configured_size.max(1)).min(entities.len());
865 batches.push(SnapshotFrame {
866 protocol_version: PROTOCOL_VERSION,
867 subscription_id: metadata.subscription_id.to_string(),
868 snapshot_id: metadata.snapshot_id.to_string(),
869 authoritative: metadata.authoritative,
870 mode: metadata.mode,
871 export: metadata.view_id.to_string(),
872 op: "snapshot",
873 key: metadata.key.map(str::to_string),
874 data: entities[offset..end].to_vec(),
875 complete: end == entities.len(),
876 });
877 offset = end;
878 }
879 batches
880}
881
882async fn send_snapshot_batches(
883 context: &SubscriptionContext,
884 subscription: &Subscription,
885 entities: &[SnapshotEntity],
886 mode: Mode,
887 batch_config: &SnapshotBatchConfig,
888) -> Result<()> {
889 let snapshot_id = Uuid::new_v4().to_string();
890 let authoritative = subscription.query.after.is_none();
891 let frames = create_snapshot_batches(
892 entities,
893 SnapshotMetadata {
894 subscription_id: &subscription.subscription_id,
895 snapshot_id: &snapshot_id,
896 authoritative,
897 mode,
898 view_id: &subscription.query.view,
899 key: subscription.query.key.as_deref(),
900 },
901 batch_config,
902 );
903
904 for frame in frames {
905 let rows = frame.data.len() as u32;
906 let json = serde_json::to_vec(&frame)?;
907 let payload = maybe_compress(&json);
908 let bytes = payload.as_bytes().len() as u64;
909 context
910 .client_manager
911 .send_compressed_async(context.client_id, payload)
912 .await
913 .map_err(|error| anyhow::anyhow!("failed to send snapshot: {error}"))?;
914 context.metrics.message_sent();
915
916 let auth_context = context.client_manager.get_auth_context(context.client_id);
917 let (metering_key, subject, _, deployment_id) = usage_identity(auth_context.as_ref());
918 emit_usage_event(
919 &context.usage_emitter,
920 WebSocketUsageEvent::SnapshotSent {
921 client_id: context.client_id.to_string(),
922 deployment_id,
923 metering_key,
924 subject,
925 view_id: subscription.query.view.clone(),
926 rows,
927 messages: 1,
928 bytes,
929 },
930 );
931 }
932 Ok(())
933}
934
935fn extract_sort_config(view_spec: &ViewSpec) -> Option<SortConfig> {
936 if let Some(sort) = view_spec
937 .pipeline
938 .as_ref()
939 .and_then(|pipeline| pipeline.sort.as_ref())
940 {
941 return Some(SortConfig {
942 field: sort.field_path.clone(),
943 order: match sort.order {
944 crate::materialized_view::SortOrder::Asc => SortOrder::Asc,
945 crate::materialized_view::SortOrder::Desc => SortOrder::Desc,
946 },
947 });
948 }
949 (view_spec.mode == Mode::List).then(|| SortConfig {
950 field: vec!["_seq".to_string()],
951 order: SortOrder::Desc,
952 })
953}
954
955fn send_control_frame<T: Serialize>(
956 context: &SubscriptionContext,
957 frame: &T,
958 view_id: &str,
959) -> Result<()> {
960 let json = serde_json::to_vec(frame)?;
961 let bytes = json.len();
962 context
963 .client_manager
964 .send_to_client(context.client_id, Arc::new(Bytes::from(json)))
965 .map_err(|error| anyhow::anyhow!("failed to send control frame: {error}"))?;
966 context.metrics.message_sent();
967 emit_update_sent_for_client(
968 &context.usage_emitter,
969 &context.client_manager,
970 context.client_id,
971 view_id,
972 bytes,
973 );
974 Ok(())
975}
976
977fn send_subscribed_frame(
978 context: &SubscriptionContext,
979 subscription: &Subscription,
980 view_spec: &ViewSpec,
981) -> Result<()> {
982 let frame = SubscribedFrame::new(
983 subscription.subscription_id.clone(),
984 subscription.query.clone(),
985 view_spec.mode,
986 extract_sort_config(view_spec),
987 );
988 send_control_frame(context, &frame, &subscription.query.view)
989}
990
991fn enforce_snapshot_limit(context: &SubscriptionContext, rows: usize) -> Result<()> {
992 context
993 .client_manager
994 .check_snapshot_allowed(context.client_id, u32::try_from(rows).unwrap_or(u32::MAX))
995 .map_err(|deny| anyhow::anyhow!(deny.reason))
996}
997
998async fn subscribe_state_then_snapshot<F, Fut, T>(
999 bus_manager: &BusManager,
1000 view_id: &str,
1001 key: &str,
1002 snapshot: F,
1003) -> (watch::Receiver<Arc<Bytes>>, T)
1004where
1005 F: FnOnce() -> Fut,
1006 Fut: Future<Output = T>,
1007{
1008 let mut receiver = bus_manager.get_or_create_state_bus(view_id, key).await;
1009 receiver.borrow_and_update();
1010 let snapshot = snapshot().await;
1011 (receiver, snapshot)
1012}
1013
1014async fn subscribe_list_then_snapshot<F, Fut, T>(
1015 bus_manager: &BusManager,
1016 view_id: &str,
1017 snapshot: F,
1018) -> (broadcast::Receiver<Arc<BusMessage>>, T)
1019where
1020 F: FnOnce() -> Fut,
1021 Fut: Future<Output = T>,
1022{
1023 let receiver = bus_manager.get_or_create_list_bus(view_id).await;
1024 let snapshot = snapshot().await;
1025 (receiver, snapshot)
1026}
1027
1028async fn attach_client_to_bus(
1029 context: &SubscriptionContext,
1030 mut subscription: Subscription,
1031 cancel_token: CancellationToken,
1032) -> Result<()> {
1033 let view_spec = context
1034 .view_index
1035 .get_view(&subscription.query.view)
1036 .cloned()
1037 .ok_or_else(|| anyhow::anyhow!("unknown view: {}", subscription.query.view))?;
1038
1039 if view_spec.mode == Mode::State && !view_spec.is_derived() && subscription.query.key.is_none()
1040 {
1041 return Err(anyhow::anyhow!("state subscriptions require query.key"));
1042 }
1043 if view_spec.is_derived() && subscription.query.take.is_none() {
1044 subscription.query.take = view_spec
1045 .pipeline
1046 .as_ref()
1047 .and_then(|pipeline| pipeline.limit);
1048 }
1049
1050 if view_spec.mode == Mode::State && !view_spec.is_derived() {
1051 attach_state_subscription(context, subscription, view_spec, cancel_token).await
1052 } else {
1053 attach_collection_subscription(context, subscription, view_spec, cancel_token).await
1054 }
1055}
1056
1057async fn attach_state_subscription(
1058 context: &SubscriptionContext,
1059 subscription: Subscription,
1060 view_spec: ViewSpec,
1061 cancel_token: CancellationToken,
1062) -> Result<()> {
1063 let view_id = subscription.query.view.clone();
1064 let key = subscription.query.key.clone().unwrap_or_default();
1065 let query = subscription.query.clone();
1066 let cache = context.entity_cache.clone();
1067 let view_spec_for_snapshot = view_spec.clone();
1068 let (mut receiver, initial) =
1069 subscribe_state_then_snapshot(&context.bus_manager, &view_id, &key, move || async move {
1070 load_query_entities(&cache, None, &view_spec_for_snapshot, &query, false).await
1071 })
1072 .await;
1073
1074 let mut snapshot_entities = initial.clone();
1075 if let Some(limit) = subscription.query.snapshot_limit {
1076 snapshot_entities.truncate(limit);
1077 }
1078 enforce_snapshot_limit(context, snapshot_entities.len())?;
1079 send_subscribed_frame(context, &subscription, &view_spec)?;
1080 if subscription.snapshot.enabled {
1081 send_snapshot_batches(
1082 context,
1083 &subscription,
1084 &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1085 view_spec.mode,
1086 &context.entity_cache.snapshot_config(),
1087 )
1088 .await?;
1089 }
1090
1091 let task_context = context.clone();
1092 let subscription_id = subscription.subscription_id.clone();
1093 let query = subscription.query.clone();
1094 let view_spec_task = view_spec.clone();
1095 let span_view = view_id.clone();
1096 let span_key = key.clone();
1097 tokio::spawn(
1098 async move {
1099 let mut member = !initial.is_empty();
1100 loop {
1101 tokio::select! {
1102 _ = cancel_token.cancelled() => break,
1103 changed = receiver.changed() => {
1104 if changed.is_err() {
1105 break;
1106 }
1107 let payload = receiver.borrow().clone();
1108 let metadata = source_frame_metadata(&payload);
1109 if metadata.op == "delete" {
1110 task_context.entity_cache.remove(&query.view, &key).await;
1111 if member && send_membership_frame(
1112 &task_context,
1113 &subscription_id,
1114 &view_spec_task,
1115 "delete",
1116 &key,
1117 Value::Null,
1118 metadata.seq,
1119 ).is_err() {
1120 break;
1121 }
1122 member = false;
1123 continue;
1124 }
1125
1126 let selected = load_query_entities(
1127 &task_context.entity_cache,
1128 None,
1129 &view_spec_task,
1130 &query,
1131 false,
1132 ).await;
1133 let is_member = !selected.is_empty();
1134 let result = match (member, is_member) {
1135 (true, true) => send_scoped_source_payload(
1136 &task_context,
1137 &subscription_id,
1138 &query.view,
1139 payload,
1140 ),
1141 (false, true) => {
1142 let (entity_key, data) = selected.into_iter().next().unwrap();
1143 send_membership_frame(
1144 &task_context,
1145 &subscription_id,
1146 &view_spec_task,
1147 "upsert",
1148 &entity_key,
1149 data,
1150 metadata.seq,
1151 )
1152 }
1153 (true, false) => send_membership_frame(
1154 &task_context,
1155 &subscription_id,
1156 &view_spec_task,
1157 "remove",
1158 &key,
1159 Value::Null,
1160 metadata.seq,
1161 ),
1162 (false, false) => Ok(()),
1163 };
1164 if result.is_err() {
1165 break;
1166 }
1167 member = is_member;
1168 }
1169 }
1170 }
1171 }
1172 .instrument(info_span!("ws.subscribe.state", client_id = %context.client_id, view = %span_view, key = %span_key)),
1173 );
1174 Ok(())
1175}
1176
1177async fn attach_collection_subscription(
1178 context: &SubscriptionContext,
1179 subscription: Subscription,
1180 view_spec: ViewSpec,
1181 cancel_token: CancellationToken,
1182) -> Result<()> {
1183 let view_id = subscription.query.view.clone();
1184 let source_view_id = view_spec
1185 .source_view
1186 .clone()
1187 .unwrap_or_else(|| view_id.clone());
1188 let query = subscription.query.clone();
1189 let cache = context.entity_cache.clone();
1190 let sorted_caches = view_spec
1191 .is_derived()
1192 .then(|| context.view_index.sorted_caches());
1193 let view_spec_for_snapshot = view_spec.clone();
1194 let (mut receiver, initial_membership) =
1195 subscribe_list_then_snapshot(&context.bus_manager, &source_view_id, move || async move {
1196 load_query_entities(
1197 &cache,
1198 sorted_caches,
1199 &view_spec_for_snapshot,
1200 &query,
1201 false,
1202 )
1203 .await
1204 })
1205 .await;
1206
1207 let mut snapshot_entities = initial_membership.clone();
1208 if let Some(limit) = subscription.query.snapshot_limit {
1209 snapshot_entities.truncate(limit);
1210 }
1211 enforce_snapshot_limit(context, snapshot_entities.len())?;
1212 send_subscribed_frame(context, &subscription, &view_spec)?;
1213 if subscription.snapshot.enabled {
1214 send_snapshot_batches(
1215 context,
1216 &subscription,
1217 &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1218 view_spec.mode,
1219 &context.entity_cache.snapshot_config(),
1220 )
1221 .await?;
1222 }
1223
1224 let task_context = context.clone();
1225 let subscription_id = subscription.subscription_id.clone();
1226 let query = subscription.query.clone();
1227 let view_spec_task = view_spec.clone();
1228 let span_view = view_id.clone();
1229 tokio::spawn(
1230 async move {
1231 let mut current = initial_membership;
1232 loop {
1233 tokio::select! {
1234 _ = cancel_token.cancelled() => break,
1235 received = receiver.recv() => {
1236 let envelope = match received {
1237 Ok(envelope) => envelope,
1238 Err(broadcast::error::RecvError::Lagged(_)) => {
1239 warn!("Subscription {} lagged; closing to preserve membership correctness", subscription_id);
1240 break;
1241 }
1242 Err(broadcast::error::RecvError::Closed) => break,
1243 };
1244 let metadata = source_frame_metadata(&envelope.payload);
1245 if metadata.op == "delete" {
1246 task_context.entity_cache.remove(&source_view_id, &envelope.key).await;
1247 if view_spec_task.is_derived() {
1248 let caches = task_context.view_index.sorted_caches();
1249 let mut guard = caches.write().await;
1250 if let Some(cache) = guard.get_mut(&query.view) {
1251 cache.remove(&envelope.key);
1252 }
1253 }
1254 }
1255
1256 let sorted_caches = view_spec_task
1257 .is_derived()
1258 .then(|| task_context.view_index.sorted_caches());
1259 let next = load_query_entities(
1260 &task_context.entity_cache,
1261 sorted_caches,
1262 &view_spec_task,
1263 &query,
1264 false,
1265 ).await;
1266 if emit_collection_delta(
1267 &task_context,
1268 &subscription_id,
1269 &view_spec_task,
1270 ¤t,
1271 &next,
1272 &envelope,
1273 &metadata,
1274 ).is_err() {
1275 break;
1276 }
1277 current = next;
1278 }
1279 }
1280 }
1281 }
1282 .instrument(info_span!("ws.subscribe.collection", client_id = %context.client_id, view = %span_view)),
1283 );
1284 Ok(())
1285}
1286
1287#[derive(Default)]
1288struct SourceFrameMetadata {
1289 op: String,
1290 seq: Option<String>,
1291}
1292
1293fn source_frame_metadata(payload: &[u8]) -> SourceFrameMetadata {
1294 serde_json::from_slice::<Value>(payload)
1295 .ok()
1296 .map(|value| SourceFrameMetadata {
1297 op: value
1298 .get("op")
1299 .and_then(Value::as_str)
1300 .unwrap_or_default()
1301 .to_string(),
1302 seq: value.get("seq").and_then(Value::as_str).map(str::to_string),
1303 })
1304 .unwrap_or_default()
1305}
1306
1307fn send_scoped_source_payload(
1308 context: &SubscriptionContext,
1309 subscription_id: &str,
1310 view_id: &str,
1311 payload: Arc<Bytes>,
1312) -> Result<()> {
1313 let mut value: Value = serde_json::from_slice(&payload)?;
1314 let object = value
1315 .as_object_mut()
1316 .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
1317 object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
1318 object.insert(
1319 "subscriptionId".to_string(),
1320 Value::String(subscription_id.to_string()),
1321 );
1322 let encoded = Arc::new(Bytes::from(serde_json::to_vec(&value)?));
1323 let bytes = encoded.len();
1324 context
1325 .client_manager
1326 .send_to_client(context.client_id, encoded)
1327 .map_err(|error| anyhow::anyhow!("failed to send live frame: {error}"))?;
1328 context.metrics.message_sent();
1329 emit_update_sent_for_client(
1330 &context.usage_emitter,
1331 &context.client_manager,
1332 context.client_id,
1333 view_id,
1334 bytes,
1335 );
1336 Ok(())
1337}
1338
1339fn send_membership_frame(
1340 context: &SubscriptionContext,
1341 subscription_id: &str,
1342 view_spec: &ViewSpec,
1343 op: &str,
1344 key: &str,
1345 mut data: Value,
1346 seq: Option<String>,
1347) -> Result<()> {
1348 apply_wire_format(&mut data, &view_spec.wire_format);
1349 let frame = Frame::scoped(
1350 subscription_id,
1351 view_spec.mode,
1352 &view_spec.id,
1353 op,
1354 key,
1355 data,
1356 seq,
1357 );
1358 let encoded = Arc::new(Bytes::from(serde_json::to_vec(&frame)?));
1359 let bytes = encoded.len();
1360 context
1361 .client_manager
1362 .send_to_client(context.client_id, encoded)
1363 .map_err(|error| anyhow::anyhow!("failed to send membership frame: {error}"))?;
1364 context.metrics.message_sent();
1365 emit_update_sent_for_client(
1366 &context.usage_emitter,
1367 &context.client_manager,
1368 context.client_id,
1369 &view_spec.id,
1370 bytes,
1371 );
1372 Ok(())
1373}
1374
1375fn emit_collection_delta(
1376 context: &SubscriptionContext,
1377 subscription_id: &str,
1378 view_spec: &ViewSpec,
1379 current: &[(String, Value)],
1380 next: &[(String, Value)],
1381 envelope: &BusMessage,
1382 metadata: &SourceFrameMetadata,
1383) -> Result<()> {
1384 let current_keys: Vec<&str> = current.iter().map(|(key, _)| key.as_str()).collect();
1385 let next_keys: Vec<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
1386 let next_set: HashSet<&str> = next_keys.iter().copied().collect();
1387
1388 for key in current_keys
1389 .iter()
1390 .copied()
1391 .filter(|key| !next_set.contains(key))
1392 {
1393 let op = if metadata.op == "delete" && key == envelope.key {
1394 "delete"
1395 } else {
1396 "remove"
1397 };
1398 send_membership_frame(
1399 context,
1400 subscription_id,
1401 view_spec,
1402 op,
1403 key,
1404 Value::Null,
1405 metadata.seq.clone(),
1406 )?;
1407 }
1408
1409 for (position, (key, data)) in next.iter().enumerate() {
1410 let previous_position = current_keys.iter().position(|candidate| *candidate == key);
1411 let changed_position = previous_position != Some(position);
1412 if previous_position.is_none() || changed_position || key == &envelope.key {
1413 let can_forward_patch = !view_spec.is_derived()
1414 && previous_position == Some(position)
1415 && key == &envelope.key
1416 && metadata.op != "delete";
1417 if can_forward_patch {
1418 send_scoped_source_payload(
1419 context,
1420 subscription_id,
1421 &view_spec.id,
1422 envelope.payload.clone(),
1423 )?;
1424 } else {
1425 let seq = metadata
1426 .seq
1427 .clone()
1428 .or_else(|| data.get("_seq").and_then(Value::as_str).map(str::to_string));
1429 send_membership_frame(
1430 context,
1431 subscription_id,
1432 view_spec,
1433 "upsert",
1434 key,
1435 data.clone(),
1436 seq,
1437 )?;
1438 }
1439 }
1440 }
1441 Ok(())
1442}
1443
1444fn to_wire_snapshot_entities(
1445 entities: Vec<(String, Value)>,
1446 view_spec: &ViewSpec,
1447) -> Vec<SnapshotEntity> {
1448 entities
1449 .into_iter()
1450 .map(|(key, mut data)| {
1451 apply_wire_format(&mut data, &view_spec.wire_format);
1452 SnapshotEntity { key, data }
1453 })
1454 .collect()
1455}
1456
1457async fn load_query_entities(
1458 entity_cache: &EntityCache,
1459 sorted_caches: Option<
1460 Arc<tokio::sync::RwLock<HashMap<String, crate::sorted_cache::SortedViewCache>>>,
1461 >,
1462 view_spec: &ViewSpec,
1463 query: &SubscriptionQuery,
1464 apply_snapshot_limit: bool,
1465) -> Vec<(String, Value)> {
1466 let (entities, preordered) = if let Some(sorted_caches) = sorted_caches {
1467 let mut caches = sorted_caches.write().await;
1468 let entities = caches
1469 .get_mut(&view_spec.id)
1470 .map(|cache| cache.get_all_ordered())
1471 .unwrap_or_default();
1472 (entities, true)
1473 } else if view_spec.mode == Mode::State {
1474 let entity = match query.key.as_deref() {
1475 Some(key) => entity_cache
1476 .get(&view_spec.id, key)
1477 .await
1478 .map(|data| vec![(key.to_string(), data)])
1479 .unwrap_or_default(),
1480 None => vec![],
1481 };
1482 (entity, true)
1483 } else {
1484 (entity_cache.get_all(&view_spec.id).await, false)
1485 };
1486 select_query_entities(entities, query, preordered, apply_snapshot_limit)
1487}
1488
1489fn select_query_entities(
1490 mut entities: Vec<(String, Value)>,
1491 query: &SubscriptionQuery,
1492 preordered: bool,
1493 apply_snapshot_limit: bool,
1494) -> Vec<(String, Value)> {
1495 entities.retain(|(key, data)| query_matches_entity(query, key, data));
1496 if !preordered {
1497 entities.sort_by(|left, right| {
1498 let left_seq = left.1.get("_seq").and_then(Value::as_str).unwrap_or("");
1499 let right_seq = right.1.get("_seq").and_then(Value::as_str).unwrap_or("");
1500 let order = if query.after.is_some() {
1501 cmp_seq(left_seq, right_seq)
1502 } else {
1503 cmp_seq(right_seq, left_seq)
1504 };
1505 order.then_with(|| left.0.cmp(&right.0))
1506 });
1507 }
1508
1509 let skip = query.skip.unwrap_or(0);
1510 let take = query.take.unwrap_or(usize::MAX);
1511 let mut selected: Vec<_> = entities.into_iter().skip(skip).take(take).collect();
1512 if apply_snapshot_limit {
1513 if let Some(limit) = query.snapshot_limit {
1514 selected.truncate(limit);
1515 }
1516 }
1517 selected
1518}
1519
1520fn query_matches_entity(query: &SubscriptionQuery, key: &str, data: &Value) -> bool {
1521 if !query.matches_key(key) {
1522 return false;
1523 }
1524 if let Some(partition) = &query.partition {
1525 if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
1526 return false;
1527 }
1528 }
1529 if let Some(after) = &query.after {
1530 let Some(seq) = data.get("_seq").and_then(Value::as_str) else {
1531 return false;
1532 };
1533 if cmp_seq(seq, after) != std::cmp::Ordering::Greater {
1534 return false;
1535 }
1536 }
1537 query
1538 .filters
1539 .iter()
1540 .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
1541}
1542
1543fn value_at_dot_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
1544 path.split('.')
1545 .try_fold(value, |current, segment| current.get(segment))
1546}
1547
1548#[cfg(test)]
1549mod tests {
1550 use super::*;
1551 use crate::cache::EntityCacheConfig;
1552 use crate::view::{Delivery, Filters, Projection};
1553 use serde_json::json;
1554 use tokio::sync::oneshot;
1555
1556 fn list_spec() -> ViewSpec {
1557 ViewSpec {
1558 id: "Thing/list".to_string(),
1559 export: "Thing".to_string(),
1560 mode: Mode::List,
1561 wire_format: Default::default(),
1562 projection: Projection::all(),
1563 filters: Filters::all(),
1564 delivery: Delivery::default(),
1565 pipeline: None,
1566 source_view: None,
1567 }
1568 }
1569
1570 #[test]
1571 fn snapshot_batches_share_identity_and_completion() {
1572 let entities = ["one", "two", "three"].map(|key| SnapshotEntity {
1573 key: key.to_string(),
1574 data: json!({"key": key}),
1575 });
1576 let batches = create_snapshot_batches(
1577 &entities,
1578 SnapshotMetadata {
1579 subscription_id: "sub-1",
1580 snapshot_id: "snapshot-1",
1581 authoritative: true,
1582 mode: Mode::List,
1583 view_id: "Thing/list",
1584 key: None,
1585 },
1586 &SnapshotBatchConfig {
1587 initial_batch_size: 2,
1588 subsequent_batch_size: 1,
1589 },
1590 );
1591 assert_eq!(batches.len(), 2);
1592 assert!(batches.iter().all(|batch| batch.subscription_id == "sub-1"));
1593 assert!(batches
1594 .iter()
1595 .all(|batch| batch.snapshot_id == "snapshot-1"));
1596 assert!(!batches[0].complete);
1597 assert!(batches[1].complete);
1598 assert!(batches.iter().all(|batch| batch.authoritative));
1599 }
1600
1601 #[test]
1602 fn empty_incremental_snapshot_is_explicitly_non_authoritative() {
1603 let batches = create_snapshot_batches(
1604 &[],
1605 SnapshotMetadata {
1606 subscription_id: "sub-1",
1607 snapshot_id: "snapshot-1",
1608 authoritative: false,
1609 mode: Mode::State,
1610 view_id: "Thing/state",
1611 key: Some("missing"),
1612 },
1613 &SnapshotBatchConfig {
1614 initial_batch_size: 1,
1615 subsequent_batch_size: 1,
1616 },
1617 );
1618 assert_eq!(batches.len(), 1);
1619 assert!(!batches[0].authoritative);
1620 assert!(batches[0].complete);
1621 assert_eq!(batches[0].key.as_deref(), Some("missing"));
1622 }
1623
1624 #[test]
1625 fn dot_path_filters_are_exact_and_type_sensitive() {
1626 let mut query = SubscriptionQuery {
1627 view: "Thing/list".to_string(),
1628 ..Default::default()
1629 };
1630 query
1631 .filters
1632 .insert("state.status".to_string(), json!("open"));
1633 query.filters.insert("metrics.count".to_string(), json!(2));
1634 assert!(query_matches_entity(
1635 &query,
1636 "one",
1637 &json!({"state": {"status": "open"}, "metrics": {"count": 2}}),
1638 ));
1639 assert!(!query_matches_entity(
1640 &query,
1641 "one",
1642 &json!({"state": {"status": "open"}, "metrics": {"count": "2"}}),
1643 ));
1644 }
1645
1646 #[test]
1647 fn take_and_skip_define_independent_deterministic_windows() {
1648 let entities: Vec<_> = (1..=6)
1649 .map(|id| {
1650 (
1651 id.to_string(),
1652 json!({"id": id, "_seq": format!("10:{id:012}")}),
1653 )
1654 })
1655 .collect();
1656 let first = SubscriptionQuery {
1657 view: "Thing/list".to_string(),
1658 take: Some(2),
1659 skip: Some(0),
1660 ..Default::default()
1661 };
1662 let second = SubscriptionQuery {
1663 skip: Some(2),
1664 ..first.clone()
1665 };
1666 let first_keys: Vec<_> = select_query_entities(entities.clone(), &first, false, false)
1667 .into_iter()
1668 .map(|(key, _)| key)
1669 .collect();
1670 let second_keys: Vec<_> = select_query_entities(entities, &second, false, false)
1671 .into_iter()
1672 .map(|(key, _)| key)
1673 .collect();
1674 assert_eq!(first_keys, ["6", "5"]);
1675 assert_eq!(second_keys, ["4", "3"]);
1676 }
1677
1678 #[tokio::test]
1679 async fn state_receiver_is_installed_before_snapshot_awaits() {
1680 let bus = BusManager::new();
1681 let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
1682 let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
1683 let bus_for_task = bus.clone();
1684 let task = tokio::spawn(async move {
1685 subscribe_state_then_snapshot(&bus_for_task, "Thing/state", "one", || async move {
1686 snapshot_started_tx.send(()).unwrap();
1687 release_snapshot_rx.await.unwrap();
1688 })
1689 .await
1690 .0
1691 });
1692 snapshot_started_rx.await.unwrap();
1693 bus.publish_state(
1694 "Thing/state",
1695 "one",
1696 Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
1697 )
1698 .await;
1699 release_snapshot_tx.send(()).unwrap();
1700 let mut receiver = task.await.unwrap();
1701 receiver.changed().await.unwrap();
1702 assert!(!receiver.borrow().is_empty());
1703 }
1704
1705 async fn assert_list_receiver_precedes_snapshot(view: &'static str) {
1706 let bus = BusManager::new();
1707 let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
1708 let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
1709 let bus_for_task = bus.clone();
1710 let task = tokio::spawn(async move {
1711 subscribe_list_then_snapshot(&bus_for_task, view, || async move {
1712 snapshot_started_tx.send(()).unwrap();
1713 release_snapshot_rx.await.unwrap();
1714 })
1715 .await
1716 .0
1717 });
1718 snapshot_started_rx.await.unwrap();
1719 bus.publish_list(
1720 view,
1721 Arc::new(BusMessage {
1722 key: "one".to_string(),
1723 entity: view.to_string(),
1724 payload: Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
1725 }),
1726 )
1727 .await;
1728 release_snapshot_tx.send(()).unwrap();
1729 let mut receiver = task.await.unwrap();
1730 assert_eq!(receiver.recv().await.unwrap().key, "one");
1731 }
1732
1733 #[tokio::test]
1734 async fn list_receiver_is_installed_before_snapshot() {
1735 assert_list_receiver_precedes_snapshot("Thing/list").await;
1736 }
1737
1738 #[tokio::test]
1739 async fn append_receiver_is_installed_before_snapshot() {
1740 assert_list_receiver_precedes_snapshot("Thing/append").await;
1741 }
1742
1743 #[tokio::test]
1744 async fn derived_source_receiver_is_installed_before_snapshot() {
1745 assert_list_receiver_precedes_snapshot("Thing/list-source").await;
1746 }
1747
1748 #[tokio::test]
1749 async fn snapshot_limit_does_not_change_live_take_skip_membership() {
1750 let cache = EntityCache::with_config(EntityCacheConfig {
1751 max_entities_per_view: 10,
1752 ..Default::default()
1753 });
1754 for id in 1..=4 {
1755 cache
1756 .upsert(
1757 "Thing/list",
1758 &id.to_string(),
1759 json!({"_seq": format!("10:{id:012}")}),
1760 )
1761 .await;
1762 }
1763 let query = SubscriptionQuery {
1764 view: "Thing/list".to_string(),
1765 take: Some(3),
1766 skip: Some(1),
1767 snapshot_limit: Some(1),
1768 ..Default::default()
1769 };
1770 let live = load_query_entities(&cache, None, &list_spec(), &query, false).await;
1771 let snapshot = load_query_entities(&cache, None, &list_spec(), &query, true).await;
1772 assert_eq!(live.len(), 3);
1773 assert_eq!(snapshot.len(), 1);
1774 }
1775
1776 #[test]
1777 fn fixture_manifest_covers_required_conformance_cases() {
1778 let manifest: Value = serde_json::from_str(include_str!(
1779 "../../../../tests/fixtures/websocket-v2/manifest.json"
1780 ))
1781 .unwrap();
1782 let names: HashSet<_> = manifest["fixtures"]
1783 .as_array()
1784 .unwrap()
1785 .iter()
1786 .filter_map(Value::as_str)
1787 .collect();
1788 for required in [
1789 "keyed-state.json",
1790 "list-windows.json",
1791 "filters.json",
1792 "multi-batch-authoritative.json",
1793 "empty-snapshot.json",
1794 "remove.json",
1795 "delete.json",
1796 "incremental-snapshot.json",
1797 "reconnect-replacement.json",
1798 "errors.json",
1799 ] {
1800 assert!(names.contains(required), "missing fixture {required}");
1801 }
1802
1803 for document in [
1804 include_str!("../../../../tests/fixtures/websocket-v2/keyed-state.json"),
1805 include_str!("../../../../tests/fixtures/websocket-v2/list-windows.json"),
1806 include_str!("../../../../tests/fixtures/websocket-v2/filters.json"),
1807 include_str!("../../../../tests/fixtures/websocket-v2/multi-batch-authoritative.json"),
1808 include_str!("../../../../tests/fixtures/websocket-v2/empty-snapshot.json"),
1809 include_str!("../../../../tests/fixtures/websocket-v2/remove.json"),
1810 include_str!("../../../../tests/fixtures/websocket-v2/delete.json"),
1811 include_str!("../../../../tests/fixtures/websocket-v2/incremental-snapshot.json"),
1812 include_str!("../../../../tests/fixtures/websocket-v2/reconnect-replacement.json"),
1813 include_str!("../../../../tests/fixtures/websocket-v2/errors.json"),
1814 ] {
1815 let fixture: Value = serde_json::from_str(document).unwrap();
1816 assert!(fixture["name"].is_string());
1817 }
1818 }
1819}