xmtp_mls/worker/
metrics.rs

1use futures::FutureExt;
2use parking_lot::Mutex;
3use std::{
4    collections::HashMap,
5    fmt::Debug,
6    future::Future,
7    hash::Hash,
8    pin::Pin,
9    sync::{
10        Arc,
11        atomic::{AtomicUsize, Ordering},
12    },
13    time::Duration,
14};
15use tokio::sync::Notify;
16use xmtp_proto::types::InstallationId;
17
18pub struct MetricInterest<Metric> {
19    fut: Pin<Box<dyn Future<Output = ()> + Send>>,
20    count: usize,
21    info: Info,
22    metric: Metric,
23}
24
25impl<Metric> MetricInterest<Metric>
26where
27    Metric: Debug,
28{
29    /// Wait for a metric to be resolved
30    pub async fn wait(self) -> Result<(), xmtp_common::time::Expired> {
31        let Self {
32            fut,
33            count,
34            info,
35            metric,
36        } = self;
37
38        let result = xmtp_common::time::timeout(Duration::from_secs(10), fut).await;
39        tracing::info!(
40            "[{}] successfully waited for {metric:?}, {:?}",
41            hex::encode(info.installation_id),
42            result
43        );
44        if info.count() >= count {
45            return Ok(());
46        }
47        tracing::error!(
48            "Timed out waiting for {:?} to be >= {}. Value: {}",
49            metric,
50            count,
51            info.count()
52        );
53        result
54    }
55}
56
57/// Information and interest in a specific metric
58#[derive(Debug, Clone)]
59struct Info {
60    count: Arc<AtomicUsize>,
61    // each Notify is a task waiting on this metric to be incremented
62    notify: Arc<Notify>,
63    installation_id: InstallationId,
64}
65
66impl Info {
67    fn new(installation_id: InstallationId) -> Self {
68        Self {
69            count: Arc::new(AtomicUsize::default()),
70            notify: Arc::new(Notify::new()),
71            installation_id,
72        }
73    }
74
75    fn count(&self) -> usize {
76        self.count.load(Ordering::SeqCst)
77    }
78
79    fn increment(&self) {
80        self.count.fetch_add(1, Ordering::Relaxed);
81    }
82
83    fn clear(&self) {
84        self.count.store(0, Ordering::SeqCst)
85    }
86
87    // Registers interest in the next time this event is fired.
88    // Returns a future that resolves when this event resolves
89    fn register_interest(&self, count: usize) -> impl Future<Output = ()> + 'static {
90        let notify = self.notify.clone();
91        let info_count = self.count.clone();
92        async move {
93            while info_count.load(Ordering::SeqCst) < count {
94                notify.notified().await;
95            }
96        }
97    }
98
99    fn fire(&self) {
100        self.notify.notify_waiters();
101    }
102}
103
104#[derive(Debug)]
105pub struct WorkerMetrics<Metric> {
106    metrics: Mutex<HashMap<Metric, Info>>,
107    installation_id: InstallationId,
108}
109
110impl<Metric> WorkerMetrics<Metric>
111where
112    Metric: PartialEq + Eq + Hash + Clone + Copy + Debug,
113{
114    pub fn new(installation_id: InstallationId) -> Self {
115        Self {
116            metrics: Mutex::default(),
117            installation_id,
118        }
119    }
120
121    fn info(&self, metric: Metric) -> Info {
122        self.metrics
123            .lock()
124            .entry(metric)
125            .or_insert(Info::new(self.installation_id))
126            .clone()
127    }
128
129    pub fn get(&self, metric: Metric) -> usize {
130        self.info(metric).count()
131    }
132
133    pub(crate) fn increment_metric(&self, metric: Metric) {
134        self.info(metric).increment();
135        tracing::trace!("[{}] firing {metric:?}", hex::encode(self.installation_id));
136        self.info(metric).fire();
137    }
138
139    pub fn reset_metrics(&self) {
140        *self.metrics.lock() = HashMap::new();
141    }
142
143    /// Register interest in a metric at a certain 'count'.
144    /// Returns a MetricInterest type that can be used to wait for the metric to reach these
145    /// parameters
146    pub fn register_interest(&self, metric_key: Metric, count: usize) -> MetricInterest<Metric> {
147        tracing::info!("registering interest in {metric_key:?}");
148        let info = self.info(metric_key);
149
150        let fut = if self
151            .metrics
152            .lock()
153            .get(&metric_key)
154            .is_some_and(|info| info.count() >= count)
155        {
156            futures::future::ready(()).boxed()
157        } else {
158            info.register_interest(count).boxed()
159        };
160
161        MetricInterest {
162            fut,
163            count,
164            info: self.info(metric_key),
165            metric: metric_key,
166        }
167    }
168
169    pub async fn do_until<F, Fut>(
170        &self,
171        metric: Metric,
172        count: usize,
173        f: F,
174    ) -> Result<(), xmtp_common::time::Expired>
175    where
176        F: Fn() -> Fut,
177        Fut: Future<Output = ()>,
178    {
179        let info = {
180            let mut m = self.metrics.lock();
181            m.entry(metric)
182                .or_insert(Info::new(self.installation_id))
183                .clone()
184        };
185        xmtp_common::time::timeout(Duration::from_secs(20), async {
186            while info.count() < count {
187                f().await;
188                xmtp_common::task::yield_now().await;
189            }
190        })
191        .await
192    }
193
194    pub fn clear_metric(&self, metric: Metric) {
195        self.metrics
196            .lock()
197            .entry(metric)
198            .or_insert(Info::new(self.installation_id))
199            .clear();
200    }
201}
xmtp_mls/worker/metrics.rs

xmtp_mls/worker/
metrics.rs
