Implement Flask server more elegant

1 files changed, 19 insertions, 5 deletions

diff --git a/session.py b/session.py
index 9fc755d..f1824d7 100644
--- a/session.py
+++ b/session.py
@@ -3,13 +3,15 @@ import random
 import secrets
 import signal
 import string
-from multiprocessing import Process
+import threading
 
 import psutil
 import requests
 
 import port_for
 from flask import Flask, abort, request
+from werkzeug.serving import make_server
+
 from vnc import run_vnc, save_password
 
 API_SERVER = "http://127.0.0.1:8000"
@@ -22,6 +24,20 @@ def random_digits(length: int):
     return "".join(random.choice(string.digits) for _ in range(length))
 
 
+class ServerThread(threading.Thread):
+    def __init__(self, app, port: int):
+        super().__init__()
+        self.srv = make_server('127.0.0.1', port, app)
+        self.ctx = app.app_context()
+        self.ctx.push()
+
+    def run(self):
+        self.srv.serve_forever()
+
+    def shutdown(self):
+        self.srv.shutdown()
+
+
 class Session:
     #: Session is running
     STATUS_RUNNING = "running"
@@ -121,9 +137,7 @@ class Session:
             else:
                 return abort(403)
 
-        self.trigger_thread = Process(
-            target=lambda: app.run("0.0.0.0", port=self.trigger_port)
-        )
+        self.trigger_thread = ServerThread(app=app, port=self.trigger_port)
         self.trigger_thread.start()
 
     def _trigger(self):
@@ -167,7 +181,7 @@ class Session:
 
         if hasattr(self, "trigger_thread"):
             print("Kill trigger service.")
-            self.trigger_thread.kill()
+            self.trigger_thread.shutdown()
 
         self.push()