utils: add DirtyFlag

easytier/src/dns/client.rs

@@ -6,7 +6,6 @@ use crate::proto::peer_rpc::RoutePeerInfo;
 use crate::proto::rpc_impl::standalone::StandAloneClient;
 use crate::proto::rpc_types::controller::BaseController;
 use crate::tunnel::tcp::TcpTunnelConnector;
-use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::task::JoinSet;
@@ -61,16 +60,15 @@ impl DnsClient {
         rpc: &mut StandAloneClient<TcpTunnelConnector>,
         heartbeat: &mut HeartbeatRequest,
     ) -> anyhow::Result<()> {
-        let request =
+        let request = if heartbeat.snapshot.is_none() || self.mgr.dirty.reset() {
-            if heartbeat.snapshot.is_none() || self.mgr.dirty.swap(false, Ordering::Acquire) {
+            heartbeat.update(self.mgr.snapshot());
-                heartbeat.update(self.mgr.snapshot());
+            heartbeat.clone().into()
-                heartbeat.clone().into()
+        } else {
-            } else {
+            let snapshot = heartbeat.snapshot.take();
-                let snapshot = heartbeat.snapshot.take();
+            let request = heartbeat.clone().into();
-                let request = heartbeat.clone().into();
+            heartbeat.snapshot = snapshot;
-                heartbeat.snapshot = snapshot;
+            request
-                request
+        };
-            };
 
         let client = rpc
             .scoped_client::<DnsServerRpcClientFactory<BaseController>>("".to_string())

easytier/src/dns/peer_mgr.rs

@@ -1,6 +1,7 @@
 use crate::common::config::ConfigLoader;
 use crate::common::PeerId;
 use crate::dns::config::{DnsExportConfig, DnsGlobalCtxExt};
+use crate::dns::utils::DirtyFlag;
 use crate::dns::zone::ZoneGroup;
 use crate::peer_center::instance::PeerCenterPeerManagerTrait;
 use crate::peers::peer_manager::PeerManager;
@@ -15,7 +16,6 @@ use anyhow::Context;
 use derive_more::Deref;
 use itertools::Itertools;
 use moka::future::Cache;
-use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -42,7 +42,7 @@ const DNS_PEER_TTL: Duration = Duration::from_secs(3);
 #[derive(Debug, Deref)]
 pub struct DnsPeerMgr {
     peers: Cache<PeerId, DnsPeerInfo>,
-    pub(super) dirty: AtomicBool,
+    pub(super) dirty: DirtyFlag,
 
     #[deref]
     mgr: Arc<PeerManager>,
@@ -53,7 +53,7 @@ impl DnsPeerMgr {
         Self {
             mgr: peer_mgr.clone(),
             peers: Cache::builder().time_to_live(DNS_PEER_TTL).build(),
-            dirty: AtomicBool::new(true),
+            dirty: DirtyFlag::new(true),
         }
     }
 
@@ -102,7 +102,7 @@ impl DnsPeerMgr {
             }
         }
 
-        self.dirty.store(true, Ordering::Release);
+        self.dirty.mark();
     }
 
     async fn fetch(&self, peer_id: PeerId) -> anyhow::Result<DnsPeerInfo> {

easytier/src/dns/server.rs

@@ -1,3 +1,13 @@
+use super::{utils::NameServerAddr, zone::Zone};
+use crate::dns::utils::DirtyFlag;
+use crate::dns::zone::ZoneGroup;
+use crate::proto::dns::DnsSnapshot;
+use crate::proto::rpc_types;
+use crate::proto::{
+    dns::{DnsServerRpc, HeartbeatRequest, HeartbeatResponse},
+    rpc_types::controller::BaseController,
+};
+use crate::utils::{DeterministicDigest, MapTryInto};
 use anyhow::Error;
 use hickory_proto::xfer::Protocol;
 use hickory_server::{
@@ -8,27 +18,13 @@ use hickory_server::{
 use itertools::Itertools;
 use moka::future::Cache;
 use std::collections::HashSet;
-use std::sync::atomic::{AtomicBool, Ordering};
 use std::{sync::Arc, time::Duration};
+use derivative::Derivative;
 use tokio::net::{TcpListener, UdpSocket};
-use tokio::sync::Notify;
+use tokio::{sync::RwLock, task::JoinHandle};
-use tokio::{
-    sync::{Mutex, RwLock},
-    task::JoinHandle,
-};
 use tokio_util::sync::CancellationToken;
 use uuid::Uuid;
 
-use super::{utils::NameServerAddr, zone::Zone};
-use crate::dns::zone::ZoneGroup;
-use crate::proto::dns::DnsSnapshot;
-use crate::proto::rpc_types;
-use crate::proto::{
-    dns::{DnsServerRpc, HeartbeatRequest, HeartbeatResponse},
-    rpc_types::controller::BaseController,
-};
-use crate::utils::{DeterministicDigest, MapTryInto};
-
 #[derive(Debug, Clone, Default)]
 pub struct DnsClientInfo {
     digest: Vec<u8>,
@@ -82,11 +78,11 @@ impl RequestHandler for DynamicCatalog {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct DnsServerDirtyState {
-    zones: AtomicBool,
+    zones: DirtyFlag,
-    addresses: AtomicBool,
+    addresses: DirtyFlag,
-    listeners: AtomicBool,
+    listeners: DirtyFlag,
 }
 
 struct DnsServerRuntime {
@@ -114,11 +110,14 @@ impl DnsServerRuntime {
     }
 }
 
-#[derive(Clone)]
+
+#[derive(Derivative)]
+#[derivative(Debug)]
 pub struct DnsServer {
     clients: Cache<Uuid, DnsClientInfo>,
-    dirty: Arc<DnsServerDirtyState>,
+    dirty: DnsServerDirtyState,
 
+    #[derivative(Debug = "ignore")]
     catalog: DynamicCatalog,
 }
 
@@ -207,18 +206,18 @@ impl DnsServer {
         let dirty = &self.dirty;
         let mut runtime = None;
         loop {
-            if dirty.zones.swap(false, Ordering::Acquire) {
+            if dirty.zones.reset() {
                 self.reload_zones().await;
             }
 
-            if dirty.addresses.swap(false, Ordering::Acquire) {
+            if dirty.addresses.reset() {
                 self.reload_addresses().await;
             }
 
-            if dirty.listeners.swap(false, Ordering::Acquire) {
+            if dirty.listeners.reset() {
                 if let Err(e) = self.reload_listeners(&mut runtime).await {
                     tracing::error!("failed to reload listeners: {:?}", e);
-                    self.dirty.listeners.store(true, Ordering::Relaxed);
+                    self.dirty.listeners.mark();
                 }
             }
 
@@ -316,13 +315,13 @@ impl DnsServerRpc for DnsServer {
             let old = self.clients.get(&id).await.unwrap_or_default();
             if new.digest != old.digest {
                 if new.zones != old.zones {
-                    self.dirty.zones.store(true, Ordering::Release);
+                    self.dirty.zones.mark();
                 }
                 if new.addresses != old.addresses {
-                    self.dirty.addresses.store(true, Ordering::Release);
+                    self.dirty.addresses.mark();
                 }
                 if new.listeners != old.listeners {
-                    self.dirty.listeners.store(true, Ordering::Release);
+                    self.dirty.listeners.mark();
                 }
 
                 self.clients.insert(id, new).await;

easytier/src/dns/utils.rs

@@ -17,6 +17,7 @@ use serde_with::{DeserializeFromStr, SerializeDisplay};
 use std::fmt::{Display, Formatter};
 use std::net::{IpAddr, SocketAddr};
 use std::str::FromStr;
+use std::sync::atomic::{AtomicBool, Ordering};
 use url::Url;
 
 pub fn sanitize(name: &str) -> String {
@@ -265,3 +266,20 @@ where
         self.0.get_nsec_records(name, lookup_options).await
     }
 }
+
+#[derive(Debug, Default)]
+pub(super) struct DirtyFlag(AtomicBool);
+
+impl DirtyFlag {
+    pub fn new(value: bool) -> Self {
+        Self(AtomicBool::new(value))
+    }
+
+    pub fn mark(&self) {
+        self.0.store(true, Ordering::Release);
+    }
+
+    pub fn reset(&self) -> bool {
+        self.0.swap(false, Ordering::Acquire)
+    }
+}