From b40b932755c70a2fe87790a774b9b66c91a30229 Mon Sep 17 00:00:00 2001
From: Fangdun Tsai <cfddream@gmail.com>
Date: Mon, 18 Sep 2023 17:25:39 +0800
Subject: [PATCH] chore(examples): graceful-shutdown

---
 Cargo.toml                             |  1 +
 examples/graceful-shutdown/Cargo.toml  | 10 ++++
 examples/graceful-shutdown/src/main.rs | 69 ++++++++++++++++++++++++++
 3 files changed, 80 insertions(+)
 create mode 100644 examples/graceful-shutdown/Cargo.toml
 create mode 100644 examples/graceful-shutdown/src/main.rs

diff --git a/Cargo.toml b/Cargo.toml
index f08fb69a..b2624247 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -27,6 +27,7 @@ members = [
   "examples/compression",
   "examples/templates/*",
   "examples/tracing",
+  "examples/graceful-shutdown",
 ]
 resolver = "2"
 
diff --git a/examples/graceful-shutdown/Cargo.toml b/examples/graceful-shutdown/Cargo.toml
new file mode 100644
index 00000000..34d8f965
--- /dev/null
+++ b/examples/graceful-shutdown/Cargo.toml
@@ -0,0 +1,10 @@
+[package]
+name = "graceful-shutdown"
+version = "0.1.0"
+edition.workspace = true
+publish = false
+
+[dependencies]
+viz.workspace = true
+
+tokio = { workspace = true, features = [ "rt-multi-thread", "macros" ] }
diff --git a/examples/graceful-shutdown/src/main.rs b/examples/graceful-shutdown/src/main.rs
new file mode 100644
index 00000000..9a38ac71
--- /dev/null
+++ b/examples/graceful-shutdown/src/main.rs
@@ -0,0 +1,69 @@
+#![deny(warnings)]
+#![allow(clippy::unused_async)]
+
+//! Graceful shutdown server.
+//!
+//! See https://github.com/hyperium/hyper/blob/master/examples/graceful_shutdown.rs
+
+use std::{net::SocketAddr, sync::Arc, time::Duration};
+use tokio::{net::TcpListener, pin};
+use viz::{server::conn::http1, Io, Request, Responder, Result, Router, Tree};
+
+async fn index(_: Request) -> Result<&'static str> {
+    Ok("Hello, World!")
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
+    let listener = TcpListener::bind(addr).await?;
+    println!("listening on {addr}");
+
+    let app = Router::new().get("/", index);
+    let tree = Arc::new(Tree::from(app));
+
+    // Use a 5 second timeout for incoming connections to the server.
+    // If a request is in progress when the 5 second timeout elapses,
+    // use a 2 second timeout for processing the final request and graceful shutdown.
+    let connection_timeouts = vec![Duration::from_secs(5), Duration::from_secs(2)];
+
+    loop {
+        // Clone the connection_timeouts so they can be passed to the new task.
+        let connection_timeouts_clone = connection_timeouts.clone();
+
+        let (stream, addr) = listener.accept().await?;
+        let tree = tree.clone();
+
+        tokio::task::spawn(async move {
+            // Pin the connection object so we can use tokio::select! below.
+            let conn = http1::Builder::new()
+                .serve_connection(Io::new(stream), Responder::new(tree, Some(addr)));
+            pin!(conn);
+
+            // Iterate the timeouts.  Use tokio::select! to wait on the
+            // result of polling the connection itself,
+            // and also on tokio::time::sleep for the current timeout duration.
+            for (iter, sleep_duration) in connection_timeouts_clone.iter().enumerate() {
+                println!("iter = {} sleep_duration = {:?}", iter, sleep_duration);
+                tokio::select! {
+                    res = conn.as_mut() => {
+                        // Polling the connection returned a result.
+                        // In this case print either the successful or error result for the connection
+                        // and break out of the loop.
+                        match res {
+                            Ok(()) => println!("after polling conn, no error"),
+                            Err(e) =>  println!("error serving connection: {:?}", e),
+                        };
+                        break;
+                    }
+                    _ = tokio::time::sleep(*sleep_duration) => {
+                        // tokio::time::sleep returned a result.
+                        // Call graceful_shutdown on the connection and continue the loop.
+                        println!("iter = {} got timeout_interval, calling conn.graceful_shutdown", iter);
+                        conn.as_mut().graceful_shutdown();
+                    }
+                }
+            }
+        });
+    }
+}