From 622b72db4b7c47832f02466c4b005e9549d3cb71 Mon Sep 17 00:00:00 2001
From: yihong <zouzou0208@gmail.com>
Date: Fri, 28 Feb 2025 14:52:40 +0800
Subject: [PATCH] Fix: add ctrl+c signal for better exit (#5469)

### What problem does this PR solve?

This patch add signal for ctrl + c that can exit the code friendly
cause code base use thread daemon can not exit friendly for being
started.

how to reproduce
1. docker-compose -f docker/docker-compose-base.yml up
2. other window `bash docker/launch_backend_service.sh`
3. stop 1 first
4. try to stop 2 then two thread can not exit which must use `kill pid`

This patch fix it
and should fix most the related issues in the `issues`

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

---------

Signed-off-by: yihong0618 <zouzou0208@gmail.com>
---
 api/ragflow_server.py    | 16 ++++++++++++++--
 rag/svr/task_executor.py | 25 +++++++++++++++++++++++--
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/api/ragflow_server.py b/api/ragflow_server.py
index d78990360..ac05dadda 100644
--- a/api/ragflow_server.py
+++ b/api/ragflow_server.py
@@ -28,6 +28,7 @@ import sys
 import time
 import traceback
 from concurrent.futures import ThreadPoolExecutor
+import threading
 
 from werkzeug.serving import run_simple
 from api import settings
@@ -42,15 +43,21 @@ from api.versions import get_ragflow_version
 from api.utils import show_configs
 from rag.settings import print_rag_settings
 
+stop_event = threading.Event()
 
 def update_progress():
-    while True:
-        time.sleep(6)
+    while not stop_event.is_set():
         try:
             DocumentService.update_progress()
+            stop_event.wait(6)
         except Exception:
             logging.exception("update_progress exception")
 
+def signal_handler(sig, frame):
+    logging.info("Received interrupt signal, shutting down...")
+    stop_event.set()
+    time.sleep(1)
+    sys.exit(0)
 
 if __name__ == '__main__':
     logging.info(r"""
@@ -96,6 +103,9 @@ if __name__ == '__main__':
     RuntimeConfig.init_env()
     RuntimeConfig.init_config(JOB_SERVER_HOST=settings.HOST_IP, HTTP_PORT=settings.HOST_PORT)
 
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+
     thread = ThreadPoolExecutor(max_workers=1)
     thread.submit(update_progress)
 
@@ -112,4 +122,6 @@ if __name__ == '__main__':
         )
     except Exception:
         traceback.print_exc()
+        stop_event.set()
+        time.sleep(1)
         os.kill(os.getpid(), signal.SIGKILL)
diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py
index 0dd239742..0d9eac99b 100644
--- a/rag/svr/task_executor.py
+++ b/rag/svr/task_executor.py
@@ -754,12 +754,33 @@ def main():
     if TRACE_MALLOC_ENABLED:
         start_tracemalloc_and_snapshot(None, None)
 
+    # Create an event to signal the background thread to exit
+    stop_event = threading.Event()
+    
     background_thread = threading.Thread(target=report_status)
     background_thread.daemon = True
     background_thread.start()
+ 
+    # Handle SIGINT (Ctrl+C)
+    def signal_handler(sig, frame):
+        logging.info("Received Ctrl+C, shutting down gracefully...")
+        stop_event.set()
+        # Give the background thread time to clean up
+        if background_thread.is_alive():
+            background_thread.join(timeout=5)
+        logging.info("Exiting...")
+        sys.exit(0)
 
-    while True:
-        handle_task()
+    signal.signal(signal.SIGINT, signal_handler)
+
+    try:
+        while not stop_event.is_set():
+            handle_task()
+    except KeyboardInterrupt:
+        logging.info("Interrupted by keyboard, shutting down...")
+        stop_event.set()
+        if background_thread.is_alive():
+            background_thread.join(timeout=5)
 
 if __name__ == "__main__":
     main()