From dba99279e7f29fb232f81528d20c89455088bade Mon Sep 17 00:00:00 2001
From: moznion <moznion@gmail.com>
Date: Mon, 23 Nov 2020 13:53:25 +0900
Subject: [PATCH] Add shutdown trigger for server

---
 src/lib.rs                     |  1 +
 src/server.rs                  | 11 +++++++++++
 src/server_shutdown_trigger.rs | 31 +++++++++++++++++++++++++++++++
 3 files changed, 43 insertions(+)
 create mode 100644 src/server_shutdown_trigger.rs

diff --git a/src/lib.rs b/src/lib.rs
index 2982981..089a66f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -9,3 +9,4 @@ pub mod server;
 pub mod secret_provider;
 pub mod request_handler;
 pub mod request;
+pub mod server_shutdown_trigger;
diff --git a/src/server.rs b/src/server.rs
index e847eb5..d784293 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -10,6 +10,7 @@ use crate::packet::Packet;
 use crate::request::Request;
 use crate::request_handler::RequestHandler;
 use crate::secret_provider::SecretProvider;
+use crate::server_shutdown_trigger::ServerShutdownTrigger;
 
 pub struct Server<T: RequestHandler, U: SecretProvider> {
     address: String,
@@ -17,6 +18,7 @@ pub struct Server<T: RequestHandler, U: SecretProvider> {
     skip_authenticity_validation: bool,
     request_handler_arc: Arc<T>,
     secret_provider_arc: Arc<U>,
+    shutdown_trigger: ServerShutdownTrigger,
 }
 
 impl<T: RequestHandler, U: SecretProvider> Server<T, U> {
@@ -27,6 +29,7 @@ impl<T: RequestHandler, U: SecretProvider> Server<T, U> {
             skip_authenticity_validation,
             request_handler_arc: Arc::new(request_handler),
             secret_provider_arc: Arc::new(secret_provider),
+            shutdown_trigger: ServerShutdownTrigger::new(),
         }
     }
 
@@ -70,10 +73,18 @@ impl<T: RequestHandler, U: SecretProvider> Server<T, U> {
                         ).await;
                     });
                 }
+                Some(_) = self.shutdown_trigger => {
+                    info!("server is shutting down");
+                    return Ok(());
+                }
             }
         }
     }
 
+    pub fn trigger_shutdown(&mut self) {
+        self.shutdown_trigger.trigger_shutdown();
+    }
+
     async fn process_request(
         conn: Arc<UdpSocket>,
         request_data: &Vec<u8>,
diff --git a/src/server_shutdown_trigger.rs b/src/server_shutdown_trigger.rs
new file mode 100644
index 0000000..7a876b7
--- /dev/null
+++ b/src/server_shutdown_trigger.rs
@@ -0,0 +1,31 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+#[derive(Copy, Clone)]
+pub struct ServerShutdownTrigger {
+    should_shutdown: bool
+}
+
+impl ServerShutdownTrigger {
+    pub(crate) fn new() -> Self {
+        ServerShutdownTrigger {
+            should_shutdown: false,
+        }
+    }
+
+    pub(crate) fn trigger_shutdown(&mut self) {
+        self.should_shutdown = true;
+    }
+}
+
+impl Future for ServerShutdownTrigger {
+    type Output = Option<()>;
+
+    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self.should_shutdown {
+            true => Poll::from(Some(())),
+            false => Poll::from(None),
+        }
+    }
+}