Add lazy P2P demand tracking and need_p2p override (#2003)

- add lazy_p2p so nodes only start background P2P for peers that actually have recent business traffic
- add need_p2p so specific peers can still request eager background P2P even when other nodes enable lazy mode
- cover the new behavior with focused connector/peer-manager tests plus three-node integration tests that verify relay-to-direct route transition
This commit is contained in:
KKRainbow
2026-03-23 09:38:57 +08:00
committed by GitHub
parent 2bfdd44759
commit 1d89ddbb16
22 changed files with 876 additions and 91 deletions
+26 -11
View File
@@ -8,7 +8,7 @@ use std::{
atomic::{AtomicBool, Ordering},
Arc,
},
time::Duration,
time::{Duration, Instant},
};
use crate::{
@@ -40,7 +40,10 @@ use rand::Rng;
use tokio::{net::UdpSocket, task::JoinSet, time::timeout};
use url::Host;
use super::{create_connector_by_url, udp_hole_punch};
use super::{
create_connector_by_url, should_background_p2p_with_peer, should_try_p2p_with_peer,
udp_hole_punch,
};
pub const DIRECT_CONNECTOR_SERVICE_ID: u32 = 1;
pub const DIRECT_CONNECTOR_BLACKLIST_TIMEOUT_SEC: u64 = 300;
@@ -58,14 +61,22 @@ impl PeerManagerForDirectConnector for PeerManager {
async fn list_peers(&self) -> Vec<PeerId> {
let mut ret = vec![];
let allow_public_server = use_global_var!(DIRECT_CONNECT_TO_PUBLIC_SERVER);
let lazy_p2p = self.get_global_ctx().get_flags().lazy_p2p;
let now = Instant::now();
let routes = self.list_routes().await;
for r in routes.iter().filter(|r| {
r.feature_flag
.map(|r| allow_public_server || !r.is_public_server)
.unwrap_or(true)
}) {
ret.push(r.peer_id);
for route in routes.iter() {
let static_allowed = should_background_p2p_with_peer(
route.feature_flag.as_ref(),
allow_public_server,
lazy_p2p,
);
let dynamic_allowed =
should_try_p2p_with_peer(route.feature_flag.as_ref(), allow_public_server)
&& self.has_recent_traffic(route.peer_id, now);
if static_allowed || dynamic_allowed {
ret.push(route.peer_id);
}
}
ret
@@ -625,7 +636,11 @@ impl DirectConnectorManager {
global_ctx.clone(),
peer_manager.clone(),
));
let client = PeerTaskManager::new(DirectConnectorLauncher(data.clone()), peer_manager);
let client = PeerTaskManager::new_with_external_signal(
DirectConnectorLauncher(data.clone()),
peer_manager.clone(),
Some(peer_manager.p2p_demand_notify()),
);
Self {
global_ctx,
data,
@@ -696,7 +711,7 @@ mod tests {
let mut f = p_a.get_global_ctx().get_flags();
f.bind_device = false;
p_a.get_global_ctx().config.set_flags(f);
p_a.get_global_ctx().set_flags(f);
p_c.get_global_ctx()
.config
@@ -765,7 +780,7 @@ mod tests {
}
let mut f = p_c.get_global_ctx().config.get_flags();
f.enable_ipv6 = ipv6;
p_c.get_global_ctx().config.set_flags(f);
p_c.get_global_ctx().set_flags(f);
let mut lis_c = ListenerManager::new(p_c.get_global_ctx(), p_c.clone());
lis_c.prepare_listeners().await.unwrap();
+1 -1
View File
@@ -334,7 +334,7 @@ mod tests {
let mut flags = global_ctx.config.get_flags();
flags.bind_device = false;
global_ctx.config.set_flags(flags);
global_ctx.set_flags(flags);
let mut connector = HttpTunnelConnector::new(test_url.clone(), global_ctx.clone());
let mut listener = TcpTunnelListener::new("tcp://0.0.0.0:25888".parse().unwrap());
+75
View File
@@ -15,6 +15,7 @@ use crate::tunnel::unix::UnixSocketTunnelConnector;
use crate::tunnel::wireguard::{WgConfig, WgTunnelConnector};
use crate::{
common::{error::Error, global_ctx::ArcGlobalCtx, idn, network::IPCollector},
proto::common::PeerFeatureFlag,
tunnel::{
check_scheme_and_get_socket_addr, ring::RingTunnelConnector, tcp::TcpTunnelConnector,
udp::UdpTunnelConnector, IpVersion, TunnelConnector,
@@ -29,6 +30,24 @@ pub mod udp_hole_punch;
pub mod dns_connector;
pub mod http_connector;
pub(crate) fn should_try_p2p_with_peer(
feature_flag: Option<&PeerFeatureFlag>,
allow_public_server: bool,
) -> bool {
feature_flag
.map(|flag| allow_public_server || !flag.is_public_server)
.unwrap_or(true)
}
pub(crate) fn should_background_p2p_with_peer(
feature_flag: Option<&PeerFeatureFlag>,
allow_public_server: bool,
lazy_p2p: bool,
) -> bool {
should_try_p2p_with_peer(feature_flag, allow_public_server)
&& (!lazy_p2p || feature_flag.map(|flag| flag.need_p2p).unwrap_or(false))
}
async fn set_bind_addr_for_peer_connector(
connector: &mut (impl TunnelConnector + ?Sized),
is_ipv4: bool,
@@ -197,3 +216,59 @@ pub async fn create_connector_by_url(
Ok(connector)
}
#[cfg(test)]
mod tests {
use crate::proto::common::PeerFeatureFlag;
use super::{should_background_p2p_with_peer, should_try_p2p_with_peer};
#[test]
fn lazy_background_p2p_requires_need_p2p() {
let no_need_p2p = PeerFeatureFlag {
need_p2p: false,
..Default::default()
};
let need_p2p = PeerFeatureFlag {
need_p2p: true,
..Default::default()
};
assert!(should_background_p2p_with_peer(
Some(&no_need_p2p),
false,
false
));
assert!(!should_background_p2p_with_peer(
Some(&no_need_p2p),
false,
true
));
assert!(should_background_p2p_with_peer(
Some(&need_p2p),
false,
true
));
}
#[test]
fn p2p_policy_respects_public_server_setting() {
let public_server = PeerFeatureFlag {
is_public_server: true,
..Default::default()
};
assert!(!should_try_p2p_with_peer(Some(&public_server), false));
assert!(should_try_p2p_with_peer(Some(&public_server), true));
assert!(!should_background_p2p_with_peer(
Some(&public_server),
false,
false
));
assert!(should_background_p2p_with_peer(
Some(&public_server),
true,
false
));
}
}
+58 -7
View File
@@ -1,7 +1,7 @@
use std::{
net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
sync::Arc,
time::Duration,
time::{Duration, Instant},
};
use anyhow::{Context, Error};
@@ -29,6 +29,8 @@ use crate::{
},
};
use crate::connector::{should_background_p2p_with_peer, should_try_p2p_with_peer};
pub const BLACKLIST_TIMEOUT_SEC: u64 = 3600;
fn handle_rpc_result<T>(
@@ -418,6 +420,7 @@ impl PeerTaskLauncher for TcpHolePunchPeerTaskLauncher {
#[tracing::instrument(skip(self, data))]
async fn collect_peers_need_task(&self, data: &Self::Data) -> Vec<Self::CollectPeerItem> {
let global_ctx = data.peer_mgr.get_global_ctx();
let lazy_p2p = global_ctx.get_flags().lazy_p2p;
let my_tcp_nat_type = NatType::try_from(
global_ctx
.get_stun_info_collector()
@@ -434,16 +437,17 @@ impl PeerTaskLauncher for TcpHolePunchPeerTaskLauncher {
}
let my_peer_id = data.peer_mgr.my_peer_id();
let now = Instant::now();
data.blacklist.cleanup();
let mut peers_to_connect = Vec::new();
for route in data.peer_mgr.list_routes().await.iter() {
if route
.feature_flag
.map(|x| x.is_public_server)
.unwrap_or(false)
{
let static_allowed =
should_background_p2p_with_peer(route.feature_flag.as_ref(), false, lazy_p2p);
let dynamic_allowed = should_try_p2p_with_peer(route.feature_flag.as_ref(), false)
&& data.peer_mgr.has_recent_traffic(route.peer_id, now);
if !static_allowed && !dynamic_allowed {
continue;
}
@@ -520,7 +524,11 @@ impl TcpHolePunchConnector {
pub fn new(peer_mgr: Arc<PeerManager>) -> Self {
Self {
server: TcpHolePunchServer::new(peer_mgr.clone()),
client: PeerTaskManager::new(TcpHolePunchPeerTaskLauncher {}, peer_mgr.clone()),
client: PeerTaskManager::new_with_external_signal(
TcpHolePunchPeerTaskLauncher {},
peer_mgr.clone(),
Some(peer_mgr.p2p_demand_notify()),
),
peer_mgr,
}
}
@@ -570,12 +578,15 @@ mod tests {
connector::tcp_hole_punch::TcpHolePunchConnector,
peers::{
peer_manager::PeerManager,
peer_task::PeerTaskLauncher,
tests::{connect_peer_manager, create_mock_peer_manager, wait_route_appear},
},
proto::common::{NatType, StunInfo},
tunnel::common::tests::wait_for_condition,
};
use super::TcpHolePunchPeerTaskLauncher;
struct MockStunInfoCollector {
udp_nat_type: NatType,
tcp_nat_type: NatType,
@@ -619,6 +630,17 @@ mod tests {
.replace_stun_info_collector(collector);
}
async fn collect_lazy_punch_peers(peer_mgr: Arc<PeerManager>) -> Vec<u32> {
let launcher = TcpHolePunchPeerTaskLauncher {};
let data = launcher.new_data(peer_mgr);
launcher
.collect_peers_need_task(&data)
.await
.into_iter()
.map(|task| task.dst_peer_id)
.collect()
}
#[tokio::test]
async fn tcp_hole_punch_connects() {
let p_a = create_mock_peer_manager().await;
@@ -701,4 +723,33 @@ mod tests {
.map(|c| c.is_empty())
.unwrap_or(true));
}
#[tokio::test]
async fn lazy_p2p_collects_tcp_hole_punch_tasks_only_after_recent_traffic() {
let p_a = create_mock_peer_manager().await;
let p_b = create_mock_peer_manager().await;
let p_c = create_mock_peer_manager().await;
replace_stun_info_collector(p_a.clone(), NatType::PortRestricted);
replace_stun_info_collector(p_b.clone(), NatType::PortRestricted);
replace_stun_info_collector(p_c.clone(), NatType::PortRestricted);
let mut flags = p_a.get_global_ctx().get_flags();
flags.lazy_p2p = true;
p_a.get_global_ctx().set_flags(flags);
connect_peer_manager(p_a.clone(), p_b.clone()).await;
connect_peer_manager(p_b.clone(), p_c.clone()).await;
wait_route_appear(p_a.clone(), p_c.clone()).await.unwrap();
assert!(!collect_lazy_punch_peers(p_a.clone())
.await
.contains(&p_c.my_peer_id()));
p_a.mark_recent_traffic(p_c.my_peer_id());
assert!(collect_lazy_punch_peers(p_a.clone())
.await
.contains(&p_c.my_peer_id()));
}
}
+53 -8
View File
@@ -1,6 +1,6 @@
use std::{
sync::{atomic::AtomicBool, Arc},
time::Duration,
time::{Duration, Instant},
};
use anyhow::{Context, Error};
@@ -32,6 +32,8 @@ use crate::{
tunnel::Tunnel,
};
use crate::connector::{should_background_p2p_with_peer, should_try_p2p_with_peer};
pub(crate) mod both_easy_sym;
pub(crate) mod common;
pub(crate) mod cone;
@@ -426,6 +428,8 @@ impl PeerTaskLauncher for UdpHolePunchPeerTaskLauncher {
}
let my_peer_id = data.peer_mgr.my_peer_id();
let lazy_p2p = data.peer_mgr.get_global_ctx().get_flags().lazy_p2p;
let now = Instant::now();
data.blacklist.cleanup();
@@ -434,11 +438,11 @@ impl PeerTaskLauncher for UdpHolePunchPeerTaskLauncher {
// 2. peers is full cone (any restricted type);
// 3. peers not in blacklist;
for route in data.peer_mgr.list_routes().await.iter() {
if route
.feature_flag
.map(|x| x.is_public_server)
.unwrap_or(false)
{
let static_allowed =
should_background_p2p_with_peer(route.feature_flag.as_ref(), false, lazy_p2p);
let dynamic_allowed = should_try_p2p_with_peer(route.feature_flag.as_ref(), false)
&& data.peer_mgr.has_recent_traffic(route.peer_id, now);
if !static_allowed && !dynamic_allowed {
continue;
}
@@ -531,7 +535,11 @@ impl UdpHolePunchConnector {
pub fn new(peer_mgr: Arc<PeerManager>) -> Self {
Self {
server: UdpHolePunchServer::new(peer_mgr.clone()),
client: PeerTaskManager::new(UdpHolePunchPeerTaskLauncher {}, peer_mgr.clone()),
client: PeerTaskManager::new_with_external_signal(
UdpHolePunchPeerTaskLauncher {},
peer_mgr.clone(),
Some(peer_mgr.p2p_demand_notify()),
),
peer_mgr,
}
}
@@ -580,12 +588,13 @@ pub mod tests {
use crate::common::stun::MockStunInfoCollector;
use crate::peers::{
peer_manager::PeerManager,
peer_task::PeerTaskLauncher,
tests::{connect_peer_manager, create_mock_peer_manager, wait_route_appear},
};
use crate::proto::common::NatType;
use crate::tunnel::common::tests::wait_for_condition;
use super::{UdpHolePunchConnector, RUN_TESTING};
use super::{UdpHolePunchConnector, UdpHolePunchPeerTaskLauncher, RUN_TESTING};
pub fn replace_stun_info_collector(peer_mgr: Arc<PeerManager>, udp_nat_type: NatType) {
let collector = Box::new(MockStunInfoCollector { udp_nat_type });
@@ -602,6 +611,17 @@ pub mod tests {
p_a
}
async fn collect_lazy_punch_peers(peer_mgr: Arc<PeerManager>) -> Vec<u32> {
let launcher = UdpHolePunchPeerTaskLauncher {};
let data = launcher.new_data(peer_mgr);
launcher
.collect_peers_need_task(&data)
.await
.into_iter()
.map(|task| task.dst_peer_id)
.collect()
}
#[rstest::rstest]
#[tokio::test]
pub async fn test_hole_punching_blacklist(
@@ -634,4 +654,29 @@ pub mod tests {
)
.await;
}
#[tokio::test]
async fn lazy_p2p_collects_udp_hole_punch_tasks_only_after_recent_traffic() {
let p_a = create_mock_peer_manager_with_mock_stun(NatType::PortRestricted).await;
let p_b = create_mock_peer_manager_with_mock_stun(NatType::PortRestricted).await;
let p_c = create_mock_peer_manager_with_mock_stun(NatType::PortRestricted).await;
let mut flags = p_a.get_global_ctx().get_flags();
flags.lazy_p2p = true;
p_a.get_global_ctx().set_flags(flags);
connect_peer_manager(p_a.clone(), p_b.clone()).await;
connect_peer_manager(p_b.clone(), p_c.clone()).await;
wait_route_appear(p_a.clone(), p_c.clone()).await.unwrap();
assert!(!collect_lazy_punch_peers(p_a.clone())
.await
.contains(&p_c.my_peer_id()));
p_a.mark_recent_traffic(p_c.my_peer_id());
assert!(collect_lazy_punch_peers(p_a.clone())
.await
.contains(&p_c.my_peer_id()));
}
}