Fix incorrect juggling of logging contexts in _PerHostRatelimiter

changelog.d/13554.misc

@@ -0,0 +1 @@
+Instrument `FederationStateIdsServlet` (`/state_ids`) for understandable traces in Jaeger.

synapse/util/ratelimitutils.py

@@ -134,12 +134,19 @@ def ratelimit(self, host: str) -> "Iterator[defer.Deferred[None]]":
 
         request_id = object()
         ret = self._on_enter(request_id)
+        ret = defer.ensureDeferred(self._trace_wait(ret))
         try:
             yield ret
         finally:
             self._on_exit(request_id)
 
     def _on_enter(self, request_id: object) -> "defer.Deferred[None]":
+        """
+        Returns:
+            A `Deferred` that resolves once the request is allowed by the rate limiter.
+            The returned `Deferred` does not follow the Synapse logging context rules
+            and must be wrapped with `make_deferred_yieldable` before awaiting it.
+        """
         time_now = self.clock.time_msec()
 
         # remove any entries from request_times which aren't within the window
@@ -222,20 +229,21 @@ def on_both(r: object) -> object:
             # Ensure that we've properly cleaned up.
             self.sleeping_requests.discard(request_id)
             self.ready_request_queue.pop(request_id, None)
-            wait_span_scope.__exit__(None, None, None)
-            wait_timer_cm.__exit__(None, None, None)
             return r
 
-        # Tracing
-        wait_span_scope = start_active_span("ratelimit wait")
-        wait_span_scope.__enter__()
-        # Metrics
-        wait_timer_cm = queue_wait_timer.time()
-        wait_timer_cm.__enter__()
-
         ret_defer.addCallbacks(on_start, on_err)
         ret_defer.addBoth(on_both)
-        return make_deferred_yieldable(ret_defer)
+        return ret_defer
+
+    async def _trace_wait(self, d: "defer.Deferred[None]") -> None:
+        """Waits for a given `Deferred` to resolve, while timing it for metrics.
+
+        Args:
+            d: The `Deferred` to wait for. Must not follow the Synapse logging context
+                rules.
+        """
+        with start_active_span("ratelimit wait"), queue_wait_timer.time():
+            await make_deferred_yieldable(d)
 
     def _on_exit(self, request_id: object) -> None:
         logger.debug("Ratelimit(%s) [%s]: Processed req", self.host, id(request_id))