Optimize presigning for replay.json (#2516)

Fixes #2515.

This PR introduces a significantly optimized logic for presigning URLs
for crawls and collections.
- For collections, the files needed from all crawls are looked up, and
then the 'presign_urls' table is merged in one pass, resulting in a
unified iterator containing files and presign urls for those files.
- For crawls, the presign URLs are also looked up once, and the same
iterator is used for a single crawl with passed in list of CrawlFiles
- URLs that are already signed are added to the return list.
- For any remaining URLs to be signed, a bulk presigning function is
added, which shares an HTTP connection and signing 8 files in parallels
(customizable via helm chart, though may not be needed). This function
is used to call the presigning API in parallel.

Author: ikreymer

backend/btrixcloud/basecrawls.py

@@ -1,7 +1,18 @@
 """base crawl type"""
 
 from datetime import datetime
-from typing import Optional, List, Union, Dict, Any, Type, TYPE_CHECKING, cast, Tuple
+from typing import (
+    Optional,
+    List,
+    Union,
+    Dict,
+    Any,
+    Type,
+    TYPE_CHECKING,
+    cast,
+    Tuple,
+    AsyncIterable,
+)
 from uuid import UUID
 import os
 import urllib.parse
@@ -76,6 +87,7 @@ def __init__(
         background_job_ops: BackgroundJobOps,
     ):
         self.crawls = mdb["crawls"]
+        self.presigned_urls = mdb["presigned_urls"]
         self.crawl_configs = crawl_configs
         self.user_manager = users
         self.orgs = orgs
@@ -464,29 +476,130 @@ async def resolve_signed_urls(
     ) -> List[CrawlFileOut]:
         """Regenerate presigned URLs for files as necessary"""
         if not files:
-            print("no files")
             return []
 
         out_files = []
 
-        for file_ in files:
-            presigned_url, expire_at = await self.storage_ops.get_presigned_url(
-                org, file_, force_update=force_update
+        cursor = self.presigned_urls.find(
+            {"_id": {"$in": [file.filename for file in files]}}
+        )
+
+        presigned = await cursor.to_list(10000)
+
+        files_dict = [file.dict() for file in files]
+
+        # need an async generator to call bulk_presigned_files
+        async def async_gen():
+            yield {"presigned": presigned, "files": files_dict, "_id": crawl_id}
+
+        out_files, _ = await self.bulk_presigned_files(async_gen(), org, force_update)
+
+        return out_files
+
+    async def get_presigned_files(
+        self, match: dict[str, Any], org: Organization
+    ) -> tuple[list[CrawlFileOut], bool]:
+        """return presigned crawl files queried as batch, merging presigns with files in one pass"""
+        cursor = self.crawls.aggregate(
+            [
+                {"$match": match},
+                {"$project": {"files": "$files", "version": 1}},
+                {
+                    "$lookup": {
+                        "from": "presigned_urls",
+                        "localField": "files.filename",
+                        "foreignField": "_id",
+                        "as": "presigned",
+                    }
+                },
+            ]
+        )
+
+        return await self.bulk_presigned_files(cursor, org)
+
+    async def bulk_presigned_files(
+        self,
+        cursor: AsyncIterable[dict[str, Any]],
+        org: Organization,
+        force_update=False,
+    ) -> tuple[list[CrawlFileOut], bool]:
+        """process presigned files in batches"""
+        resources = []
+        pages_optimized = False
+
+        sign_files = []
+
+        async for result in cursor:
+            pages_optimized = result.get("version") == 2
+
+            mapping = {}
+            # create mapping of filename -> file data
+            for file in result["files"]:
+                file["crawl_id"] = result["_id"]
+                mapping[file["filename"]] = file
+
+            if not force_update:
+                # add already presigned resources
+                for presigned in result["presigned"]:
+                    file = mapping.get(presigned["_id"])
+                    if file:
+                        file["signedAt"] = presigned["signedAt"]
+                        file["path"] = presigned["url"]
+                        resources.append(
+                            CrawlFileOut(
+                                name=os.path.basename(file["filename"]),
+                                path=presigned["url"],
+                                hash=file["hash"],
+                                size=file["size"],
+                                crawlId=file["crawl_id"],
+                                numReplicas=len(file.get("replicas") or []),
+                                expireAt=date_to_str(
+                                    presigned["signedAt"]
+                                    + self.storage_ops.signed_duration_delta
+                                ),
+                            )
+                        )
+
+                        del mapping[presigned["_id"]]
+
+            sign_files.extend(list(mapping.values()))
+
+        by_storage: dict[str, dict] = {}
+        for file in sign_files:
+            storage_ref = StorageRef(**file.get("storage"))
+            sid = str(storage_ref)
+
+            storage_group = by_storage.get(sid)
+            if not storage_group:
+                storage_group = {"ref": storage_ref, "names": [], "files": []}
+                by_storage[sid] = storage_group
+
+            storage_group["names"].append(file["filename"])
+            storage_group["files"].append(file)
+
+        for storage_group in by_storage.values():
+            s3storage = self.storage_ops.get_org_storage_by_ref(
+                org, storage_group["ref"]
             )
 
-            out_files.append(
-                CrawlFileOut(
-                    name=os.path.basename(file_.filename),
-                    path=presigned_url or "",
-                    hash=file_.hash,
-                    size=file_.size,
-                    crawlId=crawl_id,
-                    numReplicas=len(file_.replicas) if file_.replicas else 0,
-                    expireAt=date_to_str(expire_at),
-                )
+            signed_urls, expire_at = await self.storage_ops.get_presigned_urls_bulk(
+                org, s3storage, storage_group["names"]
             )
 
-        return out_files
+            for url, file in zip(signed_urls, storage_group["files"]):
+                resources.append(
+                    CrawlFileOut(
+                        name=os.path.basename(file["filename"]),
+                        path=url,
+                        hash=file["hash"],
+                        size=file["size"],
+                        crawlId=file["crawl_id"],
+                        numReplicas=len(file.get("replicas") or []),
+                        expireAt=date_to_str(expire_at),
+                    )
+                )
+
+        return resources, pages_optimized
 
     async def add_to_collection(
         self, crawl_ids: List[str], collection_id: UUID, org: Organization

backend/btrixcloud/colls.py

@@ -28,7 +28,6 @@
     UpdateColl,
     AddRemoveCrawlList,
     BaseCrawl,
-    CrawlOutWithResources,
     CrawlFileOut,
     Organization,
     PaginatedCollOutResponse,
@@ -40,6 +39,7 @@
     AddedResponse,
     DeletedResponse,
     CollectionSearchValuesResponse,
+    CollectionAllResponse,
     OrgPublicCollections,
     PublicOrgDetails,
     CollAccessType,
@@ -50,7 +50,12 @@
     MIN_UPLOAD_PART_SIZE,
     PublicCollOut,
 )
-from .utils import dt_now, slug_from_name, get_duplicate_key_error_field, get_origin
+from .utils import (
+    dt_now,
+    slug_from_name,
+    get_duplicate_key_error_field,
+    get_origin,
+)
 
 if TYPE_CHECKING:
     from .orgs import OrgOps
@@ -346,7 +351,7 @@ async def get_collection_out(
                 result["resources"],
                 crawl_ids,
                 pages_optimized,
-            ) = await self.get_collection_crawl_resources(coll_id)
+            ) = await self.get_collection_crawl_resources(coll_id, org)
 
             initial_pages, _ = await self.page_ops.list_pages(
                 crawl_ids=crawl_ids,
@@ -400,7 +405,9 @@ async def get_public_collection_out(
         if result.get("access") not in allowed_access:
             raise HTTPException(status_code=404, detail="collection_not_found")
 
-        result["resources"], _, _ = await self.get_collection_crawl_resources(coll_id)
+        result["resources"], _, _ = await self.get_collection_crawl_resources(
+            coll_id, org
+        )
 
         thumbnail = result.get("thumbnail")
         if thumbnail:
@@ -554,32 +561,24 @@ async def list_collections(
 
         return collections, total
 
+    # pylint: disable=too-many-locals
     async def get_collection_crawl_resources(
-        self, coll_id: UUID
+        self, coll_id: Optional[UUID], org: Organization
     ) -> tuple[List[CrawlFileOut], List[str], bool]:
         """Return pre-signed resources for all collection crawl files."""
-        # Ensure collection exists
-        _ = await self.get_collection_raw(coll_id)
+        match: dict[str, Any]
 
-        resources = []
-        pages_optimized = True
+        if coll_id:
+            crawl_ids = await self.get_collection_crawl_ids(coll_id)
+            match = {"_id": {"$in": crawl_ids}}
+        else:
+            crawl_ids = []
+            match = {"oid": org.id}
 
-        crawls, _ = await self.crawl_ops.list_all_base_crawls(
-            collection_id=coll_id,
-            states=list(SUCCESSFUL_STATES),
-            page_size=10_000,
-            cls_type=CrawlOutWithResources,
+        resources, pages_optimized = await self.crawl_ops.get_presigned_files(
+            match, org
         )
 
-        crawl_ids = []
-
-        for crawl in crawls:
-            crawl_ids.append(crawl.id)
-            if crawl.resources:
-                resources.extend(crawl.resources)
-            if crawl.version != 2:
-                pages_optimized = False
-
         return resources, crawl_ids, pages_optimized
 
     async def get_collection_names(self, uuids: List[UUID]):
@@ -1009,24 +1008,11 @@ async def list_collection_all(
     @app.get(
         "/orgs/{oid}/collections/$all",
         tags=["collections"],
-        response_model=Dict[str, List[CrawlFileOut]],
+        response_model=CollectionAllResponse,
     )
     async def get_collection_all(org: Organization = Depends(org_viewer_dep)):
         results = {}
-        try:
-            all_collections, _ = await colls.list_collections(org, page_size=10_000)
-            for collection in all_collections:
-                (
-                    results[collection.name],
-                    _,
-                    _,
-                ) = await colls.get_collection_crawl_resources(collection.id)
-        except Exception as exc:
-            # pylint: disable=raise-missing-from
-            raise HTTPException(
-                status_code=400, detail="Error Listing All Crawled Files: " + str(exc)
-            )
-
+        results["resources"] = await colls.get_collection_crawl_resources(None, org)
         return results
 
     @app.get(

backend/btrixcloud/models.py

@@ -1598,6 +1598,13 @@ class CollectionSearchValuesResponse(BaseModel):
     names: List[str]
 
 
+# ============================================================================
+class CollectionAllResponse(BaseModel):
+    """Response model for '$all' collection endpoint"""
+
+    resources: List[CrawlFileOut] = []
+
+
 # ============================================================================
 
 ### ORGS ###

backend/btrixcloud/pages.py

@@ -191,15 +191,22 @@ def _get_page_from_dict(
 
     async def _add_pages_to_db(self, crawl_id: str, pages: List[Page], ordered=True):
         """Add batch of pages to db in one insert"""
-        result = await self.pages.insert_many(
-            [
-                page.to_dict(
-                    exclude_unset=True, exclude_none=True, exclude_defaults=True
-                )
-                for page in pages
-            ],
-            ordered=ordered,
-        )
+        try:
+            result = await self.pages.insert_many(
+                [
+                    page.to_dict(
+                        exclude_unset=True, exclude_none=True, exclude_defaults=True
+                    )
+                    for page in pages
+                ],
+                ordered=ordered,
+            )
+        except pymongo.errors.BulkWriteError as bwe:
+            for err in bwe.details.get("writeErrors", []):
+                # ignorable duplicate key errors
+                if err.get("code") != 11000:
+                    raise
+
         if not result.inserted_ids:
             # pylint: disable=broad-exception-raised
             raise Exception("No pages inserted")