io_uring: fix: flatten threads for try_join! macro

The try join was not working as supposed to:
"Waits on multiple concurrent branches, returning when all branches
complete with Ok(_) or on the first Err(_)."
But if a single thread return with Err, try_join was still running.

The original code was from this example:
https://fasterthanli.me/articles/remote-development-with-rust-on-fly-io?ref=dailydev#a-wonderful-tcp-proxy-with-tokio-uring

But the threads was not flattened like an official tokio docs is saying:
https://docs.rs/tokio/latest/tokio/macro.try_join.html#examples
(second example: "Using try_join! with spawned tasks.")

This commit is fixing this problem.

Author: manio

src/io_uring.rs

@@ -6,6 +6,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tokio::sync::Notify;
+use tokio::task::JoinHandle;
 use tokio::time::{sleep, timeout};
 use tokio_uring::buf::BoundedBuf;
 use tokio_uring::buf::BoundedBufMut;
@@ -175,6 +176,14 @@ async fn transfer_monitor(
     }
 }
 
+async fn flatten<T>(handle: &mut JoinHandle<Result<T>>) -> Result<T> {
+    match handle.await {
+        Ok(Ok(result)) => Ok(result),
+        Ok(Err(err)) => Err(err),
+        Err(_) => Err("handling failed".into()),
+    }
+}
+
 pub async fn io_loop(
     stats_interval: Option<Duration>,
     need_restart: Arc<Notify>,
@@ -251,7 +260,11 @@ pub async fn io_loop(
         let mut monitor = tokio::spawn(transfer_monitor(stats_interval, file_bytes, stream_bytes));
 
         // Stop as soon as one of them errors
-        let res = tokio::try_join!(&mut from_file, &mut from_stream, &mut monitor);
+        let res = tokio::try_join!(
+            flatten(&mut from_file),
+            flatten(&mut from_stream),
+            flatten(&mut monitor)
+        );
         if let Err(e) = res {
             error!("{} Connection error: {}", NAME, e);
         }