xet_runtime/utils/
rw_task_lock.rs

1use std::future::Future;
2use std::mem::replace;
3use std::ops::Deref;
4
5use thiserror::Error;
6use tokio::sync::{RwLock, RwLockReadGuard};
7use tokio::task::{JoinError, JoinHandle};
8#[cfg(target_family = "wasm")]
9use tokio_with_wasm::alias as tokio;
10
11#[derive(Debug, Error)]
12#[non_exhaustive]
13pub enum RwTaskLockError {
14    #[error(transparent)]
15    JoinError(#[from] JoinError),
16
17    #[error("Attempting to access value not available due to a previously reported error.")]
18    CalledAfterError,
19}
20
21enum RwTaskLockState<T, E> {
22    Pending(JoinHandle<Result<T, E>>),
23    Ready(T),
24    Error,
25}
26
27/// Custom read guard: keeps the RwLockReadGuard alive, exposes &T.
28pub struct RwTaskLockReadGuard<'a, T, E> {
29    guard: RwLockReadGuard<'a, RwTaskLockState<T, E>>,
30}
31
32impl<T, E> Deref for RwTaskLockReadGuard<'_, T, E> {
33    type Target = T;
34    fn deref(&self) -> &T {
35        match &*self.guard {
36            RwTaskLockState::Ready(val) => val,
37            _ => unreachable!("Read guard is only constructed for Ready state"),
38        }
39    }
40}
41
42/// A one-time async-initialized, lockable value that yields a read guard after initialization.
43///
44/// `RwTaskLock<T, E>` allows you to wrap a future or ready value so the computation
45/// (e.g., background task, async function) is performed at most once, even if
46/// multiple callers invoke `.read()` concurrently. After the computation,
47/// the result is cached. If successful, all future `.read()` calls yield a
48/// read guard on the stored value. If an error occurs, all subsequent calls
49/// return the error (error value must be `Clone`).
50///
51/// # Example
52/// ```
53/// use tokio::time;
54/// use xet_runtime::utils::{RwTaskLock, RwTaskLockError};
55/// #[tokio::main]
56/// async fn main() -> Result<(), RwTaskLockError> {
57///     let lock = RwTaskLock::from_task(async {
58///         time::sleep(std::time::Duration::from_millis(50)).await;
59///         Ok::<_, RwTaskLockError>(vec![1, 2, 3])
60///     });
61///     let guard = lock.read().await?;
62///     assert_eq!(&*guard, &[1, 2, 3]);
63///     Ok(())
64/// }
65/// ```
66pub struct RwTaskLock<T, E>
67where
68    T: Send + Sync + 'static,
69    E: Send + Sync + 'static + From<RwTaskLockError>,
70{
71    state: RwLock<RwTaskLockState<T, E>>,
72}
73
74impl<T, E> RwTaskLock<T, E>
75where
76    T: Send + Sync + 'static,
77    E: Send + Sync + 'static + From<RwTaskLockError>,
78{
79    /// From a ready value.
80    pub fn from_value(val: T) -> Self {
81        Self {
82            state: RwLock::new(RwTaskLockState::Ready(val)),
83        }
84    }
85
86    /// From a future yielding Result<T, E>.
87    pub fn from_task<Fut>(fut: Fut) -> Self
88    where
89        Fut: Future<Output = Result<T, E>> + Send + 'static,
90    {
91        let task = tokio::spawn(fut);
92
93        Self {
94            state: RwLock::new(RwTaskLockState::Pending(task)),
95        }
96    }
97
98    /// Awaitable read: yields a custom read guard or error.
99    pub async fn read(&self) -> Result<RwTaskLockReadGuard<'_, T, E>, E> {
100        // Fast path
101        {
102            let state = self.state.read().await;
103            match &*state {
104                RwTaskLockState::Ready(_) => {
105                    return Ok(RwTaskLockReadGuard { guard: state });
106                },
107                RwTaskLockState::Error => return Err(E::from(RwTaskLockError::CalledAfterError)),
108                RwTaskLockState::Pending(_) => {},
109            }
110        }
111        // Acquire write lock to initialize if necessary
112        let mut state = self.state.write().await;
113
114        match replace(&mut *state, RwTaskLockState::Error) {
115            RwTaskLockState::Ready(v) => {
116                *state = RwTaskLockState::Ready(v);
117            },
118            RwTaskLockState::Error => {
119                return Err(E::from(RwTaskLockError::CalledAfterError));
120            },
121            RwTaskLockState::Pending(jh) => {
122                match jh.await.map_err(RwTaskLockError::JoinError)? {
123                    Ok(v) => {
124                        *state = RwTaskLockState::Ready(v);
125                    },
126                    Err(e) => {
127                        *state = RwTaskLockState::Error;
128                        return Err(e);
129                    },
130                };
131            },
132        };
133
134        Ok(RwTaskLockReadGuard {
135            guard: state.downgrade(),
136        })
137    }
138
139    /// Update the current value by applying an async function to it, storing the result as the new value.
140    ///
141    /// - If the current value is in the `Ready` state, the function is immediately scheduled as a background task with
142    ///   the current value, and the state becomes `Pending` until completion.
143    /// - If the value is in the `Pending` state, this chains the update: when the background task completes, the
144    ///   updater will be called on the resulting value.
145    /// - If the value is in the `Error` state, returns an error and does nothing.
146    ///
147    /// Returns `Ok(())` if the update is scheduled. Errors if the value is already in an error state.
148    ///
149    /// # Example: Chaining updates
150    /// ```
151    /// use tokio::time;
152    /// use xet_runtime::utils::{RwTaskLock, RwTaskLockError};
153    /// #[tokio::main]
154    /// async fn main() -> Result<(), RwTaskLockError> {
155    ///     let lock = RwTaskLock::from_value(10);
156    ///     lock.update(|v| async move { Ok::<_, RwTaskLockError>(v * 2) }).await?;
157    ///     assert_eq!(*lock.read().await?, 20);
158    ///
159    ///     lock.update(|v| async move { Ok::<_, RwTaskLockError>(v + 5) }).await?;
160    ///     assert_eq!(*lock.read().await?, 25);
161    ///     Ok(())
162    /// }
163    /// ```
164    ///
165    /// # Example: Chained with pending state
166    /// ```
167    /// use std::sync::Arc;
168    ///
169    /// use tokio::time;
170    /// use xet_runtime::utils::{RwTaskLock, RwTaskLockError};
171    /// #[tokio::main]
172    /// async fn main() -> Result<(), RwTaskLockError> {
173    ///     let lock = Arc::new(RwTaskLock::from_task(async {
174    ///         time::sleep(std::time::Duration::from_millis(10)).await;
175    ///         Ok::<_, RwTaskLockError>(10)
176    ///     }));
177    ///     let lock2 = lock.clone();
178    ///
179    ///     // Chain update while value is still pending
180    ///     lock2.update(|v| async move { Ok::<_, RwTaskLockError>(v + 10) }).await?;
181    ///     assert_eq!(*lock.read().await?, 20);
182    ///     Ok(())
183    /// }
184    /// ```
185    pub async fn update<Fut, Updater>(&self, updater: Updater) -> Result<(), RwTaskLockError>
186    where
187        Updater: FnOnce(T) -> Fut + Send + 'static,
188        Fut: Future<Output = Result<T, E>> + Send + 'static,
189    {
190        use RwTaskLockState::*;
191
192        let mut state_lg = self.state.write().await;
193
194        let state = replace(&mut *state_lg, RwTaskLockState::Error);
195
196        match state {
197            Pending(jh) => {
198                // Chain the old pending future, then the updater.
199                let new_task = tokio::spawn(async move {
200                    let current = jh.await.map_err(RwTaskLockError::JoinError)??;
201                    updater(current).await
202                });
203                *state_lg = Pending(new_task);
204                Ok(())
205            },
206            Ready(v) => {
207                // Start new computation from current value.
208                *state_lg = Pending(tokio::spawn(updater(v)));
209                Ok(())
210            },
211            Error => {
212                // Can't update if in error.
213                *state_lg = Error;
214                Err(RwTaskLockError::CalledAfterError)
215            },
216        }
217    }
218}
219
220#[cfg(test)]
221mod tests {
222
223    use super::*;
224
225    #[tokio::test]
226    async fn test_from_value() {
227        let lock: RwTaskLock<_, RwTaskLockError> = RwTaskLock::from_value(7);
228        let guard = lock.read().await.unwrap();
229        assert_eq!(*guard, 7);
230        let guard2 = lock.read().await.unwrap();
231        assert_eq!(*guard2, 7);
232    }
233
234    #[tokio::test]
235    async fn test_from_future_success() {
236        let lock = RwTaskLock::from_task(async {
237            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
238            Ok::<_, RwTaskLockError>(999)
239        });
240        let guard = lock.read().await.unwrap();
241        assert_eq!(*guard, 999);
242        let guard2 = lock.read().await.unwrap();
243        assert_eq!(*guard2, 999);
244    }
245
246    #[tokio::test]
247    async fn test_from_future_error() {
248        let lock = RwTaskLock::<u8, RwTaskLockError>::from_task(async { Err(RwTaskLockError::CalledAfterError) });
249        let result = lock.read().await;
250        assert!(matches!(result, Err(RwTaskLockError::CalledAfterError)));
251        let result2 = lock.read().await;
252        assert!(matches!(result2, Err(RwTaskLockError::CalledAfterError)));
253    }
254
255    #[tokio::test]
256    async fn test_concurrent_read() {
257        use std::sync::Arc;
258        let lock = Arc::new(RwTaskLock::from_task(async {
259            tokio::time::sleep(std::time::Duration::from_millis(30)).await;
260            Ok::<_, RwTaskLockError>("concurrent".to_string())
261        }));
262        let lock1 = lock.clone();
263        let lock2 = lock.clone();
264        let (a, b) = tokio::join!(lock1.read(), lock2.read());
265        assert_eq!(*a.unwrap(), "concurrent");
266        assert_eq!(*b.unwrap(), "concurrent");
267    }
268
269    #[tokio::test]
270    async fn test_error_then_retrieval() {
271        let lock = RwTaskLock::<u8, RwTaskLockError>::from_task(async { Err(RwTaskLockError::CalledAfterError) });
272        let _ = lock.read().await;
273        let result = lock.read().await;
274        assert!(matches!(result, Err(RwTaskLockError::CalledAfterError)));
275    }
276
277    #[tokio::test]
278    async fn test_update_from_ready() {
279        let lock = RwTaskLock::from_value(100);
280        lock.update(|v| async move { Ok::<_, RwTaskLockError>(v + 1) }).await.unwrap();
281        let guard = lock.read().await.unwrap();
282        assert_eq!(*guard, 101);
283    }
284
285    #[tokio::test]
286    async fn test_update_chained_pending() {
287        use std::sync::Arc;
288        let lock = Arc::new(RwTaskLock::from_task(async {
289            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
290            Ok::<_, RwTaskLockError>(5)
291        }));
292        let lock2 = lock.clone();
293        // Schedule update before initial value is ready
294        lock2.update(|v| async move { Ok::<_, RwTaskLockError>(v * 3) }).await.unwrap();
295        let guard = lock.read().await.unwrap();
296        assert_eq!(*guard, 15);
297    }
298
299    #[tokio::test]
300    async fn test_update_error_state() {
301        let lock = RwTaskLock::<i32, RwTaskLockError>::from_task(async { Err(RwTaskLockError::CalledAfterError) });
302        let _ = lock.read().await;
303        let result = lock.update(|v| async move { Ok::<_, RwTaskLockError>(v + 1) }).await;
304        assert!(matches!(result, Err(RwTaskLockError::CalledAfterError)));
305    }
306
307    #[tokio::test]
308    async fn test_update_to_error() {
309        let lock = RwTaskLock::from_value(123);
310        // Updater produces an error
311        lock.update(|_v| async move { Err(RwTaskLockError::CalledAfterError) })
312            .await
313            .unwrap();
314        let result = lock.read().await;
315        assert!(matches!(result, Err(RwTaskLockError::CalledAfterError)));
316    }
317
318    #[tokio::test]
319    async fn test_multiple_updates() {
320        let lock = RwTaskLock::from_value(1);
321        lock.update(|v| async move { Ok::<_, RwTaskLockError>(v + 10) }).await.unwrap();
322        lock.update(|v| async move { Ok::<_, RwTaskLockError>(v * 2) }).await.unwrap();
323        let guard = lock.read().await.unwrap();
324        assert_eq!(*guard, 22);
325    }
326}
xet_runtime/utils/rw_task_lock.rs

xet_runtime/utils/
rw_task_lock.rs
