1use liminal::channel::{PublishRecord, Recipient};
4use std::collections::HashMap;
5use std::num::NonZeroU32;
6use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
7use std::sync::{Arc, Mutex, OnceLock};
8use std::time::Instant;
9
10use super::error::{TapError, lock};
11use super::message::{
12 Hop, PayloadGrant, PublishFact, Reached, TapBatch, TapFilter, TapIdent, TapMessage, TapSource,
13};
14use super::ring::BoundedRing;
15
16pub const DEFAULT_TAP_CAPACITY: usize = 4096;
18
19pub const MAX_SAMPLE_ONE_IN: NonZeroU32 = match NonZeroU32::new(1 << 16) {
21 Some(most) => most,
22 None => NonZeroU32::MIN,
23};
24
25const TWO: NonZeroU32 = NonZeroU32::MIN.saturating_add(1);
26
27pub type TapWaker = Arc<dyn Fn() + Send + Sync>;
30
31#[derive(Debug)]
34struct Admitted {
35 one_in: NonZeroU32,
36 message: TapMessage,
37}
38
39#[derive(Debug)]
40struct SubscriberState {
41 ring: BoundedRing<Admitted>,
42 one_in: NonZeroU32,
43 seen: u64,
44 next_seq: u64,
45}
46
47impl SubscriberState {
48 fn admits(&mut self) -> bool {
51 let take = self.seen % u64::from(self.one_in.get()) == 0;
52 self.seen = self.seen.wrapping_add(1);
53 take
54 }
55
56 fn after_push(&mut self) {
58 if self.ring.len().saturating_mul(2) > self.ring.capacity() {
59 self.one_in = self.one_in.saturating_mul(TWO).min(MAX_SAMPLE_ONE_IN);
60 }
61 }
62}
63
64struct Subscriber {
65 id: u64,
66 filter: TapFilter,
67 grant: Option<PayloadGrant>,
68 state: Mutex<SubscriberState>,
69}
70
71impl std::fmt::Debug for Subscriber {
72 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73 formatter
74 .debug_struct("Subscriber")
75 .field("id", &self.id)
76 .field("filter", &self.filter)
77 .finish_non_exhaustive()
78 }
79}
80
81#[derive(Default)]
82struct TapShared {
83 active: AtomicUsize,
84 observed: AtomicU64,
85 #[cfg(test)]
88 observed_server: AtomicU64,
89 dropped_total: AtomicU64,
90 unsubscribe_failures: AtomicU64,
91 next_id: AtomicU64,
92 subscribers: Mutex<Vec<Arc<Subscriber>>>,
93 stations: Mutex<HashMap<u64, Arc<str>>>,
96 receivers: Mutex<HashMap<u64, u64>>,
99 buffers: Mutex<Vec<PublishRecord>>,
102 lent: AtomicU64,
104 no_stations: Arc<[Reached]>,
108 waker: OnceLock<TapWaker>,
109}
110
111impl std::fmt::Debug for TapShared {
112 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
113 formatter
114 .debug_struct("TapShared")
115 .field("active", &self.active)
116 .field("observed", &self.observed)
117 .field("dropped_total", &self.dropped_total)
118 .finish_non_exhaustive()
119 }
120}
121
122#[derive(Clone, Debug)]
124pub struct Tap {
125 shared: Arc<TapShared>,
126 node: Arc<str>,
127 depot: Option<Arc<str>>,
130 capacity: usize,
131}
132
133impl Tap {
134 #[must_use]
138 pub fn new(node: &str, capacity: usize) -> Self {
139 Self::build(node, capacity, OnceLock::new())
140 }
141
142 #[must_use]
146 pub fn with_waker(node: &str, capacity: usize, waker: TapWaker) -> Self {
147 Self::build(node, capacity, OnceLock::from(waker))
148 }
149
150 fn build(node: &str, capacity: usize, waker: OnceLock<TapWaker>) -> Self {
151 Self {
152 shared: Arc::new(TapShared {
153 waker,
154 ..TapShared::default()
155 }),
156 node: Arc::from(node),
157 depot: None,
158 capacity: capacity.max(1),
159 }
160 }
161
162 #[must_use]
166 pub fn with_depot(mut self, depot: bool) -> Self {
167 self.depot = depot.then(|| Arc::from(format!("{}.depot", self.node)));
168 self
169 }
170
171 #[must_use]
174 pub fn reserved(&self, id: &str) -> Option<&'static str> {
175 if id == &*self.node {
176 Some("exchange")
177 } else if self.depot.as_deref() == Some(id) {
178 Some("depot")
179 } else {
180 None
181 }
182 }
183
184 pub fn subscribe(&self, filter: TapFilter) -> Result<TapSubscription, TapError> {
189 self.subscribe_with(filter, None)
190 }
191
192 pub fn subscribe_with(
197 &self,
198 filter: TapFilter,
199 grant: Option<PayloadGrant>,
200 ) -> Result<TapSubscription, TapError> {
201 let id = self.shared.next_id.fetch_add(1, Ordering::Relaxed);
202 let subscriber = Arc::new(Subscriber {
203 id,
204 filter,
205 grant,
206 state: Mutex::new(SubscriberState {
207 ring: BoundedRing::new(self.capacity),
208 one_in: NonZeroU32::MIN,
209 seen: 0,
210 next_seq: 1,
211 }),
212 });
213 let mut subscribers = lock(&self.shared.subscribers, "tap subscribers")?;
214 subscribers.push(Arc::clone(&subscriber));
215 self.shared
216 .active
217 .store(subscribers.len(), Ordering::Release);
218 drop(subscribers);
219 Ok(TapSubscription {
220 shared: Arc::clone(&self.shared),
221 subscriber,
222 })
223 }
224
225 #[must_use]
229 pub fn observations(&self) -> u64 {
230 self.shared.observed.load(Ordering::Relaxed)
231 }
232
233 #[cfg(test)]
237 #[must_use]
238 pub(crate) fn server_observations(&self) -> u64 {
239 self.shared.observed_server.load(Ordering::Relaxed)
240 }
241
242 #[must_use]
244 pub fn subscriber_count(&self) -> usize {
245 self.shared.active.load(Ordering::Acquire)
246 }
247
248 #[must_use]
250 pub fn dropped_total(&self) -> u64 {
251 self.shared.dropped_total.load(Ordering::Relaxed)
252 }
253
254 #[cfg(test)]
257 pub(crate) fn poison_subscribers(&self) -> bool {
258 std::thread::scope(|scope| {
259 scope
260 .spawn(|| {
261 let _held = self.shared.subscribers.lock();
262 std::panic::resume_unwind(Box::new("poisoned for a test"));
263 })
264 .join()
265 .is_err()
266 }) && self.shared.subscribers.is_poisoned()
267 }
268
269 pub fn name_station(&self, connection: u64, station: &str) -> Result<(), TapError> {
275 lock(&self.shared.stations, "tap stations")?.insert(connection, Arc::from(station));
276 Ok(())
277 }
278
279 pub fn forget_station(&self, connection: u64) -> Result<(), TapError> {
284 lock(&self.shared.stations, "tap stations")?.remove(&connection);
285 Ok(())
286 }
287
288 pub fn name_receiver(&self, subscriber: u64, connection: u64) -> Result<(), TapError> {
294 lock(&self.shared.receivers, "tap receivers")?.insert(subscriber, connection);
295 Ok(())
296 }
297
298 pub fn forget_receiver(&self, subscriber: u64) -> Result<(), TapError> {
304 lock(&self.shared.receivers, "tap receivers")?.remove(&subscriber);
305 Ok(())
306 }
307
308 pub fn forget_receivers_of(&self, connection: u64) -> Result<(), TapError> {
313 lock(&self.shared.receivers, "tap receivers")?.retain(|_, owner| *owner != connection);
314 Ok(())
315 }
316
317 pub fn lend_recipients(&self) -> Result<Option<PublishRecord>, TapError> {
326 if self.shared.active.load(Ordering::Acquire) == 0 {
327 return Ok(None);
328 }
329 self.shared.lent.fetch_add(1, Ordering::Relaxed);
330 Ok(Some(
331 lock(&self.shared.buffers, "tap buffers")?
332 .pop()
333 .unwrap_or_default(),
334 ))
335 }
336
337 pub fn give_back(&self, mut buffer: PublishRecord) -> Result<(), TapError> {
343 buffer.clear();
344 lock(&self.shared.buffers, "tap buffers")?.push(buffer);
345 Ok(())
346 }
347
348 #[must_use]
351 pub fn recipients_lent(&self) -> u64 {
352 self.shared.lent.load(Ordering::Relaxed)
353 }
354
355 #[must_use]
358 pub fn unsubscribe_failures(&self) -> u64 {
359 self.shared.unsubscribe_failures.load(Ordering::Relaxed)
360 }
361
362 pub fn observe(&self, fact: &PublishFact<'_>) -> Result<bool, TapError> {
381 if self.shared.active.load(Ordering::Acquire) == 0 {
382 return Ok(false);
383 }
384 let observation = self.shared.observed.fetch_add(1, Ordering::Relaxed);
385 #[cfg(test)]
386 if fact.source == TapSource::Server {
387 self.shared.observed_server.fetch_add(1, Ordering::Relaxed);
388 }
389 let mut kept = false;
390 let result = self.keep(fact, observation, &mut kept);
391 if kept && let Some(waker) = self.shared.waker.get() {
392 waker();
393 }
394 result.map(|()| kept)
395 }
396
397 fn station_of(&self, source: TapSource) -> Result<Option<Arc<str>>, TapError> {
399 Ok(match source {
400 TapSource::Connection(connection) => lock(&self.shared.stations, "tap stations")?
401 .get(&connection)
402 .map(Arc::clone),
403 TapSource::Server => None,
404 })
405 }
406
407 fn stations_of(&self, recipients: &[Recipient]) -> Result<Arc<[Reached]>, TapError> {
413 if recipients.is_empty() {
414 return Ok(Arc::clone(&self.shared.no_stations));
415 }
416 let receivers = lock(&self.shared.receivers, "tap receivers")?;
417 let stations = lock(&self.shared.stations, "tap stations")?;
418 let station_at = |index: usize| {
419 recipients
420 .get(index)
421 .and_then(|(subscriber, _)| receivers.get(subscriber))
422 .and_then(|connection| stations.get(connection))
423 };
424 let first_at = |index: usize| {
427 station_at(index)
428 .filter(|station| !(0..index).any(|earlier| station_at(earlier) == Some(*station)))
429 .zip(recipients.get(index).map(|(_, queued_at)| *queued_at))
430 };
431 let named = (0..recipients.len())
432 .filter(|index| first_at(*index).is_some())
433 .count();
434 let list = gather(recipients.len(), named, first_at)?
435 .unwrap_or_else(|| Arc::clone(&self.shared.no_stations));
436 drop(stations);
437 drop(receivers);
438 Ok(list)
439 }
440
441 fn hops(
457 &self,
458 fact: &PublishFact<'_>,
459 seen_at: Instant,
460 station: Option<&Arc<str>>,
461 reached: Option<&[Reached]>,
462 failures: Option<&mut Vec<TapError>>,
463 ) -> Vec<Hop> {
464 let reached = reached.unwrap_or_default();
465 let appended_at = fact.recorded.and_then(|recorded| recorded.appended_at);
466 let depot = self.depot.as_ref().zip(appended_at);
467 let mut hops = Vec::with_capacity(
468 usize::from(station.is_some())
469 + usize::from(depot.is_some())
470 + usize::from(appended_at.is_some())
471 + reached.len(),
472 );
473 let mut early = None;
474 let mut overflow = None;
475 let mut out_of_order = None;
476 let mut place = 0;
477 let mut pass = |node: &Arc<str>, at: Option<Instant>| {
478 let hop = place;
479 place += 1;
480 let at_ms = match at {
481 None => 0,
482 Some(at) if at < seen_at => {
483 early.get_or_insert_with(|| {
484 elapsed_ms(seen_at.duration_since(at))
485 .map_or(TapError::HopOverflow { hop }, |early_ms| {
486 TapError::HopBeforeFirstSight { hop, early_ms }
487 })
488 });
489 return;
490 }
491 Some(at) => {
492 let Some(at_ms) = elapsed_ms(at.duration_since(seen_at)) else {
493 overflow.get_or_insert(TapError::HopOverflow { hop });
494 return;
495 };
496 at_ms
497 }
498 };
499 if let Some(before) = hops.last().map(|kept: &Hop| kept.at_ms)
500 && at_ms < before
501 {
502 out_of_order.get_or_insert(TapError::HopsOutOfOrder {
503 hop,
504 at_ms,
505 before_ms: before,
506 });
507 }
508 hops.push(Hop {
509 node: Arc::clone(node),
510 at_ms,
511 });
512 };
513 if let Some(station) = station {
514 pass(station, None);
515 }
516 if let Some((depot, appended_at)) = depot {
517 pass(depot, Some(appended_at));
518 }
519 if let Some(appended_at) = appended_at {
520 pass(&self.node, Some(appended_at));
521 }
522 for reached in reached {
523 pass(&reached.station, Some(reached.queued_at));
524 }
525 if let Some(failures) = failures {
526 failures.extend([early, overflow, out_of_order].into_iter().flatten());
527 }
528 hops
529 }
530
531 fn keep(
532 &self,
533 fact: &PublishFact<'_>,
534 observation: u64,
535 kept: &mut bool,
536 ) -> Result<(), TapError> {
537 let mut seen_at = None;
538 let mut station = None;
540 let mut to_stations = None;
541 let mut failures = Vec::new();
545 let mut hops_reported = false;
546 let subscribers = lock(&self.shared.subscribers, "tap subscribers")?;
547 for subscriber in subscribers.iter() {
548 if !subscriber.filter.matches(fact.subject) {
549 continue;
550 }
551 let mut state = lock(&subscriber.state, "tap subscriber")?;
552 if !state.admits() {
553 continue;
554 }
555 let one_in = state.one_in;
556 if station.is_none() {
557 station = Some(self.station_of(fact.source)?);
558 to_stations = Some(
559 match fact
560 .recorded
561 .map(|recorded| self.stations_of(recorded.recipients))
562 {
563 Some(Ok(list)) => Some(list),
564 Some(Err(error @ TapError::StationsDisagree { .. })) => {
566 failures.push(error);
567 None
568 }
569 Some(Err(error)) => return Err(error),
570 None => None,
571 },
572 );
573 }
574 let first_sight = *seen_at.get_or_insert_with(|| {
577 fact.recorded
578 .map_or_else(Instant::now, |recorded| recorded.received_at)
579 });
580 let hops = self.hops(
583 fact,
584 first_sight,
585 station.as_ref().and_then(Option::as_ref),
586 to_stations.as_ref().and_then(Option::as_deref),
587 (!hops_reported).then_some(&mut failures),
588 );
589 hops_reported = true;
590 let message = TapMessage {
591 ident: fact
592 .message_id
593 .map_or(TapIdent::Observation(observation), TapIdent::Message),
594 subject: fact.subject.to_owned(),
595 size: u64::try_from(fact.payload.len()).unwrap_or(u64::MAX),
596 outcome: fact.outcome,
597 seen_at: first_sight,
598 source: fact.source,
599 station: station.clone().flatten(),
600 to_stations: to_stations.clone().flatten(),
601 hops,
602 reply_to: fact.reply_to.map(str::to_owned),
603 held: false,
604 payload: subscriber.grant.map(|_| fact.payload.to_vec()),
605 };
606 if state.ring.push(Admitted { one_in, message }) {
607 self.shared.dropped_total.fetch_add(1, Ordering::Relaxed);
608 }
609 state.after_push();
610 drop(state);
611 *kept = true;
612 }
613 drop(subscribers);
614 match failures.len() {
615 0 => Ok(()),
616 1 => failures.pop().map_or(Ok(()), Err),
617 _ => Err(TapError::Several(failures)),
618 }
619 }
620}
621
622#[derive(Debug)]
624pub struct TapSubscription {
625 shared: Arc<TapShared>,
626 subscriber: Arc<Subscriber>,
627}
628
629impl TapSubscription {
630 #[must_use]
632 pub fn filter(&self) -> &TapFilter {
633 &self.subscriber.filter
634 }
635
636 pub fn take_batches(&self) -> Result<Vec<TapBatch>, TapError> {
645 let mut guard = lock(&self.subscriber.state, "tap subscriber")?;
646 let state = &mut *guard;
647 let (drained, dropped) = state.ring.take();
648 let mut batches: Vec<TapBatch> = Vec::new();
649 for Admitted { one_in, message } in drained {
650 match batches.last_mut() {
651 Some(open) if open.sampled == one_in => open.messages.push(message),
652 _ => batches.push(TapBatch {
653 seq: 0,
654 sampled: one_in,
655 dropped: 0,
656 messages: vec![message],
657 }),
658 }
659 }
660 match batches.first_mut() {
661 Some(first) => first.dropped = dropped,
662 None if dropped > 0 => batches.push(TapBatch {
663 seq: 0,
664 sampled: state.one_in,
665 dropped,
666 messages: Vec::new(),
667 }),
668 None => {}
669 }
670 for batch in &mut batches {
671 batch.seq = state.next_seq;
672 state.next_seq = state.next_seq.saturating_add(1);
673 }
674 state.one_in = NonZeroU32::new(state.one_in.get() / 2).unwrap_or(NonZeroU32::MIN);
676 drop(guard);
677 Ok(batches)
678 }
679
680 pub fn sample_one_in(&self) -> Result<NonZeroU32, TapError> {
685 Ok(lock(&self.subscriber.state, "tap subscriber")?.one_in)
686 }
687
688 pub fn held(&self) -> Result<usize, TapError> {
693 Ok(lock(&self.subscriber.state, "tap subscriber")?.ring.len())
694 }
695}
696
697impl Drop for TapSubscription {
698 fn drop(&mut self) {
699 let Ok(mut subscribers) = lock(&self.shared.subscribers, "tap subscribers") else {
702 self.shared
703 .unsubscribe_failures
704 .fetch_add(1, Ordering::Relaxed);
705 crate::metrics::console_failure(
706 crate::console::counts::ConsoleFailure::TapUnsubscribe,
707 1,
708 );
709 tracing::warn!("TAP_UNSUBSCRIBE_FAILED: the tap subscriber list is poisoned");
710 return;
711 };
712 let id = self.subscriber.id;
713 subscribers.retain(|subscriber| subscriber.id != id);
714 self.shared
715 .active
716 .store(subscribers.len(), Ordering::Release);
717 }
718}
719
720fn elapsed_ms(elapsed: std::time::Duration) -> Option<u64> {
723 u64::try_from(elapsed.as_millis()).ok()
724}
725
726fn gather<'station>(
734 recipients: usize,
735 named: usize,
736 first_at: impl Fn(usize) -> Option<(&'station Arc<str>, Instant)>,
737) -> Result<Option<Arc<[Reached]>>, TapError> {
738 let Some(first) = (0..recipients).find_map(&first_at) else {
739 return if named == 0 {
740 Ok(None)
741 } else {
742 Err(TapError::StationsDisagree { named, found: 0 })
743 };
744 };
745 let mut next = 0;
750 let mut found = 0_usize;
751 let reached = |(station, queued_at): (&Arc<str>, Instant)| Reached {
752 station: Arc::clone(station),
753 queued_at,
754 };
755 let list: Arc<[Reached]> = (0..named)
756 .map(|_| {
757 match (next..recipients)
758 .find_map(|index| first_at(index).map(|station| (index, station)))
759 {
760 Some((index, station)) => {
761 next = index + 1;
762 found += 1;
763 reached(station)
764 }
765 None => reached(first),
766 }
767 })
768 .collect();
769 let beyond = (next..recipients)
770 .filter(|index| first_at(*index).is_some())
771 .count();
772 if found != named || beyond != 0 {
773 return Err(TapError::StationsDisagree {
774 named,
775 found: found + beyond,
776 });
777 }
778 Ok(Some(list))
779}
780
781#[cfg(test)]
782mod gather_tests {
783 use std::sync::Arc;
784 use std::time::{Duration, Instant};
785
786 use super::{Reached, TapError, elapsed_ms, gather};
787
788 #[test]
791 fn an_elapsed_time_beyond_u64_milliseconds_has_no_reading() {
792 assert_eq!(elapsed_ms(Duration::from_millis(7)), Some(7));
793 assert_eq!(elapsed_ms(Duration::MAX), None);
794 }
795
796 #[test]
802 fn a_fill_that_disagrees_with_the_count_is_refused() {
803 let ledger: Arc<str> = Arc::from("orders.ledger");
804 let queued = Instant::now();
805 let gathered = gather(2, 2, |index| (index == 0).then_some((&ledger, queued)));
806 assert!(
807 matches!(
808 gathered,
809 Err(TapError::StationsDisagree { named: 2, found: 1 })
810 ),
811 "{gathered:?}"
812 );
813 }
814
815 #[test]
818 fn a_fill_that_agrees_builds_the_list() -> Result<(), TapError> {
819 let ledger: Arc<str> = Arc::from("orders.ledger");
820 let audit: Arc<str> = Arc::from("orders.audit");
821 let queued = Instant::now();
822 let first_at = |index: usize| match index {
823 0 => Some((&ledger, queued)),
824 2 => Some((&audit, queued)),
825 _ => None,
826 };
827 let reached = |station: &Arc<str>| Reached {
828 station: Arc::clone(station),
829 queued_at: queued,
830 };
831 assert_eq!(
832 gather(3, 2, first_at)?,
833 Some(Arc::from(vec![reached(&ledger), reached(&audit)]))
834 );
835 assert_eq!(gather(3, 0, |_| None)?, None);
836 Ok(())
837 }
838}