Add shutdown trigger for server

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;

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,9 +73,17 @@ 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>,

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),
+        }
+    }
+}