use std::{
future::Future,
sync::{Arc, Mutex, mpsc},
thread,
};
use tokio_util::sync::CancellationToken;
use super::{DiscoveryEvent, DiscoveryOptions, ServiceTypeFilter};
pub(super) enum RuntimeFlavor {
CurrentThread,
#[cfg(feature = "zeroconf")]
MultiThread,
}
impl RuntimeFlavor {
fn build(&self) -> std::io::Result<tokio::runtime::Runtime> {
let mut builder = match self {
RuntimeFlavor::CurrentThread => tokio::runtime::Builder::new_current_thread(),
#[cfg(feature = "zeroconf")]
RuntimeFlavor::MultiThread => tokio::runtime::Builder::new_multi_thread(),
};
builder.enable_all().build()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum BrowseOutcome {
Startup(String),
Stopped,
#[cfg(feature = "fake")]
Complete,
Cancelled,
}
pub(super) struct DiscoveryWorker {
shutdown: CancellationToken,
worker: Option<thread::JoinHandle<()>>,
outcome: Arc<Mutex<Option<BrowseOutcome>>>,
}
impl DiscoveryWorker {
pub(super) fn spawn<F, Fut>(
options: &DiscoveryOptions,
flavor: RuntimeFlavor,
browse: F,
) -> (Self, mpsc::Receiver<DiscoveryEvent>)
where
F: FnOnce(
String,
Option<ServiceTypeFilter>,
mpsc::Sender<DiscoveryEvent>,
CancellationToken,
) -> Fut
+ Send
+ 'static,
Fut: Future<Output = BrowseOutcome>,
{
let (tx, rx) = mpsc::channel();
let shutdown = CancellationToken::new();
let domain = options.domain().to_string();
let service_type_filter = options.service_type().cloned();
let token = shutdown.clone();
let outcome = Arc::new(Mutex::new(None));
let slot = outcome.clone();
let worker = thread::spawn(move || {
let keepalive = tx.clone();
let ended = match flavor.build() {
Ok(runtime) => runtime.block_on(browse(domain, service_type_filter, tx, token)),
Err(err) => BrowseOutcome::Startup(format!(
"failed to start mDNS runtime ({err}); try --backend fake in a build with the fake feature for sample records, or refresh to retry"
)),
};
*slot.lock().unwrap_or_else(|err| err.into_inner()) = Some(ended);
drop(keepalive);
});
(
Self {
shutdown,
worker: Some(worker),
outcome,
},
rx,
)
}
pub(super) fn outcome(&self) -> Option<BrowseOutcome> {
self.outcome
.lock()
.unwrap_or_else(|err| err.into_inner())
.clone()
}
pub(super) fn shutdown(&mut self) {
self.shutdown.cancel();
if let Some(worker) = self.worker.take() {
let _ = worker.join();
}
}
}
impl Drop for DiscoveryWorker {
fn drop(&mut self) {
self.shutdown();
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
use crate::discovery::{DiscoveryBackend, DiscoveryConfig};
fn config() -> DiscoveryConfig {
DiscoveryConfig {
backend: DiscoveryBackend::MdnsSd,
domain: "local".to_string(),
service_type: None,
}
}
fn options() -> DiscoveryOptions {
config().validate().expect("valid test options")
}
#[test]
fn browse_startup_failure_reports_the_cause_and_no_upsert() {
let (worker, rx) = DiscoveryWorker::spawn(
&options(),
RuntimeFlavor::CurrentThread,
|_domain, _service_type_filter, _tx, _shutdown| async move {
BrowseOutcome::Startup("simulated browse startup failure".to_string())
},
);
let events: Vec<_> = rx.iter().collect();
assert!(events.is_empty(), "a failed start must emit no entries");
assert_eq!(
worker.outcome(),
Some(BrowseOutcome::Startup(
"simulated browse startup failure".to_string()
))
);
}
#[test]
fn the_outcome_is_published_before_the_channel_disconnects() {
let (worker, rx) = DiscoveryWorker::spawn(
&options(),
RuntimeFlavor::CurrentThread,
|_domain, _service_type_filter, _tx, _shutdown| async move { BrowseOutcome::Stopped },
);
while rx.recv().is_ok() {}
assert_eq!(worker.outcome(), Some(BrowseOutcome::Stopped));
}
#[test]
fn the_loop_is_handed_the_validated_domain_and_filter() {
let mut config = config();
config.domain = "corp".to_string();
config.service_type = Some("_SSH._tcp".to_string());
let (_worker, rx) = DiscoveryWorker::spawn(
&config.validate().expect("valid test options"),
RuntimeFlavor::CurrentThread,
|domain, service_type_filter, tx, _shutdown| async move {
let _ = tx.send(DiscoveryEvent::Status(format!(
"{domain}/{}",
service_type_filter
.map(|filter| filter.to_string())
.unwrap_or_default()
)));
BrowseOutcome::Stopped
},
);
match rx.recv() {
Ok(DiscoveryEvent::Status(status)) => assert_eq!(status, "corp/_ssh._tcp"),
other => panic!("expected Status, got {other:?}"),
}
}
#[test]
fn shutdown_cancels_a_running_loop_and_joins_it() {
let (mut worker, rx) = DiscoveryWorker::spawn(
&options(),
RuntimeFlavor::CurrentThread,
|_domain, _service_type_filter, tx, shutdown| async move {
let _ = tx.send(DiscoveryEvent::Status("browsing".to_string()));
tokio::select! {
_ = shutdown.cancelled() => BrowseOutcome::Cancelled,
_ = tokio::time::sleep(Duration::from_secs(600)) => BrowseOutcome::Stopped,
}
},
);
assert!(matches!(rx.recv(), Ok(DiscoveryEvent::Status(_))));
worker.shutdown();
assert_eq!(worker.outcome(), Some(BrowseOutcome::Cancelled));
assert!(rx.recv().is_err());
}
#[test]
fn shutdown_is_idempotent() {
let (mut worker, _rx) = DiscoveryWorker::spawn(
&options(),
RuntimeFlavor::CurrentThread,
|_domain, _service_type_filter, _tx, shutdown| async move {
shutdown.cancelled().await;
BrowseOutcome::Cancelled
},
);
worker.shutdown();
worker.shutdown();
assert_eq!(worker.outcome(), Some(BrowseOutcome::Cancelled));
}
#[test]
fn dropping_the_worker_stops_its_producer() {
let (worker, rx) = DiscoveryWorker::spawn(
&options(),
RuntimeFlavor::CurrentThread,
|_domain, _service_type_filter, tx, shutdown| async move {
let _ = tx.send(DiscoveryEvent::Status("browsing".to_string()));
shutdown.cancelled().await;
BrowseOutcome::Cancelled
},
);
assert!(matches!(rx.recv(), Ok(DiscoveryEvent::Status(_))));
drop(worker);
assert!(rx.recv().is_err(), "the producer's sender must be gone");
}
#[test]
fn a_loop_that_never_runs_still_publishes_an_outcome() {
let (worker, rx) = DiscoveryWorker::spawn(
&options(),
RuntimeFlavor::CurrentThread,
|_domain, _service_type_filter, _tx, _shutdown| async move { BrowseOutcome::Stopped },
);
let events: Vec<_> = rx.iter().collect();
assert!(events.is_empty());
assert_eq!(worker.outcome(), Some(BrowseOutcome::Stopped));
}
}