diff --git a/src/bin/local.rs b/src/bin/local.rs index d205fc9e..bd86ed3d 100644 --- a/src/bin/local.rs +++ b/src/bin/local.rs @@ -40,6 +40,8 @@ use clap::{App, Arg}; use std::net::{SocketAddr, IpAddr}; +use coio::Scheduler; + use shadowsocks::config::{Config, ServerConfig, self}; use shadowsocks::config::DEFAULT_DNS_CACHE_CAPACITY; use shadowsocks::relay::{RelayLocal, Relay}; @@ -174,6 +176,7 @@ fn main() { let threads = matches.value_of("THREADS").unwrap_or("1").parse::() .ok().expect("`threads` should be an integer"); - RelayLocal::new(config).run(threads); - panic!("Relay stopped."); + Scheduler::with_workers(threads).run(move|| { + RelayLocal::new(config).run(); + }).unwrap(); } diff --git a/src/bin/server.rs b/src/bin/server.rs index 2bf36ac0..6b67b5c0 100644 --- a/src/bin/server.rs +++ b/src/bin/server.rs @@ -36,9 +36,12 @@ extern crate shadowsocks; extern crate log; extern crate fern; extern crate time; +extern crate coio; use clap::{App, Arg}; +use coio::Scheduler; + use shadowsocks::config::{Config, ServerConfig, self}; use shadowsocks::config::DEFAULT_DNS_CACHE_CAPACITY; use shadowsocks::relay::{RelayServer, Relay}; @@ -167,6 +170,7 @@ fn main() { let threads = matches.value_of("THREADS").unwrap_or("1").parse::() .ok().expect("`threads` should be an integer"); - RelayServer::new(config).run(threads); - panic!("Server should never ends"); + Scheduler::with_workers(threads).run(move|| { + RelayServer::new(config).run(); + }).unwrap(); } diff --git a/src/relay/local.rs b/src/relay/local.rs index 2561f693..50fbdf13 100644 --- a/src/relay/local.rs +++ b/src/relay/local.rs @@ -21,7 +21,7 @@ //! Local side -use coio::{Scheduler, Builder}; +use coio::Builder; use relay::{Relay, COROUTINE_STACK_SIZE}; use relay::tcprelay::local::TcpRelayLocal; @@ -85,29 +85,35 @@ impl RelayLocal { impl Relay for RelayLocal { #[cfg(not(feature = "enable-udp"))] - fn run(&self, threads: usize) { + fn run(&self) { if self.enable_udp { warn!("UDP relay feature is disabled, recompile with feature=\"enable-udp\" to enable this feature"); } let tcprelay = self.tcprelay.clone(); - Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || tcprelay.run()); + let fut = Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || tcprelay.run()); info!("Enabled TCP relay"); - Scheduler::run(threads); + fut.join().unwrap(); } #[cfg(feature = "enable-udp")] - fn run(&self, threads: usize) { + fn run(&self) { + let mut futs = Vec::new(); + let tcprelay = self.tcprelay.clone(); - Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || tcprelay.run()); + let tcp_fut = Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || tcprelay.run()); info!("Enabled TCP relay"); + futs.push(tcp_fut); if self.enable_udp { let udprelay = self.udprelay.clone(); - Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || udprelay.run()); + let udp_fut = Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || udprelay.run()); info!("Enabled UDP relay"); + futs.push(udp_fut); } - Scheduler::run(threads); + for fut in futs { + fut.join().unwrap(); + } } } diff --git a/src/relay/mod.rs b/src/relay/mod.rs index 06993977..c8f836ee 100644 --- a/src/relay/mod.rs +++ b/src/relay/mod.rs @@ -37,7 +37,7 @@ pub mod socks5; pub const COROUTINE_STACK_SIZE: usize = 32 * 1024; // 32KB pub trait Relay { - fn run(&self, threads: usize); + fn run(&self); } fn copy(r: &mut R, w: &mut W, prefix: &str) -> io::Result { diff --git a/src/relay/server.rs b/src/relay/server.rs index 079f1e18..63064f4e 100644 --- a/src/relay/server.rs +++ b/src/relay/server.rs @@ -21,7 +21,7 @@ //! Server side -use coio::{Scheduler, Builder}; +use coio::Builder; #[cfg(feature = "enable-udp")] use relay::udprelay::server::UdpRelayServer; @@ -85,30 +85,36 @@ impl RelayServer { impl Relay for RelayServer { #[cfg(feature = "enable-udp")] - fn run(&self, threads: usize) { + fn run(&self) { + let mut futs = Vec::new(); + let tcprelay = self.tcprelay.clone(); - Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || tcprelay.run()); + let tcp_fut = Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || tcprelay.run()); info!("Enabled TCP relay"); + futs.push(tcp_fut); if self.enable_udp { let udprelay = self.udprelay.clone(); - Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || udprelay.run()); + let udp_fut = Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || udprelay.run()); info!("Enabled UDP relay"); + futs.push(udp_fut); } - Scheduler::run(threads); + for fut in futs { + fut.join().unwrap(); + } } #[cfg(not(feature = "enable-udp"))] - fn run(&self, threads: usize) { + fn run(&self) { if self.enable_udp { warn!("UDP relay feature is disabled, recompile with feature=\"enable-udp\" to enable this feature"); } let tcprelay = self.tcprelay.clone(); - Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || tcprelay.run()); + let fut = Builder::new().stack_size(COROUTINE_STACK_SIZE).spawn(move || tcprelay.run()); info!("Enabled TCP relay"); - Scheduler::run(threads); + fut.join().unwrap(); } }