From f46904272fbfc9fd044f5fdb8e0ec445c79ba836 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <erice@element.io>
Date: Tue, 24 Jun 2025 14:02:32 -0500
Subject: [PATCH] Make sure `@measure_func` usages have the necessary
 properties

---
 synapse/federation/sender/transaction_manager.py |  3 +++
 synapse/handlers/device.py                       | 11 ++++++++++-
 synapse/handlers/message.py                      |  5 ++++-
 synapse/push/bulk_push_rule_evaluator.py         |  5 ++++-
 synapse/state/__init__.py                        |  7 +++++--
 5 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/synapse/federation/sender/transaction_manager.py b/synapse/federation/sender/transaction_manager.py
index d8a3eaa525..ca82929c18 100644
--- a/synapse/federation/sender/transaction_manager.py
+++ b/synapse/federation/sender/transaction_manager.py
@@ -60,6 +60,9 @@ class TransactionManager:
     def __init__(self, hs: "synapse.server.HomeServer"):
         self._server_name = hs.hostname
         self.clock = hs.get_clock()  # nb must be called this for @measure_func
+        self.metrics_manager = (
+            hs.metrics_manager
+        )  # nb must be called this for @measure_func
         self._store = hs.get_datastores().main
         self._transaction_actions = TransactionActions(self._store)
         self._transport_layer = hs.get_federation_transport_client()
diff --git a/synapse/handlers/device.py b/synapse/handlers/device.py
index 8f9bf92fda..97e80d9081 100644
--- a/synapse/handlers/device.py
+++ b/synapse/handlers/device.py
@@ -526,6 +526,11 @@ class DeviceHandler(DeviceWorkerHandler):
     def __init__(self, hs: "HomeServer"):
         super().__init__(hs)
 
+        self.clock = hs.get_clock()  # nb must be called this for @measure_func
+        self.metrics_manager = (
+            hs.metrics_manager
+        )  # nb must be called this for @measure_func
+
         self.federation_sender = hs.get_federation_sender()
         self._account_data_handler = hs.get_account_data_handler()
         self._storage_controllers = hs.get_storage_controllers()
@@ -1214,10 +1219,14 @@ class DeviceListUpdater(DeviceListWorkerUpdater):
     def __init__(self, hs: "HomeServer", device_handler: DeviceHandler):
         self.store = hs.get_datastores().main
         self.federation = hs.get_federation_client()
-        self.clock = hs.get_clock()
         self.device_handler = device_handler
         self._notifier = hs.get_notifier()
 
+        self.clock = hs.get_clock()  # nb must be called this for @measure_func
+        self.metrics_manager = (
+            hs.metrics_manager
+        )  # nb must be called this for @measure_func
+
         self._remote_edu_linearizer = Linearizer(name="remote_device_list")
         self._resync_linearizer = Linearizer(name="remote_device_resync")
 
diff --git a/synapse/handlers/message.py b/synapse/handlers/message.py
index 5d6ee6996f..f65e359ce3 100644
--- a/synapse/handlers/message.py
+++ b/synapse/handlers/message.py
@@ -481,7 +481,10 @@ class EventCreationHandler:
         self.store = hs.get_datastores().main
         self._storage_controllers = hs.get_storage_controllers()
         self.state = hs.get_state_handler()
-        self.clock = hs.get_clock()
+        self.clock = hs.get_clock()  # nb must be called this for @measure_func
+        self.metrics_manager = (
+            hs.metrics_manager
+        )  # nb must be called this for @measure_func
         self.validator = EventValidator()
         self.profile_handler = hs.get_profile_handler()
         self.event_builder_factory = hs.get_event_builder_factory()
diff --git a/synapse/push/bulk_push_rule_evaluator.py b/synapse/push/bulk_push_rule_evaluator.py
index 1f4f5b90c3..8e324a8ae3 100644
--- a/synapse/push/bulk_push_rule_evaluator.py
+++ b/synapse/push/bulk_push_rule_evaluator.py
@@ -129,7 +129,10 @@ class BulkPushRuleEvaluator:
     def __init__(self, hs: "HomeServer"):
         self.hs = hs
         self.store = hs.get_datastores().main
-        self.clock = hs.get_clock()
+        self.clock = hs.get_clock()  # nb must be called this for @measure_func
+        self.metrics_manager = (
+            hs.metrics_manager
+        )  # nb must be called this for @measure_func
         self._event_auth_handler = hs.get_event_auth_handler()
         self.should_calculate_push_rules = self.hs.config.push.enable_push
 
diff --git a/synapse/state/__init__.py b/synapse/state/__init__.py
index 3b24c25b81..a160634ae1 100644
--- a/synapse/state/__init__.py
+++ b/synapse/state/__init__.py
@@ -189,10 +189,13 @@ class StateHandler:
     """
 
     def __init__(self, hs: "HomeServer"):
-        self.clock = hs.get_clock()
+        self.hs = hs
         self.store = hs.get_datastores().main
         self._state_storage_controller = hs.get_storage_controllers().state
-        self.hs = hs
+        self.clock = hs.get_clock()  # nb must be called this for @measure_func
+        self.metrics_manager = (
+            hs.metrics_manager
+        )  # nb must be called this for @measure_func
         self._state_resolution_handler = hs.get_state_resolution_handler()
         self._storage_controllers = hs.get_storage_controllers()
         self._events_shard_config = hs.config.worker.events_shard_config