From 1e6bc3ba763eca665a1d99abc0d4efeaee3329e2 Mon Sep 17 00:00:00 2001
From: "Y. T. Chung" <zonyitoo@gmail.com>
Date: Thu, 17 May 2018 20:57:22 +0800
Subject: [PATCH] Optimize server listening logic, join_all

---
 src/relay/tcprelay/server.rs | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/relay/tcprelay/server.rs b/src/relay/tcprelay/server.rs
index f452cb584bd4..aedec258d244 100644
--- a/src/relay/tcprelay/server.rs
+++ b/src/relay/tcprelay/server.rs
@@ -11,13 +11,14 @@ use relay::dns_resolver::resolve;
 use relay::socks5::Address;
 use relay::tcprelay::crypto_io::{DecryptedRead, EncryptedWrite};
 
+use futures::future::join_all;
 use futures::stream::Stream;
 use futures::{self, Future};
 
 use tokio;
 use tokio::net::{TcpListener, TcpStream};
 use tokio_io::io::{ReadHalf, WriteHalf};
-use tokio_io::{AsyncRead, IoFuture};
+use tokio_io::AsyncRead;
 
 use super::{proxy_handshake, try_timeout, tunnel, DecryptedHalf, EncryptedHalfFut, TcpStreamConnect};
 
@@ -163,7 +164,7 @@ fn handle_client(server_cfg: Arc<ServerConfig>,
 
 /// Runs the server
 pub fn run(config: Arc<Config>) -> impl Future<Item = (), Error = io::Error> + Send {
-    let mut fut: Option<IoFuture<()>> = None;
+    let mut vec_fut = Vec::with_capacity(config.server.len());
 
     for svr_cfg in &config.server {
         let listener = {
@@ -190,11 +191,8 @@ pub fn run(config: Arc<Config>) -> impl Future<Item = (), Error = io::Error> + S
                                              err
                                          });
 
-        fut = Some(match fut.take() {
-                       Some(fut) => Box::new(fut.join(listening).map(|_| ())) as IoFuture<()>,
-                       None => Box::new(listening) as IoFuture<()>,
-                   })
+        vec_fut.push(boxed_future(listening));
     }
 
-    fut.expect("Must have at least one server")
+    join_all(vec_fut).map(|_| ())
 }