diff --git a/easytier/src/dns/peer_mgr.rs b/easytier/src/dns/peer_mgr.rs index cd3d1664..78fa1ff7 100644 --- a/easytier/src/dns/peer_mgr.rs +++ b/easytier/src/dns/peer_mgr.rs @@ -16,6 +16,7 @@ use crate::utils::DeterministicDigest; use anyhow::Context; use itertools::Itertools; use moka::future::Cache; +use std::ops::Deref; use std::sync::Arc; use std::time::Duration; use tracing::instrument; @@ -41,35 +42,15 @@ impl TryFrom for DnsPeerInfo { const DNS_PEER_TTL: Duration = Duration::from_secs(3); #[derive(Debug)] -pub struct DnsPeerMgr { +pub struct DnsPeerMgrInner { peers: Cache, - pub(super) dirty: DirtyFlag, + pub dirty: DirtyFlag, peer_mgr: Arc, global_ctx: ArcGlobalCtx, } -impl DnsPeerMgr { - pub fn new(peer_mgr: Arc, global_ctx: ArcGlobalCtx) -> Self { - Self { - peers: Cache::builder().time_to_live(DNS_PEER_TTL).build(), - dirty: Default::default(), - peer_mgr, - global_ctx, - } - } - - pub fn register(self: &Arc) { - self.peer_mgr - .get_peer_rpc_mgr() - .rpc_server() - .registry() - .register( - DnsPeerMgrRpcServer::new_arc(self.clone()), - &self.global_ctx.get_network_name(), - ); - } - +impl DnsPeerMgrInner { pub fn snapshot(&self) -> DnsSnapshot { let global_ctx = &self.global_ctx; @@ -149,7 +130,7 @@ impl DnsPeerMgr { } #[async_trait::async_trait] -impl DnsPeerMgrRpc for DnsPeerMgr { +impl DnsPeerMgrRpc for DnsPeerMgrInner { type Controller = BaseController; async fn get_export_config( @@ -160,3 +141,53 @@ impl DnsPeerMgrRpc for DnsPeerMgr { Ok(self.global_ctx.dns_export_config()) } } + +#[derive(Debug)] +pub struct DnsPeerMgr(Arc); + +impl DnsPeerMgr { + pub fn new(peer_mgr: Arc, global_ctx: ArcGlobalCtx) -> Self { + Self(Arc::new(DnsPeerMgrInner { + peers: Cache::builder().time_to_live(DNS_PEER_TTL).build(), + dirty: Default::default(), + peer_mgr, + global_ctx, + })) + } + + pub fn register(&self) { + self.peer_mgr + .get_peer_rpc_mgr() + .rpc_server() + .registry() + .register( + DnsPeerMgrRpcServer::new_arc(self.0.clone()), + &self.global_ctx.get_network_name(), + ); + } + + pub fn unregister(&self) -> Option<()> { + self.peer_mgr + .get_peer_rpc_mgr() + .rpc_server() + .registry() + .unregister( + DnsPeerMgrRpcServer::new_arc(self.0.clone()), + &self.global_ctx.get_network_name(), + ) + } +} + +impl Deref for DnsPeerMgr { + type Target = DnsPeerMgrInner; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl Drop for DnsPeerMgr { + fn drop(&mut self) { + self.unregister(); + } +}