Profiles: Make browser commit API call idempotent (#2728)

- Fix race condition related to browser commit time
- The profile commit request waits for browser to actual finish, and
profile saved. This can cause request to time out, resulting in a retry,
in which the browser has already been closed.
- With these changes, the commit is now 'idempotent' and returns a
waiting_for_browser until the profile is actually committed.
- On frontend, keep pinging commit endpoint with a timeout while 'waiting_for_browser' is returned, actual committed when endpoint returns profile id.

---------

Co-authored-by: sua yoo <[email protected]>

Author: ikreymer

backend/btrixcloud/crawlmanager.py

@@ -37,6 +37,7 @@ async def run_profile_browser(
         baseprofile: str = "",
         profile_filename: str = "",
         proxy_id: str = "",
+        profileid: str = "",
     ) -> str:
         """run browser for profile creation"""
 
@@ -60,6 +61,7 @@ async def run_profile_browser(
             "crawler_image": crawler_image,
             "image_pull_policy": image_pull_policy,
             "proxy_id": proxy_id or DEFAULT_PROXY_ID,
+            "profileid": profileid,
         }
 
         data = self.templates.env.get_template("profile_job.yaml").render(params)
@@ -365,14 +367,21 @@ async def get_profile_browser_metadata(self, browserid: str) -> dict[str, str]:
         except:
             return {}
 
-        return browser["metadata"]["labels"]
+        metadata = browser["metadata"]["labels"]
 
-    async def ping_profile_browser(self, browserid: str) -> None:
-        """return ping profile browser"""
+        metadata["committing"] = browser.get("spec", {}).get("committing")
+
+        return metadata
+
+    async def keep_alive_profile_browser(self, browserid: str, committing="") -> None:
+        """update profile browser to not expire"""
         expire_at = dt_now() + timedelta(seconds=30)
-        await self._patch_job(
-            browserid, {"expireTime": date_to_str(expire_at)}, "profilejobs"
-        )
+
+        update = {"expireTime": date_to_str(expire_at)}
+        if committing:
+            update["committing"] = committing
+
+        await self._patch_job(browserid, update, "profilejobs")
 
     async def rollover_restart_crawl(self, crawl_id: str) -> dict:
         """Rolling restart of crawl by updating restartTime field"""

backend/btrixcloud/profiles.py

@@ -3,6 +3,7 @@
 from typing import Optional, TYPE_CHECKING, Any, cast, Dict, List, Tuple
 from uuid import UUID, uuid4
 import os
+import asyncio
 
 from urllib.parse import urlencode
 
@@ -61,6 +62,8 @@ class ProfileOps:
     browser_fqdn_suffix: str
     router: APIRouter
 
+    bg_tasks: set
+
     def __init__(self, mdb, orgs, crawl_manager, storage_ops, background_job_ops):
         self.profiles = mdb["profiles"]
         self.orgs = orgs
@@ -79,6 +82,10 @@ def __init__(self, mdb, orgs, crawl_manager, storage_ops, background_job_ops):
 
         self.crawlconfigs = cast(CrawlConfigOps, None)
 
+        # to avoid background tasks being garbage collected
+        # see: https://stackoverflow.com/a/74059981
+        self.bg_tasks = set()
+
     def set_crawlconfigs(self, crawlconfigs):
         """set crawlconfigs ops"""
         self.crawlconfigs = crawlconfigs
@@ -128,6 +135,7 @@ async def create_new_browser(
             baseprofile=prev_profile_id,
             profile_filename=prev_profile_path,
             proxy_id=proxy_id,
+            profileid=str(uuid4()),
         )
 
         if not browserid:
@@ -166,8 +174,6 @@ async def get_profile_browser_url(
 
     async def ping_profile_browser(self, browserid: str) -> dict[str, Any]:
         """ping profile browser to keep it running"""
-        await self.crawl_manager.ping_profile_browser(browserid)
-
         json = await self._send_browser_req(browserid, "/ping")
 
         return {"success": True, "origins": json.get("origins") or []}
@@ -182,93 +188,130 @@ async def navigate_profile_browser(
 
     async def commit_to_profile(
         self,
+        metadata: dict,
         browser_commit: ProfileCreate,
         org: Organization,
         user: User,
-        metadata: dict,
         existing_profile: Optional[Profile] = None,
     ) -> dict[str, Any]:
-        """commit profile and shutdown profile browser"""
-        # pylint: disable=too-many-locals
+        """commit to profile async, returning if committed, or waiting"""
+        profileid = metadata.get("profileid")
+        if not profileid:
+            raise HTTPException(status_code=400, detail="browser_not_valid")
 
-        now = dt_now()
+        self.orgs.can_write_data(org, include_time=False)
 
-        if existing_profile:
-            profileid = existing_profile.id
-            created = existing_profile.created
-            created_by = existing_profile.createdBy
-            created_by_name = existing_profile.createdByName
-        else:
-            profileid = uuid4()
-            created = now
-            created_by = user.id
-            created_by_name = user.name if user.name else user.email
+        committing = metadata.get("committing")
+        if not committing:
+            self._run_task(
+                self.do_commit_to_profile(
+                    metadata=metadata,
+                    browser_commit=browser_commit,
+                    org=org,
+                    user=user,
+                    existing_profile=existing_profile,
+                )
+            )
 
-        filename_data = {"filename": f"profiles/profile-{profileid}.tar.gz"}
+        if committing == "done":
+            await self.crawl_manager.delete_profile_browser(browser_commit.browserid)
+            return {
+                "added": True,
+                "id": profileid,
+                "storageQuotaReached": self.orgs.storage_quota_reached(org),
+            }
 
-        json = await self._send_browser_req(
-            browser_commit.browserid, "/createProfileJS", "POST", json=filename_data
-        )
+        raise HTTPException(status_code=200, detail="waiting_for_browser")
 
+    async def do_commit_to_profile(
+        self,
+        metadata: dict,
+        browser_commit: ProfileCreate,
+        org: Organization,
+        user: User,
+        existing_profile: Optional[Profile] = None,
+    ) -> bool:
+        """commit profile and shutdown profile browser"""
+        # pylint: disable=too-many-locals
         try:
+            now = dt_now()
+
+            if existing_profile:
+                profileid = existing_profile.id
+                created = existing_profile.created
+                created_by = existing_profile.createdBy
+                created_by_name = existing_profile.createdByName
+            else:
+                profileid = UUID(metadata["profileid"])
+                created = now
+                created_by = user.id
+                created_by_name = user.name if user.name else user.email
+
+            filename_data = {"filename": f"profiles/profile-{profileid}.tar.gz"}
+
+            json = await self._send_browser_req(
+                browser_commit.browserid,
+                "/createProfileJS",
+                "POST",
+                json=filename_data,
+                committing="committing",
+            )
             resource = json["resource"]
-        except:
-            # pylint: disable=raise-missing-from
-            raise HTTPException(status_code=400, detail="browser_not_valid")
-
-        await self.crawl_manager.delete_profile_browser(browser_commit.browserid)
 
-        # backwards compatibility
-        file_size = resource.get("size") or resource.get("bytes")
+            # backwards compatibility
+            file_size = resource.get("size") or resource.get("bytes")
 
-        profile_file = ProfileFile(
-            hash=resource["hash"],
-            size=file_size,
-            filename=resource["path"],
-            storage=org.storage,
-        )
+            profile_file = ProfileFile(
+                hash=resource["hash"],
+                size=file_size,
+                filename=resource["path"],
+                storage=org.storage,
+            )
 
-        baseid = metadata.get("btrix.baseprofile")
-        if baseid:
-            print("baseid", baseid)
-            baseid = UUID(baseid)
+            baseid = metadata.get("btrix.baseprofile")
+            if baseid:
+                print("baseid", baseid)
+                baseid = UUID(baseid)
+
+            profile = Profile(
+                id=profileid,
+                name=browser_commit.name,
+                description=browser_commit.description,
+                created=created,
+                createdBy=created_by,
+                createdByName=created_by_name,
+                modified=now,
+                modifiedBy=user.id,
+                modifiedByName=user.name if user.name else user.email,
+                origins=json["origins"],
+                resource=profile_file,
+                userid=UUID(metadata.get("btrix.user")),
+                oid=org.id,
+                baseid=baseid,
+                crawlerChannel=browser_commit.crawlerChannel,
+                proxyId=browser_commit.proxyId,
+            )
 
-        self.orgs.can_write_data(org, include_time=False)
+            await self.profiles.find_one_and_update(
+                {"_id": profile.id}, {"$set": profile.to_dict()}, upsert=True
+            )
 
-        profile = Profile(
-            id=profileid,
-            name=browser_commit.name,
-            description=browser_commit.description,
-            created=created,
-            createdBy=created_by,
-            createdByName=created_by_name,
-            modified=now,
-            modifiedBy=user.id,
-            modifiedByName=user.name if user.name else user.email,
-            origins=json["origins"],
-            resource=profile_file,
-            userid=UUID(metadata.get("btrix.user")),
-            oid=org.id,
-            baseid=baseid,
-            crawlerChannel=browser_commit.crawlerChannel,
-            proxyId=browser_commit.proxyId,
-        )
+            await self.background_job_ops.create_replica_jobs(
+                org.id, profile_file, str(profileid), "profile"
+            )
 
-        await self.profiles.find_one_and_update(
-            {"_id": profile.id}, {"$set": profile.to_dict()}, upsert=True
-        )
+            await self.orgs.inc_org_bytes_stored(org.id, file_size, "profile")
 
-        await self.background_job_ops.create_replica_jobs(
-            org.id, profile_file, str(profileid), "profile"
-        )
+            await self.crawl_manager.keep_alive_profile_browser(
+                browser_commit.browserid, committing="done"
+            )
 
-        await self.orgs.inc_org_bytes_stored(org.id, file_size, "profile")
+        # pylint: disable=broad-except
+        except Exception as e:
+            print("Profile commit failed", e)
+            return False
 
-        return {
-            "added": True,
-            "id": str(profile.id),
-            "storageQuotaReached": self.orgs.storage_quota_reached(org),
-        }
+        return True
 
     async def update_profile_metadata(
         self, profileid: UUID, update: ProfileUpdate, user: User
@@ -432,8 +475,13 @@ async def _send_browser_req(
         path: str,
         method: str = "GET",
         json: Optional[dict[str, Any]] = None,
+        committing="",
     ) -> dict[str, Any]:
         """make request to browser api to get state"""
+        await self.crawl_manager.keep_alive_profile_browser(
+            browserid, committing=committing
+        )
+
         try:
             async with aiohttp.ClientSession() as session:
                 async with session.request(
@@ -443,7 +491,8 @@ async def _send_browser_req(
                 ) as resp:
                     json = await resp.json()
 
-        except Exception:
+        except Exception as e:
+            print(e)
             # pylint: disable=raise-missing-from
             raise HTTPException(status_code=200, detail="waiting_for_browser")
 
@@ -470,6 +519,12 @@ async def calculate_org_profile_file_storage(self, oid: UUID) -> int:
 
         return total_size
 
+    def _run_task(self, func) -> None:
+        """add bg tasks to set to avoid premature garbage collection"""
+        task = asyncio.create_task(func)
+        self.bg_tasks.add(task)
+        task.add_done_callback(self.bg_tasks.discard)
+
 
 # ============================================================================
 # pylint: disable=redefined-builtin,invalid-name,too-many-locals,too-many-arguments
@@ -529,7 +584,9 @@ async def commit_browser_to_new(
     ):
         metadata = await browser_get_metadata(browser_commit.browserid, org)
 
-        return await ops.commit_to_profile(browser_commit, org, user, metadata)
+        return await ops.commit_to_profile(
+            browser_commit=browser_commit, org=org, user=user, metadata=metadata
+        )
 
     @router.patch("/{profileid}", response_model=UpdatedResponse)
     async def commit_browser_to_existing(

backend/test/conftest.py

@@ -702,7 +702,7 @@ def profile_id(admin_auth_headers, default_org_id, profile_browser_id):
 
     # Create profile
     start_time = time.monotonic()
-    time_limit = 300
+    time_limit = 30
     while True:
         try:
             r = requests.post(
@@ -793,7 +793,7 @@ def profile_2_id(admin_auth_headers, default_org_id, profile_browser_2_id):
 
     # Create profile
     start_time = time.monotonic()
-    time_limit = 300
+    time_limit = 30
     while True:
         try:
             r = requests.post(