pbs-tools/src/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 type SourceFuture<T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send>>;
  66
  67 struct BroadCastFutureBinding<T> {
  68     broadcast: BroadcastData<T>,
  69     future: Option<SourceFuture<T>>,
  70 }
  71
  72 /// Broadcast future results to registered listeners
  73 pub struct BroadcastFuture<T> {
  74     inner: Arc<Mutex<BroadCastFutureBinding<T>>>,
  75 }
  76
  77 impl<T: Clone + Send + 'static> BroadcastFuture<T> {
  78     /// Create instance for specified source future.
  79     ///
  80     /// The result of the future is sent to all registered listeners.
  81     pub fn new(source: Box<dyn Future<Output = Result<T, Error>> + Send>) -> Self {
  82         let inner = BroadCastFutureBinding {
  83             broadcast: BroadcastData::new(),
  84             future: Some(Pin::from(source)),
  85         };
  86         Self { inner: Arc::new(Mutex::new(inner)) }
  87     }
  88
  89     /// Creates a new instance with a oneshot channel as trigger
  90     pub fn new_oneshot() -> (Self, oneshot::Sender<Result<T, Error>>) {
  91
  92         let (tx, rx) = oneshot::channel::<Result<T, Error>>();
  93         let rx = rx
  94             .map_err(Error::from)
  95             .and_then(futures::future::ready);
  96
  97         (Self::new(Box::new(rx)), tx)
  98     }
  99
 100     fn notify_listeners(
 101         inner: Arc<Mutex<BroadCastFutureBinding<T>>>,
 102         result: Result<T, String>,
 103     ) {
 104         let mut data = inner.lock().unwrap();
 105         data.broadcast.notify_listeners(result);
 106     }
 107
 108     fn spawn(inner: Arc<Mutex<BroadCastFutureBinding<T>>>) -> impl Future<Output = Result<T, Error>> {
 109         let mut data = inner.lock().unwrap();
 110
 111         if let Some(source) = data.future.take() {
 112
 113             let inner1 = inner.clone();
 114
 115             let task = source.map(move |value| {
 116                 match value {
 117                     Ok(value) => Self::notify_listeners(inner1, Ok(value)),
 118                     Err(err) => Self::notify_listeners(inner1, Err(err.to_string())),
 119                 }
 120             });
 121             tokio::spawn(task);
 122         }
 123
 124         data.broadcast.listen()
 125     }
 126
 127     /// Register a listener
 128     pub fn listen(&self) -> impl Future<Output = Result<T, Error>> {
 129         let inner2 = self.inner.clone();
 130         async move { Self::spawn(inner2).await }
 131     }
 132 }
 133
 134 #[test]
 135 fn test_broadcast_future() {
 136     use std::sync::atomic::{AtomicUsize, Ordering};
 137
 138     static CHECKSUM: AtomicUsize  = AtomicUsize::new(0);
 139
 140     let (sender, trigger) = BroadcastFuture::new_oneshot();
 141
 142     let receiver1 = sender.listen()
 143         .map_ok(|res| {
 144             CHECKSUM.fetch_add(res, Ordering::SeqCst);
 145         })
 146         .map_err(|err| { panic!("got error {}", err); })
 147         .map(|_| ());
 148
 149     let receiver2 = sender.listen()
 150         .map_ok(|res| {
 151             CHECKSUM.fetch_add(res*2, Ordering::SeqCst);
 152         })
 153         .map_err(|err| { panic!("got error {}", err); })
 154         .map(|_| ());
 155
 156     let rt = tokio::runtime::Runtime::new().unwrap();
 157     rt.block_on(async move {
 158         let r1 = tokio::spawn(receiver1);
 159         let r2 = tokio::spawn(receiver2);
 160
 161         trigger.send(Ok(1)).unwrap();
 162         let _ = r1.await;
 163         let _ = r2.await;
 164     });
 165
 166     let result = CHECKSUM.load(Ordering::SeqCst);
 167
 168     assert_eq!(result, 3);
 169
 170     // the result stays available until the BroadcastFuture is dropped
 171     rt.block_on(sender.listen()
 172         .map_ok(|res| {
 173             CHECKSUM.fetch_add(res*4, Ordering::SeqCst);
 174         })
 175         .map_err(|err| { panic!("got error {}", err); })
 176         .map(|_| ()));
 177
 178     let result = CHECKSUM.load(Ordering::SeqCst);
 179     assert_eq!(result, 7);
 180 }