Add endpoint to get counts of running crawls (per-org, or across all orgs) (#3549)

Part of #3541 

This PR adds:

- New `/orgs/all/crawlconfigs/running` (superadmin) and
`/orgs/{oid}/crawlconfigs/running` endpoints that provide counts of the
number of workflows in all running, paused, and waiting states, as well
as totals.
- A new index to the crawls mongo collection to help facilitate the
query
- Some basic tests

Author: tw4l

backend/btrixcloud/crawlconfigs.py

@@ -27,6 +27,7 @@
 )
 
 from .models import (
+    ALL_CRAWL_STATES,
     SUCCESSFUL_STATES,
     TYPE_ALL_CRAWL_STATES,
     ConfigRevision,
@@ -35,6 +36,7 @@
     CrawlConfigDeletedResponse,
     CrawlConfigIn,
     CrawlConfigOut,
+    CrawlConfigRunningCountsResponse,
     CrawlConfigSearchValues,
     CrawlConfigUpdateResponse,
     CrawlerChannel,
@@ -1689,6 +1691,103 @@ async def validate_custom_behavior(self, url: str) -> dict[str, bool]:
 
         return {"success": True}
 
+    async def get_running_counts(
+        self, org: Organization | None = None
+    ) -> CrawlConfigRunningCountsResponse:
+        """Return counts of running workflows, total and status, optionally by org"""
+
+        state_count_logger = logger.bind(oid=org.id if org else None)
+
+        try:
+            match_query: dict[str, UUID | str] = {}
+            if org:
+                match_query["oid"] = org.id
+
+            res = await self.crawls.aggregate(
+                [
+                    {"$match": match_query},
+                    {"$group": {"_id": "$state", "count": {"$count": {}}}},
+                ]
+            ).to_list()
+
+            state_counts: dict[str, int] = {}
+
+            for state_dict in res:
+                state = state_dict["_id"]
+                count = state_dict.get("count", 0)
+                if state not in ALL_CRAWL_STATES:
+                    state_count_logger.error(
+                        "unexpected_crawl_state_found", state=state, count=count
+                    )
+                else:
+                    state_counts[state] = count
+
+            # Running states
+            running = state_counts.get("running", 0)
+            pending_wait = state_counts.get("pending-wait", 0)
+            generate_wacz = state_counts.get("generate-wacz", 0)
+            uploading_wacz = state_counts.get("uploading-wacz", 0)
+            rate_limited = state_counts.get("rate-limited", 0)
+            total_running = (
+                running + pending_wait + generate_wacz + uploading_wacz + rate_limited
+            )
+
+            # Paused states
+            paused = state_counts.get("paused", 0)
+            paused_storage = state_counts.get("paused_storage_quota_reached", 0)
+            paused_time = state_counts.get("paused_time_quota_reached", 0)
+            paused_read_only = state_counts.get("paused_org_readonly", 0)
+            paused_rate_limit = state_counts.get("paused_rate_limit_time_reached", 0)
+
+            total_paused = (
+                paused
+                + paused_storage
+                + paused_time
+                + paused_read_only
+                + paused_rate_limit
+            )
+
+            # Waiting states
+            starting = state_counts.get("starting", 0)
+            waiting_capacity = state_counts.get("waiting_capacity", 0)
+            waiting_org_limit = state_counts.get("waiting_org_limit", 0)
+            waiting_dedupe = state_counts.get("waiting_dedupe", 0)
+            total_waiting = (
+                starting + waiting_capacity + waiting_org_limit + waiting_dedupe
+            )
+
+            total = total_running + total_paused + total_waiting
+
+            return CrawlConfigRunningCountsResponse(
+                totalRunningPausedWaiting=total,
+                totalRunning=total_running,
+                totalPaused=total_paused,
+                totalWaiting=total_waiting,
+                # Running states
+                running=running,
+                pendingWait=pending_wait,
+                generateWACZ=generate_wacz,
+                uploadingWACZ=uploading_wacz,
+                rateLimited=rate_limited,
+                # Paused states
+                paused=paused,
+                pausedStorageQuotaReached=paused_storage,
+                pausedTimeQuotaReached=paused_time,
+                pausedOrgReadOnly=paused_read_only,
+                pausedRateLimitTimeReached=paused_rate_limit,
+                # Waiting states
+                starting=starting,
+                waitingCapacity=waiting_capacity,
+                waitingOrgLimit=waiting_org_limit,
+                waitingDedupeIndex=waiting_dedupe,
+            )
+        except Exception:
+            state_count_logger.exception(
+                "running_workflow_counts_calculation_failed",
+            )
+            # pylint: disable=raise-missing-from
+            raise HTTPException(status_code=400, detail="calculation_failure")
+
 
 # ============================================================================
 # pylint: disable=too-many-locals
@@ -1945,6 +2044,24 @@ async def get_all_crawler_proxies(
 
         return ops.get_crawler_proxies()
 
+    @router.get("/running", response_model=CrawlConfigRunningCountsResponse)
+    async def get_org_crawl_config_running_counts(
+        org: Organization = Depends(org_viewer_dep),
+    ):
+        return await ops.get_running_counts(org)
+
+    @app.get(
+        "/orgs/all/crawlconfigs/running",
+        response_model=CrawlConfigRunningCountsResponse,
+    )
+    async def get_all_crawl_config_running_counts(
+        user: User = Depends(user_dep),
+    ):
+        if not user.is_superuser:
+            raise HTTPException(status_code=403, detail="Not Allowed")
+
+        return await ops.get_running_counts()
+
     @app.get(
         "/orgs/{oid}/crawlconfigs/{cid}/public/replay.json",
         response_model=CrawlOutWithResources,

backend/btrixcloud/crawls.py

@@ -129,6 +129,9 @@ async def init_index(self):
                 ("started", pymongo.ASCENDING),
             ]
         )
+        await self.crawls.create_index(
+            [("oid", pymongo.HASHED), ("state", pymongo.DESCENDING)]
+        )
         await self.crawls.create_index([("finished", pymongo.DESCENDING)])
         await self.crawls.create_index([("oid", pymongo.HASHED)])
         await self.crawls.create_index([("cid", pymongo.HASHED)])

backend/btrixcloud/models.py

@@ -709,6 +709,36 @@ class TagsResponse(BaseModel):
     tags: list[TagCount]
 
 
+# ============================================================================
+class CrawlConfigRunningCountsResponse(BaseModel):
+    """Response model for counts of running workflows (total and by status)"""
+
+    totalRunningPausedWaiting: int = 0
+    totalRunning: int = 0
+    totalPaused: int = 0
+    totalWaiting: int = 0
+
+    # Running states
+    running: int = 0
+    pendingWait: int = 0
+    generateWACZ: int = 0
+    uploadingWACZ: int = 0
+    rateLimited: int = 0
+
+    # Paused states
+    paused: int = 0
+    pausedStorageQuotaReached: int = 0
+    pausedTimeQuotaReached: int = 0
+    pausedOrgReadOnly: int = 0
+    pausedRateLimitTimeReached: int = 0
+
+    # Waiting states
+    starting: int = 0
+    waitingCapacity: int = 0
+    waitingOrgLimit: int = 0
+    waitingDedupeIndex: int = 0
+
+
 # ============================================================================
 class CrawlConfigSearchValues(BaseModel):
     """Response model for adding crawlconfigs"""

backend/test/test_run_crawl.py

@@ -176,6 +176,43 @@ def test_remove_exclusion(admin_auth_headers, default_org_id):
     assert r.json()["success"] == True
 
 
+def test_running_workflow_counts(
+    admin_auth_headers, crawler_auth_headers, default_org_id
+):
+    # Verify running workflow counts are updated
+    r = requests.get(
+        f"{API_PREFIX}/orgs/{default_org_id}/crawlconfigs/running",
+        headers=admin_auth_headers,
+    )
+    assert r.status_code == 200
+    data = r.json()
+    assert data["totalRunningPausedWaiting"] >= 1
+    assert data["totalRunning"] >= 1
+    assert (
+        data["running"] >= 1 or data["generateWACZ"] >= 1 or data["uploadingWACZ"] >= 1
+    )
+
+    # Verify again but from non-org-specific endpoint
+    r = requests.get(
+        f"{API_PREFIX}/orgs/all/crawlconfigs/running",
+        headers=admin_auth_headers,
+    )
+    assert r.status_code == 200
+    data = r.json()
+    assert data["totalRunningPausedWaiting"] >= 1
+    assert data["totalRunning"] >= 1
+    assert (
+        data["running"] >= 1 or data["generateWACZ"] >= 1 or data["uploadingWACZ"] >= 1
+    )
+
+    # Check that non-org-specific endpoint is only available to superadmins
+    r = requests.get(
+        f"{API_PREFIX}/orgs/all/crawlconfigs/running",
+        headers=crawler_auth_headers,
+    )
+    assert r.status_code == 403
+
+
 def test_wait_for_complete(admin_auth_headers, default_org_id):
     state = None
     data = None