mysten_common/sync/
high_water_mark.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::future::Future;
5
6use tokio::sync::watch;
7
8/// A value that only moves forward, with waiters that resolve once it reaches a target.
9///
10/// `advance_to` is the only mutator and ignores anything at or below the current value, so
11/// writers that observe progress out of order cannot move the mark backwards.
12#[derive(Debug)]
13pub struct HighWaterMark<T> {
14    sender: watch::Sender<T>,
15}
16
17impl<T: Copy + Ord + Send + Sync + 'static> HighWaterMark<T> {
18    pub fn new(initial: T) -> Self {
19        Self {
20            sender: watch::Sender::new(initial),
21        }
22    }
23
24    pub fn current(&self) -> T {
25        *self.sender.borrow()
26    }
27
28    /// Returns true if the mark moved.
29    pub fn advance_to(&self, value: T) -> bool {
30        self.sender.send_if_modified(|current| {
31            if value > *current {
32                *current = value;
33                true
34            } else {
35                false
36            }
37        })
38    }
39
40    /// Resolves to the first observed value at or above `target`, or to `None` if the mark is
41    /// dropped before reaching it. The returned future does not borrow `self`.
42    pub fn wait_for(&self, target: T) -> impl Future<Output = Option<T>> + Send + 'static {
43        let mut receiver = self.sender.subscribe();
44        async move {
45            receiver
46                .wait_for(|current| *current >= target)
47                .await
48                .ok()
49                .map(|current| *current)
50        }
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57
58    #[tokio::test]
59    async fn advance_only_moves_forward() {
60        let mark = HighWaterMark::new(0u64);
61        assert!(mark.advance_to(2));
62        assert!(!mark.advance_to(1));
63        assert!(!mark.advance_to(2));
64        assert_eq!(mark.current(), 2);
65    }
66
67    #[tokio::test]
68    async fn wait_for_already_reached() {
69        let mark = HighWaterMark::new(5u64);
70        assert_eq!(mark.wait_for(3).await, Some(5));
71    }
72
73    #[tokio::test]
74    async fn wait_for_resolves_on_advance() {
75        let mark = HighWaterMark::new(0u64);
76        let wait = mark.wait_for(2);
77        mark.advance_to(1);
78        mark.advance_to(3);
79        assert_eq!(wait.await, Some(3));
80    }
81
82    #[tokio::test]
83    async fn wait_for_none_when_dropped() {
84        let mark = HighWaterMark::new(0u64);
85        let wait = mark.wait_for(1);
86        drop(mark);
87        assert_eq!(wait.await, None);
88    }
89}
mysten_common/sync/high_water_mark.rs

mysten_common/sync/
high_water_mark.rs
