use the latest coio

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::<usize>()
         .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();
 }

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::<usize>()
         .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();
 }

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();
+        }
     }
 }

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: Read, W: Write>(r: &mut R, w: &mut W, prefix: &str) -> io::Result<u64> {

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();
     }
 }