Be able to shutdown homeserver that hasn't setup (#19187)

For example, a homeserver can fail to `setup` if it fails to connect to
the database.

Fix #19188

Follow-up to #18828


### Background

As part of Element's plan to support a light form of vhosting (virtual
host) (multiple instances of Synapse in the same Python process) (c.f
Synapse Pro for small hosts), we're currently diving into the details
and implications of running multiple instances of Synapse in the same
Python process.

"Clean tenant deprovisioning" tracked internally by
element-hq/synapse-small-hosts#50

Author: MadLittleMods

changelog.d/19187.misc

@@ -0,0 +1 @@
+Fix `HomeServer.shutdown()` failing if the homeserver hasn't been setup yet.

synapse/api/errors.py

@@ -856,6 +856,12 @@ def to_synapse_error(self) -> SynapseError:
         return ProxiedRequestError(self.code, errmsg, errcode, j)
 
 
+class HomeServerNotSetupException(Exception):
+    """
+    Raised when an operation is attempted on the HomeServer before setup() has been called.
+    """
+
+
 class ShadowBanError(Exception):
     """
     Raised when a shadow-banned user attempts to perform an action.

synapse/crypto/keyring.py

@@ -21,6 +21,7 @@
 
 import abc
 import logging
+from contextlib import ExitStack
 from typing import TYPE_CHECKING, Callable, Iterable
 
 import attr
@@ -150,57 +151,81 @@ class Keyring:
     """
 
     def __init__(
-        self, hs: "HomeServer", key_fetchers: "Iterable[KeyFetcher] | None" = None
+        self,
+        hs: "HomeServer",
+        test_only_key_fetchers: "list[KeyFetcher] | None" = None,
     ):
-        self.server_name = hs.hostname
+        """
+        Args:
+            hs: The HomeServer instance
+            test_only_key_fetchers: Dependency injection for tests only. If provided,
+                these key fetchers will be used instead of the default ones.
+        """
+        # Clean-up to avoid partial initialization leaving behind references.
+        with ExitStack() as exit:
+            self.server_name = hs.hostname
+
+            self._key_fetchers: list[KeyFetcher] = []
+            if test_only_key_fetchers is None:
+                # Always fetch keys from the database.
+                store_key_fetcher = StoreKeyFetcher(hs)
+                exit.callback(store_key_fetcher.shutdown)
+                self._key_fetchers.append(store_key_fetcher)
+
+                # Fetch keys from configured trusted key servers, if any exist.
+                key_servers = hs.config.key.key_servers
+                if key_servers:
+                    perspectives_key_fetcher = PerspectivesKeyFetcher(hs)
+                    exit.callback(perspectives_key_fetcher.shutdown)
+                    self._key_fetchers.append(perspectives_key_fetcher)
+
+                # Finally, fetch keys from the origin server directly.
+                server_key_fetcher = ServerKeyFetcher(hs)
+                exit.callback(server_key_fetcher.shutdown)
+                self._key_fetchers.append(server_key_fetcher)
+            else:
+                self._key_fetchers = test_only_key_fetchers
+
+            self._fetch_keys_queue: BatchingQueue[
+                _FetchKeyRequest, dict[str, dict[str, FetchKeyResult]]
+            ] = BatchingQueue(
+                name="keyring_server",
+                hs=hs,
+                clock=hs.get_clock(),
+                # The method called to fetch each key
+                process_batch_callback=self._inner_fetch_key_requests,
+            )
+            exit.callback(self._fetch_keys_queue.shutdown)
 
-        if key_fetchers is None:
-            # Always fetch keys from the database.
-            mutable_key_fetchers: list[KeyFetcher] = [StoreKeyFetcher(hs)]
-            # Fetch keys from configured trusted key servers, if any exist.
-            key_servers = hs.config.key.key_servers
-            if key_servers:
-                mutable_key_fetchers.append(PerspectivesKeyFetcher(hs))
-            # Finally, fetch keys from the origin server directly.
-            mutable_key_fetchers.append(ServerKeyFetcher(hs))
-
-            self._key_fetchers: Iterable[KeyFetcher] = tuple(mutable_key_fetchers)
-        else:
-            self._key_fetchers = key_fetchers
-
-        self._fetch_keys_queue: BatchingQueue[
-            _FetchKeyRequest, dict[str, dict[str, FetchKeyResult]]
-        ] = BatchingQueue(
-            name="keyring_server",
-            hs=hs,
-            clock=hs.get_clock(),
-            # The method called to fetch each key
-            process_batch_callback=self._inner_fetch_key_requests,
-        )
+            self._is_mine_server_name = hs.is_mine_server_name
 
-        self._is_mine_server_name = hs.is_mine_server_name
+            # build a FetchKeyResult for each of our own keys, to shortcircuit the
+            # fetcher.
+            self._local_verify_keys: dict[str, FetchKeyResult] = {}
+            for key_id, key in hs.config.key.old_signing_keys.items():
+                self._local_verify_keys[key_id] = FetchKeyResult(
+                    verify_key=key, valid_until_ts=key.expired
+                )
 
-        # build a FetchKeyResult for each of our own keys, to shortcircuit the
-        # fetcher.
-        self._local_verify_keys: dict[str, FetchKeyResult] = {}
-        for key_id, key in hs.config.key.old_signing_keys.items():
-            self._local_verify_keys[key_id] = FetchKeyResult(
-                verify_key=key, valid_until_ts=key.expired
+            vk = get_verify_key(hs.signing_key)
+            self._local_verify_keys[f"{vk.alg}:{vk.version}"] = FetchKeyResult(
+                verify_key=vk,
+                valid_until_ts=2**63,  # fake future timestamp
             )
 
-        vk = get_verify_key(hs.signing_key)
-        self._local_verify_keys[f"{vk.alg}:{vk.version}"] = FetchKeyResult(
-            verify_key=vk,
-            valid_until_ts=2**63,  # fake future timestamp
-        )
+            # We reached the end of the block which means everything was successful, so
+            # no exit handlers are needed (remove them all).
+            exit.pop_all()
 
     def shutdown(self) -> None:
         """
         Prepares the KeyRing for garbage collection by shutting down it's queues.
         """
         self._fetch_keys_queue.shutdown()
+
         for key_fetcher in self._key_fetchers:
             key_fetcher.shutdown()
+        self._key_fetchers.clear()
 
     async def verify_json_for_server(
         self,
@@ -521,9 +546,21 @@ class StoreKeyFetcher(KeyFetcher):
     """KeyFetcher impl which fetches keys from our data store"""
 
     def __init__(self, hs: "HomeServer"):
-        super().__init__(hs)
-
-        self.store = hs.get_datastores().main
+        # Clean-up to avoid partial initialization leaving behind references.
+        with ExitStack() as exit:
+            super().__init__(hs)
+            # `KeyFetcher` keeps a reference to `hs` which we need to clean up if
+            # something goes wrong so we can cleanly shutdown the homeserver.
+            exit.callback(super().shutdown)
+
+            # An error can be raised here if someone tried to create a `StoreKeyFetcher`
+            # before the homeserver is fully set up (`HomeServerNotSetupException:
+            # HomeServer.setup must be called before getting datastores`).
+            self.store = hs.get_datastores().main
+
+            # We reached the end of the block which means everything was successful, so
+            # no exit handlers are needed (remove them all).
+            exit.pop_all()
 
     async def _fetch_keys(
         self, keys_to_fetch: list[_FetchKeyRequest]
@@ -543,9 +580,21 @@ async def _fetch_keys(
 
 class BaseV2KeyFetcher(KeyFetcher):
     def __init__(self, hs: "HomeServer"):
-        super().__init__(hs)
-
-        self.store = hs.get_datastores().main
+        # Clean-up to avoid partial initialization leaving behind references.
+        with ExitStack() as exit:
+            super().__init__(hs)
+            # `KeyFetcher` keeps a reference to `hs` which we need to clean up if
+            # something goes wrong so we can cleanly shutdown the homeserver.
+            exit.callback(super().shutdown)
+
+            # An error can be raised here if someone tried to create a `StoreKeyFetcher`
+            # before the homeserver is fully set up (`HomeServerNotSetupException:
+            # HomeServer.setup must be called before getting datastores`).
+            self.store = hs.get_datastores().main
+
+            # We reached the end of the block which means everything was successful, so
+            # no exit handlers are needed (remove them all).
+            exit.pop_all()
 
     async def process_v2_response(
         self, from_server: str, response_json: JsonDict, time_added_ms: int

synapse/server.py

@@ -54,6 +54,7 @@
 from synapse.api.auth.internal import InternalAuth
 from synapse.api.auth.mas import MasDelegatedAuth
 from synapse.api.auth_blocking import AuthBlocking
+from synapse.api.errors import HomeServerNotSetupException
 from synapse.api.filtering import Filtering
 from synapse.api.ratelimiting import Ratelimiter, RequestRatelimiter
 from synapse.app._base import unregister_sighups
@@ -399,7 +400,7 @@ def run_as_background_process(
         """
         if self._is_shutdown:
             raise Exception(
-                f"Cannot start background process. HomeServer has been shutdown {len(self._background_processes)} {len(self.get_clock()._looping_calls)} {len(self.get_clock()._call_id_to_delayed_call)}"
+                "Cannot start background process. HomeServer has been shutdown"
             )
 
         # Ignore linter error as this is the one location this should be called.
@@ -466,7 +467,17 @@ async def shutdown(self) -> None:
 
         # TODO: Cleanup replication pieces
 
-        self.get_keyring().shutdown()
+        keyring: Keyring | None = None
+        try:
+            keyring = self.get_keyring()
+        except HomeServerNotSetupException:
+            # If the homeserver wasn't fully setup, keyring won't have existed before
+            # this and will fail to be initialized but it cleans itself up for any
+            # partial initialization problem.
+            pass
+
+        if keyring:
+            keyring.shutdown()
 
         # Cleanup metrics associated with the homeserver
         for later_gauge in all_later_gauges_to_clean_up_on_shutdown.values():
@@ -478,8 +489,12 @@ async def shutdown(self) -> None:
             self.config.server.server_name
         )
 
-        for db in self.get_datastores().databases:
-            db.stop_background_updates()
+        try:
+            for db in self.get_datastores().databases:
+                db.stop_background_updates()
+        except HomeServerNotSetupException:
+            # If the homeserver wasn't fully setup, the datastores won't exist
+            pass
 
         if self.should_send_federation():
             try:
@@ -513,8 +528,12 @@ async def shutdown(self) -> None:
                 pass
         self._background_processes.clear()
 
-        for db in self.get_datastores().databases:
-            db._db_pool.close()
+        try:
+            for db in self.get_datastores().databases:
+                db._db_pool.close()
+        except HomeServerNotSetupException:
+            # If the homeserver wasn't fully setup, the datastores won't exist
+            pass
 
     def register_async_shutdown_handler(
         self,
@@ -677,7 +696,9 @@ def get_clock(self) -> Clock:
 
     def get_datastores(self) -> Databases:
         if not self.datastores:
-            raise Exception("HomeServer.setup must be called before getting datastores")
+            raise HomeServerNotSetupException(
+                "HomeServer.setup must be called before getting datastores"
+            )
 
         return self.datastores