change_stream.rs (18944B)
1 use std::collections::BTreeMap; 2 use std::num::{NonZeroU64, NonZeroUsize}; 3 4 use tokio::sync::{mpsc, watch}; 5 6 use crate::{AppSnapshot, SnapshotRevision}; 7 8 #[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)] 9 pub struct ChangeSubscriptionId(NonZeroU64); 10 11 impl ChangeSubscriptionId { 12 #[must_use] 13 pub const fn value(self) -> u64 { 14 self.0.get() 15 } 16 } 17 18 #[derive(Clone, Debug, Eq, PartialEq)] 19 pub struct SnapshotChange { 20 snapshot: AppSnapshot, 21 previous_revision: Option<SnapshotRevision>, 22 } 23 24 impl SnapshotChange { 25 #[must_use] 26 pub const fn revision(&self) -> SnapshotRevision { 27 self.snapshot.revision() 28 } 29 30 #[must_use] 31 pub const fn snapshot(&self) -> &AppSnapshot { 32 &self.snapshot 33 } 34 35 #[must_use] 36 pub fn into_snapshot(self) -> AppSnapshot { 37 self.snapshot 38 } 39 40 #[must_use] 41 pub const fn previous_revision(&self) -> Option<SnapshotRevision> { 42 self.previous_revision 43 } 44 45 #[must_use] 46 pub fn recovers_gap_after(&self, observed: SnapshotRevision) -> bool { 47 self.previous_revision 48 .is_some_and(|previous| previous != observed) 49 } 50 } 51 52 pub struct SnapshotChangeReceiver { 53 receiver: mpsc::Receiver<SnapshotChange>, 54 latest: watch::Receiver<SnapshotChange>, 55 last_delivered_revision: Option<SnapshotRevision>, 56 } 57 58 impl SnapshotChangeReceiver { 59 pub async fn receive(&mut self) -> Option<SnapshotChange> { 60 loop { 61 let queue_open = match self.receiver.try_recv() { 62 Ok(change) => { 63 if let Some(change) = self.deliver_newer(change) { 64 return Some(change); 65 } 66 continue; 67 } 68 Err(mpsc::error::TryRecvError::Empty) => true, 69 Err(mpsc::error::TryRecvError::Disconnected) => false, 70 }; 71 let retained = self.latest.borrow().clone(); 72 if let Some(change) = self.deliver_newer(retained) { 73 return Some(change); 74 } 75 if !queue_open { 76 return None; 77 } 78 if let Some(change) = self.receiver.recv().await 79 && let Some(change) = self.deliver_newer(change) 80 { 81 return Some(change); 82 } 83 } 84 } 85 86 fn deliver_newer(&mut self, change: SnapshotChange) -> Option<SnapshotChange> { 87 let revision = change.revision(); 88 if self 89 .last_delivered_revision 90 .is_some_and(|delivered| revision <= delivered) 91 { 92 return None; 93 } 94 self.last_delivered_revision = Some(revision); 95 Some(change) 96 } 97 } 98 99 struct SnapshotChangeSubscriber { 100 sender: mpsc::Sender<SnapshotChange>, 101 latest: watch::Sender<SnapshotChange>, 102 } 103 104 pub struct OrderedSnapshotChanges { 105 latest: AppSnapshot, 106 next_subscription: u64, 107 subscribers: BTreeMap<ChangeSubscriptionId, SnapshotChangeSubscriber>, 108 closed: bool, 109 } 110 111 impl OrderedSnapshotChanges { 112 #[must_use] 113 pub fn new(initial_snapshot: AppSnapshot) -> Self { 114 Self { 115 latest: initial_snapshot, 116 next_subscription: 1, 117 subscribers: BTreeMap::new(), 118 closed: false, 119 } 120 } 121 122 #[must_use] 123 pub const fn last_revision(&self) -> SnapshotRevision { 124 self.latest.revision() 125 } 126 127 /// Registers a bounded consumer for future changes. 128 /// 129 /// # Errors 130 /// 131 /// Returns `None` if the subscription identifier space is exhausted. 132 pub fn subscribe( 133 &mut self, 134 capacity: NonZeroUsize, 135 ) -> Option<(ChangeSubscriptionId, SnapshotChangeReceiver)> { 136 if self.closed { 137 return None; 138 } 139 let id = ChangeSubscriptionId(NonZeroU64::new(self.next_subscription)?); 140 self.next_subscription = self.next_subscription.checked_add(1)?; 141 let initial = SnapshotChange { 142 snapshot: self.latest.clone(), 143 previous_revision: None, 144 }; 145 let (sender, receiver) = mpsc::channel(capacity.get()); 146 let (latest, retained) = watch::channel(initial.clone()); 147 sender.try_send(initial).ok()?; 148 self.subscribers 149 .insert(id, SnapshotChangeSubscriber { sender, latest }); 150 Some(( 151 id, 152 SnapshotChangeReceiver { 153 receiver, 154 latest: retained, 155 last_delivered_revision: None, 156 }, 157 )) 158 } 159 160 #[must_use] 161 pub fn unsubscribe(&mut self, id: ChangeSubscriptionId) -> bool { 162 self.subscribers.remove(&id).is_some() 163 } 164 165 pub fn publish(&mut self, snapshot: AppSnapshot) { 166 self.subscribers 167 .retain(|_, subscriber| !subscriber.sender.is_closed()); 168 if self.closed || snapshot.revision() <= self.latest.revision() { 169 return; 170 } 171 let change = SnapshotChange { 172 previous_revision: Some(self.latest.revision()), 173 snapshot, 174 }; 175 self.latest = change.snapshot.clone(); 176 self.subscribers.retain(|_, subscriber| { 177 // Retain before enqueue so saturation and sender closure preserve the tail. 178 drop(subscriber.latest.send_replace(change.clone())); 179 match subscriber.sender.try_send(change.clone()) { 180 Ok(()) | Err(mpsc::error::TrySendError::Full(_)) => true, 181 Err(mpsc::error::TrySendError::Closed(_)) => false, 182 } 183 }); 184 } 185 186 pub fn close(&mut self) { 187 self.closed = true; 188 self.subscribers.clear(); 189 } 190 } 191 192 #[cfg(test)] 193 mod tests { 194 use std::future::Future; 195 use std::num::NonZeroUsize; 196 use std::task::Poll; 197 use std::time::Duration; 198 199 use crate::{ 200 AppSnapshot, OrderedSnapshotChanges, RelayConfiguration, SessionState, SnapshotChange, 201 SnapshotChangeReceiver, SnapshotRevision, 202 }; 203 204 const RECEIVE_TIMEOUT: Duration = Duration::from_secs(1); 205 206 #[tokio::test] 207 async fn change_stream_publishes_monotonic_revisions_to_multiple_consumers() { 208 let mut changes = OrderedSnapshotChanges::new(snapshot(0)); 209 let (_, mut first) = changes 210 .subscribe(NonZeroUsize::new(4).expect("capacity")) 211 .expect("first subscription"); 212 let (_, mut second) = changes 213 .subscribe(NonZeroUsize::new(4).expect("capacity")) 214 .expect("second subscription"); 215 216 assert_eq!( 217 first.receive().await.expect("initial").revision(), 218 revision(0) 219 ); 220 assert_eq!( 221 second.receive().await.expect("initial").revision(), 222 revision(0) 223 ); 224 changes.publish(snapshot(1)); 225 changes.publish(snapshot(1)); 226 changes.publish(snapshot(2)); 227 228 for receiver in [&mut first, &mut second] { 229 assert_eq!( 230 receiver.receive().await.expect("revision 1").revision(), 231 revision(1) 232 ); 233 assert_eq!( 234 receiver.receive().await.expect("revision 2").revision(), 235 revision(2) 236 ); 237 } 238 assert_eq!(changes.last_revision(), revision(2)); 239 } 240 241 #[tokio::test] 242 async fn slow_consumers_expose_a_revision_gap_without_blocking_publication() { 243 let mut changes = OrderedSnapshotChanges::new(snapshot(0)); 244 let (_, mut receiver) = changes 245 .subscribe(NonZeroUsize::new(1).expect("capacity")) 246 .expect("subscription"); 247 248 assert_eq!( 249 receiver.receive().await.expect("initial").revision(), 250 revision(0) 251 ); 252 changes.publish(snapshot(1)); 253 changes.publish(snapshot(2)); 254 let first = receiver.receive().await.expect("first"); 255 assert_eq!(first.revision(), revision(1)); 256 changes.publish(snapshot(3)); 257 let recovered = receiver.receive().await.expect("gap recovery"); 258 assert_eq!(recovered.revision(), revision(3)); 259 assert!(recovered.recovers_gap_after(first.revision())); 260 } 261 262 #[tokio::test] 263 async fn close_terminates_consumers_and_rejects_later_subscriptions() { 264 let mut changes = OrderedSnapshotChanges::new(snapshot(0)); 265 let (_, mut receiver) = changes 266 .subscribe(NonZeroUsize::new(1).expect("capacity")) 267 .expect("subscription"); 268 receiver.receive().await.expect("initial"); 269 270 changes.close(); 271 changes.publish(snapshot(1)); 272 assert!(receiver.receive().await.is_none()); 273 assert!( 274 changes 275 .subscribe(NonZeroUsize::new(1).expect("capacity")) 276 .is_none() 277 ); 278 } 279 280 #[tokio::test] 281 async fn capacity_one_recovers_final_tail_without_later_publication() { 282 let mut changes = OrderedSnapshotChanges::new(snapshot(0)); 283 let (_, mut receiver) = changes 284 .subscribe(NonZeroUsize::new(1).expect("capacity")) 285 .expect("subscription"); 286 let initial = next_change(&mut receiver, "initial").await; 287 assert_eq!(initial.revision(), revision(0)); 288 289 changes.publish(snapshot(1)); 290 changes.publish(snapshot(2)); 291 let queued = next_change(&mut receiver, "queued revision 1").await; 292 assert_eq!(queued.revision(), revision(1)); 293 assert_eq!(queued.previous_revision(), Some(revision(0))); 294 let tail = next_change(&mut receiver, "final revision 2 without another publish").await; 295 assert_eq!(tail.revision(), revision(2)); 296 assert_eq!(tail.snapshot(), &snapshot(2)); 297 assert_eq!(tail.previous_revision(), Some(revision(1))); 298 assert!(!tail.recovers_gap_after(queued.revision())); 299 assert_eq!(changes.last_revision(), revision(2)); 300 } 301 302 #[tokio::test] 303 async fn sustained_overflow_coalesces_latest_tail_with_original_gap_metadata() { 304 let mut changes = OrderedSnapshotChanges::new(snapshot(0)); 305 let (_, mut receiver) = changes 306 .subscribe(NonZeroUsize::new(1).expect("capacity")) 307 .expect("subscription"); 308 next_change(&mut receiver, "initial").await; 309 310 for value in 1..=128 { 311 changes.publish(snapshot(value)); 312 } 313 assert_eq!(changes.last_revision(), revision(128)); 314 let queued = next_change(&mut receiver, "queued revision 1").await; 315 assert_eq!(queued.revision(), revision(1)); 316 let latest = next_change(&mut receiver, "coalesced revision 128").await; 317 assert_eq!(latest.revision(), revision(128)); 318 assert_eq!(latest.snapshot(), &snapshot(128)); 319 assert_eq!(latest.previous_revision(), Some(revision(127))); 320 assert!(latest.recovers_gap_after(queued.revision())); 321 322 changes.publish(snapshot(128)); 323 changes.publish(snapshot(127)); 324 changes.close(); 325 assert!(next_delivery(&mut receiver).await.is_none()); 326 } 327 328 #[tokio::test] 329 async fn queued_initial_snapshot_precedes_coalesced_latest_tail() { 330 let mut changes = OrderedSnapshotChanges::new(snapshot(0)); 331 let (_, mut receiver) = changes 332 .subscribe(NonZeroUsize::new(1).expect("capacity")) 333 .expect("subscription"); 334 changes.publish(snapshot(1)); 335 changes.publish(snapshot(2)); 336 337 let initial = next_change(&mut receiver, "queued initial snapshot").await; 338 assert_eq!(initial.revision(), revision(0)); 339 assert_eq!(initial.previous_revision(), None); 340 let latest = next_change(&mut receiver, "latest after queued initial").await; 341 assert_eq!(latest.revision(), revision(2)); 342 assert_eq!(latest.previous_revision(), Some(revision(1))); 343 assert!(latest.recovers_gap_after(initial.revision())); 344 changes.close(); 345 assert!(next_delivery(&mut receiver).await.is_none()); 346 } 347 348 #[tokio::test] 349 async fn independently_paced_consumers_recover_their_own_latest_tail() { 350 let mut changes = OrderedSnapshotChanges::new(snapshot(0)); 351 let (_, mut slow) = changes 352 .subscribe(NonZeroUsize::new(1).expect("capacity")) 353 .expect("slow subscription"); 354 let (_, mut fast) = changes 355 .subscribe(NonZeroUsize::new(1).expect("capacity")) 356 .expect("fast subscription"); 357 next_change(&mut slow, "slow initial").await; 358 next_change(&mut fast, "fast initial").await; 359 360 for value in 1..=3 { 361 changes.publish(snapshot(value)); 362 let delivered = next_change(&mut fast, "fast consumer revision").await; 363 assert_eq!(delivered.revision(), revision(value)); 364 assert_eq!(delivered.previous_revision(), Some(revision(value - 1))); 365 assert!(!delivered.recovers_gap_after(revision(value - 1))); 366 } 367 let queued = next_change(&mut slow, "slow queued revision 1").await; 368 assert_eq!(queued.revision(), revision(1)); 369 let tail = next_change(&mut slow, "slow retained revision 3").await; 370 assert_eq!(tail.revision(), revision(3)); 371 assert_eq!(tail.previous_revision(), Some(revision(2))); 372 assert!(tail.recovers_gap_after(queued.revision())); 373 changes.close(); 374 assert!(next_delivery(&mut slow).await.is_none()); 375 assert!(next_delivery(&mut fast).await.is_none()); 376 } 377 378 #[tokio::test] 379 async fn full_close_drains_queued_change_and_final_tail_before_none() { 380 let mut changes = OrderedSnapshotChanges::new(snapshot(0)); 381 let (_, mut receiver) = changes 382 .subscribe(NonZeroUsize::new(1).expect("capacity")) 383 .expect("subscription"); 384 next_change(&mut receiver, "initial").await; 385 changes.publish(snapshot(1)); 386 changes.publish(snapshot(2)); 387 changes.close(); 388 changes.publish(snapshot(3)); 389 390 assert_eq!( 391 next_change(&mut receiver, "queued revision before close") 392 .await 393 .revision(), 394 revision(1) 395 ); 396 let tail = next_change(&mut receiver, "retained final revision before close").await; 397 assert_eq!(tail.revision(), revision(2)); 398 assert_eq!(tail.previous_revision(), Some(revision(1))); 399 assert!(next_delivery(&mut receiver).await.is_none()); 400 assert!(next_delivery(&mut receiver).await.is_none()); 401 assert_eq!(changes.last_revision(), revision(2)); 402 } 403 404 #[tokio::test] 405 async fn unsubscribe_drains_retained_tail_without_future_updates() { 406 let mut changes = OrderedSnapshotChanges::new(snapshot(0)); 407 let (id, mut receiver) = changes 408 .subscribe(NonZeroUsize::new(1).expect("capacity")) 409 .expect("subscription"); 410 next_change(&mut receiver, "initial").await; 411 changes.publish(snapshot(1)); 412 changes.publish(snapshot(2)); 413 assert!(changes.unsubscribe(id)); 414 assert!(!changes.unsubscribe(id)); 415 changes.publish(snapshot(3)); 416 417 assert_eq!( 418 next_change(&mut receiver, "queued revision before unsubscribe") 419 .await 420 .revision(), 421 revision(1) 422 ); 423 let tail = next_change(&mut receiver, "retained revision before unsubscribe").await; 424 assert_eq!(tail.revision(), revision(2)); 425 assert_eq!(tail.previous_revision(), Some(revision(1))); 426 assert!(next_delivery(&mut receiver).await.is_none()); 427 assert_eq!(changes.last_revision(), revision(3)); 428 } 429 430 #[tokio::test] 431 async fn cancelled_pending_receive_preserves_future_queued_and_retained_tail() { 432 let mut changes = OrderedSnapshotChanges::new(snapshot(0)); 433 let (_, mut receiver) = changes 434 .subscribe(NonZeroUsize::new(1).expect("capacity")) 435 .expect("subscription"); 436 next_change(&mut receiver, "initial").await; 437 { 438 let pending = receiver.receive(); 439 tokio::pin!(pending); 440 tokio::time::timeout( 441 RECEIVE_TIMEOUT, 442 std::future::poll_fn(|context| { 443 assert!(pending.as_mut().poll(context).is_pending()); 444 Poll::Ready(()) 445 }), 446 ) 447 .await 448 .expect("empty receive was polled before cancellation"); 449 } 450 451 changes.publish(snapshot(1)); 452 changes.publish(snapshot(2)); 453 assert_eq!( 454 next_change(&mut receiver, "queued revision after cancellation") 455 .await 456 .revision(), 457 revision(1) 458 ); 459 let tail = next_change(&mut receiver, "retained revision after cancellation").await; 460 assert_eq!(tail.revision(), revision(2)); 461 assert_eq!(tail.previous_revision(), Some(revision(1))); 462 changes.publish(snapshot(3)); 463 assert_eq!( 464 next_change(&mut receiver, "future revision after recovered tail") 465 .await 466 .revision(), 467 revision(3) 468 ); 469 changes.close(); 470 assert!(next_delivery(&mut receiver).await.is_none()); 471 } 472 473 #[tokio::test] 474 async fn unchanged_publish_reclaims_abandoned_receiver_registration() { 475 let mut changes = OrderedSnapshotChanges::new(snapshot(0)); 476 let (id, receiver) = changes.subscribe(NonZeroUsize::MIN).expect("subscription"); 477 drop(receiver); 478 479 changes.publish(snapshot(0)); 480 481 assert!( 482 !changes.unsubscribe(id), 483 "unchanged observation must reclaim a closed receiver" 484 ); 485 assert_eq!(changes.last_revision(), revision(0)); 486 let (_, mut replacement) = changes.subscribe(NonZeroUsize::MIN).expect("replacement"); 487 let initial = next_change(&mut replacement, "replacement initial snapshot").await; 488 assert_eq!(initial.revision(), revision(0)); 489 assert!(initial.previous_revision().is_none()); 490 changes.close(); 491 assert!(next_delivery(&mut replacement).await.is_none()); 492 } 493 494 async fn next_delivery(receiver: &mut SnapshotChangeReceiver) -> Option<SnapshotChange> { 495 tokio::time::timeout(RECEIVE_TIMEOUT, receiver.receive()) 496 .await 497 .expect("snapshot receive completed within its bound") 498 } 499 500 async fn next_change(receiver: &mut SnapshotChangeReceiver, label: &str) -> SnapshotChange { 501 next_delivery(receiver).await.expect(label) 502 } 503 504 fn revision(value: u64) -> SnapshotRevision { 505 SnapshotRevision::from_value(value) 506 } 507 508 fn snapshot(value: u64) -> AppSnapshot { 509 if value == 0 { 510 AppSnapshot::booting() 511 } else { 512 AppSnapshot::ready( 513 revision(value), 514 RelayConfiguration::default(), 515 Vec::new(), 516 None, 517 SessionState::SignedOut, 518 None, 519 None, 520 ) 521 .expect("snapshot") 522 } 523 } 524 }