xmtp_events/
bus.rs

1use std::{
2    cell::Cell,
3    collections::VecDeque,
4    sync::atomic::{AtomicBool, Ordering},
5    sync::{Arc, Weak},
6};
7
8use parking_lot::{Mutex, ReentrantMutex};
9use tokio::sync::Notify;
10use xmtp_common::{MaybeSend, MaybeSync};
11
12use crate::{ClientEvent, EventFilter, EventKind, Lagged};
13
14const UNBOUNDED_QUEUE_WARNING_DEPTH: usize = 10_000;
15const APP_QUEUE_DEPTH: usize = 1024;
16
17thread_local! {
18    static FILTER_DEPTH: Cell<usize> = const { Cell::new(0) };
19}
20
21struct FilterGuard;
22
23impl FilterGuard {
24    fn enter() -> Self {
25        FILTER_DEPTH.with(|depth| depth.set(depth.get() + 1));
26        Self
27    }
28}
29
30impl Drop for FilterGuard {
31    fn drop(&mut self) {
32        FILTER_DEPTH.with(|depth| depth.set(depth.get() - 1));
33    }
34}
35
36/// Extra facts used by filters. A DM identifier is attached by the emitter.
37#[derive(Clone, Debug, Default, PartialEq, Eq)]
38pub struct EventContext {
39    pub dm_identifier: Option<Vec<u8>>,
40    /// The emitter sets this only after checking the stored reference and sender.
41    pub references_own_messages: bool,
42    /// Used to decide whether a received message is live when the write commits.
43    pub message_expires_at_ns: Option<i64>,
44}
45
46/// One bus item. A write may carry a public event, an internal fact, or both.
47#[derive(Clone, Debug, PartialEq, Eq)]
48pub struct EventEnvelope<I> {
49    pub client: Option<ClientEvent>,
50    pub internal: Option<I>,
51    pub context: EventContext,
52}
53
54impl<I> EventEnvelope<I> {
55    pub fn new(client: Option<ClientEvent>, internal: Option<I>, context: EventContext) -> Self {
56        Self {
57            client,
58            internal,
59            context,
60        }
61    }
62}
63
64/// The only interface an emitter needs. A public-only crate uses an adapter.
65pub trait EventWriter<I = ()>: MaybeSend + MaybeSync {
66    fn emit(&self, client: Option<ClientEvent>, internal: Option<I>) {
67        self.emit_with_context(client, internal, EventContext::default());
68    }
69
70    fn emit_with_context(
71        &self,
72        client: Option<ClientEvent>,
73        internal: Option<I>,
74        context: EventContext,
75    );
76}
77
78struct BusInner<I> {
79    subscriptions: Mutex<Vec<Weak<SubscriptionInner<I>>>>,
80    dispatch_lock: ReentrantMutex<()>,
81    buffer_lock: ReentrantMutex<()>,
82    app_closed: AtomicBool,
83    internal_closed: AtomicBool,
84}
85
86/// One client's live event bus. The client owns this value and passes writers down.
87pub struct EventBus<I> {
88    inner: Arc<BusInner<I>>,
89}
90
91impl<I> Clone for EventBus<I> {
92    fn clone(&self) -> Self {
93        Self {
94            inner: self.inner.clone(),
95        }
96    }
97}
98
99impl<I> Default for EventBus<I> {
100    fn default() -> Self {
101        Self::new()
102    }
103}
104
105impl<I> EventBus<I> {
106    pub fn new() -> Self {
107        Self {
108            inner: Arc::new(BusInner {
109                subscriptions: Mutex::new(Vec::new()),
110                dispatch_lock: ReentrantMutex::new(()),
111                buffer_lock: ReentrantMutex::new(()),
112                app_closed: AtomicBool::new(false),
113                internal_closed: AtomicBool::new(false),
114            }),
115        }
116    }
117
118    /// Registers before returning. `None` gives an internal worker an unbounded queue.
119    pub fn subscribe(&self, filter: EventFilter<I>, queue_depth: Option<usize>) -> Subscription<I> {
120        self.subscribe_with_role(filter, queue_depth, false)
121    }
122
123    /// Register an app listener while the client is open.
124    pub fn subscribe_app(&self, filter: EventFilter<I>) -> Option<Subscription<I>> {
125        if self.inner.app_closed.load(Ordering::Acquire) {
126            return None;
127        }
128        let subscription = self.subscribe_with_role(filter, Some(APP_QUEUE_DEPTH), true);
129        (!subscription.is_closed()).then_some(subscription)
130    }
131
132    fn subscribe_with_role(
133        &self,
134        filter: EventFilter<I>,
135        queue_depth: Option<usize>,
136        app: bool,
137    ) -> Subscription<I> {
138        self.assert_not_dispatching();
139        let _dispatch = self.inner.dispatch_lock.lock();
140        let inner = Arc::new(SubscriptionInner {
141            filter,
142            queue_depth,
143            app,
144            state: Mutex::new(QueueState::default()),
145            changed: Notify::new(),
146        });
147        if (app && self.inner.app_closed.load(Ordering::Acquire))
148            || (!app && self.inner.internal_closed.load(Ordering::Acquire))
149        {
150            inner.close();
151        }
152        self.inner.subscriptions.lock().push(Arc::downgrade(&inner));
153        Subscription { inner }
154    }
155
156    /// End app listeners before worker shutdown begins.
157    pub fn close_app_subscriptions(&self) {
158        self.inner.app_closed.store(true, Ordering::Release);
159        self.close_role(true);
160    }
161
162    /// End legacy and other internal listeners after worker shutdown.
163    pub fn close_internal_subscriptions(&self) {
164        self.inner.internal_closed.store(true, Ordering::Release);
165        self.close_role(false);
166    }
167
168    fn close_role(&self, app: bool) {
169        let _dispatch = self.inner.dispatch_lock.lock();
170        let mut subscriptions = self.inner.subscriptions.lock();
171        subscriptions.retain(|weak| {
172            if let Some(inner) = weak.upgrade() {
173                if inner.app == app {
174                    inner.close();
175                }
176                true
177            } else {
178                false
179            }
180        });
181    }
182
183    /// Runs one synchronous emitting write. Flush follows a successful commit only.
184    pub fn with_buffer<R, E>(
185        &self,
186        f: impl FnOnce(&EventBuffer<'_, I>) -> Result<R, E>,
187    ) -> Result<R, E>
188    where
189        I: Clone,
190    {
191        self.assert_not_dispatching();
192        assert!(
193            !self.inner.buffer_lock.is_owned_by_current_thread(),
194            "nested event buffer on the same bus"
195        );
196        let _guard = self.inner.buffer_lock.lock();
197        let buffer = EventBuffer {
198            bus: self,
199            pending: Mutex::new(Vec::new()),
200        };
201        let result = f(&buffer);
202        if result.is_ok() {
203            buffer.flush();
204        } else {
205            buffer.clear();
206        }
207        result
208    }
209
210    fn publish(&self, mut event: EventEnvelope<I>)
211    where
212        I: Clone,
213    {
214        if event.client.is_none() && event.internal.is_none() {
215            return;
216        }
217        self.assert_not_dispatching();
218        let _dispatch = self.inner.dispatch_lock.lock();
219        if matches!(event.client.as_ref(), Some(ClientEvent::MessageReceived(_)))
220            && event
221                .context
222                .message_expires_at_ns
223                .is_some_and(|expiry| expiry <= xmtp_common::time::now_ns())
224        {
225            event.client = None;
226            if event.internal.is_none() {
227                return;
228            }
229        }
230        if let Some(client) = &event.client {
231            tracing::debug!(kind = client.kind().name(), "client event emitted");
232        }
233        let mut entries = self.inner.subscriptions.lock();
234        entries.retain(|weak| {
235            if let Some(subscription) = weak.upgrade() {
236                subscription.deliver(&event);
237                true
238            } else {
239                false
240            }
241        });
242    }
243
244    fn assert_not_dispatching(&self) {
245        assert!(
246            FILTER_DEPTH.with(|depth| depth.get() == 0),
247            "an internal event filter must not call an event bus"
248        );
249        assert!(
250            !self.inner.dispatch_lock.is_owned_by_current_thread(),
251            "an internal event filter must not call its event bus"
252        );
253    }
254}
255
256impl<I> EventWriter<I> for EventBus<I>
257where
258    I: Clone + MaybeSend + MaybeSync + 'static,
259{
260    fn emit_with_context(
261        &self,
262        client: Option<ClientEvent>,
263        internal: Option<I>,
264        context: EventContext,
265    ) {
266        self.publish(EventEnvelope::new(client, internal, context));
267    }
268}
269
270/// A transaction's pending events. Only `EventBus::with_buffer` can create it.
271pub struct EventBuffer<'a, I> {
272    bus: &'a EventBus<I>,
273    pending: Mutex<Vec<EventEnvelope<I>>>,
274}
275
276impl<I: Clone> EventBuffer<'_, I> {
277    /// Publish after the matching state write committed. Events of one write use kind-table order.
278    fn flush(&self) {
279        let mut pending = std::mem::take(&mut *self.pending.lock());
280        pending.sort_by_key(|event| event.client.as_ref().map(|client| client.kind()));
281        for event in pending {
282            self.bus.publish(event);
283        }
284    }
285
286    /// Discard after rollback.
287    fn clear(&self) {
288        self.pending.lock().clear();
289    }
290
291    /// Keep the outer write's events when one storage savepoint rolls back.
292    pub fn savepoint<R, E>(&self, f: impl FnOnce(&Self) -> Result<R, E>) -> Result<R, E> {
293        let checkpoint = self.pending.lock().len();
294        let result = f(self);
295        if result.is_err() {
296            self.pending.lock().truncate(checkpoint);
297        }
298        result
299    }
300}
301
302impl<I> EventWriter<I> for EventBuffer<'_, I>
303where
304    I: Clone + MaybeSend + MaybeSync + 'static,
305{
306    fn emit_with_context(
307        &self,
308        client: Option<ClientEvent>,
309        internal: Option<I>,
310        context: EventContext,
311    ) {
312        if client.is_some() || internal.is_some() {
313            self.pending
314                .lock()
315                .push(EventEnvelope::new(client, internal, context));
316        }
317    }
318}
319
320/// A weak, public-only handle for lower crates and shared resources.
321pub struct PublicBusWriter<I> {
322    bus: Weak<BusInner<I>>,
323}
324
325impl<I> PublicBusWriter<I> {
326    pub fn new(bus: &EventBus<I>) -> Self {
327        Self {
328            bus: Arc::downgrade(&bus.inner),
329        }
330    }
331}
332
333impl<I> EventWriter<()> for PublicBusWriter<I>
334where
335    I: Clone + MaybeSend + MaybeSync + 'static,
336{
337    fn emit_with_context(
338        &self,
339        client: Option<ClientEvent>,
340        _internal: Option<()>,
341        context: EventContext,
342    ) {
343        if let Some(inner) = self.bus.upgrade() {
344            EventBus { inner }.publish(EventEnvelope::new(client, None, context));
345        }
346    }
347}
348
349/// Public-only view of a transaction buffer.
350pub struct PublicBufferWriter<'b, 'bus, I> {
351    buffer: &'b EventBuffer<'bus, I>,
352}
353
354impl<'b, 'bus, I> PublicBufferWriter<'b, 'bus, I> {
355    pub fn new(buffer: &'b EventBuffer<'bus, I>) -> Self {
356        Self { buffer }
357    }
358}
359
360impl<I> EventWriter<()> for PublicBufferWriter<'_, '_, I>
361where
362    I: Clone + MaybeSend + MaybeSync + 'static,
363{
364    fn emit_with_context(
365        &self,
366        client: Option<ClientEvent>,
367        _internal: Option<()>,
368        context: EventContext,
369    ) {
370        self.buffer.emit_with_context(client, None, context);
371    }
372}
373
374struct QueueState<I> {
375    items: VecDeque<EventEnvelope<I>>,
376    discarded: u64,
377    lagged_after: usize,
378    in_flight: usize,
379    warned: bool,
380    closed: bool,
381}
382
383impl<I> Default for QueueState<I> {
384    fn default() -> Self {
385        Self {
386            items: VecDeque::new(),
387            discarded: 0,
388            lagged_after: 0,
389            in_flight: 0,
390            warned: false,
391            closed: false,
392        }
393    }
394}
395
396struct SubscriptionInner<I> {
397    filter: EventFilter<I>,
398    queue_depth: Option<usize>,
399    app: bool,
400    state: Mutex<QueueState<I>>,
401    changed: Notify,
402}
403
404impl<I> SubscriptionInner<I> {
405    fn close(&self) {
406        let mut state = self.state.lock();
407        state.closed = true;
408        state.items.clear();
409        state.discarded = 0;
410        drop(state);
411        self.changed.notify_waiters();
412    }
413}
414
415impl<I: Clone> SubscriptionInner<I> {
416    fn deliver(&self, event: &EventEnvelope<I>) {
417        let public = event
418            .client
419            .as_ref()
420            .filter(|client| self.filter.matches_public(client, &event.context));
421        let internal = event.internal.as_ref().filter(|internal| {
422            let _filter_guard = FilterGuard::enter();
423            self.filter.matches_internal(internal)
424        });
425        if public.is_none() && internal.is_none() {
426            return;
427        }
428        let selected =
429            EventEnvelope::new(public.cloned(), internal.cloned(), event.context.clone());
430        let mut state = self.state.lock();
431        if state.closed {
432            return;
433        }
434        if self
435            .queue_depth
436            .is_some_and(|depth| state.items.len() + state.in_flight >= depth)
437        {
438            if state.discarded == 0 {
439                state.lagged_after = state.items.len();
440            }
441            state.discarded = state.discarded.saturating_add(1);
442        } else {
443            state.items.push_back(selected);
444            if self.queue_depth.is_none()
445                && state.items.len() > UNBOUNDED_QUEUE_WARNING_DEPTH
446                && !state.warned
447            {
448                state.warned = true;
449                tracing::warn!(
450                    depth = state.items.len(),
451                    "unbounded event subscription queue is growing"
452                );
453            }
454        }
455        drop(state);
456        self.changed.notify_one();
457    }
458
459    fn take(&self, lease: bool) -> Option<Option<EventEnvelope<I>>> {
460        let mut state = self.state.lock();
461        if state.closed {
462            return Some(None);
463        }
464        if state.discarded > 0 && state.lagged_after == 0 {
465            let discarded = std::mem::take(&mut state.discarded);
466            let event = EventEnvelope::new(
467                Some(ClientEvent::Lagged(Lagged { discarded })),
468                None,
469                EventContext::default(),
470            );
471            return Some(Some(event));
472        }
473        if let Some(event) = state.items.pop_front() {
474            if state.lagged_after > 0 {
475                state.lagged_after -= 1;
476            }
477            if lease {
478                state.in_flight += 1;
479            }
480            return Some(Some(event));
481        }
482        None
483    }
484}
485
486/// One independent, live subscription. Dropping it also closes it.
487pub struct Subscription<I> {
488    inner: Arc<SubscriptionInner<I>>,
489}
490
491impl<I> Subscription<I> {
492    pub fn close(&self) {
493        self.inner.close();
494    }
495
496    pub fn is_closed(&self) -> bool {
497        self.inner.state.lock().closed
498    }
499}
500
501impl<I: Clone> Subscription<I> {
502    /// Return the next item, or `None` after close. No runtime is needed to emit.
503    pub async fn next(&self) -> Option<EventEnvelope<I>> {
504        loop {
505            let notified = self.inner.changed.notified();
506            tokio::pin!(notified);
507            notified.as_mut().enable();
508            if let Some(item) = self.inner.take(false) {
509                return item;
510            }
511            notified.await;
512        }
513    }
514
515    /// Keep a callback's item inside the queue bound until its returned value completes.
516    pub async fn next_for_callback(&self) -> Option<EventLease<I>> {
517        loop {
518            let notified = self.inner.changed.notified();
519            tokio::pin!(notified);
520            notified.as_mut().enable();
521            if let Some(item) = self.inner.take(true) {
522                return item.map(|event| {
523                    let counted = event
524                        .client
525                        .as_ref()
526                        .is_none_or(|client| client.kind() != EventKind::Lagged);
527                    EventLease {
528                        event,
529                        subscription: self.inner.clone(),
530                        counted,
531                    }
532                });
533            }
534            notified.await;
535        }
536    }
537
538    /// Drain ready items after a worker wake. This includes a pending `lagged` item.
539    pub fn drain(&self) -> Vec<EventEnvelope<I>> {
540        let mut items = Vec::new();
541        while let Some(Some(event)) = self.inner.take(false) {
542            items.push(event);
543        }
544        items
545    }
546}
547
548impl<I> Drop for Subscription<I> {
549    fn drop(&mut self) {
550        self.close();
551    }
552}
553
554/// A callback handoff. Drop it after the callback's future completes.
555pub struct EventLease<I> {
556    pub event: EventEnvelope<I>,
557    subscription: Arc<SubscriptionInner<I>>,
558    counted: bool,
559}
560
561impl<I> Drop for EventLease<I> {
562    fn drop(&mut self) {
563        let mut state = self.subscription.state.lock();
564        if self.counted {
565            state.in_flight -= 1;
566        }
567        drop(state);
568        self.subscription.changed.notify_one();
569    }
570}
571
572#[cfg(test)]
573mod tests;
xmtp_events/bus.rs

xmtp_events/
bus.rs
