More shutdown cleanup

Author: devonh

synapse/api/auth_blocking.py

@@ -36,7 +36,7 @@
 
 class AuthBlocking:
     def __init__(self, hs: "HomeServer"):
-        self.store = hs.get_datastores().main
+        self.store = weakref.proxy(hs.get_datastores().main)
 
         self._server_notices_mxid = hs.config.servernotices.server_notices_mxid
         self._hs_disabled = hs.config.server.hs_disabled

synapse/api/ratelimiting.py

@@ -21,6 +21,7 @@
 #
 
 from typing import TYPE_CHECKING, Dict, Hashable, Optional, Tuple
+import weakref
 
 from synapse.api.errors import LimitExceededError
 from synapse.config.ratelimiting import RatelimitSettings
@@ -33,6 +34,7 @@
     from synapse.module_api.callbacks.ratelimit_callbacks import (
         RatelimitModuleApiCallbacks,
     )
+    from synapse.server import HomeServer
 
 
 class Ratelimiter:
@@ -75,6 +77,7 @@ class Ratelimiter:
 
     def __init__(
         self,
+        hs: "HomeServer",
         store: DataStore,
         clock: Clock,
         cfg: RatelimitSettings,
@@ -83,7 +86,7 @@ def __init__(
         self.clock = clock
         self.rate_hz = cfg.per_second
         self.burst_count = cfg.burst_count
-        self.store = store
+        self.store = weakref.proxy(store)
         self._limiter_name = cfg.key
         self._ratelimit_callbacks = ratelimit_callbacks
 
@@ -94,7 +97,7 @@ def __init__(
         #   * The rate_hz (leak rate) of this particular bucket.
         self.actions: Dict[Hashable, Tuple[float, float, float]] = {}
 
-        self.clock.looping_call(self._prune_message_counts, 60 * 1000)
+        hs.register_looping_call(self.clock.looping_call(self._prune_message_counts, 60 * 1000))
 
     def _get_key(
         self, requester: Optional[Requester], key: Optional[Hashable]
@@ -348,6 +351,7 @@ async def ratelimit(
 class RequestRatelimiter:
     def __init__(
         self,
+        hs: "HomeServer",
         store: DataStore,
         clock: Clock,
         rc_message: RatelimitSettings,
@@ -358,6 +362,7 @@ def __init__(
 
         # The rate_hz and burst_count are overridden on a per-user basis
         self.request_ratelimiter = Ratelimiter(
+            hs=hs,
             store=self.store,
             clock=self.clock,
             cfg=RatelimitSettings(key=rc_message.key, per_second=0, burst_count=0),
@@ -368,6 +373,7 @@ def __init__(
         # by the presence of rate limits in the config
         if rc_admin_redaction:
             self.admin_redaction_ratelimiter: Optional[Ratelimiter] = Ratelimiter(
+                hs=hs,
                 store=self.store,
                 clock=self.clock,
                 cfg=rc_admin_redaction,

synapse/app/homeserver.py

@@ -69,6 +69,7 @@
 from synapse.rest.well_known import well_known_resource
 from synapse.server import HomeServer
 from synapse.storage import DataStore
+from synapse.types import ISynapseReactor
 from synapse.util.check_dependencies import VERSION, check_requirements
 from synapse.util.httpresourcetree import create_resource_tree
 from synapse.util.module_loader import load_module
@@ -83,6 +84,14 @@ def gz_wrap(r: Resource) -> Resource:
 class SynapseHomeServer(HomeServer):
     DATASTORE_CLASS = DataStore
 
+    def shutdown(self) -> None:
+        super().shutdown()
+        for listener in self._listening_services:
+            listener.loseConnection()
+        self._listening_services.clear()
+
+        self._reactor.stop()
+
     def _listener_http(
         self,
         config: HomeServerConfig,
@@ -308,7 +317,8 @@ def start_listening(self) -> None:
                 logger.warning("Unrecognized listener type: %s", listener.type)
 
 
-def setup(config_options: List[str]) -> SynapseHomeServer:
+from typing import Optional
+def setup(config_options: List[str], reactor: Optional[ISynapseReactor]=None) -> SynapseHomeServer:
     """
     Args:
         config_options_options: The options passed to Synapse. Usually `sys.argv[1:]`.
@@ -365,6 +375,7 @@ def setup(config_options: List[str]) -> SynapseHomeServer:
         config.server.server_name,
         config=config,
         version_string=f"Synapse/{VERSION}",
+        reactor=reactor,
     )
 
     synapse.config.logger.setup_logging(hs, config, use_worker_options=False)

synapse/app/phone_stats_home.py

@@ -206,11 +206,12 @@ def performance_stats_init() -> None:
     )
 
     # monthly active user limiting functionality
-    clock.looping_call(
+    hs.register_looping_call(clock.looping_call(
         hs.get_datastores().main.reap_monthly_active_users,
         ONE_HOUR_SECONDS * MILLISECONDS_PER_SECOND,
-    )
-    hs.get_datastores().main.reap_monthly_active_users()
+    ))
+    # TODO: (devon) how does this hold onto a DB pool reference?
+    #hs.register_background_process(hs.get_datastores().main.reap_monthly_active_users())
 
     @wrap_as_background_process("generate_monthly_active_users")
     async def generate_monthly_active_users() -> None:
@@ -234,7 +235,7 @@ async def generate_monthly_active_users() -> None:
 
     if hs.config.server.limit_usage_by_mau or hs.config.server.mau_stats_only:
         generate_monthly_active_users()
-        clock.looping_call(generate_monthly_active_users, 5 * 60 * 1000)
+        hs.register_looping_call(clock.looping_call(generate_monthly_active_users, 5 * 60 * 1000))
     # End of monthly active user settings
 
     if hs.config.metrics.report_stats:

synapse/appservice/scheduler.py

@@ -523,7 +523,7 @@ def force_retry(self) -> None:
         if self.scheduled_recovery:
             self.clock.cancel_call_later(self.scheduled_recovery)
         # Run a retry, which will resechedule a recovery if it fails.
-        run_as_background_process(
+        deferred = run_as_background_process(
             "retry",
             self.retry,
         )

synapse/federation/federation_client.py

@@ -133,7 +133,7 @@ def __init__(self, hs: "HomeServer"):
         super().__init__(hs)
 
         self.pdu_destination_tried: Dict[str, Dict[str, int]] = {}
-        self._clock.looping_call(self._clear_tried_cache, 60 * 1000)
+        hs.register_looping_call(self._clock.looping_call(self._clear_tried_cache, 60 * 1000))
         self.state = hs.get_state_handler()
         self.transport_layer = hs.get_federation_transport_client()
 
@@ -143,6 +143,7 @@ def __init__(self, hs: "HomeServer"):
         # Cache mapping `event_id` to a tuple of the event itself and the `pull_origin`
         # (which server we pulled the event from)
         self._get_pdu_cache: ExpiringCache[str, Tuple[EventBase, str]] = ExpiringCache(
+            hs=hs,
             cache_name="get_pdu_cache",
             server_name=self.server_name,
             clock=self._clock,
@@ -162,6 +163,7 @@ def __init__(self, hs: "HomeServer"):
             Tuple[str, bool],
             Tuple[JsonDict, Sequence[JsonDict], Sequence[JsonDict], Sequence[str]],
         ] = ExpiringCache(
+            hs=hs,
             cache_name="get_room_hierarchy_cache",
             server_name=self.server_name,
             clock=self._clock,

synapse/federation/federation_server.py

@@ -300,7 +300,7 @@ async def on_incoming_transaction(
             # Start a periodic check for old staged events. This is to handle
             # the case where locks time out, e.g. if another process gets killed
             # without dropping its locks.
-            self._clock.looping_call(self._handle_old_staged_events, 60 * 1000)
+            self.hs.register_looping_call(self._clock.looping_call(self._handle_old_staged_events, 60 * 1000))
 
         # keep this as early as possible to make the calculated origin ts as
         # accurate as possible.

synapse/federation/persistence.py

@@ -28,6 +28,7 @@
 
 import logging
 from typing import Optional, Tuple
+import weakref
 
 from synapse.federation.units import Transaction
 from synapse.storage.databases.main import DataStore
@@ -40,7 +41,7 @@ class TransactionActions:
     """Defines persistence actions that relate to handling Transactions."""
 
     def __init__(self, datastore: DataStore):
-        self.store = datastore
+        self.store = weakref.proxy(datastore)
 
     async def have_responded(
         self, origin: str, transaction: Transaction

synapse/federation/send_queue.py

@@ -112,12 +112,12 @@ def __init__(self, hs: "HomeServer"):
         # lambda binds to the queue rather than to the name of the queue which
         # changes. ARGH.
         def register(name: str, queue: Sized) -> None:
-            LaterGauge(
+            hs.register_later_gauge(LaterGauge(
                 "synapse_federation_send_queue_%s_size" % (queue_name,),
                 "",
                 [],
                 lambda: len(queue),
-            )
+            ))
 
         for queue_name in [
             "presence_map",
@@ -129,7 +129,7 @@ def register(name: str, queue: Sized) -> None:
         ]:
             register(queue_name, getattr(self, queue_name))
 
-        self.clock.looping_call(self._clear_queue, 30 * 1000)
+        hs.register_looping_call(self.clock.looping_call(self._clear_queue, 30 * 1000))
 
     def _next_pos(self) -> int:
         pos = self.pos

synapse/federation/sender/__init__.py

@@ -389,7 +389,7 @@ def __init__(self, hs: "HomeServer"):
         # map from destination to PerDestinationQueue
         self._per_destination_queues: Dict[str, PerDestinationQueue] = {}
 
-        LaterGauge(
+        hs.register_later_gauge(LaterGauge(
             "synapse_federation_transaction_queue_pending_destinations",
             "",
             [],
@@ -398,24 +398,24 @@ def __init__(self, hs: "HomeServer"):
                 for d in self._per_destination_queues.values()
                 if d.transmission_loop_running
             ),
-        )
+        ))
 
-        LaterGauge(
+        hs.register_later_gauge(LaterGauge(
             "synapse_federation_transaction_queue_pending_pdus",
             "",
             [],
             lambda: sum(
                 d.pending_pdu_count() for d in self._per_destination_queues.values()
             ),
-        )
-        LaterGauge(
+        ))
+        hs.register_later_gauge(LaterGauge(
             "synapse_federation_transaction_queue_pending_edus",
             "",
             [],
             lambda: sum(
                 d.pending_edu_count() for d in self._per_destination_queues.values()
             ),
-        )
+        ))
 
         self._is_processing = False
         self._last_poked_id = -1
@@ -426,16 +426,16 @@ def __init__(self, hs: "HomeServer"):
             1.0 / hs.config.ratelimiting.federation_rr_transactions_per_room_per_second
         )
         self._destination_wakeup_queue = _DestinationWakeupQueue(
-            self, self.clock, max_delay_s=rr_txn_interval_per_room_s
+            weakref.proxy(self), self.clock, max_delay_s=rr_txn_interval_per_room_s
         )
 
         # Regularly wake up destinations that have outstanding PDUs to be caught up
-        self.clock.looping_call_now(
+        hs.register_looping_call(self.clock.looping_call_now(
             run_as_background_process,
             WAKEUP_RETRY_PERIOD_SEC * 1000.0,
             "wake_destinations_needing_catchup",
             self._wake_destinations_needing_catchup,
-        )
+        ))
 
     def _get_per_destination_queue(
         self, destination: str