diff --git a/synapse/storage/database.py b/synapse/storage/database.py
index 57cc1d76e0..fc8dff8d17 100644
--- a/synapse/storage/database.py
+++ b/synapse/storage/database.py
@@ -41,6 +41,7 @@ from prometheus_client import Histogram
from typing_extensions import Literal
from twisted.enterprise import adbapi
+from twisted.python import reflect
from synapse.api.errors import StoreError
from synapse.config.database import DatabaseConnectionConfig
@@ -90,6 +91,20 @@ UNIQUE_INDEX_BACKGROUND_UPDATES = {
}
+class NastyConnectionWrapper:
+ def __init__(self, connection):
+ self._connection = connection
+ self._synapse_parent_context = None
+
+ def commit(self, *args, **kwargs):
+ with LoggingContext("db_commit", parent_context = self._synapse_parent_context):
+ with opentracing.start_active_span("db.conn.commit"):
+ self._connection.commit(*args, **kwargs)
+
+ def __getattr__(self, item):
+ return getattr(self._connection, item)
+
+
def make_pool(
reactor, db_config: DatabaseConnectionConfig, engine: BaseDatabaseEngine
) -> adbapi.ConnectionPool:
@@ -105,9 +120,28 @@ def make_pool(
# etc.
with LoggingContext("db.on_new_connection"):
engine.on_new_connection(
- LoggingDatabaseConnection(conn, engine, "on_new_connection")
+ LoggingDatabaseConnection(
+ conn, engine, "on_new_connection"
+ )
)
+ # HACK Patch the connection's commit function so that we can see
+ # how long it's taking from Jaeger. To do that, we need to patch the
+ # dbapi module's 'connect' method so that it returns a wrapped 'Connection'
+ # object to the connection pool. (psycopg2's Connection class is a C thing
+ # which we can't monkey-patch directly).
+ dbapiname = db_config.config["name"]
+ dbapi = reflect.namedModule(dbapiname)
+ if not getattr(dbapi, "_synapse_wrapped_dbapi", False):
+ real_connect = dbapi.connect
+
+ def wrapped_connect(*args, **kwargs):
+ conn = real_connect(*args, **kwargs)
+ return NastyConnectionWrapper(conn)
+
+ dbapi.connect = wrapped_connect
+ dbapi._synapse_wrapped_dbapi = True
+
connection_pool = adbapi.ConnectionPool(
db_config.config["name"],
cp_reactor=reactor,
@@ -800,6 +834,10 @@ class DatabasePool:
# pool).
assert not self.engine.in_transaction(conn)
+ # HACK: record the parent context in 'conn' so that we can tie later commits
+ # back to it
+ conn._connection._synapse_parent_context = parent_context
+
with LoggingContext(
str(curr_context), parent_context=parent_context
) as context:
|
