diff --git a/easytier/src/dns/client.rs b/easytier/src/dns/client.rs index 5ac8b3f0..31fe3f43 100644 --- a/easytier/src/dns/client.rs +++ b/easytier/src/dns/client.rs @@ -48,6 +48,7 @@ impl DnsClient { ..Default::default() }; loop { + self.mgr.dirty.notified().await; if let Err(e) = self.heartbeat(&mut rpc, &mut heartbeat).await { tracing::error!("DnsClient heartbeat failed: {:?}", e); } @@ -60,7 +61,7 @@ impl DnsClient { rpc: &mut StandAloneClient, heartbeat: &mut HeartbeatRequest, ) -> anyhow::Result<()> { - let request = if heartbeat.snapshot.is_none() || self.mgr.dirty.reset() { + let request = if heartbeat.snapshot.is_none() || self.mgr.dirty.peers.reset() { heartbeat.update(self.mgr.snapshot()); heartbeat.clone().into() } else { diff --git a/easytier/src/dns/peer_mgr.rs b/easytier/src/dns/peer_mgr.rs index db3d5073..d07d3a51 100644 --- a/easytier/src/dns/peer_mgr.rs +++ b/easytier/src/dns/peer_mgr.rs @@ -13,11 +13,12 @@ use crate::proto::rpc_types; use crate::proto::rpc_types::controller::BaseController; use crate::utils::DeterministicDigest; use anyhow::Context; -use derive_more::Deref; +use derive_more::{Deref, DerefMut}; use itertools::Itertools; use moka::future::Cache; use std::sync::Arc; use std::time::Duration; +use tokio::sync::Notify; #[derive(Debug, Clone)] pub struct DnsPeerInfo { @@ -39,10 +40,18 @@ impl TryFrom for DnsPeerInfo { const DNS_PEER_TTL: Duration = Duration::from_secs(3); +#[derive(Debug, Default, Deref, DerefMut)] +pub struct DnsPeerMgrDirtyState { + pub(crate) peers: DirtyFlag, + #[deref] + #[deref_mut] + notify: Notify, +} + #[derive(Debug, Deref)] pub struct DnsPeerMgr { peers: Cache, - pub(super) dirty: DirtyFlag, + pub(super) dirty: DnsPeerMgrDirtyState, #[deref] mgr: Arc, @@ -53,7 +62,7 @@ impl DnsPeerMgr { Self { mgr: peer_mgr.clone(), peers: Cache::builder().time_to_live(DNS_PEER_TTL).build(), - dirty: DirtyFlag::new(true), + dirty: Default::default(), } } @@ -88,6 +97,8 @@ impl DnsPeerMgr { } }; + self.dirty.peers.mark(); + match self.fetch(peer_id).await { Ok(info) => { self.peers.insert(peer_id, info).await; @@ -101,8 +112,8 @@ impl DnsPeerMgr { self.peers.invalidate(&peer_id).await; } } - - self.dirty.mark(); + + self.dirty.notify_one(); } async fn fetch(&self, peer_id: PeerId) -> anyhow::Result { diff --git a/easytier/src/dns/server.rs b/easytier/src/dns/server.rs index e93c951f..63903425 100644 --- a/easytier/src/dns/server.rs +++ b/easytier/src/dns/server.rs @@ -20,8 +20,10 @@ use moka::future::Cache; use std::collections::HashSet; use std::{sync::Arc, time::Duration}; use derivative::Derivative; +use derive_more::{Deref, DerefMut}; use tokio::net::{TcpListener, UdpSocket}; use tokio::{sync::RwLock, task::JoinHandle}; +use tokio::sync::Notify; use tokio_util::sync::CancellationToken; use uuid::Uuid; @@ -78,11 +80,15 @@ impl RequestHandler for DynamicCatalog { } } -#[derive(Debug, Default)] +// TODO: same as DnsPeerMgrDirtyState +#[derive(Debug, Default, Deref, DerefMut)] pub struct DnsServerDirtyState { zones: DirtyFlag, addresses: DirtyFlag, listeners: DirtyFlag, + #[deref] + #[deref_mut] + notify: Notify, } struct DnsServerRuntime { @@ -206,6 +212,8 @@ impl DnsServer { let dirty = &self.dirty; let mut runtime = None; loop { + dirty.notified().await; + if dirty.zones.reset() { self.reload_zones().await; } @@ -218,6 +226,7 @@ impl DnsServer { if let Err(e) = self.reload_listeners(&mut runtime).await { tracing::error!("failed to reload listeners: {:?}", e); self.dirty.listeners.mark(); + self.dirty.notify_one(); } } @@ -325,6 +334,7 @@ impl DnsServerRpc for DnsServer { } self.clients.insert(id, new).await; + self.dirty.notify_one(); } false } else { diff --git a/easytier/src/dns/utils.rs b/easytier/src/dns/utils.rs index cbbbc467..728b385a 100644 --- a/easytier/src/dns/utils.rs +++ b/easytier/src/dns/utils.rs @@ -267,7 +267,7 @@ where } } -#[derive(Debug, Default)] +#[derive(Debug)] pub(super) struct DirtyFlag(AtomicBool); impl DirtyFlag { @@ -283,3 +283,9 @@ impl DirtyFlag { self.0.swap(false, Ordering::Acquire) } } + +impl Default for DirtyFlag { + fn default() -> Self { + Self::new(true) + } +}