chore: Increase graceful timeout to 30 seconds and handle server stop…

…ping signal in web UI
DitZii · Jun 18, 2024 · 8bba38d · 8bba38d
1 parent c05668e
commit 8bba38d
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 5 deletions.
diff --git a/src/ui/gunicorn.conf.py b/src/ui/gunicorn.conf.py
@@ -1,8 +1,11 @@
 from contextlib import suppress
 from hashlib import sha256
+from json import JSONDecodeError, dumps, loads
 from os import cpu_count, getenv, getpid, sep, urandom
 from os.path import join
 from pathlib import Path
+from signal import SIGINT, SIGTERM, signal
+from threading import Lock
 from regex import compile as re_compile
 from sys import path as sys_path
 from time import sleep
@@ -36,7 +39,7 @@
 worker_class = "gthread"
 threads = int(getenv("MAX_THREADS", MAX_WORKERS * 2))
 max_requests_jitter = min(8, MAX_WORKERS)
-graceful_timeout = 5
+graceful_timeout = 30
 
 DEBUG = getenv("DEBUG", False)
 
@@ -46,12 +49,16 @@
     reload = True
     reload_extra_files = [file.as_posix() for file in Path(sep, "usr", "share", "bunkerweb", "ui", "templates").iterdir()]
 
+LOCK = Lock()
+
 
 def on_starting(server):
+    TMP_DIR.mkdir(parents=True, exist_ok=True)
     if not getenv("FLASK_SECRET") and not TMP_DIR.joinpath(".flask_secret").is_file():
-        TMP_DIR.mkdir(parents=True, exist_ok=True)
         TMP_DIR.joinpath(".flask_secret").write_text(sha256(urandom(32)).hexdigest(), encoding="utf-8")
 
+    TMP_DIR.joinpath(".ui.json").write_text("{}", encoding="utf-8")
+
     LOGGER = setup_logger("UI")
 
     db = Database(LOGGER, ui=True)
@@ -125,6 +132,29 @@ def on_starting(server):
     LOGGER.info("UI is ready")
 
 
+def handle_stop(signum=None, frame=None):
+    if not TMP_DIR.joinpath(".ui.json").is_file():
+        return
+
+    ui_data = "Error"
+    while ui_data == "Error":
+        with suppress(JSONDecodeError):
+            ui_data = loads(TMP_DIR.joinpath(".ui.json").read_text(encoding="utf-8"))
+
+    ui_data["SERVER_STOPPING"] = True
+
+    with LOCK:
+        TMP_DIR.joinpath(".ui.json").write_text(dumps(ui_data), encoding="utf-8")
+
+
+signal(SIGINT, handle_stop)
+signal(SIGTERM, handle_stop)
+
+
+def on_reload(server):
+    handle_stop()
+
+
 def when_ready(server):
     RUN_DIR.mkdir(parents=True, exist_ok=True)
     RUN_DIR.joinpath("ui.pid").write_text(str(getpid()), encoding="utf-8")
@@ -135,3 +165,4 @@ def on_exit(server):
     RUN_DIR.joinpath("ui.pid").unlink(missing_ok=True)
     TMP_DIR.joinpath("ui.healthy").unlink(missing_ok=True)
     TMP_DIR.joinpath(".flask_secret").unlink(missing_ok=True)
+    TMP_DIR.joinpath(".ui.json").unlink(missing_ok=True)
diff --git a/src/ui/main.py b/src/ui/main.py
@@ -21,7 +21,7 @@
 from dateutil.parser import parse as dateutil_parse
 from docker import DockerClient
 from docker.errors import NotFound as docker_NotFound, APIError as docker_APIError, DockerException
-from flask import Flask, Response, flash, jsonify, redirect, render_template, request, send_file, session, url_for
+from flask import Flask, Response, flash, jsonify, make_response, redirect, render_template, request, send_file, session, url_for
 from flask_login import current_user, LoginManager, login_required, login_user, logout_user
 from flask_wtf.csrf import CSRFProtect, CSRFError
 from hashlib import sha256
@@ -458,11 +458,16 @@ def handle_csrf_error(_):
 
 @app.before_request
 def before_request():
+    ui_data = get_ui_data()
+
+    if ui_data.get("SERVER_STOPPING", False):
+        response = make_response(jsonify({"message": "Server is shutting down, try again later."}), 503)
+        response.headers["Retry-After"] = 30  # Clients should retry after 30 seconds # type: ignore
+        return response
+
     app.config["SCRIPT_NONCE"] = sha256(urandom(32)).hexdigest()
 
     if not request.path.startswith(("/css", "/images", "/js", "/json", "/webfonts")):
-        ui_data = get_ui_data()
-
         if (
             app.config["DB"].database_uri
             and app.config["DB"].readonly
@@ -645,10 +650,12 @@ def setup():
         random_url=f"/{''.join(choice(ascii_letters + digits) for _ in range(10))}",
     )
 
+
 @app.route("/setup/loading", methods=["GET"])
 def setup_loading():
     return render_template("setup_loading.html")
 
+
 @app.route("/totp", methods=["GET", "POST"])
 @login_required
 def totp():