From e8d04cff7c8f8a90be8d80d2e797714e07c0e4a4 Mon Sep 17 00:00:00 2001
From: Michael Pilosov <consistentbayes@gmail.com>
Date: Mon, 20 Mar 2023 00:33:07 -0600
Subject: [PATCH] working with two ctrl-c, previous commit does not

---
 src/announce_server/client.py    | 39 +++++++++++++++++++++-----------
 src/announce_server/decorator.py | 33 +++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/src/announce_server/client.py b/src/announce_server/client.py
index fe7602e..09b5059 100644
--- a/src/announce_server/client.py
+++ b/src/announce_server/client.py
@@ -1,6 +1,7 @@
 import argparse
 import asyncio
 import signal
+import sys
 import threading
 from http.server import HTTPServer, SimpleHTTPRequestHandler
 
@@ -9,6 +10,7 @@ import socketio
 from .decorator import register_service
 from .get_ip import get_ip_address
 
+http_server_thread = None
 
 def start_client(host_ip="0.0.0.0", host_port=4999):
     @register_service(
@@ -19,26 +21,37 @@ def start_client(host_ip="0.0.0.0", host_port=4999):
         host_port=host_port,
     )
     def server(port=13373):
-        httpd = HTTPServer(("", port), SimpleHTTPRequestHandler)
-        print(f"Serving HTTP on 0.0.0.0 port {port} (http://0.0.0.0:{port}/) ...")
-
-        def serve_forever_nonblocking():
+        def start_server():
+            global http_server_thread
+            print(f"Serving HTTP on 0.0.0.0 port {port} (http://0.0.0.0:{port}/) ...")
+            httpd = HTTPServer(("", port), SimpleHTTPRequestHandler)
             httpd.serve_forever()
-
-        server_thread = threading.Thread(target=serve_forever_nonblocking)
+        server_thread = threading.Thread(target=start_server)
+        server_thread.daemon = True
         server_thread.start()
+        server_thread.join()
+
+    def signal_handler(signal, frame):
+        print("Cleaning up and shutting down...")
+
+        # If the HTTP server thread is running, shut it down
+        if http_server_thread is not None and http_server_thread.is_alive():
+            http_server_thread.shutdown()
 
-        return server_thread
+        sys.exit(0)
 
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
     try:
-        server_thread = server()
-        while True:
-            pass
-    except (KeyboardInterrupt, asyncio.exceptions.CancelledError):
-        print("Cleaning up and shutting down...")
-        # server_thread.join()
+        server()
+        signal.pause()
+    except asyncio.exceptions.CancelledError:
+        print("CancelledError")
+        # signal_handler(signal.SIGINT, None)
+        pass
     finally:
         print("Shutting down test client...")
+        sys.exit(0)
 
 
 if __name__ == "__main__":
diff --git a/src/announce_server/decorator.py b/src/announce_server/decorator.py
index 1224ec1..a366016 100644
--- a/src/announce_server/decorator.py
+++ b/src/announce_server/decorator.py
@@ -6,6 +6,7 @@ import socketio
 
 sio = socketio.AsyncClient()
 
+import signal
 
 async def _announce_server(**kwargs):
     SERVER_NAME = kwargs.get("name", "server_1")
@@ -68,6 +69,38 @@ def register_service(task=None, **outer_kwargs):
 
     return wrapper
 
+# def register_service(task=None, **outer_kwargs):
+#     if task is None:
+#         return lambda f: register_service(f, **outer_kwargs)
+
+#     @wraps(task)
+#     def wrapper(*args, **kwargs):
+#         async def main(*args, **kwargs):
+#             loop = asyncio.get_event_loop()
+#             host_block_thread = loop.run_in_executor(None, task)
+
+#             # Announce the server to the host
+#             await _announce_server(**outer_kwargs)
+
+#             # Set up signal handlers to clean up properly
+#             def signal_handler(signum, frame):
+#                 print(f"Received signal {signum}. Cleaning up and shutting down...")
+#                 host_block_thread.cancel()
+#                 import sys
+#                 sys.exit(0)
+
+#             signal.signal(signal.SIGINT, signal_handler)
+#             signal.signal(signal.SIGTERM, signal_handler)
+
+#             # Wait for host_block to finish or to be cancelled
+#             try:
+#                 await host_block_thread
+#             except asyncio.CancelledError:
+#                 pass
+
+#         return asyncio.run(main())
+
+#     return wrapper
 
 def announce_server(*args, **kwargs):
     """Wrapper for register_service"""