Refactor Measure in synapse/http/federation/well_known_resolver.py

Author: MadLittleMods

synapse/http/federation/matrix_federation_agent.py

@@ -48,6 +48,7 @@
 from synapse.http.federation.well_known_resolver import WellKnownResolver
 from synapse.http.proxyagent import ProxyAgent
 from synapse.logging.context import make_deferred_yieldable, run_in_background
+from synapse.metrics.homeserver_metrics_manager import HomeserverMetricsManager
 from synapse.types import ISynapseReactor
 from synapse.util import Clock
 
@@ -93,6 +94,7 @@ class MatrixFederationAgent:
     def __init__(
         self,
         reactor: ISynapseReactor,
+        metrics_manager: HomeserverMetricsManager,
         tls_client_options_factory: Optional[FederationPolicyForHTTPS],
         user_agent: bytes,
         ip_allowlist: Optional[IPSet],
@@ -128,6 +130,7 @@ def __init__(
         if _well_known_resolver is None:
             _well_known_resolver = WellKnownResolver(
                 reactor,
+                metrics_manager,
                 agent=BlocklistingAgentWrapper(
                     ProxyAgent(
                         reactor,

synapse/http/federation/well_known_resolver.py

@@ -35,6 +35,7 @@
 
 from synapse.http.client import BodyExceededMaxSize, read_body_with_max_size
 from synapse.logging.context import make_deferred_yieldable
+from synapse.metrics.homeserver_metrics_manager import HomeserverMetricsManager
 from synapse.util import Clock, json_decoder
 from synapse.util.caches.ttlcache import TTLCache
 from synapse.util.metrics import Measure
@@ -92,13 +93,15 @@ class WellKnownResolver:
     def __init__(
         self,
         reactor: IReactorTime,
+        metrics_manager: HomeserverMetricsManager,
         agent: IAgent,
         user_agent: bytes,
         well_known_cache: Optional[TTLCache[bytes, Optional[bytes]]] = None,
         had_well_known_cache: Optional[TTLCache[bytes, bool]] = None,
     ):
         self._reactor = reactor
         self._clock = Clock(reactor)
+        self.metrics_manager = metrics_manager
 
         if well_known_cache is None:
             well_known_cache = _well_known_cache

synapse/http/matrixfederationclient.py

@@ -406,6 +406,7 @@ def __init__(
         self.server_name = hs.hostname
 
         self.reactor = hs.get_reactor()
+        self.metrics_manager = hs.metrics_manager
 
         user_agent = hs.version_string
         if hs.config.server.user_agent_suffix:
@@ -418,6 +419,7 @@ def __init__(
             # Talk to federation directly
             federation_agent: IAgent = MatrixFederationAgent(
                 self.reactor,
+                self.metrics_manager,
                 tls_client_options_factory,
                 user_agent.encode("ascii"),
                 hs.config.server.federation_ip_range_allowlist,

tests/handlers/test_typing.py

@@ -84,14 +84,6 @@ def make_homeserver(
 
         # we mock out the federation client too
         self.mock_federation_client = AsyncMock(spec=["put_json"])
-        self.mock_federation_client.put_json.return_value = (200, "OK")
-        self.mock_federation_client.agent = MatrixFederationAgent(
-            reactor,
-            tls_client_options_factory=None,
-            user_agent=b"SynapseInTrialTest/0.0.0",
-            ip_allowlist=None,
-            ip_blocklist=IPSet(),
-        )
 
         # the tests assume that we are starting at unix time 1000
         reactor.pump((1000,))
@@ -104,6 +96,18 @@ def make_homeserver(
             replication_streams={},
         )
 
+        # Finish off mocking the federation client
+        self.mock_federation_client.put_json.return_value = (200, "OK")
+        self.mock_federation_client.agent = MatrixFederationAgent(
+            reactor=reactor,
+            # After we get access to the `hs` homeserver instance, we can replace the federation agent
+            metrics_manager=hs.metrics_manager,
+            tls_client_options_factory=None,
+            user_agent=b"SynapseInTrialTest/0.0.0",
+            ip_allowlist=None,
+            ip_blocklist=IPSet(),
+        )
+
         return hs
 
     def create_resource_dict(self) -> Dict[str, Resource]:

tests/http/federation/test_matrix_federation_agent.py

@@ -59,6 +59,7 @@
     LoggingContextOrSentinel,
     current_context,
 )
+from synapse.metrics.homeserver_metrics_manager import HomeserverMetricsManager
 from synapse.types import ISynapseReactor
 from synapse.util.caches.ttlcache import TTLCache
 
@@ -82,6 +83,8 @@ def setUp(self) -> None:
         self._config = config = HomeServerConfig()
         config.parse_config_dict(config_dict, "", "")
 
+        self.metrics_manager = HomeserverMetricsManager()
+
         self.tls_factory = FederationPolicyForHTTPS(config)
 
         self.well_known_cache: TTLCache[bytes, Optional[bytes]] = TTLCache(
@@ -92,6 +95,7 @@ def setUp(self) -> None:
         )
         self.well_known_resolver = WellKnownResolver(
             self.reactor,
+            self.metrics_manager,
             Agent(self.reactor, contextFactory=self.tls_factory),
             b"test-agent",
             well_known_cache=self.well_known_cache,
@@ -270,6 +274,7 @@ def _make_agent(self) -> MatrixFederationAgent:
         """
         return MatrixFederationAgent(
             reactor=cast(ISynapseReactor, self.reactor),
+            metrics_manager=self.metrics_manager,
             tls_client_options_factory=self.tls_factory,
             user_agent=b"test-agent",  # Note that this is unused since _well_known_resolver is provided.
             ip_allowlist=IPSet(),
@@ -1012,13 +1017,15 @@ def test_get_well_known_unsigned_cert(self) -> None:
         tls_factory = FederationPolicyForHTTPS(config)
         agent = MatrixFederationAgent(
             reactor=self.reactor,
+            metrics_manager=self.metrics_manager,
             tls_client_options_factory=tls_factory,
             user_agent=b"test-agent",  # This is unused since _well_known_resolver is passed below.
             ip_allowlist=IPSet(),
             ip_blocklist=IPSet(),
             _srv_resolver=self.mock_resolver,
             _well_known_resolver=WellKnownResolver(
                 cast(ISynapseReactor, self.reactor),
+                self.metrics_manager,
                 Agent(self.reactor, contextFactory=tls_factory),
                 b"test-agent",
                 well_known_cache=self.well_known_cache,

tests/replication/test_federation_sender_shard.py

@@ -36,6 +36,7 @@
 from synapse.events import EventBase, make_event_from_dict
 from synapse.handlers.typing import TypingWriterHandler
 from synapse.http.federation.matrix_federation_agent import MatrixFederationAgent
+from synapse.metrics.homeserver_metrics_manager import HomeserverMetricsManager
 from synapse.rest.admin import register_servlets_for_client_rest_resource
 from synapse.rest.client import login, room
 from synapse.server import HomeServer
@@ -69,6 +70,7 @@ def setUp(self) -> None:
         reactor, _ = get_clock()
         self.matrix_federation_agent = MatrixFederationAgent(
             reactor,
+            metrics_manager=HomeserverMetricsManager(),
             tls_client_options_factory=None,
             user_agent=b"SynapseInTrialTest/0.0.0",
             ip_allowlist=None,