diff --git a/synapse/util/metrics.py b/synapse/util/metrics.py
index c51b641125..4ea930d3e8 100644
--- a/synapse/util/metrics.py
+++ b/synapse/util/metrics.py
@@ -13,10 +13,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from twisted.internet import defer
from synapse.util.logcontext import LoggingContext
import synapse.metrics
+from functools import wraps
import logging
@@ -47,10 +49,22 @@ block_db_txn_duration = metrics.register_distribution(
)
+def measure_func(name):
+ def wrapper(func):
+ @wraps(func)
+ @defer.inlineCallbacks
+ def measured_func(self, *args, **kwargs):
+ with Measure(self.clock, name):
+ r = yield func(self, *args, **kwargs)
+ defer.returnValue(r)
+ return measured_func
+ return wrapper
+
+
class Measure(object):
__slots__ = [
"clock", "name", "start_context", "start", "new_context", "ru_utime",
- "ru_stime", "db_txn_count", "db_txn_duration"
+ "ru_stime", "db_txn_count", "db_txn_duration", "created_context"
]
def __init__(self, clock, name):
@@ -58,17 +72,22 @@ class Measure(object):
self.name = name
self.start_context = None
self.start = None
+ self.created_context = False
def __enter__(self):
self.start = self.clock.time_msec()
self.start_context = LoggingContext.current_context()
- if self.start_context:
- self.ru_utime, self.ru_stime = self.start_context.get_resource_usage()
- self.db_txn_count = self.start_context.db_txn_count
- self.db_txn_duration = self.start_context.db_txn_duration
+ if not self.start_context:
+ self.start_context = LoggingContext("Measure")
+ self.start_context.__enter__()
+ self.created_context = True
+
+ self.ru_utime, self.ru_stime = self.start_context.get_resource_usage()
+ self.db_txn_count = self.start_context.db_txn_count
+ self.db_txn_duration = self.start_context.db_txn_duration
def __exit__(self, exc_type, exc_val, exc_tb):
- if exc_type is not None or not self.start_context:
+ if isinstance(exc_type, Exception) or not self.start_context:
return
duration = self.clock.time_msec() - self.start
@@ -78,8 +97,8 @@ class Measure(object):
if context != self.start_context:
logger.warn(
- "Context have unexpectedly changed from '%s' to '%s'. (%r)",
- context, self.start_context, self.name
+ "Context has unexpectedly changed from '%s' to '%s'. (%r)",
+ self.start_context, context, self.name
)
return
@@ -91,7 +110,12 @@ class Measure(object):
block_ru_utime.inc_by(ru_utime - self.ru_utime, self.name)
block_ru_stime.inc_by(ru_stime - self.ru_stime, self.name)
- block_db_txn_count.inc_by(context.db_txn_count - self.db_txn_count, self.name)
+ block_db_txn_count.inc_by(
+ context.db_txn_count - self.db_txn_count, self.name
+ )
block_db_txn_duration.inc_by(
context.db_txn_duration - self.db_txn_duration, self.name
)
+
+ if self.created_context:
+ self.start_context.__exit__(exc_type, exc_val, exc_tb)
|
