Skip to content

Commit 2b5c052

Browse files
rikaaa0928rikaaa0928
committed
Revert "Preload connector cache at object startup"
connect no pre create This reverts commit 833e6e5
1 parent 5f67bd9 commit 2b5c052

File tree

2 files changed

+83
-39
lines changed

2 files changed

+83
-39
lines changed

‎src/object.rs‎

Lines changed: 56 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use std::io;
77
use std::io::Error;
88
use std::sync::Arc;
99
use tokio::spawn;
10+
use tokio::sync::Mutex; // Already present, but ensure it's used for cache
1011

1112
pub mod config;
1213
pub mod raw_udp;
@@ -16,20 +17,16 @@ pub mod udp;
1617
pub struct Object {
1718
config: Arc<ObjectConfig>,
1819
router: Arc<dyn RouterSet>,
20+
connector_cache: Arc<Mutex<HashMap<String, Arc<Box<dyn RunConnector>>>>>, // New field
1921
}
2022

2123
impl Object {
2224
pub fn new(config: Arc<ObjectConfig>, router: Arc<dyn RouterSet>) -> Self {
23-
Self { config, router }
24-
}
25-
26-
async fn init_connectors(&self) -> io::Result<HashMap<String, Arc<Box<dyn RunConnector>>>> {
27-
let mut cache = HashMap::with_capacity(self.config.connector.len());
28-
for (name, conn_conf) in &self.config.connector {
29-
let connector = connector::create(conn_conf).await?;
30-
cache.insert(name.clone(), Arc::new(connector));
25+
Self {
26+
config,
27+
router,
28+
connector_cache: Arc::new(Mutex::new(HashMap::new())), // Initialize cache
3129
}
32-
Ok(cache)
3330
}
3431

3532
pub async fn start(&self) -> io::Result<()> {
@@ -42,7 +39,8 @@ impl Object {
4239
e
4340
})?;
4441
let main_acceptor = Arc::new(acc);
45-
let connector_cache_outer = Arc::new(self.init_connectors().await?);
42+
let connector_cache_outer = self.connector_cache.clone(); // Clone cache Arc for the loop
43+
4644
loop {
4745
let (mut acc_stream, _) = main_acceptor.accept().await.map_err(|e| {
4846
error!("Failed to accept connection: {}", e);
@@ -51,7 +49,7 @@ impl Object {
5149
let main_acceptor_clone = Arc::clone(&main_acceptor);
5250
let router_clone = Arc::clone(&router_outer);
5351
let config_clone = Arc::clone(&config_outer);
54-
let connector_cache_clone = Arc::clone(&connector_cache_outer);
52+
let connector_cache_clone = Arc::clone(&connector_cache_outer); // Clone cache Arc for the spawned task
5553

5654
spawn(async move {
5755
match acc_stream {
@@ -85,20 +83,53 @@ impl Object {
8583
addr_ref,
8684
)
8785
.await;
88-
let connector_obj =
89-
match connector_cache_clone.get(client_name.as_str()) {
90-
Some(cached_connector) => {
91-
debug!(
92-
"Reusing cached connector for client: {}",
93-
client_name
94-
);
95-
Arc::clone(cached_connector)
96-
}
97-
None => {
98-
error!("Connector '{}' not preloaded", client_name);
99-
return Ok(());
100-
}
101-
};
86+
let conn_conf = match config_clone
87+
.connector
88+
.get(client_name.as_str())
89+
{
90+
Some(c) => c,
91+
None => {
92+
error!("Connector config '{}' not found", client_name);
93+
return Ok(()); // Exit the task for this connection
94+
}
95+
};
96+
97+
let connector_obj: Arc<Box<dyn RunConnector>>;
98+
{
99+
let mut connector_cache_guard =
100+
connector_cache_clone.lock().await;
101+
if let Some(cached_connector) =
102+
connector_cache_guard.get(client_name.as_str())
103+
{
104+
connector_obj = Arc::clone(cached_connector);
105+
debug!(
106+
"Reusing cached connector for client: {}",
107+
client_name
108+
);
109+
} else {
110+
debug!(
111+
"Creating new connector for client: {}",
112+
client_name
113+
);
114+
let new_connector =
115+
match connector::create(conn_conf).await {
116+
Ok(c) => c,
117+
Err(e) => {
118+
error!(
119+
"Failed to create connector '{}': {}",
120+
client_name, e
121+
);
122+
return Ok(());
123+
}
124+
};
125+
let new_connector_arc = Arc::new(new_connector);
126+
connector_cache_guard.insert(
127+
client_name.clone(),
128+
Arc::clone(&new_connector_arc),
129+
);
130+
connector_obj = new_connector_arc;
131+
}
132+
}
102133

103134
debug!("Handshake successful {:?}", addr_ref);
104135
let client_stream_res = Arc::clone(&connector_obj)

‎src/object/raw_udp.rs‎

Lines changed: 27 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,20 @@
1-
use crate::def::{RouterSet, RunConnector, RunUdpReader, RunUdpWriter, UDPPacket};
1+
use crate::connector;
2+
use crate::def::{RouterSet, RunConnector, RunUdpReader, RunUdpWriter, UDPPacket}; // Added RunConnector
23
use crate::object::config::ObjectConfig;
34
use crate::util::RunAddr;
45
use log::{debug, warn};
5-
use std::collections::HashMap;
6-
use std::io::{self, Error, ErrorKind, Result};
6+
use std::collections::HashMap; // Added HashMap
7+
use std::io::{self, Error, ErrorKind, Result}; // Ensure io is imported for io::Error::new
78
use std::sync::Arc;
8-
use tokio::sync::Notify;
9+
use tokio::sync::{Mutex, Notify}; // Added Mutex
910
use tokio::{select, spawn};
1011

1112
pub async fn handle_raw_udp(
1213
mut r: Box<dyn RunUdpReader>,
1314
w: Box<dyn RunUdpWriter>,
1415
config: Arc<ObjectConfig>,
1516
router: Arc<dyn RouterSet>,
16-
connector_cache: Arc<HashMap<String, Arc<Box<dyn RunConnector>>>>,
17+
connector_cache: Arc<Mutex<HashMap<String, Arc<Box<dyn RunConnector>>>>>, // New argument
1718
) -> Result<()> {
1819
debug!("raw udp, route based on the first packet");
1920
let first_packet = r.read().await?;
@@ -30,15 +31,27 @@ pub async fn handle_raw_udp(
3031
)
3132
.await;
3233

33-
let connector_obj = connector_cache
34-
.get(client_name.as_str())
35-
.cloned()
36-
.ok_or_else(|| {
37-
io::Error::new(
38-
io::ErrorKind::NotFound,
39-
format!("Connector '{}' not preloaded for UDP", client_name),
40-
)
41-
})?;
34+
let conn_conf = config.connector.get(client_name.as_str()).ok_or_else(|| {
35+
io::Error::new(
36+
io::ErrorKind::NotFound,
37+
format!("Connector config '{}' not found for UDP", client_name),
38+
)
39+
})?;
40+
41+
let connector_obj: Arc<Box<dyn RunConnector>>;
42+
{
43+
let mut cache_guard = connector_cache.lock().await;
44+
if let Some(cached_connector) = cache_guard.get(client_name.as_str()) {
45+
connector_obj = Arc::clone(cached_connector);
46+
debug!("Reusing cached connector for UDP: {}", client_name);
47+
} else {
48+
debug!("Creating new connector for UDP: {}", client_name);
49+
let new_connector = connector::create(conn_conf).await?;
50+
let new_connector_arc = Arc::new(new_connector);
51+
cache_guard.insert(client_name.clone(), Arc::clone(&new_connector_arc));
52+
connector_obj = new_connector_arc;
53+
}
54+
}
4255

4356
let (mut udp_reader, udp_writer) = connector_obj
4457
.udp_tunnel(format!(

0 commit comments

Comments
 (0)