Fix background job email failures (#3045)

Fixes #3044 

- Use `asyncio.create_task` to send background job failure emails
(tracking tasks in set in ops class to avoid garbage collection)
- Update job in database before sending emails
- Add try/except around email sending so failures are logged instead of
throwing uncaught exception
- Rename `handle_replica_job_finished` to `handle_replica_job_succeeded`
to better reflect intent
- Fix failed background job email template and ensure that job sent to
it is JSON-serializable with `model_dump(mode="json")

Author: tw4l

backend/btrixcloud/background_jobs.py

@@ -81,17 +81,27 @@ def __init__(self, mdb, email, user_manager, org_ops, crawl_manager, storage_ops
             responses={404: {"description": "Not found"}},
         )
 
+        # to avoid background tasks being garbage collected
+        # see: https://stackoverflow.com/a/74059981
+        self.bg_tasks = set()
+
     def set_ops(self, base_crawl_ops: BaseCrawlOps, profile_ops: ProfileOps) -> None:
         """basecrawlops and profileops for updating files"""
         self.base_crawl_ops = base_crawl_ops
         self.profile_ops = profile_ops
 
+    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)
+
     def strip_bucket(self, endpoint_url: str) -> tuple[str, str]:
         """split the endpoint_url into the origin and return rest of endpoint as bucket path"""
         parts = urlsplit(endpoint_url)
         return parts.scheme + "://" + parts.netloc + "/", parts.path[1:]
 
-    async def handle_replica_job_finished(self, job: CreateReplicaJob) -> None:
+    async def handle_replica_job_succeeded(self, job: CreateReplicaJob) -> None:
         """Update replicas in corresponding file objects, based on type"""
         res = None
         if job.object_type in ("crawl", "upload"):
@@ -516,36 +526,41 @@ async def job_finished(
             raise HTTPException(status_code=400, detail="invalid_job_type")
 
         if success and job_type == BgJobType.CREATE_REPLICA:
-            await self.handle_replica_job_finished(cast(CreateReplicaJob, job))
+            await self.handle_replica_job_succeeded(cast(CreateReplicaJob, job))
 
         if job_type == BgJobType.DELETE_REPLICA:
             await self.handle_delete_replica_job_finished(cast(DeleteReplicaJob, job))
 
-        if not success:
-            await self._send_bg_job_failure_email(job, finished)
-
         await self.jobs.find_one_and_update(
             {"_id": job_id, "oid": oid},
             {"$set": {"success": success, "finished": finished}},
         )
 
+        if not success:
+            await self._send_bg_job_failure_email(job, finished)
+
     async def _send_bg_job_failure_email(self, job: BackgroundJob, finished: datetime):
         print(
             f"Background job {job.id} failed, sending email to superuser",
             flush=True,
         )
-        superuser = await self.user_manager.get_superuser()
-        org = None
-        if job.oid:
-            org = await self.org_ops.get_org_by_id(job.oid)
-        await asyncio.get_event_loop().run_in_executor(
-            None,
-            self.email.send_background_job_failed,
-            job,
-            finished,
-            superuser.email,
-            org,
-        )
+        try:
+            superuser = await self.user_manager.get_superuser()
+            org = None
+            if job.oid:
+                org = await self.org_ops.get_org_by_id(job.oid)
+
+            self._run_task(
+                self.email.send_background_job_failed(
+                    job, finished, superuser.email, org
+                )
+            )
+        # pylint: disable=broad-exception-caught
+        except Exception as err:
+            print(
+                f"Error sending bg job failure email for job {job.id}: {err}",
+                flush=True,
+            )
 
     async def get_background_job(
         self, job_id: str, oid: Optional[UUID] = None
@@ -780,11 +795,8 @@ async def retry_failed_org_background_jobs(
         Keep track of tasks in set to prevent them from being garbage collected
         See: https://stackoverflow.com/a/74059981
         """
-        bg_tasks = set()
         async for job in self.jobs.find({"oid": org.id, "success": False}):
-            task = asyncio.create_task(self.retry_background_job(job["_id"], org))
-            bg_tasks.add(task)
-            task.add_done_callback(bg_tasks.discard)
+            self._run_task(self.retry_background_job(job["_id"], org))
         return {"success": True}
 
     async def retry_all_failed_background_jobs(
@@ -795,14 +807,11 @@ async def retry_all_failed_background_jobs(
         Keep track of tasks in set to prevent them from being garbage collected
         See: https://stackoverflow.com/a/74059981
         """
-        bg_tasks = set()
         async for job in self.jobs.find({"success": False}):
             org = None
             if job.get("oid"):
                 org = await self.org_ops.get_org_by_id(job["oid"])
-            task = asyncio.create_task(self.retry_background_job(job["_id"], org))
-            bg_tasks.add(task)
-            task.add_done_callback(bg_tasks.discard)
+            self._run_task(self.retry_background_job(job["_id"], org))
         return {"success": True}
 
 

backend/btrixcloud/emailsender.py

@@ -197,7 +197,7 @@ async def send_background_job_failed(
         await self._send_encrypted(
             receiver_email,
             "failedBgJob",
-            job=job,
+            job=job.model_dump(mode="json"),
             org=str(org.id) if org else None,
             finished=finished.isoformat(),
         )

backend/btrixcloud/operator/bgjobs.py

@@ -43,7 +43,7 @@ async def finalize_background_job(self, data: MCDecoratorSyncData) -> dict:
         success = status.get("succeeded") == spec.get("parallelism")
         if not success:
             print(
-                "Succeeded: {status.get('succeeded')}, Num Pods: {spec.get('parallelism')}"
+                f"Succeeded: {status.get('succeeded')}, Num Pods: {spec.get('parallelism')}"
             )
         start_time = status.get("startTime")
         completion_time = status.get("completionTime")
@@ -75,10 +75,6 @@ async def finalize_background_job(self, data: MCDecoratorSyncData) -> dict:
                 finished=finished,
                 oid=org_id,
             )
-            # print(
-            #    f"{job_type} background job completed: success: {success}, {job_id}",
-            #    flush=True,
-            # )
 
         # pylint: disable=broad-except
         except Exception:

backend/test/test_emails.py

@@ -10,7 +10,14 @@
 import pytest
 
 from btrixcloud.emailsender import EmailSender
-from btrixcloud.models import Organization, InvitePending, EmailStr, StorageRef
+from btrixcloud.models import (
+    Organization,
+    InvitePending,
+    EmailStr,
+    StorageRef,
+    CreateReplicaJob,
+)
+from btrixcloud.utils import dt_now
 
 EMAILS_HOST_PREFIX = (
     os.environ.get("EMAIL_TEMPLATE_ENDPOINT") or "http://127.0.0.1:30872"
@@ -193,17 +200,19 @@ async def test_send_password_reset(email_sender, capsys):
 @pytest.mark.asyncio
 async def test_send_background_job_failed(email_sender, sample_org, capsys):
     """Test sending background job failure notification"""
-    # Create a mock job
-    job = {
-        "id": str(uuid.uuid4()),
-        "type": "create_replica",
-        "crawl_id": "test_crawl_123",
-        "started": datetime.now().isoformat(),
-    }
-    finished = datetime.now()
+    job = CreateReplicaJob(
+        id="fake-create-replica-job",
+        oid=sample_org.id,
+        success=False,
+        started=dt_now(),
+        file_path="path/to/file.wacz",
+        object_type="crawl",
+        object_id="sample-crawl-id",
+        replica_storage=StorageRef(name="test-storage"),
+    )
 
     await email_sender.send_background_job_failed(
-        job=job, finished=finished, receiver_email="[email protected]", org=sample_org
+        job=job, finished=dt_now(), receiver_email="[email protected]", org=sample_org
     )
 
     # Check log output

emails/emails/failed-bg-job.tsx

@@ -13,7 +13,10 @@ export const schema = z.object({
     object_type: z.string().optional(),
     object_id: z.string().optional(),
     file_path: z.string().optional(),
-    replica_storage: z.string().optional(),
+    replica_storage: z.object({
+      name: z.string(),
+      custom: z.boolean(),
+    }).optional(),
   }),
   finished: z.coerce.date(),
 });
@@ -84,9 +87,9 @@ export const FailedBgJobEmail = ({
             <Code>{job.file_path}</Code>
           </DataRow>
         )}
-        {job.replica_storage && (
+        {job.replica_storage && job.replica_storage.name && (
           <DataRow label="Replica Storage">
-            <Code>{job.replica_storage}</Code>
+            <Code>{job.replica_storage.name}</Code>
           </DataRow>
         )}
       </table>
@@ -104,7 +107,7 @@ FailedBgJobEmail.PreviewProps = {
     object_type: "object_type",
     object_id: "object_id",
     file_path: "file_path",
-    replica_storage: "replica_storage",
+    replica_storage: {name: "replica_storage", custom: false},
   },
   finished: new Date(),
 } satisfies FailedBgJobEmailProps;