chore: add extra dimension to celery concurrency limit hits (PostHog#28783)

orian · web-flow · commit 512fbf6ea3f1 · 2025-02-17T11:17:59.000Z
diff --git a/posthog/clickhouse/client/limit.py b/posthog/clickhouse/client/limit.py
@@ -11,7 +11,7 @@
 CONCURRENT_TASKS_LIMIT_EXCEEDED_COUNTER = Counter(
     "posthog_celery_task_concurrency_limit_exceeded",
     "Number of times a Celery task exceeded the concurrency limit",
-    ["task_name", "limit"],
+    ["task_name", "limit", "limit_name"],
 )
 
 # Lua script for atomic check, remove expired if limit hit, and increment with TTL
@@ -44,7 +44,9 @@ class CeleryConcurrencyLimitExceeded(Exception):
     pass
 
 
-def limit_concurrency(max_concurrent_tasks: int, key: Optional[Callable] = None, ttl: int = 60 * 15) -> Callable:
+def limit_concurrency(
+    max_concurrent_tasks: int, key: Optional[Callable] = None, ttl: int = 60 * 15, limit_name: str = ""
+) -> Callable:
     def decorator(task_func):
         @wraps(task_func)
         def wrapper(*args, **kwargs):
@@ -64,7 +66,9 @@ def wrapper(*args, **kwargs):
                 redis_client.eval(lua_script, 1, running_tasks_key, current_time, task_id, max_concurrent_tasks, ttl)
                 == 0
             ):
-                CONCURRENT_TASKS_LIMIT_EXCEEDED_COUNTER.labels(task_name=task_name, limit=max_concurrent_tasks).inc()
+                CONCURRENT_TASKS_LIMIT_EXCEEDED_COUNTER.labels(
+                    task_name=task_name, limit=max_concurrent_tasks, limit_name=limit_name
+                ).inc()
 
                 raise CeleryConcurrencyLimitExceeded(
                     f"Exceeded maximum concurrent tasks limit: {max_concurrent_tasks} for key: {dynamic_key}"
diff --git a/posthog/tasks/tasks.py b/posthog/tasks/tasks.py
@@ -52,9 +52,11 @@ def redis_heartbeat() -> None:
     expires=60 * 10,  # Do not run queries that got stuck for more than this
     reject_on_worker_lost=True,
 )
-@limit_concurrency(150)  # Do not go above what CH can handle (max_concurrent_queries)
+@limit_concurrency(150, limit_name="global")  # Do not go above what CH can handle (max_concurrent_queries)
 @limit_concurrency(
-    10, key=lambda *args, **kwargs: kwargs.get("team_id") or args[0]
+    10,
+    key=lambda *args, **kwargs: kwargs.get("team_id") or args[0],
+    limit_name="per_team",
 )  # Do not run too many queries at once for the same team
 def process_query_task(
     team_id: int,
