From da6f5161fe34cea5c9fb10d889c8149f3325ee62 Mon Sep 17 00:00:00 2001
From: Bulat Kurbanov <kurbanovbul@gmail.com>
Date: Mon, 8 Jan 2024 00:58:08 +0100
Subject: [PATCH] Add gracefull shutdown

---
 src/bots_manager/axum_server.rs | 17 +++++++++++++++--
 src/bots_manager/mod.rs         |  2 +-
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/bots_manager/axum_server.rs b/src/bots_manager/axum_server.rs
index 5bd63ad..0c622bd 100644
--- a/src/bots_manager/axum_server.rs
+++ b/src/bots_manager/axum_server.rs
@@ -4,7 +4,13 @@ use axum::routing::post;
 
 use reqwest::StatusCode;
 
-use std::net::SocketAddr;
+use std::{
+    net::SocketAddr,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+};
 
 use teloxide::types::{Update, UpdateKind};
 
@@ -14,7 +20,7 @@ use tracing::log;
 
 use crate::bots_manager::{internal::start_bot, BOTS_DATA, BOTS_ROUTES, SERVER_PORT};
 
-pub async fn start_axum_server() {
+pub async fn start_axum_server(stop_signal: Arc<AtomicBool>) {
     async fn telegram_request(Path(token): Path<String>, input: String) -> impl IntoResponse {
         let (_, r_tx) = match BOTS_ROUTES.get(&token).await {
             Some(tx) => tx,
@@ -80,6 +86,13 @@ pub async fn start_axum_server() {
 
         axum::Server::bind(&addr)
             .serve(router.into_make_service())
+            .with_graceful_shutdown(async move {
+                loop {
+                    if !stop_signal.load(Ordering::SeqCst) {
+                        break;
+                    };
+                }
+            })
             .await
             .expect("Axum server error");
 
diff --git a/src/bots_manager/mod.rs b/src/bots_manager/mod.rs
index 7de07c6..8806672 100644
--- a/src/bots_manager/mod.rs
+++ b/src/bots_manager/mod.rs
@@ -98,7 +98,7 @@ impl BotsManager {
     }
 
     pub async fn start(running: Arc<AtomicBool>) {
-        start_axum_server().await;
+        start_axum_server(running.clone()).await;
 
         let mut interval = time::interval(Duration::from_secs(5));