workers: create abstract Worker class, use asyncio (bug 1744327)

This is a work in progress, do not land.

Author: zzzeid

landoapi/cli.py

@@ -1,6 +1,7 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
+import asyncio
 import logging
 import os
 import sys
@@ -81,10 +82,10 @@ def landing_worker():
     for system in get_subsystems(exclude=exclusions):
         system.ensure_ready()
 
-    from landoapi.landing_worker import LandingWorker
+    from landoapi.workers.landing_worker import LandingWorker
 
     worker = LandingWorker()
-    worker.start()
+    asyncio.run(worker.start())
 
 
 @cli.command(name="revision-worker")
@@ -95,10 +96,10 @@ def revision_worker():
     for system in get_subsystems(exclude=exclusions):
         system.ensure_ready()
 
-    from landoapi.revision_worker import RevisionWorker
+    from landoapi.workers.revision_worker import RevisionWorker
 
     worker = RevisionWorker()
-    worker.start()
+    asyncio.run(worker.start())
 
 
 @cli.command(name="run-pre-deploy-sequence")

landoapi/workers/__init__.py

@@ -0,0 +1,149 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+"""This module contains the abstract repo worker implementation."""
+
+import logging
+import os
+import subprocess
+import re
+import asyncio
+import signal
+from flask import current_app
+from landoapi.repos import repo_clone_subsystem
+from landoapi.treestatus import treestatus_subsystem
+from landoapi.models.configuration import ConfigurationVariable
+
+
+logger = logging.getLogger(__name__)
+
+
+class Worker:
+    def __init__(self, pause_key, sleep_seconds=5, with_ssh=True, with_aws=True):
+        SSH_PRIVATE_KEY_ENV_KEY = "SSH_PRIVATE_KEY"
+
+        # ConfigurationVariable key used to control the paused state of this worker.
+        self.PAUSE_KEY = pause_key
+        self.sleep_seconds = sleep_seconds
+
+        if with_aws:
+            config_keys = ["AWS_SECRET_KEY", "AWS_ACCESS_KEY", "PATCH_BUCKET_NAME"]
+            self.config = {k: current_app.config[k] for k in config_keys}
+
+        # The list of all repos that are enabled for this worker
+        self.applicable_repos = (
+            list(repo_clone_subsystem.repos)
+            if hasattr(repo_clone_subsystem, "repos")
+            else []
+        )
+
+        # The list of all repos that have open trees; refreshed when needed via
+        # `self.refresh_enabled_repos`.
+        self.enabled_repos = []
+
+        # This is True when a worker active, and False when it is shut down
+        self.running = False
+
+        # This is True when the worker is busy processing a job
+        self.busy = False
+
+        if with_ssh:
+            # Fetch ssh private key from the environment. Note that this key should be
+            # stored in standard format including all new lines and new line at the end
+            # of the file.
+            self.ssh_private_key = os.environ.get(SSH_PRIVATE_KEY_ENV_KEY)
+            if not self.ssh_private_key:
+                logger.warning(f"No {SSH_PRIVATE_KEY_ENV_KEY} present in environment.")
+
+    @staticmethod
+    def _setup_ssh(ssh_private_key):
+        """Add a given private ssh key to ssh agent.
+
+        SSH keys are needed in order to push to repositories that have an ssh
+        push path.
+
+        The private key should be passed as it is in the key file, including all
+        new line characters and the new line character at the end.
+
+        Args:
+            ssh_private_key (str): A string representing the private SSH key file.
+        """
+        # Set all the correct environment variables
+        agent_process = subprocess.run(
+            ["ssh-agent", "-s"], capture_output=True, universal_newlines=True
+        )
+
+        # This pattern will match keys and values, and ignore everything after the
+        # semicolon. For example, the output of `agent_process` is of the form:
+        #     SSH_AUTH_SOCK=/tmp/ssh-c850kLXXOS5e/agent.120801; export SSH_AUTH_SOCK;
+        #     SSH_AGENT_PID=120802; export SSH_AGENT_PID;
+        #     echo Agent pid 120802;
+        pattern = re.compile("(.+)=([^;]*)")
+        for key, value in pattern.findall(agent_process.stdout):
+            logger.info(f"_setup_ssh: setting {key} to {value}")
+            os.environ[key] = value
+
+        # Add private SSH key to agent
+        # NOTE: ssh-add seems to output everything to stderr, including upon exit 0.
+        add_process = subprocess.run(
+            ["ssh-add", "-"],
+            input=ssh_private_key,
+            capture_output=True,
+            universal_newlines=True,
+        )
+        if add_process.returncode != 0:
+            raise Exception(add_process.stderr)
+        logger.info("Added private SSH key from environment.")
+
+    @property
+    def _paused(self):
+        return ConfigurationVariable.get(self.PAUSE_KEY, False)
+
+    def _setup(self):
+        loop = asyncio.get_event_loop()
+        loop.add_signal_handler(
+            signal.SIGINT,
+            lambda *args, **kwargs: asyncio.create_task(
+                self.exit_gracefully(*args, **kwargs)
+            ),
+        )
+        loop.add_signal_handler(
+            signal.SIGTERM,
+            lambda *args, **kwargs: asyncio.create_task(
+                self.exit_gracefully(*args, **kwargs)
+            ),
+        )
+
+        if hasattr(self, "ssh_private_key"):
+            self._setup_ssh(self.ssh_private_key)
+
+    async def _start(self, *args, **kwargs):
+        self.running = True
+        while self.running:
+            while self._paused:
+                await asyncio.sleep(self.sleep_seconds)
+            await self.loop(*args, **kwargs)
+
+    async def sleep(self, sleep_seconds):
+        await asyncio.sleep(self.sleep_seconds)
+
+    def refresh_enabled_repos(self):
+        self.enabled_repos = [
+            r
+            for r in self.applicable_repos
+            if treestatus_subsystem.client.is_open(repo_clone_subsystem.repos[r].tree)
+        ]
+        logger.info(f"{len(self.enabled_repos)} enabled repos: {self.enabled_repos}")
+
+    async def start(self):
+        self._setup()
+        await self._start()
+
+    async def loop(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    async def exit_gracefully(self, *args):
+        logger.info(f"Worker exiting gracefully {args}")
+        while self.busy:
+            await asyncio.sleep(self.sleep_seconds)
+        self.running = False