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 self.observe_from(fact, None)
382 }
383
384 pub(crate) fn observe_embedded(
385 &self,
386 fact: &PublishFact<'_>,
387 host: &str,
388 ) -> Result<bool, TapError> {
389 self.observe_from(fact, Some(host))
390 }
391
392 fn observe_from(&self, fact: &PublishFact<'_>, host: Option<&str>) -> Result<bool, TapError> {
393 if self.shared.active.load(Ordering::Acquire) == 0 {
394 return Ok(false);
395 }
396 let observation = self.shared.observed.fetch_add(1, Ordering::Relaxed);
397 #[cfg(test)]
398 if fact.source == TapSource::Server && host.is_none() {
399 self.shared.observed_server.fetch_add(1, Ordering::Relaxed);
400 }
401 let mut kept = false;
402 let result = self.keep(fact, observation, &mut kept, host);
403 if kept && let Some(waker) = self.shared.waker.get() {
404 waker();
405 }
406 result.map(|()| kept)
407 }
408
409 fn station_of(&self, source: TapSource) -> Result<Option<Arc<str>>, TapError> {
411 Ok(match source {
412 TapSource::Connection(connection) => lock(&self.shared.stations, "tap stations")?
413 .get(&connection)
414 .map(Arc::clone),
415 TapSource::Server => None,
416 })
417 }
418
419 fn stations_of(&self, recipients: &[Recipient]) -> Result<Arc<[Reached]>, TapError> {
425 if recipients.is_empty() {
426 return Ok(Arc::clone(&self.shared.no_stations));
427 }
428 let receivers = lock(&self.shared.receivers, "tap receivers")?;
429 let stations = lock(&self.shared.stations, "tap stations")?;
430 let station_at = |index: usize| {
431 recipients
432 .get(index)
433 .and_then(|(subscriber, _)| receivers.get(subscriber))
434 .and_then(|connection| stations.get(connection))
435 };
436 let first_at = |index: usize| {
439 station_at(index)
440 .filter(|station| !(0..index).any(|earlier| station_at(earlier) == Some(*station)))
441 .zip(recipients.get(index).map(|(_, queued_at)| *queued_at))
442 };
443 let named = (0..recipients.len())
444 .filter(|index| first_at(*index).is_some())
445 .count();
446 let list = gather(recipients.len(), named, first_at)?
447 .unwrap_or_else(|| Arc::clone(&self.shared.no_stations));
448 drop(stations);
449 drop(receivers);
450 Ok(list)
451 }
452
453 fn hops(
469 &self,
470 fact: &PublishFact<'_>,
471 seen_at: Instant,
472 station: Option<&Arc<str>>,
473 reached: Option<&[Reached]>,
474 failures: Option<&mut Vec<TapError>>,
475 ) -> Vec<Hop> {
476 let reached = reached.unwrap_or_default();
477 let appended_at = fact.recorded.and_then(|recorded| recorded.appended_at);
478 let depot = self.depot.as_ref().zip(appended_at);
479 let mut hops = Vec::with_capacity(
480 usize::from(station.is_some())
481 + usize::from(depot.is_some())
482 + usize::from(appended_at.is_some())
483 + reached.len(),
484 );
485 let mut early = None;
486 let mut overflow = None;
487 let mut out_of_order = None;
488 let mut place = 0;
489 let mut pass = |node: &Arc<str>, at: Option<Instant>| {
490 let hop = place;
491 place += 1;
492 let at_ms = match at {
493 None => 0,
494 Some(at) if at < seen_at => {
495 early.get_or_insert_with(|| {
496 elapsed_ms(seen_at.duration_since(at))
497 .map_or(TapError::HopOverflow { hop }, |early_ms| {
498 TapError::HopBeforeFirstSight { hop, early_ms }
499 })
500 });
501 return;
502 }
503 Some(at) => {
504 let Some(at_ms) = elapsed_ms(at.duration_since(seen_at)) else {
505 overflow.get_or_insert(TapError::HopOverflow { hop });
506 return;
507 };
508 at_ms
509 }
510 };
511 if let Some(before) = hops.last().map(|kept: &Hop| kept.at_ms)
512 && at_ms < before
513 {
514 out_of_order.get_or_insert(TapError::HopsOutOfOrder {
515 hop,
516 at_ms,
517 before_ms: before,
518 });
519 }
520 hops.push(Hop {
521 node: Arc::clone(node),
522 at_ms,
523 });
524 };
525 if let Some(station) = station {
526 pass(station, None);
527 }
528 if let Some((depot, appended_at)) = depot {
529 pass(depot, Some(appended_at));
530 }
531 if let Some(appended_at) = appended_at {
532 pass(&self.node, Some(appended_at));
533 }
534 for reached in reached {
535 pass(&reached.station, Some(reached.queued_at));
536 }
537 if let Some(failures) = failures {
538 failures.extend([early, overflow, out_of_order].into_iter().flatten());
539 }
540 hops
541 }
542
543 fn keep(
544 &self,
545 fact: &PublishFact<'_>,
546 observation: u64,
547 kept: &mut bool,
548 host: Option<&str>,
549 ) -> Result<(), TapError> {
550 let mut seen_at = None;
551 let mut station = None;
553 let mut to_stations = None;
554 let mut failures = Vec::new();
558 let mut hops_reported = false;
559 let subscribers = lock(&self.shared.subscribers, "tap subscribers")?;
560 for subscriber in subscribers.iter() {
561 if !subscriber.filter.matches(fact.subject) {
562 continue;
563 }
564 let mut state = lock(&subscriber.state, "tap subscriber")?;
565 if !state.admits() {
566 continue;
567 }
568 let one_in = state.one_in;
569 if station.is_none() {
570 station = Some(match host {
571 Some(host) => Some(Arc::from(format!("liminal.{host}"))),
572 None => self.station_of(fact.source)?,
573 });
574 to_stations = Some(
575 match fact
576 .recorded
577 .map(|recorded| self.stations_of(recorded.recipients))
578 {
579 Some(Ok(list)) => Some(list),
580 Some(Err(error @ TapError::StationsDisagree { .. })) => {
582 failures.push(error);
583 None
584 }
585 Some(Err(error)) => return Err(error),
586 None => None,
587 },
588 );
589 }
590 let first_sight = *seen_at.get_or_insert_with(|| {
593 fact.recorded
594 .map_or_else(Instant::now, |recorded| recorded.received_at)
595 });
596 let hops = self.hops(
599 fact,
600 first_sight,
601 station.as_ref().and_then(Option::as_ref),
602 to_stations.as_ref().and_then(Option::as_deref),
603 (!hops_reported).then_some(&mut failures),
604 );
605 hops_reported = true;
606 let message = TapMessage {
607 ident: fact
608 .message_id
609 .map_or(TapIdent::Observation(observation), TapIdent::Message),
610 subject: fact.subject.to_owned(),
611 size: u64::try_from(fact.payload.len()).unwrap_or(u64::MAX),
612 outcome: fact.outcome,
613 seen_at: first_sight,
614 source: fact.source,
615 station: station.clone().flatten(),
616 to_stations: to_stations.clone().flatten(),
617 hops,
618 reply_to: fact.reply_to.map(str::to_owned),
619 held: false,
620 payload: subscriber.grant.map(|_| fact.payload.to_vec()),
621 };
622 if state.ring.push(Admitted { one_in, message }) {
623 self.shared.dropped_total.fetch_add(1, Ordering::Relaxed);
624 }
625 state.after_push();
626 drop(state);
627 *kept = true;
628 }
629 drop(subscribers);
630 match failures.len() {
631 0 => Ok(()),
632 1 => failures.pop().map_or(Ok(()), Err),
633 _ => Err(TapError::Several(failures)),
634 }
635 }
636}
637
638#[derive(Debug)]
640pub struct TapSubscription {
641 shared: Arc<TapShared>,
642 subscriber: Arc<Subscriber>,
643}
644
645impl TapSubscription {
646 #[must_use]
648 pub fn filter(&self) -> &TapFilter {
649 &self.subscriber.filter
650 }
651
652 pub fn take_batches(&self) -> Result<Vec<TapBatch>, TapError> {
661 let mut guard = lock(&self.subscriber.state, "tap subscriber")?;
662 let state = &mut *guard;
663 let (drained, dropped) = state.ring.take();
664 let mut batches: Vec<TapBatch> = Vec::new();
665 for Admitted { one_in, message } in drained {
666 match batches.last_mut() {
667 Some(open) if open.sampled == one_in => open.messages.push(message),
668 _ => batches.push(TapBatch {
669 seq: 0,
670 sampled: one_in,
671 dropped: 0,
672 messages: vec![message],
673 }),
674 }
675 }
676 match batches.first_mut() {
677 Some(first) => first.dropped = dropped,
678 None if dropped > 0 => batches.push(TapBatch {
679 seq: 0,
680 sampled: state.one_in,
681 dropped,
682 messages: Vec::new(),
683 }),
684 None => {}
685 }
686 for batch in &mut batches {
687 batch.seq = state.next_seq;
688 state.next_seq = state.next_seq.saturating_add(1);
689 }
690 state.one_in = NonZeroU32::new(state.one_in.get() / 2).unwrap_or(NonZeroU32::MIN);
692 drop(guard);
693 Ok(batches)
694 }
695
696 pub fn sample_one_in(&self) -> Result<NonZeroU32, TapError> {
701 Ok(lock(&self.subscriber.state, "tap subscriber")?.one_in)
702 }
703
704 pub fn held(&self) -> Result<usize, TapError> {
709 Ok(lock(&self.subscriber.state, "tap subscriber")?.ring.len())
710 }
711}
712
713impl Drop for TapSubscription {
714 fn drop(&mut self) {
715 let Ok(mut subscribers) = lock(&self.shared.subscribers, "tap subscribers") else {
718 self.shared
719 .unsubscribe_failures
720 .fetch_add(1, Ordering::Relaxed);
721 crate::metrics::console_failure(
722 crate::console::counts::ConsoleFailure::TapUnsubscribe,
723 1,
724 );
725 tracing::warn!("TAP_UNSUBSCRIBE_FAILED: the tap subscriber list is poisoned");
726 return;
727 };
728 let id = self.subscriber.id;
729 subscribers.retain(|subscriber| subscriber.id != id);
730 self.shared
731 .active
732 .store(subscribers.len(), Ordering::Release);
733 }
734}
735
736fn elapsed_ms(elapsed: std::time::Duration) -> Option<u64> {
739 u64::try_from(elapsed.as_millis()).ok()
740}
741
742fn gather<'station>(
750 recipients: usize,
751 named: usize,
752 first_at: impl Fn(usize) -> Option<(&'station Arc<str>, Instant)>,
753) -> Result<Option<Arc<[Reached]>>, TapError> {
754 let Some(first) = (0..recipients).find_map(&first_at) else {
755 return if named == 0 {
756 Ok(None)
757 } else {
758 Err(TapError::StationsDisagree { named, found: 0 })
759 };
760 };
761 let mut next = 0;
766 let mut found = 0_usize;
767 let reached = |(station, queued_at): (&Arc<str>, Instant)| Reached {
768 station: Arc::clone(station),
769 queued_at,
770 };
771 let list: Arc<[Reached]> = (0..named)
772 .map(|_| {
773 match (next..recipients)
774 .find_map(|index| first_at(index).map(|station| (index, station)))
775 {
776 Some((index, station)) => {
777 next = index + 1;
778 found += 1;
779 reached(station)
780 }
781 None => reached(first),
782 }
783 })
784 .collect();
785 let beyond = (next..recipients)
786 .filter(|index| first_at(*index).is_some())
787 .count();
788 if found != named || beyond != 0 {
789 return Err(TapError::StationsDisagree {
790 named,
791 found: found + beyond,
792 });
793 }
794 Ok(Some(list))
795}
796
797#[cfg(test)]
798mod gather_tests {
799 use std::sync::Arc;
800 use std::time::{Duration, Instant};
801
802 use super::{Reached, TapError, elapsed_ms, gather};
803
804 #[test]
807 fn an_elapsed_time_beyond_u64_milliseconds_has_no_reading() {
808 assert_eq!(elapsed_ms(Duration::from_millis(7)), Some(7));
809 assert_eq!(elapsed_ms(Duration::MAX), None);
810 }
811
812 #[test]
818 fn a_fill_that_disagrees_with_the_count_is_refused() {
819 let ledger: Arc<str> = Arc::from("orders.ledger");
820 let queued = Instant::now();
821 let gathered = gather(2, 2, |index| (index == 0).then_some((&ledger, queued)));
822 assert!(
823 matches!(
824 gathered,
825 Err(TapError::StationsDisagree { named: 2, found: 1 })
826 ),
827 "{gathered:?}"
828 );
829 }
830
831 #[test]
834 fn a_fill_that_agrees_builds_the_list() -> Result<(), TapError> {
835 let ledger: Arc<str> = Arc::from("orders.ledger");
836 let audit: Arc<str> = Arc::from("orders.audit");
837 let queued = Instant::now();
838 let first_at = |index: usize| match index {
839 0 => Some((&ledger, queued)),
840 2 => Some((&audit, queued)),
841 _ => None,
842 };
843 let reached = |station: &Arc<str>| Reached {
844 station: Arc::clone(station),
845 queued_at: queued,
846 };
847 assert_eq!(
848 gather(3, 2, first_at)?,
849 Some(Arc::from(vec![reached(&ledger), reached(&audit)]))
850 );
851 assert_eq!(gather(3, 0, |_| None)?, None);
852 Ok(())
853 }
854}