feat(backend): add database connection throttling

Also add metrics related to in use connection.
First attempt was to use metric callback so that we grab the actual
usage but was hard since we require access to the request context that has the
DB driver object referenced.
A simpler version is to keep track of the last usage value seen and
update the metric accordingly.

Note that the max_concurrent_queries setting should match the query
thread pool on the server side (defaults to 400) and be distributed
among server workers.

Signed-off-by: Fatih Acar <[email protected]>

Author: fatih-acar

backend/infrahub/config.py

@@ -249,6 +249,12 @@ class DatabaseSettings(BaseSettings):
     retry_limit: int = Field(
         default=3, description="Maximum number of times a transient issue in a transaction should be retried."
     )
+    max_concurrent_queries: int = Field(
+        default=0, ge=0, description="Maximum number of concurrent queries that can run (0 means unlimited)."
+    )
+    max_concurrent_queries_delay: float = Field(
+        default=0.01, ge=0, description="Delay to add when max_concurrent_queries is reached."
+    )
 
     @property
     def database_name(self) -> str:

backend/infrahub/database/__init__.py

@@ -34,7 +34,7 @@
 
 from .constants import DatabaseType, Neo4jRuntime
 from .memgraph import DatabaseManagerMemgraph
-from .metrics import QUERY_EXECUTION_METRICS, TRANSACTION_RETRIES
+from .metrics import CONNECTION_POOL_USAGE, QUERY_EXECUTION_METRICS, TRANSACTION_RETRIES
 from .neo4j import DatabaseManagerNeo4j
 
 if TYPE_CHECKING:
@@ -335,6 +335,14 @@ async def execute_query_with_metadata(
         context: dict[str, str] | None = None,
         type: QueryType | None = None,  # pylint: disable=redefined-builtin
     ) -> tuple[list[Record], dict[str, Any]]:
+        connpool_usage = self._driver._pool.in_use_connection_count(self._driver._pool.address)
+        CONNECTION_POOL_USAGE.labels(self._driver._pool.address).set(float(connpool_usage))
+
+        if config.SETTINGS.database.max_concurrent_queries:
+            while connpool_usage > config.SETTINGS.database.max_concurrent_queries:  # noqa: ASYNC110
+                await asyncio.sleep(config.SETTINGS.database.max_concurrent_queries_delay)
+                connpool_usage = self._driver._pool.in_use_connection_count(self._driver._pool.address)
+
         with trace.get_tracer(__name__).start_as_current_span("execute_db_query_with_metadata") as span:
             span.set_attribute("query", query)
             if name:

backend/infrahub/database/metrics.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from prometheus_client import Counter, Histogram
+from prometheus_client import Counter, Gauge, Histogram
 
 METRIC_PREFIX = "infrahub_db"
 
@@ -16,3 +16,9 @@
     "Number of transaction that have been retried due to transcient error",
     labelnames=["name"],
 )
+
+CONNECTION_POOL_USAGE = Gauge(
+    f"{METRIC_PREFIX}_last_connection_pool_usage",
+    "Number of last known active connections in the pool",
+    labelnames=["address"],
+)