Time CPU time spent in each log context

2 files changed, 49 insertions, 7 deletions

diff --git a/synapse/http/server.py b/synapse/http/server.py
index b60e905a62..05ecff9b24 100644
--- a/synapse/http/server.py
+++ b/synapse/http/server.py
@@ -77,7 +77,7 @@ def request_handler(request_handler):
         global _next_request_id
         request_id = "%s-%s" % (request.method, _next_request_id)
         _next_request_id += 1
-        with LoggingContext(request_id) as request_context:
+        with LoggingContext(request_id, True) as request_context:
             request_context.request = request_id
             with request.processing():
                 try:
diff --git a/synapse/util/logcontext.py b/synapse/util/logcontext.py
index 7e6062c1b8..8a36bbf32d 100644
--- a/synapse/util/logcontext.py
+++ b/synapse/util/logcontext.py
@@ -16,10 +16,16 @@ from twisted.internet import defer
 
 import threading
 import logging
+from resource import getrusage, RUSAGE_SELF
+
 
 logger = logging.getLogger(__name__)
 
 
+def _get_cpu_time():
+    return getrusage(RUSAGE_SELF).ru_utime * 1000
+
+
 class LoggingContext(object):
     """Additional context for log formatting. Contexts are scoped within a
     "with" block. Contexts inherit the state of their parent contexts.
@@ -27,7 +33,10 @@ class LoggingContext(object):
         name (str): Name for the context for debugging.
     """
 
-    __slots__ = ["parent_context", "name", "__dict__"]
+    __slots__ = [
+        "parent_context", "name", "__dict__", "cpu_time", "_current_cpu_time",
+        "_log_time"
+    ]
 
     thread_local = threading.local()
 
@@ -42,11 +51,20 @@ class LoggingContext(object):
         def copy_to(self, record):
             pass
 
+        def resume(self):
+            pass
+
+        def pause(self):
+            pass
+
     sentinel = Sentinel()
 
-    def __init__(self, name=None):
+    def __init__(self, name=None, log_time=False):
         self.parent_context = None
         self.name = name
+        self.cpu_time = 0
+        self._current_cpu_time = 0
+        self._log_time = log_time
 
     def __str__(self):
         return "%s@%x" % (self.name, id(self))
@@ -62,6 +80,7 @@ class LoggingContext(object):
             raise Exception("Attempt to enter logging context multiple times")
         self.parent_context = self.current_context()
         self.thread_local.current_context = self
+        self.resume()
         return self
 
     def __exit__(self, type, value, traceback):
@@ -79,9 +98,24 @@ class LoggingContext(object):
                     self.thread_local.current_context,
                     self
                 )
+        self.pause()
+
+        if self._log_time:
+            logger.info("Elapsed CPU time: %d %r", self.cpu_time, self.name)
+
         self.thread_local.current_context = self.parent_context
         self.parent_context = None
 
+    def resume(self):
+        if self._log_time:
+            self._current_cpu_time = _get_cpu_time()
+
+    def pause(self):
+        if self._log_time:
+            now = _get_cpu_time()
+            self.cpu_time += now - self._current_cpu_time
+            self._current_cpu_time = now
+
     def __getattr__(self, name):
         """Delegate member lookup to parent context"""
         return getattr(self.parent_context, name)
@@ -121,15 +155,24 @@ class PreserveLoggingContext(object):
     exited. Used to restore the context after a function using
     @defer.inlineCallbacks is resumed by a callback from the reactor."""
 
-    __slots__ = ["current_context"]
+    __slots__ = ["current_context", "inner_context"]
+
+    def __init__(self, inner_context=LoggingContext.sentinel):
+        object.__setattr__(self, "inner_context", inner_context)
 
     def __enter__(self):
         """Captures the current logging context"""
         self.current_context = LoggingContext.current_context()
-        LoggingContext.thread_local.current_context = LoggingContext.sentinel
+        LoggingContext.thread_local.current_context = self.inner_context
+
+        self.current_context.pause()
+        self.inner_context.resume()
 
     def __exit__(self, type, value, traceback):
         """Restores the current logging context"""
+        self.inner_context.pause()
+        self.current_context.resume()
+
         LoggingContext.thread_local.current_context = self.current_context
 
         if self.current_context is not LoggingContext.sentinel:
@@ -164,8 +207,7 @@ class _PreservingContextDeferred(defer.Deferred):
 
     def _wrap_callback(self, f):
         def g(res, *args, **kwargs):
-            with PreserveLoggingContext():
-                LoggingContext.thread_local.current_context = self._log_context
+            with PreserveLoggingContext(self._log_context):
                 res = f(res, *args, **kwargs)
             return res
         return g