mysten_common/sync/
notify_read.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::ZipDebugEqIteratorExt;
5use crate::debug_fatal;
6
7use futures::future::{Either, join_all};
8use parking_lot::Mutex;
9use parking_lot::MutexGuard;
10use std::collections::HashMap;
11use std::collections::hash_map::DefaultHasher;
12use std::future::Future;
13use std::hash::{Hash, Hasher};
14use std::mem;
15use std::ops::Deref;
16use std::pin::Pin;
17use std::sync::Arc;
18use std::sync::atomic::AtomicUsize;
19use std::sync::atomic::Ordering;
20use std::task::{Context, Poll};
21use std::time::Duration;
22use tokio::sync::oneshot;
23use tokio::time::Instant;
24use tokio::time::interval_at;
25use tracing::warn;
26
27use crate::sync::oneshot as blocking_oneshot;
28
29/// A registered waiter: async waiters hold a tokio oneshot, blocking waiters (see
30/// [`NotifyRead::register_one_blocking`]) hold a [`blocking_oneshot`] whose receiver
31/// blocks the OS thread.
32enum NotifySender<V> {
33    Async(oneshot::Sender<V>),
34    Blocking(blocking_oneshot::Sender<V>),
35}
36
37impl<V> NotifySender<V> {
38    fn send(self, value: V) {
39        match self {
40            NotifySender::Async(sender) => {
41                sender.send(value).ok();
42            }
43            NotifySender::Blocking(sender) => {
44                sender.send(value).ok();
45            }
46        }
47    }
48
49    fn is_closed(&self) -> bool {
50        match self {
51            NotifySender::Async(sender) => sender.is_closed(),
52            NotifySender::Blocking(sender) => sender.is_closed(),
53        }
54    }
55}
56
57type Registrations<V> = Vec<NotifySender<V>>;
58
59/// Interval duration for logging waiting keys when reads take too long
60const LONG_WAIT_LOG_INTERVAL_SECS: u64 = 10;
61
62/// Minimum interval between stall reports for a given task name, across every
63/// read blocked on it.
64const STALL_LOG_INTERVAL_SECS: u64 = 30;
65
66/// Number of this read's own keys included in a stall report.
67const MAX_SAMPLED_KEYS: usize = 32;
68
69pub const CHECKPOINT_BUILDER_NOTIFY_READ_TASK_NAME: &str =
70    "CheckpointBuilder::notify_read_executed_effects";
71
72pub struct NotifyRead<K, V> {
73    pending: Vec<Mutex<HashMap<K, Registrations<V>>>>,
74    count_pending: AtomicUsize,
75    // Last stall report per task name.
76    last_stall_log: Mutex<HashMap<&'static str, Instant>>,
77}
78
79impl<K: Eq + Hash + Clone, V: Clone> NotifyRead<K, V> {
80    pub fn new() -> Self {
81        let pending = (0..255).map(|_| Default::default()).collect();
82        let count_pending = Default::default();
83        Self {
84            pending,
85            count_pending,
86            last_stall_log: Default::default(),
87        }
88    }
89
90    /// Returns true if this caller should emit the stall report for `task_name`.
91    /// Any number of reads may be blocked at once; only one of them logs.
92    fn throttle_stall_log(&self, task_name: &'static str) -> bool {
93        let now = Instant::now();
94        let mut last_log = self.last_stall_log.lock();
95        match last_log.get(task_name) {
96            Some(last)
97                if now.duration_since(*last) < Duration::from_secs(STALL_LOG_INTERVAL_SECS) =>
98            {
99                false
100            }
101            _ => {
102                last_log.insert(task_name, now);
103                true
104            }
105        }
106    }
107
108    /// Asynchronously notifies waiters and return number of remaining pending registration
109    pub fn notify(&self, key: &K, value: &V) -> usize {
110        let registrations = self.pending(key).remove(key);
111        let Some(registrations) = registrations else {
112            return self.count_pending.load(Ordering::Relaxed);
113        };
114        let rem = self
115            .count_pending
116            .fetch_sub(registrations.len(), Ordering::Relaxed);
117        for registration in registrations {
118            registration.send(value.clone());
119        }
120        rem
121    }
122
123    pub fn register_one(&self, key: &K) -> Registration<&NotifyRead<K, V>, K, V> {
124        self.register_with(key, self)
125    }
126
127    /// Like [`NotifyRead::register_one`], but the returned registration holds the
128    /// registry by `Arc` instead of borrowing it, so it can be stored in long-lived
129    /// structures.
130    pub fn register_one_owned(self: &Arc<Self>, key: &K) -> OwnedRegistration<K, V> {
131        self.register_with(key, self.clone())
132    }
133
134    fn register_with<R>(&self, key: &K, this: R) -> Registration<R, K, V>
135    where
136        R: Deref<Target = NotifyRead<K, V>>,
137    {
138        self.count_pending.fetch_add(1, Ordering::Relaxed);
139        let (sender, receiver) = oneshot::channel();
140        self.register(key, NotifySender::Async(sender));
141        Registration {
142            this,
143            registration: Some((key.clone(), receiver)),
144        }
145    }
146
147    /// Register a waiter whose receiver blocks the OS thread (see
148    /// [`BlockingRegistration::wait`]). Must not be awaited from async code.
149    pub fn register_one_blocking(&self, key: &K) -> BlockingRegistration<'_, K, V> {
150        self.count_pending.fetch_add(1, Ordering::Relaxed);
151        let (sender, receiver) = blocking_oneshot::channel();
152        self.register(key, NotifySender::Blocking(sender));
153        BlockingRegistration {
154            this: self,
155            registration: Some((key.clone(), receiver)),
156        }
157    }
158
159    pub fn register_all(&self, keys: &[K]) -> Vec<Registration<&NotifyRead<K, V>, K, V>> {
160        keys.iter().map(|key| self.register_one(key)).collect()
161    }
162
163    fn register(&self, key: &K, sender: NotifySender<V>) {
164        self.pending(key)
165            .entry(key.clone())
166            .or_default()
167            .push(sender);
168    }
169
170    fn pending(&self, key: &K) -> MutexGuard<'_, HashMap<K, Registrations<V>>> {
171        let mut state = DefaultHasher::new();
172        key.hash(&mut state);
173        let hash = state.finish();
174        let pending = self
175            .pending
176            .get((hash % self.pending.len() as u64) as usize)
177            .unwrap();
178        pending.lock()
179    }
180
181    pub fn num_pending(&self) -> usize {
182        self.count_pending.load(Ordering::Relaxed)
183    }
184
185    /// Blocking version of [`Self::read`] for a single key: returns `fetch(key)` if the
186    /// value is already available, and otherwise blocks the calling OS thread until the
187    /// key is notified.
188    ///
189    /// Must not be called from an async context. Under msim it may only be called from
190    /// a blocking-pool thread (e.g. inside `spawn_blocking`), where the wait yields the
191    /// thread's quantum between readiness checks.
192    pub fn read_one_blocking(
193        &self,
194        task_name: &'static str,
195        key: &K,
196        fetch: impl FnOnce(&K) -> Option<V>,
197    ) -> V {
198        let _metrics_scope = mysten_metrics::monitored_scope(task_name);
199        let registration = self.register_one_blocking(key);
200        // As in `read`, fetch after registering so that a concurrent notify cannot be
201        // missed. If the value is already available the registration is dropped, which
202        // de-registers it.
203        if let Some(value) = fetch(key) {
204            return value;
205        }
206        registration.wait()
207    }
208
209    fn cleanup(&self, key: &K) {
210        let mut pending = self.pending(key);
211        // it is possible that registration was fulfilled before we get here
212        let Some(registrations) = pending.get_mut(key) else {
213            return;
214        };
215        let mut count_deleted = 0usize;
216        registrations.retain(|s| {
217            let delete = s.is_closed();
218            if delete {
219                count_deleted += 1;
220            }
221            !delete
222        });
223        self.count_pending
224            .fetch_sub(count_deleted, Ordering::Relaxed);
225        if registrations.is_empty() {
226            pending.remove(key);
227        }
228    }
229}
230
231impl<K: Eq + Hash + Clone + Unpin + std::fmt::Debug + Send + Sync + 'static, V: Clone + Unpin>
232    NotifyRead<K, V>
233{
234    pub async fn read(
235        &self,
236        task_name: &'static str,
237        keys: &[K],
238        fetch: impl FnOnce(&[K]) -> Vec<Option<V>>,
239    ) -> Vec<V> {
240        let _metrics_scope = mysten_metrics::monitored_scope(task_name);
241        let registrations = self.register_all(keys);
242
243        let results = fetch(keys);
244        // Snapshot of what was missing at fetch time.
245        let waiting_keys: Vec<K> = keys
246            .iter()
247            .zip_debug_eq(results.iter())
248            .filter(|(_key, result)| result.is_none())
249            .map(|(key, _result)| key.clone())
250            .collect();
251
252        let results = results
253            .into_iter()
254            .zip_debug_eq(registrations)
255            .map(|(a, r)| match a {
256                // Note that Some() clause also drops registration that is already fulfilled
257                Some(ready) => Either::Left(futures::future::ready(ready)),
258                None => Either::Right(r),
259            });
260
261        let join = join_all(results);
262        if waiting_keys.is_empty() {
263            return join.await;
264        }
265
266        tokio::pin!(join);
267        let start_time = Instant::now();
268        let mut interval = interval_at(
269            start_time + Duration::from_secs(LONG_WAIT_LOG_INTERVAL_SECS),
270            Duration::from_secs(LONG_WAIT_LOG_INTERVAL_SECS),
271        );
272
273        loop {
274            tokio::select! {
275                values = &mut join => return values,
276                _ = interval.tick() => {
277                    let elapsed_secs = start_time.elapsed().as_secs();
278
279                    // Deduplicate logging by task name. When many reads are blocked,
280                    // logs will sample blocked reads per task and the keys we are
281                    // waiting on for those reads.
282                    if self.throttle_stall_log(task_name) {
283                        let mut sample: Vec<&K> = Vec::with_capacity(MAX_SAMPLED_KEYS);
284                        let mut outstanding = 0usize;
285                        for key in &waiting_keys {
286                            if self.pending(key).contains_key(key) {
287                                outstanding += 1;
288                                if sample.len() < MAX_SAMPLED_KEYS {
289                                    sample.push(key);
290                                }
291                            }
292                        }
293
294                        warn!(
295                            "[{task_name}] Still waiting {elapsed_secs}s. {} registrations pending, this read still blocked on {outstanding} of {} key(s): {sample:?}",
296                            self.num_pending(),
297                            waiting_keys.len(),
298                        );
299                    }
300
301                    if task_name == CHECKPOINT_BUILDER_NOTIFY_READ_TASK_NAME && elapsed_secs >= 60 {
302                        debug_fatal!("{} is stuck", task_name);
303                    }
304                }
305            }
306        }
307    }
308}
309
310/// Registration resolves to the value but also provides safe cancellation
311/// When Registration is dropped before it is resolved, we de-register from the pending list
312///
313/// Generic over how it holds the registry: borrowed for the await-in-place pattern,
314/// or by `Arc` ([`OwnedRegistration`]) so it can be stored in long-lived structures.
315pub struct Registration<R, K: Eq + Hash + Clone, V: Clone>
316where
317    R: Deref<Target = NotifyRead<K, V>>,
318{
319    this: R,
320    registration: Option<(K, oneshot::Receiver<V>)>,
321}
322
323pub type OwnedRegistration<K, V> = Registration<Arc<NotifyRead<K, V>>, K, V>;
324
325impl<R, K: Eq + Hash + Clone, V: Clone> Registration<R, K, V>
326where
327    R: Deref<Target = NotifyRead<K, V>>,
328{
329    pub fn key(&self) -> &K {
330        &self
331            .registration
332            .as_ref()
333            .expect("registration is only taken on drop")
334            .0
335    }
336
337    pub fn try_recv(&mut self) -> Result<V, oneshot::error::TryRecvError> {
338        self.registration
339            .as_mut()
340            .expect("registration is only taken on drop")
341            .1
342            .try_recv()
343    }
344}
345
346impl<R, K: Eq + Hash + Clone + Unpin, V: Clone + Unpin> Future for Registration<R, K, V>
347where
348    R: Deref<Target = NotifyRead<K, V>> + Unpin,
349{
350    type Output = V;
351
352    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
353        let receiver = self
354            .registration
355            .as_mut()
356            .map(|(_key, receiver)| receiver)
357            .expect("poll can not be called after drop");
358        let poll = Pin::new(receiver).poll(cx);
359        if poll.is_ready() {
360            // When polling complete we no longer need to cancel
361            self.registration.take();
362        }
363        poll.map(|r| r.expect("Sender never drops when registration is pending"))
364    }
365}
366
367impl<R, K: Eq + Hash + Clone, V: Clone> Drop for Registration<R, K, V>
368where
369    R: Deref<Target = NotifyRead<K, V>>,
370{
371    fn drop(&mut self) {
372        if let Some((key, receiver)) = self.registration.take() {
373            mem::drop(receiver);
374            // Receiver is dropped before cleanup
375            self.this.cleanup(&key)
376        }
377    }
378}
379
380/// Blocking counterpart of [`Registration`]: resolved via [`Self::wait`], which blocks
381/// the calling OS thread. Dropping it before waiting de-registers from the pending
382/// list.
383pub struct BlockingRegistration<'a, K: Eq + Hash + Clone, V: Clone> {
384    this: &'a NotifyRead<K, V>,
385    registration: Option<(K, blocking_oneshot::Receiver<V>)>,
386}
387
388impl<K: Eq + Hash + Clone, V: Clone> BlockingRegistration<'_, K, V> {
389    /// Block the calling thread until the key is notified. See
390    /// [`NotifyRead::read_one_blocking`] for the msim constraints.
391    pub fn wait(mut self) -> V {
392        let (_key, receiver) = self
393            .registration
394            .take()
395            .expect("registration is only taken here, and wait consumes self");
396        // No cleanup needed after this point: a successful recv means notify() removed
397        // the registration, and on panic the sender is already gone.
398        receiver
399            .blocking_recv()
400            .expect("Sender never drops when registration is pending")
401    }
402}
403
404impl<K: Eq + Hash + Clone, V: Clone> Drop for BlockingRegistration<'_, K, V> {
405    fn drop(&mut self) {
406        if let Some((key, receiver)) = self.registration.take() {
407            mem::drop(receiver);
408            // Receiver is dropped before cleanup
409            self.this.cleanup(&key)
410        }
411    }
412}
413
414impl<K: Eq + Hash + Clone, V: Clone> Default for NotifyRead<K, V> {
415    fn default() -> Self {
416        Self::new()
417    }
418}
419
420#[cfg(test)]
421mod tests {
422    use super::*;
423    use futures::future::join_all;
424    use std::sync::Arc;
425    use tokio::time::timeout;
426
427    #[tokio::test]
428    pub async fn test_notify_read() {
429        let notify_read = NotifyRead::<u64, u64>::new();
430        let mut registrations = notify_read.register_all(&[1, 2, 3]);
431        assert_eq!(3, notify_read.count_pending.load(Ordering::Relaxed));
432        registrations.pop();
433        assert_eq!(2, notify_read.count_pending.load(Ordering::Relaxed));
434        notify_read.notify(&2, &2);
435        notify_read.notify(&1, &1);
436        let reads = join_all(registrations).await;
437        assert_eq!(0, notify_read.count_pending.load(Ordering::Relaxed));
438        assert_eq!(reads, vec![1, 2]);
439        // ensure cleanup is done correctly
440        for pending in &notify_read.pending {
441            assert!(pending.lock().is_empty());
442        }
443    }
444
445    #[tokio::test]
446    pub async fn test_register_one_owned() {
447        let notify_read = Arc::new(NotifyRead::<u64, u64>::new());
448
449        let mut fired = notify_read.register_one_owned(&1);
450        let dropped = notify_read.register_one_owned(&2);
451        assert_eq!(2, notify_read.num_pending());
452        assert_eq!(&1, fired.key());
453        assert_eq!(Err(oneshot::error::TryRecvError::Empty), fired.try_recv());
454
455        // A notified value is observable via try_recv without the registration ever
456        // being polled as a future, which is what lets an owner test readiness
457        // synchronously.
458        notify_read.notify(&1, &7);
459        assert_eq!(Ok(7), fired.try_recv());
460        assert_eq!(1, notify_read.num_pending());
461
462        // Dropping deregisters; notifying a departed key is a no-op.
463        drop(dropped);
464        assert_eq!(0, notify_read.num_pending());
465        notify_read.notify(&2, &9);
466        drop(fired);
467        assert_eq!(0, notify_read.num_pending());
468        for pending in &notify_read.pending {
469            assert!(pending.lock().is_empty());
470        }
471    }
472
473    #[tokio::test]
474    pub async fn test_notify_read_cancellation() {
475        let notify_read = Arc::new(NotifyRead::<u64, u64>::new());
476
477        // Start a read that will wait indefinitely
478        let read_future = notify_read.read(
479            "test_task",
480            &[1, 2, 3],
481            |_keys| vec![None, None, None], // All keys will wait
482        );
483
484        // Use timeout to cancel the read after a short duration
485        let result = timeout(Duration::from_millis(100), read_future).await;
486
487        // Verify the read was cancelled
488        assert!(result.is_err());
489
490        // Give some time for cleanup to complete
491        tokio::time::sleep(Duration::from_millis(50)).await;
492
493        // When the read is cancelled, the registrations are cleaned up
494        // so the pending count should be 0
495        assert_eq!(0, notify_read.count_pending.load(Ordering::Relaxed));
496
497        // Verify all pending maps are empty (cleanup was performed)
498        for pending in &notify_read.pending {
499            assert!(pending.lock().is_empty());
500        }
501    }
502
503    #[tokio::test(start_paused = true)]
504    pub async fn test_stall_log_throttle() {
505        let notify_read = NotifyRead::<u64, u64>::new();
506
507        assert!(notify_read.throttle_stall_log("task_a"));
508        assert!(!notify_read.throttle_stall_log("task_a"));
509
510        // A report for one task name must not silence a different one.
511        assert!(notify_read.throttle_stall_log("task_b"));
512        assert!(!notify_read.throttle_stall_log("task_a"));
513
514        tokio::time::advance(Duration::from_secs(STALL_LOG_INTERVAL_SECS + 1)).await;
515        assert!(notify_read.throttle_stall_log("task_a"));
516    }
517
518    #[tokio::test(start_paused = true)]
519    pub async fn test_read_blocked_past_log_interval() {
520        let notify_read = Arc::new(NotifyRead::<u64, u64>::new());
521
522        let reader = notify_read.clone();
523        let handle = tokio::spawn(async move {
524            reader
525                .read("test_task", &[1, 2, 3], |_keys| vec![Some(10), None, None])
526                .await
527        });
528
529        // Outlive several ticks so the read exercises the stall reporting path.
530        tokio::time::advance(Duration::from_secs(LONG_WAIT_LOG_INTERVAL_SECS * 4)).await;
531        assert!(!handle.is_finished());
532
533        notify_read.notify(&2, &20);
534        notify_read.notify(&3, &30);
535
536        // Values are returned in key order, not completion order.
537        assert_eq!(handle.await.unwrap(), vec![10, 20, 30]);
538        assert_eq!(0, notify_read.count_pending.load(Ordering::Relaxed));
539        for pending in &notify_read.pending {
540            assert!(pending.lock().is_empty());
541        }
542    }
543}
mysten_common/sync/notify_read.rs

mysten_common/sync/
notify_read.rs
