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          421 : #[derive(Clone)]
       7              : pub struct Completion(TaskTrackerToken);
       8              : 
       9              : /// Barrier will wait until all clones of [`Completion`] have been dropped.
      10         5166 : #[derive(Clone)]
      11              : pub struct Barrier(TaskTracker);
      12              : 
      13              : impl Default for Barrier {
      14            8 :     fn default() -> Self {
      15            8 :         let (_, rx) = channel();
      16            8 :         rx
      17            8 :     }
      18              : }
      19              : 
      20              : impl Barrier {
      21         4568 :     pub async fn wait(self) {
      22         4568 :         self.0.wait().await;
      23         4546 :     }
      24              : 
      25         2895 :     pub async fn maybe_wait(barrier: Option<Barrier>) {
      26         2895 :         if let Some(b) = barrier {
      27            0 :             b.wait().await
      28         2895 :         }
      29         2895 :     }
      30              : }
      31              : 
      32              : impl PartialEq for Barrier {
      33            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         3539 : pub fn channel() -> (Completion, Barrier) {
      42         3539 :     let tracker = TaskTracker::new();
      43         3539 :     // otherwise wait never exits
      44         3539 :     tracker.close();
      45         3539 : 
      46         3539 :     let token = tracker.token();
      47         3539 :     (Completion(token), Barrier(tracker))
      48         3539 : }