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        let checkpoint_builder_stall_threshold_secs = if crate::in_antithesis() {
273            // Antithesis faults can delay checkpoint effects throughout node recovery.
274            120
275        } else {
276            60
277        };
278
279        loop {
280            tokio::select! {
281                values = &mut join => return values,
282                _ = interval.tick() => {
283                    let elapsed_secs = start_time.elapsed().as_secs();
284
285                    // Deduplicate logging by task name. When many reads are blocked,
286                    // logs will sample blocked reads per task and the keys we are
287                    // waiting on for those reads.
288                    if self.throttle_stall_log(task_name) {
289                        let mut sample: Vec<&K> = Vec::with_capacity(MAX_SAMPLED_KEYS);
290                        let mut outstanding = 0usize;
291                        for key in &waiting_keys {
292                            if self.pending(key).contains_key(key) {
293                                outstanding += 1;
294                                if sample.len() < MAX_SAMPLED_KEYS {
295                                    sample.push(key);
296                                }
297                            }
298                        }
299
300                        warn!(
301                            "[{task_name}] Still waiting {elapsed_secs}s. {} registrations pending, this read still blocked on {outstanding} of {} key(s): {sample:?}",
302                            self.num_pending(),
303                            waiting_keys.len(),
304                        );
305                    }
306
307                    if task_name == CHECKPOINT_BUILDER_NOTIFY_READ_TASK_NAME
308                        && elapsed_secs >= checkpoint_builder_stall_threshold_secs
309                    {
310                        debug_fatal!("{} is stuck", task_name);
311                    }
312                }
313            }
314        }
315    }
316}
317
318/// Registration resolves to the value but also provides safe cancellation
319/// When Registration is dropped before it is resolved, we de-register from the pending list
320///
321/// Generic over how it holds the registry: borrowed for the await-in-place pattern,
322/// or by `Arc` ([`OwnedRegistration`]) so it can be stored in long-lived structures.
323pub struct Registration<R, K: Eq + Hash + Clone, V: Clone>
324where
325    R: Deref<Target = NotifyRead<K, V>>,
326{
327    this: R,
328    registration: Option<(K, oneshot::Receiver<V>)>,
329}
330
331pub type OwnedRegistration<K, V> = Registration<Arc<NotifyRead<K, V>>, K, V>;
332
333impl<R, K: Eq + Hash + Clone, V: Clone> Registration<R, K, V>
334where
335    R: Deref<Target = NotifyRead<K, V>>,
336{
337    pub fn key(&self) -> &K {
338        &self
339            .registration
340            .as_ref()
341            .expect("registration is only taken on drop")
342            .0
343    }
344
345    pub fn try_recv(&mut self) -> Result<V, oneshot::error::TryRecvError> {
346        self.registration
347            .as_mut()
348            .expect("registration is only taken on drop")
349            .1
350            .try_recv()
351    }
352}
353
354impl<R, K: Eq + Hash + Clone + Unpin, V: Clone + Unpin> Future for Registration<R, K, V>
355where
356    R: Deref<Target = NotifyRead<K, V>> + Unpin,
357{
358    type Output = V;
359
360    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
361        let receiver = self
362            .registration
363            .as_mut()
364            .map(|(_key, receiver)| receiver)
365            .expect("poll can not be called after drop");
366        let poll = Pin::new(receiver).poll(cx);
367        if poll.is_ready() {
368            // When polling complete we no longer need to cancel
369            self.registration.take();
370        }
371        poll.map(|r| r.expect("Sender never drops when registration is pending"))
372    }
373}
374
375impl<R, K: Eq + Hash + Clone, V: Clone> Drop for Registration<R, K, V>
376where
377    R: Deref<Target = NotifyRead<K, V>>,
378{
379    fn drop(&mut self) {
380        if let Some((key, receiver)) = self.registration.take() {
381            mem::drop(receiver);
382            // Receiver is dropped before cleanup
383            self.this.cleanup(&key)
384        }
385    }
386}
387
388/// Blocking counterpart of [`Registration`]: resolved via [`Self::wait`], which blocks
389/// the calling OS thread. Dropping it before waiting de-registers from the pending
390/// list.
391pub struct BlockingRegistration<'a, K: Eq + Hash + Clone, V: Clone> {
392    this: &'a NotifyRead<K, V>,
393    registration: Option<(K, blocking_oneshot::Receiver<V>)>,
394}
395
396impl<K: Eq + Hash + Clone, V: Clone> BlockingRegistration<'_, K, V> {
397    /// Block the calling thread until the key is notified. See
398    /// [`NotifyRead::read_one_blocking`] for the msim constraints.
399    pub fn wait(mut self) -> V {
400        let (_key, receiver) = self
401            .registration
402            .take()
403            .expect("registration is only taken here, and wait consumes self");
404        // No cleanup needed after this point: a successful recv means notify() removed
405        // the registration, and on panic the sender is already gone.
406        receiver
407            .blocking_recv()
408            .expect("Sender never drops when registration is pending")
409    }
410}
411
412impl<K: Eq + Hash + Clone, V: Clone> Drop for BlockingRegistration<'_, K, V> {
413    fn drop(&mut self) {
414        if let Some((key, receiver)) = self.registration.take() {
415            mem::drop(receiver);
416            // Receiver is dropped before cleanup
417            self.this.cleanup(&key)
418        }
419    }
420}
421
422impl<K: Eq + Hash + Clone, V: Clone> Default for NotifyRead<K, V> {
423    fn default() -> Self {
424        Self::new()
425    }
426}
427
428#[cfg(test)]
429mod tests {
430    use super::*;
431    use futures::future::join_all;
432    use std::sync::Arc;
433    use tokio::time::timeout;
434
435    #[tokio::test]
436    pub async fn test_notify_read() {
437        let notify_read = NotifyRead::<u64, u64>::new();
438        let mut registrations = notify_read.register_all(&[1, 2, 3]);
439        assert_eq!(3, notify_read.count_pending.load(Ordering::Relaxed));
440        registrations.pop();
441        assert_eq!(2, notify_read.count_pending.load(Ordering::Relaxed));
442        notify_read.notify(&2, &2);
443        notify_read.notify(&1, &1);
444        let reads = join_all(registrations).await;
445        assert_eq!(0, notify_read.count_pending.load(Ordering::Relaxed));
446        assert_eq!(reads, vec![1, 2]);
447        // ensure cleanup is done correctly
448        for pending in &notify_read.pending {
449            assert!(pending.lock().is_empty());
450        }
451    }
452
453    #[tokio::test]
454    pub async fn test_register_one_owned() {
455        let notify_read = Arc::new(NotifyRead::<u64, u64>::new());
456
457        let mut fired = notify_read.register_one_owned(&1);
458        let dropped = notify_read.register_one_owned(&2);
459        assert_eq!(2, notify_read.num_pending());
460        assert_eq!(&1, fired.key());
461        assert_eq!(Err(oneshot::error::TryRecvError::Empty), fired.try_recv());
462
463        // A notified value is observable via try_recv without the registration ever
464        // being polled as a future, which is what lets an owner test readiness
465        // synchronously.
466        notify_read.notify(&1, &7);
467        assert_eq!(Ok(7), fired.try_recv());
468        assert_eq!(1, notify_read.num_pending());
469
470        // Dropping deregisters; notifying a departed key is a no-op.
471        drop(dropped);
472        assert_eq!(0, notify_read.num_pending());
473        notify_read.notify(&2, &9);
474        drop(fired);
475        assert_eq!(0, notify_read.num_pending());
476        for pending in &notify_read.pending {
477            assert!(pending.lock().is_empty());
478        }
479    }
480
481    #[tokio::test]
482    pub async fn test_notify_read_cancellation() {
483        let notify_read = Arc::new(NotifyRead::<u64, u64>::new());
484
485        // Start a read that will wait indefinitely
486        let read_future = notify_read.read(
487            "test_task",
488            &[1, 2, 3],
489            |_keys| vec![None, None, None], // All keys will wait
490        );
491
492        // Use timeout to cancel the read after a short duration
493        let result = timeout(Duration::from_millis(100), read_future).await;
494
495        // Verify the read was cancelled
496        assert!(result.is_err());
497
498        // Give some time for cleanup to complete
499        tokio::time::sleep(Duration::from_millis(50)).await;
500
501        // When the read is cancelled, the registrations are cleaned up
502        // so the pending count should be 0
503        assert_eq!(0, notify_read.count_pending.load(Ordering::Relaxed));
504
505        // Verify all pending maps are empty (cleanup was performed)
506        for pending in &notify_read.pending {
507            assert!(pending.lock().is_empty());
508        }
509    }
510
511    #[tokio::test(start_paused = true)]
512    pub async fn test_stall_log_throttle() {
513        let notify_read = NotifyRead::<u64, u64>::new();
514
515        assert!(notify_read.throttle_stall_log("task_a"));
516        assert!(!notify_read.throttle_stall_log("task_a"));
517
518        // A report for one task name must not silence a different one.
519        assert!(notify_read.throttle_stall_log("task_b"));
520        assert!(!notify_read.throttle_stall_log("task_a"));
521
522        tokio::time::advance(Duration::from_secs(STALL_LOG_INTERVAL_SECS + 1)).await;
523        assert!(notify_read.throttle_stall_log("task_a"));
524    }
525
526    #[tokio::test(start_paused = true)]
527    pub async fn test_read_blocked_past_log_interval() {
528        let notify_read = Arc::new(NotifyRead::<u64, u64>::new());
529
530        let reader = notify_read.clone();
531        let handle = tokio::spawn(async move {
532            reader
533                .read("test_task", &[1, 2, 3], |_keys| vec![Some(10), None, None])
534                .await
535        });
536
537        // Outlive several ticks so the read exercises the stall reporting path.
538        tokio::time::advance(Duration::from_secs(LONG_WAIT_LOG_INTERVAL_SECS * 4)).await;
539        assert!(!handle.is_finished());
540
541        notify_read.notify(&2, &20);
542        notify_read.notify(&3, &30);
543
544        // Values are returned in key order, not completion order.
545        assert_eq!(handle.await.unwrap(), vec![10, 20, 30]);
546        assert_eq!(0, notify_read.count_pending.load(Ordering::Relaxed));
547        for pending in &notify_read.pending {
548            assert!(pending.lock().is_empty());
549        }
550    }
551}
mysten_common/sync/notify_read.rs

mysten_common/sync/
notify_read.rs
