From e9781094101af349e13558cfd27bb594ee116da9 Mon Sep 17 00:00:00 2001
From: black-sliver <59490463+black-sliver@users.noreply.github.com>
Date: Sun, 19 May 2024 20:40:36 +0200
Subject: [PATCH] WebHost: properly stop worker threads (#3340)

* WebHost: properly stop worker threads

* Less jank

* Forgot the try-catch around the while true
---
 WebHost.py                 | 10 +++++++++-
 WebHostLib/autolauncher.py | 28 ++++++++++++++++++----------
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/WebHost.py b/WebHost.py
index 8ccf6b68..9b5edd32 100644
--- a/WebHost.py
+++ b/WebHost.py
@@ -117,7 +117,7 @@ if __name__ == "__main__":
     logging.basicConfig(format='[%(asctime)s] %(message)s', level=logging.INFO)
 
     from WebHostLib.lttpsprites import update_sprites_lttp
-    from WebHostLib.autolauncher import autohost, autogen
+    from WebHostLib.autolauncher import autohost, autogen, stop
     from WebHostLib.options import create as create_options_files
 
     try:
@@ -138,3 +138,11 @@ if __name__ == "__main__":
         else:
             from waitress import serve
             serve(app, port=app.config["PORT"], threads=app.config["WAITRESS_THREADS"])
+    else:
+        from time import sleep
+        try:
+            while True:
+                sleep(1)  # wait for process to be killed
+        except (SystemExit, KeyboardInterrupt):
+            pass
+    stop()  # stop worker threads
diff --git a/WebHostLib/autolauncher.py b/WebHostLib/autolauncher.py
index e70a31bc..08a1309e 100644
--- a/WebHostLib/autolauncher.py
+++ b/WebHostLib/autolauncher.py
@@ -3,16 +3,26 @@ from __future__ import annotations
 import json
 import logging
 import multiprocessing
-import time
 import typing
-from uuid import UUID
 from datetime import timedelta, datetime
+from threading import Event, Thread
+from uuid import UUID
 
 from pony.orm import db_session, select, commit
 
 from Utils import restricted_loads
 from .locker import Locker, AlreadyRunningException
 
+_stop_event = Event()
+
+
+def stop():
+    """Stops previously launched threads"""
+    global _stop_event
+    stop_event = _stop_event
+    _stop_event = Event()  # new event for new threads
+    stop_event.set()
+
 
 def handle_generation_success(seed_id):
     logging.info(f"Generation finished for seed {seed_id}")
@@ -63,6 +73,7 @@ def cleanup():
 
 def autohost(config: dict):
     def keep_running():
+        stop_event = _stop_event
         try:
             with Locker("autohost"):
                 cleanup()
@@ -72,8 +83,7 @@ def autohost(config: dict):
                     hosters.append(hoster)
                     hoster.start()
 
-                while 1:
-                    time.sleep(0.1)
+                while not stop_event.wait(0.1):
                     with db_session:
                         rooms = select(
                             room for room in Room if
@@ -86,12 +96,12 @@ def autohost(config: dict):
         except AlreadyRunningException:
             logging.info("Autohost reports as already running, not starting another.")
 
-    import threading
-    threading.Thread(target=keep_running, name="AP_Autohost").start()
+    Thread(target=keep_running, name="AP_Autohost").start()
 
 
 def autogen(config: dict):
     def keep_running():
+        stop_event = _stop_event
         try:
             with Locker("autogen"):
 
@@ -112,8 +122,7 @@ def autogen(config: dict):
                             commit()
                         select(generation for generation in Generation if generation.state == STATE_ERROR).delete()
 
-                    while 1:
-                        time.sleep(0.1)
+                    while not stop_event.wait(0.1):
                         with db_session:
                             # for update locks the database row(s) during transaction, preventing writes from elsewhere
                             to_start = select(
@@ -124,8 +133,7 @@ def autogen(config: dict):
         except AlreadyRunningException:
             logging.info("Autogen reports as already running, not starting another.")
 
-    import threading
-    threading.Thread(target=keep_running, name="AP_Autogen").start()
+    Thread(target=keep_running, name="AP_Autogen").start()
 
 
 multiworlds: typing.Dict[type(Room.id), MultiworldInstance] = {}