iota_common/sync/
async_once_cell.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::sync::Arc;
6
7use parking_lot::Mutex;
8use tokio::sync::{OwnedRwLockWriteGuard, RwLock};
9
10/// This structure contains a cell for a single value.
11/// The cell can be written only once, and can be read many times.
12/// Readers are provided with async API, that waits for write to happen.
13/// This is similar to tokio::sync::watch, except one difference:
14/// * tokio::sync::watch requires existing receiver to work. If no subscriber is registered, and the
15///   value is sent to channel, the value is dropped
16/// * Unlike with tokio::sync::watch, it is possible to write to AsyncOnceCell when no readers are
17///   registered, and value will be available later when AsyncOnceCell::get is called
18pub struct AsyncOnceCell<T> {
19    value: Arc<RwLock<Option<T>>>,
20    writer: Mutex<Option<OwnedRwLockWriteGuard<Option<T>>>>,
21}
22
23impl<T: Send + Clone> AsyncOnceCell<T> {
24    pub fn new() -> Self {
25        let value = Arc::new(RwLock::new(None));
26        let writer = value
27            .clone()
28            .try_write_owned()
29            .expect("Write lock can not fail here");
30        let writer = Mutex::new(Some(writer));
31        Self { value, writer }
32    }
33
34    pub async fn get(&self) -> T {
35        self.value
36            .read()
37            .await
38            .as_ref()
39            .cloned()
40            .expect("Value is available when writer is dropped")
41    }
42
43    /// Sets the value and notifies waiters. Return error if called twice
44    #[expect(clippy::result_unit_err)]
45    pub fn set(&self, value: T) -> Result<(), ()> {
46        let mut writer = self.writer.lock();
47        match writer.take() {
48            None => Err(()),
49            Some(mut writer) => {
50                *writer = Some(value);
51                Ok(())
52            }
53        }
54    }
55}
56
57impl<T: Send + Clone> Default for AsyncOnceCell<T> {
58    fn default() -> Self {
59        Self::new()
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[tokio::test]
68    async fn async_once_cell_test() {
69        let cell = Arc::new(AsyncOnceCell::<u64>::new());
70        let cell2 = cell.clone();
71        let wait = tokio::spawn(async move { cell2.get().await });
72        tokio::task::yield_now().await;
73        cell.set(15).unwrap();
74        assert!(cell.set(16).is_err());
75        assert_eq!(15, cell.get().await);
76        assert_eq!(15, wait.await.unwrap());
77    }
78}
iota_common/sync/async_once_cell.rs

iota_common/sync/
async_once_cell.rs
