Commit ugly code that crashed rustc.

Author: Jonathan Bastien-Filiatrault

Cargo.toml

@@ -12,7 +12,7 @@ repository = "https://github.com/zerospam/smtpbis.git"
 
 [dependencies]
 rustyknife = {version="0.2", features=["quoted-string-rfc2047"]}
-tokio = {version="0.2", features=["signal", "io-util", "sync", "signal", "rt-core", "tcp", "dns", "rt-threaded"]}
+tokio = {version="0.2", features=["signal", "io-util", "sync", "signal", "rt-core", "tcp", "dns", "rt-threaded", "macros"]}
 tokio-util = {version="0.2", features=["codec"]}
 bytes = "0.5"
 futures = "0.3"

src/bin/smtpbis-server/main.rs

@@ -15,7 +15,6 @@ use futures_util::stream::TryStreamExt;
 
 use tokio::net::{TcpListener, TcpStream};
 use tokio::prelude::*;
-use tokio::runtime::Runtime;
 use tokio::sync::oneshot::Receiver;
 
 use tokio_rustls::rustls::{
@@ -160,18 +159,16 @@ impl DummyHandler {
     }
 }
 
-fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let mut rt = Runtime::new()?;
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let (listen_shutdown_tx, listen_shutdown_rx) = tokio::sync::oneshot::channel();
+    let join = tokio::spawn(listen_loop(listen_shutdown_rx));
 
-    rt.block_on(async {
-        let (listen_shutdown_tx, listen_shutdown_rx) = tokio::sync::oneshot::channel();
-        tokio::spawn(listen_loop(listen_shutdown_rx));
-
-        tokio::signal::ctrl_c().await.unwrap();
-        listen_shutdown_tx.send(()).unwrap();
-        println!("Waiting for tasks to finish...");
-        // FIXME: actually wait on tasks here.
-    });
+    tokio::signal::ctrl_c().await.unwrap();
+    listen_shutdown_tx.send(()).unwrap();
+    println!("Waiting for tasks to finish...");
+    join.await?;
+    println!("Done !");
 
     Ok(())
 }
@@ -187,7 +184,7 @@ async fn listen_loop(mut shutdown: Receiver<()>) {
     let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
     let shutdown_rx = shutdown_rx.map_err(|_| ()).shared();
 
-    loop {
+    smtpbis::taskjoin::join_tasks(move |spawn| Box::pin(async move { loop {
         let accept = listener.accept();
         pin_mut!(accept);
 
@@ -197,18 +194,18 @@ async fn listen_loop(mut shutdown: Receiver<()>) {
                 let mut shutdown_rx = shutdown_rx.clone();
                 let tls_config = tls_config.clone();
 
-                tokio::spawn(async move {
+                spawn.read().unwrap().push(tokio::spawn(async move {
                     let smtp_res = serve_smtp(socket, addr, tls_config, &mut shutdown_rx).await;
                     println!("SMTP task done: {:?}", smtp_res);
-                })
+                }))
             }
             Either::Right(..) => {
                 println!("socket listening loop stopping");
                 shutdown_tx.send(()).unwrap();
                 break;
             }
         };
-    }
+    }})).await;
 }
 
 async fn serve_smtp(

src/lib.rs

@@ -4,6 +4,7 @@ mod codecs;
 mod reply;
 mod server;
 mod syntax;
+pub mod taskjoin;
 
 pub use codecs::{LineCodec, LineError};
 pub use reply::*;

src/taskjoin.rs

@@ -0,0 +1,41 @@
+use std::sync::{Arc, RwLock};
+
+use futures::stream::futures_unordered::FuturesUnordered;
+use futures::Future;
+use futures::future::{Either};
+use futures_util::stream::StreamExt;
+
+pub async fn join_tasks<Fut, Fut2, FutO, Task>(fut: Fut) -> FutO
+where Fut: FnOnce(Arc<RwLock<FuturesUnordered<Task>>>) -> Fut2,
+      Fut2: Future<Output=FutO> + Unpin,
+      FutO: Sized,
+      Task: Future,
+{
+    let fo = Arc::new(RwLock::new(FuturesUnordered::new()));
+    let output;
+
+    let mut task = fut(fo.clone());
+
+    loop {
+        match futures::future::select(&mut task, fo.write().unwrap().next()).await {
+            Either::Left((spawn_output, _fo_next)) => {
+                output = spawn_output;
+                eprintln!("spawner ended");
+                break;
+            },
+            Either::Right((_task_output, _task)) => {
+                eprintln!("task awaited (spawner alive)");
+            }
+        }
+    }
+
+    let mut fo = fo.write().unwrap();
+
+    eprintln!("waiting for remaining tasks");
+    while !fo.is_empty() {
+        fo.next().await;
+        eprintln!("task awaited (join)");
+    }
+    eprintln!("done");
+    output
+}