fix multi client

client/src/main.rs

@@ -1,15 +1,17 @@
 use crate::config::{ClientConfig, Profile};
 use crate::task::run;
-use communicator::{Identity, TunnelEndpoint, TunnelListener, connect_tunnel};
+use communicator::{ClientManager, Identity, TunnelEndpoint, TunnelListener, connect_tunnel};
 use lazy_static::lazy_static;
 use log::{LevelFilter, error, info};
 use simplelog::{ColorChoice, Config, TermLogger, TerminalMode};
+use std::collections::{HashMap, VecDeque};
 use std::fs;
 use std::str::FromStr;
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 use std::time::Duration;
 use structopt::StructOpt;
-use tokio::sync::{OwnedSemaphorePermit, Semaphore};
+use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc};
+use tokio::task::JoinSet;
 use tokio::time::sleep;
 use tokio_util::sync::CancellationToken;
 
@@ -60,7 +62,59 @@ async fn main() -> anyhow::Result<()> {
         })
         .collect::<Result<Vec<_>, _>>()?;
 
-    let mut tasks = vec![];
+    let mut join_set = JoinSet::new();
+
+    let mut senders: HashMap<String, Arc<Sender<Arc<ClientManager>>>> = HashMap::new();
+    for profile in &profiles {
+        senders.insert(profile.0.clone(), Arc::new(Sender::new()));
+    }
+    let (liveness_tx, mut liveness_rx) = mpsc::channel::<()>(1);
+    let senders = Arc::new(senders);
+    let cancel_token = CancellationToken::new();
+    {
+        let cancel_token = cancel_token.clone();
+        join_set.spawn(async move {
+            if liveness_rx.recv().await.is_none() {
+                cancel_token.cancel();
+            }
+        });
+    }
+
+    for server_conf in &CONFIG.connect.server {
+        let server_conf = server_conf.clone().into_tunnel_config(Identity::Client)?;
+        let url = server_conf.url.clone();
+        let server = TunnelListener::bind(server_conf).await?;
+        let senders = senders.clone();
+        let cancel_token = cancel_token.clone();
+        info!("tcp server started on {}", url);
+        join_set.spawn(async move {
+            loop {
+                tokio::select! {
+                    _ = cancel_token.cancelled() => {
+                        info!("tcp server {} stopped", url);
+                        break;
+                    }
+
+                    endpoint_res = server.accept() => {
+                        let endpoint = match endpoint_res {
+                            Ok(e) => e,
+                            Err(e) => {
+                                error!("Failed to accept connection: {}", e);
+                                continue;
+                            }
+                        };
+
+                        if let TunnelEndpoint::Client(client) = endpoint {
+                            for (_, sender) in senders.iter() {
+                                    sender.send(client.clone());
+                            }
+                        }
+                    }
+                }
+            }
+        });
+    }
+
     for profile in profiles {
         let profile = Arc::new(profile.clone());
         let semaphore = Arc::new(Semaphore::new(1));
@@ -80,48 +134,45 @@ async fn main() -> anyhow::Result<()> {
                 drop(permit);
             }
         };
-        for server_conf in &CONFIG.connect.server {
-            let server_conf = server_conf.clone().into_tunnel_config(Identity::Client)?;
-            let url = server_conf.url.clone();
-            let server = TunnelListener::bind(server_conf).await?;
+
+        {
+            let sender = senders.get(&profile.0).unwrap().clone();
             let semaphore = semaphore.clone();
             let cancel_token = cancel_token.clone();
             let profile = profile.clone();
-            info!("tcp server started on {}", url);
-            tasks.push(tokio::task::spawn(async move {
+            let liveness_tx = liveness_tx.clone();
+            join_set.spawn(async move {
+                let _guard = liveness_tx;
                 loop {
-                    let endpoint = server
-                        .accept()
-                        .await
-                        .map_err(|e| error!("Failed to accept connection: {}", e));
-                    let endpoint = if let Ok(endpoint) = endpoint {
-                        endpoint
-                    } else {
+                    let client = sender.recv();
+                    if client.is_none() {
                         continue;
-                    };
-                    if let TunnelEndpoint::Client(client) = endpoint {
-                        loop {
-                            if client.get_client().await.is_err() {
-                                break;
+                    }
+                    let client = client.unwrap();
+                    if cancel_token.is_cancelled() {
+                        return;
+                    }
+                    loop {
+                        if client.get_client().await.is_err() {
+                            break;
+                        }
+                        if cancel_token.is_cancelled() {
+                            return;
+                        }
+                        let permit = semaphore.clone().acquire_owned().await.unwrap();
+                        let result = run(client.clone(), profile.clone()).await;
+                        match result {
+                            Ok(true) => {
+                                post_task(profile.clone(), permit, cancel_token.clone()).await;
                             }
-                            if cancel_token.is_cancelled() {
-                                return;
-                            }
-                            let permit = semaphore.clone().acquire_owned().await.unwrap();
-                            let result = run(client.clone(), profile.clone()).await;
-                            match result {
-                                Ok(true) => {
-                                    post_task(profile.clone(), permit, cancel_token.clone()).await;
-                                }
-                                Err(error) => {
-                                    error!("{}", error);
-                                }
-                                _ => {}
+                            Err(error) => {
+                                error!("{}", error);
                             }
+                            _ => {}
                         }
                     }
                 }
-            }));
+            });
         }
 
         for client_conf in &CONFIG.connect.client {
@@ -130,7 +181,7 @@ async fn main() -> anyhow::Result<()> {
             let cancel_token = cancel_token.clone();
             let profile = profile.clone();
             info!("tcp client started for {}", client_conf.url);
-            tasks.push(tokio::task::spawn(async move {
+            join_set.spawn(async move {
                 loop {
                     if cancel_token.is_cancelled() {
                         return;
@@ -164,18 +215,46 @@ async fn main() -> anyhow::Result<()> {
                             }
                         }
                     }
+                    if cancel_token.is_cancelled() {
+                        return;
+                    }
                     sleep(Duration::from_secs(10)).await;
                 }
-            }));
+            });
         }
     }
 
-    for task in tasks {
-        let _ = task.await.map_err(|e| error!("{}", e));
+    drop(liveness_tx);
+
+    while let Some(res) = join_set.join_next().await {
+        if let Err(e) = res {
+            error!("{}", e);
+        }
     }
 
     Ok(())
 }
+
+struct Sender<T> {
+    inner: RwLock<VecDeque<T>>,
+}
+
+impl<T> Sender<T> {
+    pub fn new() -> Self {
+        Self {
+            inner: RwLock::new(VecDeque::new()),
+        }
+    }
+
+    pub fn send(&self, item: T) {
+        self.inner.write().unwrap().push_back(item);
+    }
+
+    pub fn recv(&self) -> Option<T> {
+        self.inner.write().unwrap().pop_front()
+    }
+}
+
 lazy_static! {
     pub static ref CONFIG: ClientConfig = {
         let raw = fs::read_to_string("sekai-unpacker-client.yaml").unwrap();

client/src/task.rs

@@ -9,12 +9,14 @@ use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use tokio::sync::{RwLock, Semaphore};
+use tokio::task::JoinSet;
 
 pub async fn run(
     client: Arc<ClientManager>,
     profile: Arc<(String, Profile)>,
 ) -> anyhow::Result<bool> {
     info!("[{}]: Starting sync", profile.0);
+    tokio::fs::create_dir_all(&profile.1.path).await?;
     let sync_resp = sync(&mut client.get_client().await?, &profile.1).await?;
     let id = sync_resp.id;
     let local_manifest = Arc::new(
@@ -41,7 +43,7 @@ pub async fn run(
     info!("[{}]: Collected {} tasks", profile.0, tasks.len());
     let n = 5;
     let semaphore = Arc::new(Semaphore::new(n));
-    let mut handles = Vec::new();
+    let mut join_set = JoinSet::new();
     for task in tasks {
         let permit = semaphore.clone().acquire_owned().await?;
         let client = client.clone();
@@ -49,7 +51,7 @@ pub async fn run(
         let local_manifest = local_manifest.clone();
         let profile = profile.clone();
 
-        handles.push(tokio::task::spawn(async move {
+        join_set.spawn(async move {
             let req = DownloadRequest {
                 id: id.clone(),
                 task: task.clone(),
@@ -68,12 +70,11 @@ pub async fn run(
                 .await
                 .unwrap();
             drop(permit);
-        }));
+        });
     }
     let mut succeed = 0;
     let mut failed = 0;
-    for handle in handles {
-        let r = handle.await;
+    while let Some(r) = join_set.join_next().await {
         if let Err(e) = r {
             error!("{}", e);
             failed += 1;