use tokio::sync::watch;
use crate::errors::Error;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum State {
Connected,
Disconnected,
ShutDown,
}
#[derive(Debug)]
pub(crate) struct ConnectionSignal {
sender: watch::Sender<State>,
}
impl Default for ConnectionSignal {
fn default() -> Self {
Self {
sender: watch::channel(State::Connected).0,
}
}
}
impl ConnectionSignal {
pub(crate) fn is_connected(&self) -> bool {
*self.sender.borrow() == State::Connected
}
pub(crate) fn set_connected(&self) {
self.set(State::Connected);
}
pub(crate) fn set_disconnected(&self) {
self.set(State::Disconnected);
}
pub(crate) fn shutdown(&self) {
self.sender.send_replace(State::ShutDown);
}
pub(crate) async fn wait_connected(&self) -> Result<(), Error> {
let mut receiver = self.sender.subscribe();
loop {
let state = *receiver.borrow_and_update();
match state {
State::Connected => return Ok(()),
State::ShutDown => return Err(Error::Shutdown),
State::Disconnected if receiver.changed().await.is_err() => return Err(Error::Shutdown),
State::Disconnected => {}
}
}
}
fn set(&self, new_state: State) {
self.sender.send_if_modified(|state| {
if *state == State::ShutDown || *state == new_state {
return false;
}
*state = new_state;
true
});
}
}
#[cfg(test)]
#[path = "connection_signal_tests.rs"]
mod tests;