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          690 : #[derive(Clone)]
       9              : pub struct Completion(mpsc::Sender<()>);
      10              : 
      11              : /// Barrier will wait until all clones of [`Completion`] have been dropped.
      12         3249 : #[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         1876 :     pub async fn wait(self) {
      24         1876 :         self.0.lock().await.recv().await;
      25         1557 :     }
      26              : 
      27         2956 :     pub async fn maybe_wait(barrier: Option<Barrier>) {
      28         2956 :         if let Some(b) = barrier {
      29          685 :             b.wait().await
      30         2271 :         }
      31         2734 :     }
      32              : }
      33              : 
      34              : impl PartialEq for Barrier {
      35            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         2096 : pub fn channel() -> (Completion, Barrier) {
      45         2096 :     let (tx, rx) = mpsc::channel::<()>(1);
      46         2096 :     let rx = Mutex::new(rx);
      47         2096 :     let rx = Arc::new(rx);
      48         2096 :     (Completion(tx), Barrier(rx))
      49         2096 : }