diff --git a/easytier/src/dns/client.rs b/easytier/src/dns/client.rs index c3251283..5ac8b3f0 100644 --- a/easytier/src/dns/client.rs +++ b/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, heartbeat: &mut HeartbeatRequest, ) -> anyhow::Result<()> { - let request = - if heartbeat.snapshot.is_none() || self.mgr.dirty.swap(false, Ordering::Acquire) { - heartbeat.update(self.mgr.snapshot()); - heartbeat.clone().into() - } else { - let snapshot = heartbeat.snapshot.take(); - let request = heartbeat.clone().into(); - heartbeat.snapshot = snapshot; - request - }; + let request = if heartbeat.snapshot.is_none() || self.mgr.dirty.reset() { + heartbeat.update(self.mgr.snapshot()); + heartbeat.clone().into() + } else { + let snapshot = heartbeat.snapshot.take(); + let request = heartbeat.clone().into(); + heartbeat.snapshot = snapshot; + request + }; let client = rpc .scoped_client::>("".to_string()) diff --git a/easytier/src/dns/peer_mgr.rs b/easytier/src/dns/peer_mgr.rs index 3418e205..db3d5073 100644 --- a/easytier/src/dns/peer_mgr.rs +++ b/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, - pub(super) dirty: AtomicBool, + pub(super) dirty: DirtyFlag, #[deref] mgr: Arc, @@ -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 { diff --git a/easytier/src/dns/server.rs b/easytier/src/dns/server.rs index 747f8b77..e93c951f 100644 --- a/easytier/src/dns/server.rs +++ b/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::{Mutex, RwLock}, - task::JoinHandle, -}; +use tokio::{sync::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, @@ -82,11 +78,11 @@ impl RequestHandler for DynamicCatalog { } } -#[derive(Debug)] +#[derive(Debug, Default)] pub struct DnsServerDirtyState { - zones: AtomicBool, - addresses: AtomicBool, - listeners: AtomicBool, + zones: DirtyFlag, + addresses: DirtyFlag, + listeners: DirtyFlag, } struct DnsServerRuntime { @@ -114,11 +110,14 @@ impl DnsServerRuntime { } } -#[derive(Clone)] + +#[derive(Derivative)] +#[derivative(Debug)] pub struct DnsServer { clients: Cache, - dirty: Arc, + 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; diff --git a/easytier/src/dns/utils.rs b/easytier/src/dns/utils.rs index 77e570a4..cbbbc467 100644 --- a/easytier/src/dns/utils.rs +++ b/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) + } +}