use std::{future::Future, sync::mpsc, thread};
use tokio_util::sync::CancellationToken;
use super::{Discovery, DiscoveryConfig, DiscoveryEvent, fake};
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()
}
}
pub(super) struct DiscoveryWorker {
receiver: Option<mpsc::Receiver<DiscoveryEvent>>,
shutdown: CancellationToken,
worker: Option<thread::JoinHandle<()>>,
}
impl DiscoveryWorker {
pub(super) fn spawn<F, Fut>(config: &DiscoveryConfig, flavor: RuntimeFlavor, browse: F) -> Self
where
F: FnOnce(String, Option<String>, mpsc::Sender<DiscoveryEvent>, CancellationToken) -> Fut
+ Send
+ 'static,
Fut: Future<Output = ()>,
{
let (tx, rx) = mpsc::channel();
let shutdown = CancellationToken::new();
let domain = config.domain.clone();
let service_type_filter = config.service_type.clone();
let token = shutdown.clone();
let worker = thread::spawn(move || {
let runtime = match flavor.build() {
Ok(runtime) => runtime,
Err(err) => {
let _ = tx.send(DiscoveryEvent::Status(format!(
"failed to start mDNS runtime ({err}); using sample records"
)));
fake::spawn(domain, service_type_filter, tx);
return;
}
};
runtime.block_on(browse(domain, service_type_filter, tx, token));
});
Self {
receiver: Some(rx),
shutdown,
worker: Some(worker),
}
}
}
impl Discovery for DiscoveryWorker {
fn events(&mut self) -> mpsc::Receiver<DiscoveryEvent> {
self.receiver
.take()
.expect("discovery receiver can only be taken once")
}
}
impl Drop for DiscoveryWorker {
fn drop(&mut self) {
self.shutdown.cancel();
if let Some(worker) = self.worker.take() {
let _ = worker.join();
}
}
}