TLA  Line data    Source code

       1                 : use tokio_util::task::{task_tracker::TaskTrackerToken, TaskTracker};
       2                 : 
       3                 : /// While a reference is kept around, the associated [`Barrier::wait`] will wait.
       4                 : ///
       5                 : /// Can be cloned, moved and kept around in futures as "guard objects".
       6 CBC         404 : #[derive(Clone)]
       7                 : pub struct Completion(TaskTrackerToken);
       8                 : 
       9                 : /// Barrier will wait until all clones of [`Completion`] have been dropped.
      10            2850 : #[derive(Clone)]
      11                 : pub struct Barrier(TaskTracker);
      12                 : 
      13                 : impl Default for Barrier {
      14               4 :     fn default() -> Self {
      15               4 :         let (_, rx) = channel();
      16               4 :         rx
      17               4 :     }
      18                 : }
      19                 : 
      20                 : impl Barrier {
      21            2263 :     pub async fn wait(self) {
      22            2263 :         self.0.wait().await;
      23            2238 :     }
      24                 : 
      25            2520 :     pub async fn maybe_wait(barrier: Option<Barrier>) {
      26            2520 :         if let Some(b) = barrier {
      27 UBC           0 :             b.wait().await
      28 CBC        2520 :         }
      29            2520 :     }
      30                 : }
      31                 : 
      32                 : impl PartialEq for Barrier {
      33 UBC           0 :     fn eq(&self, other: &Self) -> bool {
      34               0 :         TaskTracker::ptr_eq(&self.0, &other.0)
      35               0 :     }
      36                 : }
      37                 : 
      38                 : impl Eq for Barrier {}
      39                 : 
      40                 : /// Create new Guard and Barrier pair.
      41 CBC        2802 : pub fn channel() -> (Completion, Barrier) {
      42            2802 :     let tracker = TaskTracker::new();
      43            2802 :     // otherwise wait never exits
      44            2802 :     tracker.close();
      45            2802 : 
      46            2802 :     let token = tracker.token();
      47            2802 :     (Completion(token), Barrier(tracker))
      48            2802 : }