iota_common/sync/
notify_read.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    collections::{HashMap, HashSet, hash_map::DefaultHasher},
7    error::Error,
8    future::Future,
9    hash::{Hash, Hasher},
10    mem,
11    pin::Pin,
12    sync::{
13        Arc,
14        atomic::{AtomicUsize, Ordering},
15    },
16    task::{Context, Poll},
17    time::Duration,
18};
19
20use futures::future::{Either, join_all};
21use iota_metrics::spawn_monitored_task;
22use parking_lot::{Mutex, MutexGuard};
23use tokio::{
24    sync::oneshot,
25    time::{Instant, interval_at},
26};
27use tracing::warn;
28
29use crate::debug_fatal;
30
31type Registrations<V> = Vec<oneshot::Sender<V>>;
32
33/// Wrapper that ensures a spawned task is aborted when dropped
34struct TaskAbortOnDrop {
35    handle: Option<tokio::task::JoinHandle<()>>,
36}
37
38impl TaskAbortOnDrop {
39    fn new(handle: tokio::task::JoinHandle<()>) -> Self {
40        Self {
41            handle: Some(handle),
42        }
43    }
44}
45
46impl Drop for TaskAbortOnDrop {
47    fn drop(&mut self) {
48        if let Some(handle) = self.handle.take() {
49            handle.abort();
50        }
51    }
52}
53
54/// Interval duration for logging waiting keys when reads take too long
55const LONG_WAIT_LOG_INTERVAL_SECS: u64 = 10;
56
57pub const CHECKPOINT_BUILDER_NOTIFY_READ_TASK_NAME: &str =
58    "CheckpointBuilder::notify_read_executed_effects";
59
60pub struct NotifyRead<K, V> {
61    pending: Vec<Mutex<HashMap<K, Registrations<V>>>>,
62    count_pending: AtomicUsize,
63}
64
65impl<K: Eq + Hash + Clone, V: Clone> NotifyRead<K, V> {
66    pub fn new() -> Self {
67        let pending = (0..255).map(|_| Default::default()).collect();
68        let count_pending = Default::default();
69        Self {
70            pending,
71            count_pending,
72        }
73    }
74
75    /// Asynchronously notifies waiters and return number of remaining pending
76    /// registration
77    pub fn notify(&self, key: &K, value: &V) -> usize {
78        let registrations = self.pending(key).remove(key);
79        let Some(registrations) = registrations else {
80            return self.count_pending.load(Ordering::Relaxed);
81        };
82        let rem = self
83            .count_pending
84            .fetch_sub(registrations.len(), Ordering::Relaxed);
85        for registration in registrations {
86            registration.send(value.clone()).ok();
87        }
88        rem
89    }
90
91    pub fn register_one(&self, key: &K) -> Registration<'_, K, V> {
92        self.count_pending.fetch_add(1, Ordering::Relaxed);
93        let (sender, receiver) = oneshot::channel();
94        self.register(key, sender);
95        Registration {
96            this: self,
97            registration: Some((key.clone(), receiver)),
98        }
99    }
100
101    pub fn register_all(&self, keys: &[K]) -> Vec<Registration<'_, K, V>> {
102        self.count_pending.fetch_add(keys.len(), Ordering::Relaxed);
103        let mut registrations = vec![];
104        for key in keys.iter() {
105            let (sender, receiver) = oneshot::channel();
106            self.register(key, sender);
107            let registration = Registration {
108                this: self,
109                registration: Some((key.clone(), receiver)),
110            };
111            registrations.push(registration);
112        }
113        registrations
114    }
115
116    fn register(&self, key: &K, sender: oneshot::Sender<V>) {
117        self.pending(key)
118            .entry(key.clone())
119            .or_default()
120            .push(sender);
121    }
122
123    fn pending(&self, key: &K) -> MutexGuard<'_, HashMap<K, Registrations<V>>> {
124        let mut state = DefaultHasher::new();
125        key.hash(&mut state);
126        let hash = state.finish();
127        let pending = self
128            .pending
129            .get((hash % self.pending.len() as u64) as usize)
130            .unwrap();
131        pending.lock()
132    }
133
134    pub fn num_pending(&self) -> usize {
135        self.count_pending.load(Ordering::Relaxed)
136    }
137
138    fn cleanup(&self, key: &K) {
139        let mut pending = self.pending(key);
140        // it is possible that registration was fulfilled before we get here
141        let Some(registrations) = pending.get_mut(key) else {
142            return;
143        };
144        let mut count_deleted = 0usize;
145        registrations.retain(|s| {
146            let delete = s.is_closed();
147            if delete {
148                count_deleted += 1;
149            }
150            !delete
151        });
152        self.count_pending
153            .fetch_sub(count_deleted, Ordering::Relaxed);
154        if registrations.is_empty() {
155            pending.remove(key);
156        }
157    }
158}
159
160impl<K: Eq + Hash + Clone + Unpin + std::fmt::Debug + Send + Sync + 'static, V: Clone + Unpin>
161    NotifyRead<K, V>
162{
163    pub async fn read<E: Error>(
164        &self,
165        task_name: &'static str,
166        keys: &[K],
167        fetch: impl FnOnce(&[K]) -> Result<Vec<Option<V>>, E>,
168    ) -> Result<Vec<V>, E> {
169        let _metrics_scope = iota_metrics::monitored_scope(task_name);
170        let registrations = self.register_all(keys);
171
172        let results = fetch(keys)?;
173
174        // Track which keys are still waiting
175        let waiting_keys: HashSet<K> = keys
176            .iter()
177            .zip(results.iter())
178            .filter(|&(_key, result)| result.is_none())
179            .map(|(key, _result)| key.clone())
180            .collect();
181        let has_waiting_keys = !waiting_keys.is_empty();
182        let waiting_keys = Arc::new(Mutex::new(waiting_keys));
183
184        // Spawn logging task if there are waiting keys
185        let _log_handle_guard = if has_waiting_keys {
186            let waiting_keys_clone = waiting_keys.clone();
187            let start_time = Instant::now();
188            let task_name = task_name.to_string();
189
190            let handle = spawn_monitored_task!(async move {
191                // Only start logging after the first interval.
192                let start = Instant::now() + Duration::from_secs(LONG_WAIT_LOG_INTERVAL_SECS);
193                let mut interval =
194                    interval_at(start, Duration::from_secs(LONG_WAIT_LOG_INTERVAL_SECS));
195
196                loop {
197                    interval.tick().await;
198                    let current_waiting = waiting_keys_clone.lock();
199                    if current_waiting.is_empty() {
200                        break;
201                    }
202                    let keys_vec: Vec<_> = current_waiting.iter().cloned().collect();
203                    drop(current_waiting); // Release lock before logging
204
205                    let elapsed_secs = start_time.elapsed().as_secs();
206
207                    warn!(
208                        "[{task_name}] Still waiting for {elapsed_secs}s for {} keys: {keys_vec:?}",
209                        keys_vec.len(),
210                    );
211
212                    if task_name == CHECKPOINT_BUILDER_NOTIFY_READ_TASK_NAME && elapsed_secs >= 60 {
213                        debug_fatal!("{task_name} is stuck");
214                    }
215                }
216            });
217            Some(TaskAbortOnDrop::new(handle))
218        } else {
219            None
220        };
221
222        let results =
223            results
224                .into_iter()
225                .zip(registrations)
226                .zip(keys.iter())
227                .map(|((a, r), key)| match a {
228                    // Note that Some() clause also drops registration that is already fulfilled
229                    Some(ready) => Either::Left(futures::future::ready(ready)),
230                    None => {
231                        let waiting_keys = waiting_keys.clone();
232                        let key = key.clone();
233                        Either::Right(async move {
234                            let result = r.await;
235                            // Remove this key from the waiting set
236                            waiting_keys.lock().remove(&key);
237                            result
238                        })
239                    }
240                });
241
242        // The logging task will be automatically aborted when _log_handle_guard is
243        // dropped
244
245        Ok(join_all(results).await)
246    }
247}
248
249/// Registration resolves to the value but also provides safe cancellation
250/// When Registration is dropped before it is resolved, we de-register from the
251/// pending list
252pub struct Registration<'a, K: Eq + Hash + Clone, V: Clone> {
253    this: &'a NotifyRead<K, V>,
254    registration: Option<(K, oneshot::Receiver<V>)>,
255}
256
257impl<K: Eq + Hash + Clone + Unpin, V: Clone + Unpin> Future for Registration<'_, K, V> {
258    type Output = V;
259
260    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
261        let receiver = self
262            .registration
263            .as_mut()
264            .map(|(_key, receiver)| receiver)
265            .expect("poll can not be called after drop");
266        let poll = Pin::new(receiver).poll(cx);
267        if poll.is_ready() {
268            // When polling complete we no longer need to cancel
269            self.registration.take();
270        }
271        poll.map(|r| r.expect("Sender never drops when registration is pending"))
272    }
273}
274
275impl<K: Eq + Hash + Clone, V: Clone> Drop for Registration<'_, K, V> {
276    fn drop(&mut self) {
277        if let Some((key, receiver)) = self.registration.take() {
278            mem::drop(receiver);
279            // Receiver is dropped before cleanup
280            self.this.cleanup(&key)
281        }
282    }
283}
284impl<K: Eq + Hash + Clone, V: Clone> Default for NotifyRead<K, V> {
285    fn default() -> Self {
286        Self::new()
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use std::{convert::Infallible, sync::Arc};
293
294    use futures::future::join_all;
295    use tokio::time::timeout;
296
297    use super::*;
298
299    #[tokio::test]
300    pub async fn test_notify_read() {
301        let notify_read = NotifyRead::<u64, u64>::new();
302        let mut registrations = notify_read.register_all(&[1, 2, 3]);
303        assert_eq!(3, notify_read.count_pending.load(Ordering::Relaxed));
304        registrations.pop();
305        assert_eq!(2, notify_read.count_pending.load(Ordering::Relaxed));
306        notify_read.notify(&2, &2);
307        notify_read.notify(&1, &1);
308        let reads = join_all(registrations).await;
309        assert_eq!(0, notify_read.count_pending.load(Ordering::Relaxed));
310        assert_eq!(reads, vec![1, 2]);
311        // ensure cleanup is done correctly
312        for pending in &notify_read.pending {
313            assert!(pending.lock().is_empty());
314        }
315    }
316
317    #[tokio::test]
318    pub async fn test_notify_read_cancellation() {
319        let notify_read = Arc::new(NotifyRead::<u64, u64>::new());
320
321        // Start a read that will wait indefinitely
322        let read_future = notify_read.read::<Infallible>(
323            "test_task",
324            &[1, 2, 3],
325            |_keys| Ok(vec![None, None, None]), // All keys will wait
326        );
327
328        // Use timeout to cancel the read after a short duration
329        let result = timeout(Duration::from_millis(100), read_future).await;
330
331        // Verify the read was cancelled
332        assert!(result.is_err());
333
334        // Give some time for cleanup to complete
335        tokio::time::sleep(Duration::from_millis(50)).await;
336
337        // When the read is cancelled, the registrations are cleaned up
338        // so the pending count should be 0
339        assert_eq!(0, notify_read.count_pending.load(Ordering::Relaxed));
340
341        // Verify all pending maps are empty (cleanup was performed)
342        for pending in &notify_read.pending {
343            assert!(pending.lock().is_empty());
344        }
345    }
346}
iota_common/sync/notify_read.rs

iota_common/sync/
notify_read.rs
