use notify in dirty states

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<TcpTunnelConnector>,
         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 {

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<DnsExportConfig> 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<PeerId, DnsPeerInfo>,
-    pub(super) dirty: DirtyFlag,
+    pub(super) dirty: DnsPeerMgrDirtyState,
 
     #[deref]
     mgr: Arc<PeerManager>,
@@ -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<DnsPeerInfo> {

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 {

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)
+    }
+}