Merge pull request #13514 from mvdbeek/support_preload

[22.01] Support ``--preload`` and worker_id for server processes

lib/galaxy/model/orm/engine_factory.py

@@ -8,6 +8,7 @@ from multiprocessing.util import register_after_fork
 from sqlalchemy import (
     create_engine,
     event,
+    exc,
 )
 from sqlalchemy.engine import Engine
 
@@ -93,5 +94,22 @@ def build_engine(url, engine_options, database_query_profiling_proxy=False, trac
         connect_args['check_same_thread'] = False
     # Create the database engine
     engine = create_engine(url, connect_args=connect_args, **engine_options)
+
+    # Prevent sharing connection across fork: https://docs.sqlalchemy.org/en/14/core/pooling.html#using-connection-pools-with-multiprocessing-or-os-fork
     register_after_fork(engine, lambda e: e.dispose())
+
+    @event.listens_for(engine, "connect")
+    def connect(dbapi_connection, connection_record):
+        connection_record.info["pid"] = os.getpid()
+
+    @event.listens_for(engine, "checkout")
+    def checkout(dbapi_connection, connection_record, connection_proxy):
+        pid = os.getpid()
+        if connection_record.info["pid"] != pid:
+            connection_record.dbapi_connection = connection_proxy.dbapi_connection = None
+            raise exc.DisconnectionError(
+                "Connection record belongs to pid %s, "
+                "attempting to check out in pid %s" % (connection_record.info["pid"], pid)
+            )
+
     return engine

lib/galaxy/model/unittest_utils/data_app.py

@@ -55,6 +55,7 @@ class GalaxyDataTestConfig(Bunch):
         self.jobs_directory = os.path.join(self.data_dir, 'jobs_directory')
         self.new_file_path = os.path.join(self.data_dir, 'tmp')
         self.file_path = os.path.join(self.data_dir, 'files')
+        self.server_name = "main"
 
     def __del__(self):
         if self._remove_root:

lib/galaxy/web_stack/__init__.py

@@ -5,7 +5,17 @@ import json
 import logging
 import multiprocessing
 import os
-from typing import Callable, Dict, FrozenSet, List, Optional, Tuple, Type
+import sys
+import threading
+from typing import (
+    Callable,
+    Dict,
+    FrozenSet,
+    List,
+    Optional,
+    Tuple,
+    Type,
+)
 from urllib.request import install_opener
 
 # The uwsgi module is automatically injected by the parent uwsgi process and only exists that way.  If anything works,
@@ -199,6 +209,8 @@ class ApplicationStack:
 
     def set_postfork_server_name(self, app):
         new_server_name = self.server_name_template.format(**self.facts)
+        if "GUNICORN_WORKER_ID" in os.environ:
+            new_server_name = f"{new_server_name}.{os.environ['GUNICORN_WORKER_ID']}"
         multiprocessing.current_process().name = app.config.server_name = new_server_name
         log.debug('server_name set to: %s', new_server_name)
 
@@ -557,6 +569,45 @@ class PasteApplicationStack(ApplicationStack):
     name = 'Python Paste'
 
 
+class GunicornApplicationStack(ApplicationStack):
+    name = "Gunicorn"
+    do_post_fork = "--preload" in os.environ.get("GUNICORN_CMD_ARGS", "") or "--preload" in sys.argv
+    postfork_functions: List[Callable] = []
+    # Will be set to True by external hook
+    late_postfork_event = threading.Event()
+
+    @classmethod
+    def register_postfork_function(cls, f, *args, **kwargs):
+        # do_post_fork determines if we need to run postfork functions
+        if cls.do_post_fork:
+            # if so, we call ApplicationStack.late_postfork once after forking ...
+            if not cls.postfork_functions:
+                os.register_at_fork(after_in_child=cls.late_postfork)
+            # ... and store everything we need to run in ApplicationStack.postfork_functions
+            cls.postfork_functions.append(lambda: f(*args, **kwargs))
+        else:
+            f(*args, **kwargs)
+
+    @classmethod
+    def run_postfork(cls):
+        cls.late_postfork_event.wait(1)
+        for f in cls.postfork_functions:
+            f()
+
+    @classmethod
+    def late_postfork(cls):
+        # We can't run postfork functions immediately, because this is before the gunicorn `post_fork` hook runs,
+        # and we depend on the `post_fork` hook to set a worker id.
+        t = threading.Thread(target=cls.run_postfork)
+        t.start()
+
+    def log_startup(self):
+        msg = [f"Galaxy server instance '{self.config.server_name}' is running"]
+        if "GUNICORN_LISTENERS" in os.environ:
+            msg.append(f'serving on {os.environ["GUNICORN_LISTENERS"]}')
+        log.info("\n".join(msg))
+
+
 class WeblessApplicationStack(ApplicationStack):
     name = 'Webless'
 
@@ -602,7 +653,9 @@ def application_stack_class() -> Type[ApplicationStack]:
     """Returns the correct ApplicationStack class for the stack under which
     this Galaxy process is running.
     """
-    if uwsgi is not None and hasattr(uwsgi, 'numproc'):
+    if "gunicorn" in os.environ.get("SERVER_SOFTWARE", ""):
+        return GunicornApplicationStack
+    elif uwsgi is not None and hasattr(uwsgi, "numproc"):
         return UWSGIApplicationStack
     else:
         # cleverer ideas welcome

lib/galaxy/web_stack/gunicorn_config.py

@@ -0,0 +1,67 @@
+"""
+Gunicorn config file based on https://gist.github.com/hynek/ba655c8756924a5febc5285c712a7946
+"""
+import logging
+import os
+import sys
+
+log = logging.getLogger(__name__)
+
+
+def on_starting(server):
+    """
+    Attach a set of IDs that can be temporarily re-used.
+
+    Used on reloads when each worker exists twice.
+    """
+    server._worker_id_overload = set()
+
+
+def nworkers_changed(server, new_value, old_value):
+    """
+    Gets called on startup too.
+
+    Set the current number of workers. Required if we raise the worker count
+    temporarily using TTIN because server.cfg.workers won't be updated and if
+    one of those workers dies, we wouldn't know the ids go that far.
+    """
+    server._worker_id_current_workers = new_value
+
+
+def _next_worker_id(server):
+    """
+    If there are IDs open for re-use, take one. Else look for a free one.
+    """
+    if server._worker_id_overload:
+        return server._worker_id_overload.pop()
+
+    in_use = set(w._worker_id for w in tuple(server.WORKERS.values()) if w.alive)
+    free = set(range(1, server._worker_id_current_workers + 1)) - in_use
+
+    return free.pop()
+
+
+def on_reload(server):
+    """
+    Add a full set of ids into overload so it can be re-used once.
+    """
+    server._worker_id_overload = set(range(1, server.cfg.workers + 1))
+
+
+def pre_fork(server, worker):
+    """
+    Attach the next free worker_id before forking off.
+    """
+    worker._worker_id = _next_worker_id(server)
+
+
+def post_fork(server, worker):
+    """
+    Put worker_id and listeners into an env variable for further use within the app.
+    """
+    os.environ["GUNICORN_WORKER_ID"] = str(worker._worker_id)
+    os.environ["GUNICORN_LISTENERS"] = ",".join(str(bind) for bind in server.LISTENERS)
+    if "--preload" in os.environ.get("GUNICORN_CMD_ARGS", "") or "--preload" in sys.argv:
+        from galaxy.web_stack import GunicornApplicationStack
+
+        GunicornApplicationStack.late_postfork_event.set()

lib/galaxy/webapps/galaxy/fast_factory.py

@@ -37,7 +37,7 @@ with the following command-line.
 
 ::
 
-    gunicorn 'galaxy.webapps.galaxy.fast_factory:factory()' --env GALAXY_CONFIG_FILE=config/galaxy.ini --pythonpath lib -w 4 -k uvicorn.workers.UvicornWorker
+    gunicorn 'galaxy.webapps.galaxy.fast_factory:factory()' --env GALAXY_CONFIG_FILE=config/galaxy.ini --pythonpath lib -w 4 -k uvicorn.workers.UvicornWorker --config lib/galaxy/web_stack/gunicorn_config.py
 
 """
 

lib/tool_shed/webapp/fast_factory.py

@@ -30,7 +30,7 @@ with the following command-line.
 
 ::
 
-    gunicorn 'tool_shed.webapp.fast_factory:factory()' --env TOOL_SHED_CONFIG_FILE=config/tool_shed.yml --pythonpath lib -w 4 -k uvicorn.workers.UvicornWorker
+    gunicorn 'tool_shed.webapp.fast_factory:factory()' --env TOOL_SHED_CONFIG_FILE=config/tool_shed.yml --pythonpath lib -w 4 -k uvicorn.workers.UvicornWorker --config lib/galaxy/web_stack/gunicorn_config.py
 
 """
 

packages/web_stack/test-requirements.txt

@@ -1 +1,2 @@
+gunicorn
 pytest

scripts/common_startup_functions.sh

@@ -163,10 +163,10 @@ find_server() {
             ;;
         reports)
             # TODO: is this really the only way to configure the port?
-            GUNICORN_CMD_ARGS=${GUNICORN_CMD_ARGS:-"--bind=localhost:9001"}
+            GUNICORN_CMD_ARGS=${GUNICORN_CMD_ARGS:-"--bind=localhost:9001 --config lib/galaxy/web_stack/gunicorn_config.py"}
             ;;
         tool_shed)
-            GUNICORN_CMD_ARGS=${GUNICORN_CMD_ARGS:-"--bind=localhost:9009"}
+            GUNICORN_CMD_ARGS=${GUNICORN_CMD_ARGS:-"--bind=localhost:9009 --config lib/galaxy/web_stack/gunicorn_config.py"}
     esac
 
     APP_WEBSERVER=${APP_WEBSERVER:-$default_webserver}

test/unit/web_stack/test_application_stack.py

@@ -0,0 +1,38 @@
+from gunicorn import SERVER_SOFTWARE
+
+from galaxy.model.unittest_utils import GalaxyDataTestApp
+from galaxy.tool_util.verify.wait import wait_on
+from galaxy.web_stack import (
+    application_stack_class,
+    application_stack_instance,
+    GunicornApplicationStack,
+    WeblessApplicationStack,
+)
+
+
+def test_application_stack_class_factory(monkeypatch):
+    assert application_stack_class() == WeblessApplicationStack
+    monkeypatch.setenv("SERVER_SOFTWARE", SERVER_SOFTWARE)
+    assert application_stack_class() == GunicornApplicationStack
+
+
+def test_application_stack_instance_factory(monkeypatch):
+    app = GalaxyDataTestApp()
+    monkeypatch.setenv("SERVER_SOFTWARE", SERVER_SOFTWARE)
+    stack_instance = application_stack_instance(app=app, config=app.config)
+    assert isinstance(stack_instance, GunicornApplicationStack)
+    assert stack_instance.name == "Gunicorn"
+
+
+def test_application_stack_post_fork(monkeypatch):
+    app = GalaxyDataTestApp()
+    monkeypatch.setenv("SERVER_SOFTWARE", SERVER_SOFTWARE)
+    monkeypatch.setenv("GUNICORN_WORKER_ID", "100")
+    monkeypatch.setattr(GunicornApplicationStack, "do_post_fork", True)
+    stack_instance = application_stack_instance(app=app, config=app.config)
+    assert isinstance(stack_instance, GunicornApplicationStack)
+    stack_instance.register_postfork_function(lambda: stack_instance.set_postfork_server_name(app))
+    assert not app.config.server_name.endswith(".100")
+    stack_instance.late_postfork()
+    stack_instance.late_postfork_event.set()
+    wait_on(lambda: app.config.server_name.endswith(".100") or None, desc="server name to change", timeout=1, delta=0.1)