Merge pull request #3528 from antcybersec/fix-smoke-test-timeout

sgoggins · web-flow · commit 7a46497aaed6 · 2026-02-16T11:08:19.000-06:00
Fix flaky smoke tests by adding SIGTERM signal handling
diff --git a/augur/application/cli/backend.py b/augur/application/cli/backend.py
@@ -14,18 +14,17 @@
 import traceback
 import requests
 from redis.exceptions import ConnectionError as RedisConnectionError
-from urllib.parse import urlparse
 
 from augur.tasks.start_tasks import augur_collection_monitor, create_collection_status_records
 from augur.tasks.git.facade_tasks import clone_repos
 from augur.tasks.github.contributors import process_contributors
 from augur.tasks.github.util.github_api_key_handler import GithubApiKeyHandler
 from augur.tasks.gitlab.gitlab_api_key_handler import GitlabApiKeyHandler
 from augur.tasks.data_analysis.contributor_breadth_worker.contributor_breadth_worker import contributor_breadth_model
-from augur.tasks.init.redis_connection import get_redis_connection 
 from augur.application.db.models import UserRepo
 from augur.application.db.session import DatabaseSession
 from augur.application.logs import AugurLogger
+from augur.application.service_manager import AugurServiceManager, cleanup_collection_status_and_rabbit, clean_collection_status
 from augur.application.db.lib import get_value
 from augur.application.cli import test_connection, test_db_connection, with_database, DatabaseContext
 import sqlalchemy as s
@@ -36,7 +35,6 @@
 
 logger = AugurLogger("augur", reset_logfiles=reset_logs).get_logger()
 
-
 @click.group('server', short_help='Commands for controlling the backend API server & data collection workers')
 @click.pass_context
 def cli(ctx):
@@ -55,17 +53,23 @@ def start(ctx, disable_collection, development, pidfile, port):
     """Start Augur's backend server."""
     with open(pidfile, "w") as pidfile_io:
         pidfile_io.write(str(os.getpid()))
-        
+
+    manager = AugurServiceManager(ctx, pidfile, disable_collection)
+
+    # Register signal handlers for graceful shutdown
+    signal.signal(signal.SIGTERM, manager.shutdown_signal_handler)
+    signal.signal(signal.SIGINT, manager.shutdown_signal_handler)
+
     try:
         if os.environ.get('AUGUR_DOCKER_DEPLOY') != "1":
             raise_open_file_limit(100000)
-    except Exception as e: 
+    except Exception as e:
         logger.error(
                     ''.join(traceback.format_exception(None, e, e.__traceback__)))
-        
+
         logger.error("Failed to raise open file limit!")
         raise e
-    
+
     if development:
         os.environ["AUGUR_DEV"] = "1"
         logger.info("Starting in development mode")
@@ -101,6 +105,7 @@ def start(ctx, disable_collection, development, pidfile, port):
 
     gunicorn_command = f"gunicorn -c {gunicorn_location} -b {host}:{port} augur.api.server:app --log-file {gunicorn_log_file}"
     server = subprocess.Popen(gunicorn_command.split(" "))
+    manager.server = server
 
     logger.info("awaiting Gunicorn start")
     while not server.poll():
@@ -123,6 +128,7 @@ def start(ctx, disable_collection, development, pidfile, port):
     logger.info(f"The API is available at '{api_response.json()['route']}'")
 
     processes = start_celery_worker_processes((core_worker_count, secondary_worker_count, facade_worker_count), disable_collection)
+    manager.processes = processes
 
     celery_beat_schedule_db = os.getenv("CELERYBEAT_SCHEDULE_DB", "celerybeat-schedule.db")
     if os.path.exists(celery_beat_schedule_db):
@@ -132,8 +138,10 @@ def start(ctx, disable_collection, development, pidfile, port):
     log_level = get_value("Logging", "log_level")
     celery_beat_process = None
     celery_command = f"celery -A augur.tasks.init.celery_app.celery_app beat -l {log_level.lower()} -s {celery_beat_schedule_db}"
-    celery_beat_process = subprocess.Popen(celery_command.split(" "))    
+    celery_beat_process = subprocess.Popen(celery_command.split(" "))
+    manager.celery_beat_process = celery_beat_process    
     keypub = KeyPublisher()
+    manager.keypub = keypub
     
     if not disable_collection:
         if os.environ.get('AUGUR_DOCKER_DEPLOY') != "1":
@@ -180,29 +188,15 @@ def start(ctx, disable_collection, development, pidfile, port):
     try:
         server.wait()
     except KeyboardInterrupt:
-        
-        if server:
-            logger.info("Shutting down server")
-            server.terminate()
-
-        logger.info("Shutting down all celery worker processes")
-        for p in processes:
-            if p:
-                p.terminate()
-
-        if celery_beat_process:
-            logger.info("Shutting down celery beat process")
-            celery_beat_process.terminate()
-
-        if not disable_collection:
-
+        # Signal handler will take care of cleanup
+        pass
+    finally:
+        # Ensure pidfile is cleaned up if we exit normally
+        if os.path.exists(pidfile):
             try:
-                keypub.shutdown()
-                cleanup_collection_status_and_rabbit(logger, ctx.obj.engine)
-            except RedisConnectionError:
-                pass
-            
-    os.unlink(pidfile)
+                os.unlink(pidfile)
+            except OSError as e:
+                logger.error(f"Could not remove pidfile {pidfile}: {e}")
 
 def start_celery_worker_processes(worker_counts: tuple[int, int, int], disable_collection=False):
     """
@@ -344,84 +338,6 @@ def augur_stop(signal, logger, engine):
         cleanup_collection_status_and_rabbit(logger, engine)
 
 
-def cleanup_collection_status_and_rabbit(logger, engine):
-    clear_redis_caches()
-
-    connection_string = get_value("RabbitMQ", "connection_string")
-
-    with DatabaseSession(logger, engine=engine) as session:
-
-        clean_collection_status(session)
-
-    clear_rabbitmq_messages(connection_string)
-
-def clear_redis_caches():
-    """Clears the redis databases that celery and redis use."""
-
-    logger.info("Flushing all redis databases this instance was using")
-    celery_purge_command = "celery -A augur.tasks.init.celery_app.celery_app purge -f"
-    subprocess.call(celery_purge_command.split(" "))
-
-    redis_connection = get_redis_connection()
-    redis_connection.flushdb()
-
-def clear_all_message_queues(connection_string):
-    queues = ['celery','secondary','scheduling','facade']
-
-    virtual_host_string = connection_string.split("/")[-1]
-
-    #Parse username and password with urllib
-    parsed = urlparse(connection_string)
-
-    for q in queues:
-        curl_cmd = f"curl -i -u {parsed.username}:{parsed.password} -XDELETE http://localhost:15672/api/queues/{virtual_host_string}/{q}"
-        subprocess.call(curl_cmd.split(" "),stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-
-
-def clear_rabbitmq_messages(connection_string):
-    #virtual_host_string = connection_string.split("/")[-1]
-
-    logger.info("Clearing all messages from celery queue in rabbitmq")
-    from augur.tasks.init.celery_app import celery_app
-    celery_app.control.purge()
-
-    clear_all_message_queues(connection_string)
-    #rabbitmq_purge_command = f"sudo rabbitmqctl purge_queue celery -p {virtual_host_string}"
-    #subprocess.call(rabbitmq_purge_command.split(" "))
-
-#Make sure that database reflects collection status when processes are killed/stopped.
-def clean_collection_status(session):
-    session.execute_sql(s.sql.text("""
-        UPDATE augur_operations.collection_status 
-        SET core_status='Pending',core_task_id = NULL
-        WHERE core_status='Collecting' AND core_data_last_collected IS NULL;
-
-        UPDATE augur_operations.collection_status
-        SET core_status='Success',core_task_id = NULL
-        WHERE core_status='Collecting' AND core_data_last_collected IS NOT NULL;
-
-        UPDATE augur_operations.collection_status 
-        SET secondary_status='Pending',secondary_task_id = NULL
-        WHERE secondary_status='Collecting' AND secondary_data_last_collected IS NULL;
-
-        UPDATE augur_operations.collection_status 
-        SET secondary_status='Success',secondary_task_id = NULL
-        WHERE secondary_status='Collecting' AND secondary_data_last_collected IS NOT NULL;
-
-        UPDATE augur_operations.collection_status 
-        SET facade_status='Update', facade_task_id=NULL
-        WHERE facade_status LIKE '%Collecting%' and facade_data_last_collected IS NULL;
-
-        UPDATE augur_operations.collection_status 
-        SET facade_status='Success', facade_task_id=NULL
-        WHERE facade_status LIKE '%Collecting%' and facade_data_last_collected IS NOT NULL;
-
-        UPDATE augur_operations.collection_status
-        SET facade_status='Pending', facade_task_id=NULL
-        WHERE facade_status='Failed Clone' OR facade_status='Initializing';
-    """))
-    #TODO: write timestamp for currently running repos.
-
 def assign_orphan_repos_to_default_user(session):
     query = s.sql.text("""
         SELECT repo_id FROM repo WHERE repo_id NOT IN (SELECT repo_id FROM augur_operations.user_repos)
diff --git a/augur/application/service_manager.py b/augur/application/service_manager.py
@@ -0,0 +1,166 @@
+import sys
+import os
+import subprocess
+import sqlalchemy as s
+from augur.application.logs import AugurLogger
+from augur.application.db.session import DatabaseSession
+from augur.application.db.lib import get_value
+from augur.tasks.init.redis_connection import get_redis_connection
+from urllib.parse import urlparse
+
+logger = AugurLogger("augur_servicemanager").get_logger()
+
+
+class AugurServiceManager:
+    """ Provides a storage space for references to the various components of augur
+    This enables them to all be properly shut down in the event a shutdown signal (SIGINT - AKA ctrl-c, or SIGTERM) is received 
+    """
+    def __init__(self, ctx, pidfile, disable_collection):
+        self.ctx = ctx
+        self.pidfile = pidfile
+        self.disable_collection = disable_collection
+        self.server = None
+        self.processes = []
+        self.celery_beat_process = None
+        self.keypub = None
+        self.shutting_down = False
+
+    def shutdown_signal_handler(self, signum, frame):
+        if self.shutting_down:
+            return
+        
+        self.shutting_down = True
+        logger.info(f"Received signal {signum}, shutting down gracefully")
+
+        # Stop server
+        if self.server:
+            logger.info("Stopping server")
+            self.server.terminate()
+            try:
+                self.server.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                logger.warning("Server did not terminate in time, killing")
+                self.server.kill()
+
+        # Stop celery workers
+        logger.info("Stopping celery workers")
+        for p in self.processes:
+            if p and p.poll() is None:
+                p.terminate()
+        
+        # Wait for workers to terminate
+        for p in self.processes:
+            if p:
+                try:
+                    p.wait(timeout=3)
+                except subprocess.TimeoutExpired:
+                    logger.warning(f"Worker {p.pid} did not terminate in time, killing")
+                    p.kill()
+
+        # Stop celery beat
+        if self.celery_beat_process:
+            logger.info("Stopping celery beat")
+            self.celery_beat_process.terminate()
+            try:
+                self.celery_beat_process.wait(timeout=3)
+            except subprocess.TimeoutExpired:
+                logger.warning("Celery beat did not terminate in time, killing")
+                self.celery_beat_process.kill()
+
+        # Cleanup collection resources
+        if not self.disable_collection:
+            try:
+                if self.keypub:
+                    self.keypub.shutdown()
+                cleanup_collection_status_and_rabbit(logger, self.ctx.obj.engine)
+            except Exception as e:
+                logger.debug(f"Error during collection cleanup: {e}")
+
+        # Remove pidfile
+        if os.path.exists(self.pidfile):
+            try:
+                os.unlink(self.pidfile)
+            except OSError as e:
+                logger.error(f"Could not remove pidfile {self.pidfile}: {e}")
+
+        sys.exit(0)
+
+def cleanup_collection_status_and_rabbit(logger, engine):
+    # TODO: tech debt: this should probbaly be in a helper function but its so tightly coupled with other stuff
+    clear_redis_caches()
+
+    connection_string = get_value("RabbitMQ", "connection_string")
+
+    with DatabaseSession(logger, engine=engine) as session:
+
+        clean_collection_status(session)
+
+    clear_rabbitmq_messages(connection_string)
+
+def clear_redis_caches():
+    """Clears the redis databases that celery and redis use."""
+
+    logger.info("Flushing all redis databases this instance was using")
+    celery_purge_command = "celery -A augur.tasks.init.celery_app.celery_app purge -f"
+    subprocess.call(celery_purge_command.split(" "))
+
+    redis_connection = get_redis_connection()
+    redis_connection.flushdb()
+
+#Make sure that database reflects collection status when processes are killed/stopped.
+def clean_collection_status(session):
+    session.execute_sql(s.sql.text("""
+        UPDATE augur_operations.collection_status 
+        SET core_status='Pending',core_task_id = NULL
+        WHERE core_status='Collecting' AND core_data_last_collected IS NULL;
+
+        UPDATE augur_operations.collection_status
+        SET core_status='Success',core_task_id = NULL
+        WHERE core_status='Collecting' AND core_data_last_collected IS NOT NULL;
+
+        UPDATE augur_operations.collection_status 
+        SET secondary_status='Pending',secondary_task_id = NULL
+        WHERE secondary_status='Collecting' AND secondary_data_last_collected IS NULL;
+
+        UPDATE augur_operations.collection_status 
+        SET secondary_status='Success',secondary_task_id = NULL
+        WHERE secondary_status='Collecting' AND secondary_data_last_collected IS NOT NULL;
+
+        UPDATE augur_operations.collection_status 
+        SET facade_status='Update', facade_task_id=NULL
+        WHERE facade_status LIKE '%Collecting%' and facade_data_last_collected IS NULL;
+
+        UPDATE augur_operations.collection_status 
+        SET facade_status='Success', facade_task_id=NULL
+        WHERE facade_status LIKE '%Collecting%' and facade_data_last_collected IS NOT NULL;
+
+        UPDATE augur_operations.collection_status
+        SET facade_status='Pending', facade_task_id=NULL
+        WHERE facade_status='Failed Clone' OR facade_status='Initializing';
+    """))
+    #TODO: write timestamp for currently running repos.
+
+
+def clear_rabbitmq_messages(connection_string):
+    #virtual_host_string = connection_string.split("/")[-1]
+
+    logger.info("Clearing all messages from celery queue in rabbitmq")
+    from augur.tasks.init.celery_app import celery_app
+    celery_app.control.purge()
+
+    clear_all_message_queues(connection_string)
+    #rabbitmq_purge_command = f"sudo rabbitmqctl purge_queue celery -p {virtual_host_string}"
+    #subprocess.call(rabbitmq_purge_command.split(" "))
+
+
+def clear_all_message_queues(connection_string):
+    queues = ['celery','secondary','scheduling','facade']
+
+    virtual_host_string = connection_string.split("/")[-1]
+
+    #Parse username and password with urllib
+    parsed = urlparse(connection_string)
+
+    for q in queues:
+        curl_cmd = f"curl -i -u {parsed.username}:{parsed.password} -XDELETE http://localhost:15672/api/queues/{virtual_host_string}/{q}"
+        subprocess.call(curl_cmd.split(" "),stdout=subprocess.PIPE, stderr=subprocess.PIPE)
