use std::future::{Future, poll_fn};
use std::pin::{Pin, pin};
use std::task::Poll;
use futures_core::Stream;
use tokio::sync::watch;
use crate::config::ProxyConfig;
use crate::watch::{ProxyWatcher, WatchEvent};
#[must_use]
pub fn watch_channel(watcher: ProxyWatcher) -> watch::Receiver<ProxyConfig> {
bridge(watcher.current(), watcher)
}
fn bridge<S>(initial: ProxyConfig, stream: S) -> watch::Receiver<ProxyConfig>
where
S: Stream<Item = WatchEvent> + Unpin + Send + 'static,
{
let (sender, receiver) = watch::channel(initial);
let mut stream = stream;
tokio::spawn(async move {
let mut closed = pin!(sender.closed());
loop {
let next = poll_fn(|cx| {
if closed.as_mut().poll(cx).is_ready() {
return Poll::Ready(None);
}
Pin::new(&mut stream).poll_next(cx)
})
.await;
let Some(item) = next else {
break;
};
if let WatchEvent::Snapshot { state } = item {
let config = state.config;
sender.send_if_modified(|current| {
if *current == config {
return false;
}
*current = config;
true
});
}
}
});
receiver
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::Context;
use super::*;
use crate::config::ProxyConfigSource;
use crate::error::Error;
use crate::mode::ProxyMode;
use crate::watch::{WatchHealth, WatchState};
fn live() -> WatchHealth {
WatchHealth {
degraded: Vec::new(),
has_live_notifications: true,
poll_interval: None,
stopped: false,
}
}
fn snapshot(config: ProxyConfig) -> WatchEvent {
WatchEvent::Snapshot {
state: WatchState {
config,
health: live(),
},
}
}
fn failure(error: Error) -> WatchEvent {
WatchEvent::Error {
error,
state: WatchState {
config: ProxyConfig::from_source(ProxyConfigSource::Env, ProxyMode::Direct),
health: live(),
},
}
}
struct Fed(tokio::sync::mpsc::UnboundedReceiver<WatchEvent>);
impl Stream for Fed {
type Item = WatchEvent;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.0.poll_recv(cx)
}
}
fn runtime() -> tokio::runtime::Runtime {
tokio::runtime::Builder::new_current_thread()
.build()
.expect("current-thread runtime")
}
#[test]
fn changes_reach_every_receiver_and_noise_is_filtered() {
let initial = ProxyConfig::direct();
let changed =
ProxyConfig::from_source(ProxyConfigSource::Registry, ProxyMode::WpadAutoDetect);
runtime().block_on(async {
let (events, stream) = tokio::sync::mpsc::unbounded_channel();
let mut receiver = bridge(initial.clone(), Fed(stream));
let clone = receiver.clone();
assert_eq!(*receiver.borrow(), initial);
events.send(snapshot(initial.clone())).unwrap();
events.send(snapshot(changed.clone())).unwrap();
receiver.changed().await.expect("a change is published");
assert_eq!(*receiver.borrow_and_update(), changed);
assert_eq!(*clone.borrow(), changed);
events.send(failure(Error::Unsupported)).unwrap();
events
.send(WatchEvent::Snapshot {
state: WatchState {
config: changed.clone(),
health: WatchHealth {
degraded: vec![ProxyConfigSource::GroupPolicy],
..live()
},
},
})
.unwrap();
tokio::task::yield_now().await;
assert!(
!receiver
.has_changed()
.expect("the stream is still open, so the channel is too"),
"neither an error nor a health-only snapshot may publish a configuration"
);
drop(events);
assert!(receiver.changed().await.is_err());
});
}
struct Idle(Arc<AtomicBool>);
impl Stream for Idle {
type Item = WatchEvent;
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Poll::Pending
}
}
impl Drop for Idle {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
#[test]
fn dropping_the_last_receiver_releases_the_watcher() {
let dropped = Arc::new(AtomicBool::new(false));
runtime().block_on(async {
let receiver = bridge(ProxyConfig::direct(), Idle(Arc::clone(&dropped)));
drop(receiver);
tokio::task::yield_now().await;
tokio::task::yield_now().await;
assert!(
dropped.load(Ordering::SeqCst),
"the bridge task must end once the last receiver is gone"
);
});
}
}