src/tools/broadcast_future.rs

   1 use std::future::Future;
   2 use std::pin::Pin;
   3 use std::sync::{Arc, Mutex};
   4
   5 use anyhow::{format_err, Error};
   6 use futures::future::{FutureExt, TryFutureExt};
   7 use tokio::sync::oneshot;
   8
   9 /// Broadcast results to registered listeners using asnyc oneshot channels
  10 #[derive(Default)]
  11 pub struct BroadcastData<T> {
  12     result: Option<Result<T, String>>,
  13     listeners: Vec<oneshot::Sender<Result<T, Error>>>,
  14 }
  15
  16 impl <T: Clone> BroadcastData<T> {
  17
  18     pub fn new() -> Self {
  19         Self {
  20             result: None,
  21             listeners: vec![],
  22         }
  23     }
  24
  25     pub fn notify_listeners(&mut self, result: Result<T, String>) {
  26
  27         self.result = Some(result.clone());
  28
  29         loop {
  30             match self.listeners.pop() {
  31                 None => { break; },
  32                 Some(ch) => {
  33                     match &result {
  34                         Ok(result) => { let _ = ch.send(Ok(result.clone())); },
  35                         Err(err) => { let _ = ch.send(Err(format_err!("{}", err))); },
  36                     }
  37                 },
  38             }
  39         }
  40     }
  41
  42     pub fn listen(&mut self) -> impl Future<Output = Result<T, Error>> {
  43         use futures::future::{ok, Either};
  44
  45         match &self.result {
  46             None => {},
  47             Some(Ok(result)) => return Either::Left(ok(result.clone())),
  48             Some(Err(err)) => return Either::Left(futures::future::err(format_err!("{}", err))),
  49         }
  50
  51         let (tx, rx) = oneshot::channel::<Result<T, Error>>();
  52
  53         self.listeners.push(tx);
  54
  55         Either::Right(rx
  56             .map(|res| match res {
  57                 Ok(Ok(t)) => Ok(t),
  58                 Ok(Err(e)) => Err(e),
  59                 Err(e) => Err(Error::from(e)),
  60             })
  61         )
  62     }
  63 }
  64
  65 /// Broadcast future results to registered listeners
  66 pub struct BroadcastFuture<T> {
  67     inner: Arc<
  68         Mutex<(
  69             BroadcastData<T>,
  70             Option<Pin<Box<dyn Future<Output = Result<T, Error>> + Send>>>,
  71         )>,
  72     >,
  73 }
  74
  75 impl<T: Clone + Send + 'static> BroadcastFuture<T> {
  76     /// Create instance for specified source future.
  77     ///
  78     /// The result of the future is sent to all registered listeners.
  79     pub fn new(source: Box<dyn Future<Output = Result<T, Error>> + Send>) -> Self {
  80         Self { inner: Arc::new(Mutex::new((BroadcastData::new(), Some(Pin::from(source))))) }
  81     }
  82
  83     /// Creates a new instance with a oneshot channel as trigger
  84     pub fn new_oneshot() -> (Self, oneshot::Sender<Result<T, Error>>) {
  85
  86         let (tx, rx) = oneshot::channel::<Result<T, Error>>();
  87         let rx = rx
  88             .map_err(Error::from)
  89             .and_then(futures::future::ready);
  90
  91         (Self::new(Box::new(rx)), tx)
  92     }
  93
  94     fn notify_listeners(
  95         inner: Arc<
  96             Mutex<(
  97                 BroadcastData<T>,
  98                 Option<Pin<Box<dyn Future<Output = Result<T, Error>> + Send>>>,
  99             )>,
 100         >,
 101         result: Result<T, String>,
 102     ) {
 103         let mut data = inner.lock().unwrap();
 104         data.0.notify_listeners(result);
 105     }
 106
 107     fn spawn(
 108         inner: Arc<
 109             Mutex<(
 110                 BroadcastData<T>,
 111                 Option<Pin<Box<dyn Future<Output = Result<T, Error>> + Send>>>,
 112             )>,
 113         >,
 114     ) -> impl Future<Output = Result<T, Error>> {
 115         let mut data = inner.lock().unwrap();
 116
 117         if let Some(source) = data.1.take() {
 118
 119             let inner1 = inner.clone();
 120
 121             let task = source.map(move |value| {
 122                 match value {
 123                     Ok(value) => Self::notify_listeners(inner1, Ok(value)),
 124                     Err(err) => Self::notify_listeners(inner1, Err(err.to_string())),
 125                 }
 126             });
 127             tokio::spawn(task);
 128         }
 129
 130         data.0.listen()
 131     }
 132
 133     /// Register a listener
 134     pub fn listen(&self) -> impl Future<Output = Result<T, Error>> {
 135         let inner2 = self.inner.clone();
 136         async move { Self::spawn(inner2).await }
 137     }
 138 }
 139
 140 #[test]
 141 fn test_broadcast_future() {
 142     use std::sync::atomic::{AtomicUsize, Ordering};
 143
 144     static CHECKSUM: AtomicUsize  = AtomicUsize::new(0);
 145
 146     let (sender, trigger) = BroadcastFuture::new_oneshot();
 147
 148     let receiver1 = sender.listen()
 149         .map_ok(|res| {
 150             CHECKSUM.fetch_add(res, Ordering::SeqCst);
 151         })
 152         .map_err(|err| { panic!("got error {}", err); })
 153         .map(|_| ());
 154
 155     let receiver2 = sender.listen()
 156         .map_ok(|res| {
 157             CHECKSUM.fetch_add(res*2, Ordering::SeqCst);
 158         })
 159         .map_err(|err| { panic!("got error {}", err); })
 160         .map(|_| ());
 161
 162     let rt = tokio::runtime::Runtime::new().unwrap();
 163     rt.block_on(async move {
 164         let r1 = tokio::spawn(receiver1);
 165         let r2 = tokio::spawn(receiver2);
 166
 167         trigger.send(Ok(1)).unwrap();
 168         let _ = r1.await;
 169         let _ = r2.await;
 170     });
 171
 172     let result = CHECKSUM.load(Ordering::SeqCst);
 173
 174     assert_eq!(result, 3);
 175 }