TLA  Line data    Source code

       1                 : use std::sync::Arc;
       2                 : 
       3                 : use tokio::sync::{mpsc, Mutex};
       4                 : 
       5                 : /// While a reference is kept around, the associated [`Barrier::wait`] will wait.
       6                 : ///
       7                 : /// Can be cloned, moved and kept around in futures as "guard objects".
       8 CBC         878 : #[derive(Clone)]
       9                 : pub struct Completion(mpsc::Sender<()>);
      10                 : 
      11                 : /// Barrier will wait until all clones of [`Completion`] have been dropped.
      12            3654 : #[derive(Clone)]
      13                 : pub struct Barrier(Arc<Mutex<mpsc::Receiver<()>>>);
      14                 : 
      15                 : impl Default for Barrier {
      16               6 :     fn default() -> Self {
      17               6 :         let (_, rx) = channel();
      18               6 :         rx
      19               6 :     }
      20                 : }
      21                 : 
      22                 : impl Barrier {
      23            2406 :     pub async fn wait(self) {
      24            2406 :         self.0.lock().await.recv().await;
      25            2128 :     }
      26                 : 
      27            2820 :     pub async fn maybe_wait(barrier: Option<Barrier>) {
      28            2820 :         if let Some(b) = barrier {
      29             691 :             b.wait().await
      30            2129 :         }
      31            2632 :     }
      32                 : }
      33                 : 
      34                 : impl PartialEq for Barrier {
      35 UBC           0 :     fn eq(&self, other: &Self) -> bool {
      36               0 :         // we don't use dyn so this is good
      37               0 :         Arc::ptr_eq(&self.0, &other.0)
      38               0 :     }
      39                 : }
      40                 : 
      41                 : impl Eq for Barrier {}
      42                 : 
      43                 : /// Create new Guard and Barrier pair.
      44 CBC        2554 : pub fn channel() -> (Completion, Barrier) {
      45            2554 :     let (tx, rx) = mpsc::channel::<()>(1);
      46            2554 :     let rx = Mutex::new(rx);
      47            2554 :     let rx = Arc::new(rx);
      48            2554 :     (Completion(tx), Barrier(rx))
      49            2554 : }