chore(examples): use consistent serving method

Signed-off-by: Roman Volosatovs <rvolosatovs@riseup.net>

Author: rvolosatovs

examples/rust/hello-nats-server/src/main.rs

@@ -3,9 +3,10 @@ use core::pin::pin;
 use anyhow::Context as _;
 use clap::Parser;
 use futures::stream::select_all;
-use futures::{StreamExt as _, TryStreamExt as _};
+use futures::StreamExt as _;
+use tokio::task::JoinSet;
 use tokio::{select, signal};
-use tracing::{info, warn};
+use tracing::{debug, error, info, warn};
 use url::Url;
 
 mod bindings {
@@ -60,28 +61,45 @@ async fn main() -> anyhow::Result<()> {
         .context("failed to serve `wrpc-examples.hello/handler.hello`")?;
     // NOTE: This will conflate all invocation streams into a single stream via `futures::stream::SelectAll`,
     // to customize this, iterate over the returned `invocations` and set up custom handling per export
-    let mut invocations = select_all(invocations.into_iter().map(
-        |(instance, name, invocations)| {
-            invocations
-                .try_buffer_unordered(16) // handle up to 16 invocations concurrently
-                .map(move |res| (instance, name, res))
-        },
-    ));
+    let mut invocations = select_all(
+        invocations
+            .into_iter()
+            .map(|(instance, name, invocations)| invocations.map(move |res| (instance, name, res))),
+    );
     let shutdown = signal::ctrl_c();
     let mut shutdown = pin!(shutdown);
+    let mut tasks = JoinSet::new();
     loop {
         select! {
             Some((instance, name, res)) = invocations.next() => {
                 match res {
-                    Ok(()) => {
-                        info!(instance, name, "invocation successfully handled");
+                    Ok(fut) => {
+                        debug!(instance, name, "invocation accepted");
+                        tasks.spawn(async move {
+                            if let Err(err) = fut.await {
+                                warn!(?err, "failed to handle invocation");
+                            } else {
+                                info!(instance, name, "invocation successfully handled");
+                            }
+                        });
                     }
                     Err(err) => {
                         warn!(?err, instance, name, "failed to accept invocation");
                     }
                 }
             }
+            Some(res) = tasks.join_next() => {
+                if let Err(err) = res {
+                    error!(?err, "failed to join task")
+                }
+            }
             res = &mut shutdown => {
+                // wait for all invocations to complete
+                while let Some(res) = tasks.join_next().await {
+                    if let Err(err) = res {
+                        error!(?err, "failed to join task")
+                    }
+                }
                 return res.context("failed to listen for ^C")
             }
         }

examples/rust/hello-tcp-server/src/main.rs

@@ -6,9 +6,10 @@ use std::sync::Arc;
 use anyhow::Context as _;
 use clap::Parser;
 use futures::stream::select_all;
-use futures::{StreamExt as _, TryStreamExt as _};
+use futures::StreamExt as _;
+use tokio::task::JoinSet;
 use tokio::{select, signal};
-use tracing::{error, info, warn};
+use tracing::{debug, error, info, warn};
 
 mod bindings {
     wit_bindgen_wrpc::generate!({
@@ -61,29 +62,46 @@ async fn main() -> anyhow::Result<()> {
         .context("failed to serve `wrpc-examples.hello/handler.hello`")?;
     // NOTE: This will conflate all invocation streams into a single stream via `futures::stream::SelectAll`,
     // to customize this, iterate over the returned `invocations` and set up custom handling per export
-    let mut invocations = select_all(invocations.into_iter().map(
-        |(instance, name, invocations)| {
-            invocations
-                .try_buffer_unordered(16) // handle up to 16 invocations concurrently
-                .map(move |res| (instance, name, res))
-        },
-    ));
+    let mut invocations = select_all(
+        invocations
+            .into_iter()
+            .map(|(instance, name, invocations)| invocations.map(move |res| (instance, name, res))),
+    );
     let shutdown = signal::ctrl_c();
     let mut shutdown = pin!(shutdown);
+    let mut tasks = JoinSet::new();
     loop {
         select! {
             Some((instance, name, res)) = invocations.next() => {
                 match res {
-                    Ok(()) => {
-                        info!(instance, name, "invocation successfully handled");
+                    Ok(fut) => {
+                        debug!(instance, name, "invocation accepted");
+                        tasks.spawn(async move {
+                            if let Err(err) = fut.await {
+                                warn!(?err, "failed to handle invocation");
+                            } else {
+                                info!(instance, name, "invocation successfully handled");
+                            }
+                        });
                     }
                     Err(err) => {
                         warn!(?err, instance, name, "failed to accept invocation");
                     }
                 }
             }
+            Some(res) = tasks.join_next() => {
+                if let Err(err) = res {
+                    error!(?err, "failed to join task")
+                }
+            }
             res = &mut shutdown => {
                 accept.abort();
+                // wait for all invocations to complete
+                while let Some(res) = tasks.join_next().await {
+                    if let Err(err) = res {
+                        error!(?err, "failed to join task")
+                    }
+                }
                 return res.context("failed to listen for ^C")
             }
         }

examples/rust/hello-web-server/src/main.rs

@@ -105,9 +105,9 @@ async fn main() -> anyhow::Result<()> {
             .into_iter()
             .map(|(instance, name, invocations)| invocations.map(move |res| (instance, name, res))),
     );
-    let mut tasks = JoinSet::new();
     let shutdown = signal::ctrl_c();
     let mut shutdown = pin!(shutdown);
+    let mut tasks = JoinSet::new();
     loop {
         select! {
             Some((instance, name, res)) = invocations.next() => {
@@ -134,6 +134,7 @@ async fn main() -> anyhow::Result<()> {
             }
             res = &mut shutdown => {
                 accept.abort();
+                // wait for all invocations to complete
                 while let Some(res) = tasks.join_next().await {
                     if let Err(err) = res {
                         error!(?err, "failed to join task")

examples/rust/streams-nats-server/src/main.rs

@@ -4,9 +4,10 @@ use anyhow::Context as _;
 use bytes::Bytes;
 use clap::Parser;
 use futures::stream::select_all;
-use futures::{Stream, StreamExt as _, TryStreamExt as _};
+use futures::{Stream, StreamExt as _};
+use tokio::task::JoinSet;
 use tokio::{select, signal};
-use tracing::{info, warn};
+use tracing::{debug, error, info, warn};
 use url::Url;
 
 mod bindings {
@@ -70,28 +71,45 @@ async fn main() -> anyhow::Result<()> {
         .context("failed to serve `wrpc-examples:streams/handler.echo`")?;
     // NOTE: This will conflate all invocation streams into a single stream via `futures::stream::SelectAll`,
     // to customize this, iterate over the returned `invocations` and set up custom handling per export
-    let mut invocations = select_all(invocations.into_iter().map(
-        |(instance, name, invocations)| {
-            invocations
-                .try_buffer_unordered(16) // handle up to 16 invocations concurrently
-                .map(move |res| (instance, name, res))
-        },
-    ));
+    let mut invocations = select_all(
+        invocations
+            .into_iter()
+            .map(|(instance, name, invocations)| invocations.map(move |res| (instance, name, res))),
+    );
     let shutdown = signal::ctrl_c();
     let mut shutdown = pin!(shutdown);
+    let mut tasks = JoinSet::new();
     loop {
         select! {
             Some((instance, name, res)) = invocations.next() => {
                 match res {
-                    Ok(()) => {
-                        info!(instance, name, "invocation successfully handled");
+                    Ok(fut) => {
+                        debug!(instance, name, "invocation accepted");
+                        tasks.spawn(async move {
+                            if let Err(err) = fut.await {
+                                warn!(?err, "failed to handle invocation");
+                            } else {
+                                info!(instance, name, "invocation successfully handled");
+                            }
+                        });
                     }
                     Err(err) => {
                         warn!(?err, instance, name, "failed to accept invocation");
                     }
                 }
             }
+            Some(res) = tasks.join_next() => {
+                if let Err(err) = res {
+                    error!(?err, "failed to join task")
+                }
+            }
             res = &mut shutdown => {
+                // wait for all invocations to complete
+                while let Some(res) = tasks.join_next().await {
+                    if let Err(err) = res {
+                        error!(?err, "failed to join task")
+                    }
+                }
                 return res.context("failed to listen for ^C")
             }
         }

examples/rust/wasi-keyvalue-nats-server/src/main.rs

@@ -3,9 +3,10 @@ use core::pin::pin;
 use anyhow::Context as _;
 use clap::Parser;
 use futures::stream::select_all;
-use futures::{StreamExt as _, TryStreamExt as _};
+use futures::StreamExt as _;
+use tokio::task::JoinSet;
 use tokio::{select, signal};
-use tracing::{info, warn};
+use tracing::{debug, error, info, warn};
 use url::Url;
 
 #[derive(Parser, Debug)]
@@ -41,28 +42,45 @@ async fn main() -> anyhow::Result<()> {
         .context("failed to serve `wasi:keyvalue`")?;
     // NOTE: This will conflate all invocation streams into a single stream via `futures::stream::SelectAll`,
     // to customize this, iterate over the returned `invocations` and set up custom handling per export
-    let mut invocations = select_all(invocations.into_iter().map(
-        |(instance, name, invocations)| {
-            invocations
-                .try_buffer_unordered(16) // handle up to 16 invocations concurrently
-                .map(move |res| (instance, name, res))
-        },
-    ));
+    let mut invocations = select_all(
+        invocations
+            .into_iter()
+            .map(|(instance, name, invocations)| invocations.map(move |res| (instance, name, res))),
+    );
     let shutdown = signal::ctrl_c();
     let mut shutdown = pin!(shutdown);
+    let mut tasks = JoinSet::new();
     loop {
         select! {
             Some((instance, name, res)) = invocations.next() => {
                 match res {
-                    Ok(()) => {
-                        info!(instance, name, "invocation successfully handled");
+                    Ok(fut) => {
+                        debug!(instance, name, "invocation accepted");
+                        tasks.spawn(async move {
+                            if let Err(err) = fut.await {
+                                warn!(?err, "failed to handle invocation");
+                            } else {
+                                info!(instance, name, "invocation successfully handled");
+                            }
+                        });
                     }
                     Err(err) => {
                         warn!(?err, instance, name, "failed to accept invocation");
                     }
                 }
             }
+            Some(res) = tasks.join_next() => {
+                if let Err(err) = res {
+                    error!(?err, "failed to join task")
+                }
+            }
             res = &mut shutdown => {
+                // wait for all invocations to complete
+                while let Some(res) = tasks.join_next().await {
+                    if let Err(err) = res {
+                        error!(?err, "failed to join task")
+                    }
+                }
                 return res.context("failed to listen for ^C")
             }
         }

examples/rust/wasi-keyvalue-tcp-server/src/main.rs

@@ -4,9 +4,10 @@ use std::sync::Arc;
 use anyhow::Context as _;
 use clap::Parser;
 use futures::stream::select_all;
-use futures::{StreamExt as _, TryStreamExt as _};
+use futures::StreamExt as _;
+use tokio::task::JoinSet;
 use tokio::{select, signal};
-use tracing::{error, info, warn};
+use tracing::{debug, error, info, warn};
 
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
@@ -42,29 +43,46 @@ async fn main() -> anyhow::Result<()> {
             .context("failed to serve `wasi:keyvalue`")?;
     // NOTE: This will conflate all invocation streams into a single stream via `futures::stream::SelectAll`,
     // to customize this, iterate over the returned `invocations` and set up custom handling per export
-    let mut invocations = select_all(invocations.into_iter().map(
-        |(instance, name, invocations)| {
-            invocations
-                .try_buffer_unordered(16) // handle up to 16 invocations concurrently
-                .map(move |res| (instance, name, res))
-        },
-    ));
+    let mut invocations = select_all(
+        invocations
+            .into_iter()
+            .map(|(instance, name, invocations)| invocations.map(move |res| (instance, name, res))),
+    );
     let shutdown = signal::ctrl_c();
     let mut shutdown = pin!(shutdown);
+    let mut tasks = JoinSet::new();
     loop {
         select! {
             Some((instance, name, res)) = invocations.next() => {
                 match res {
-                    Ok(()) => {
-                        info!(instance, name, "invocation successfully handled");
+                    Ok(fut) => {
+                        debug!(instance, name, "invocation accepted");
+                        tasks.spawn(async move {
+                            if let Err(err) = fut.await {
+                                warn!(?err, "failed to handle invocation");
+                            } else {
+                                info!(instance, name, "invocation successfully handled");
+                            }
+                        });
                     }
                     Err(err) => {
                         warn!(?err, instance, name, "failed to accept invocation");
                     }
                 }
             }
+            Some(res) = tasks.join_next() => {
+                if let Err(err) = res {
+                    error!(?err, "failed to join task")
+                }
+            }
             res = &mut shutdown => {
                 accept.abort();
+                // wait for all invocations to complete
+                while let Some(res) = tasks.join_next().await {
+                    if let Err(err) = res {
+                        error!(?err, "failed to join task")
+                    }
+                }
                 return res.context("failed to listen for ^C")
             }
         }