Merge pull request #358 from phenobarbital/dev

fix callback and complete execution on Background Service

Author: phenobarbital

navigator/background/service/__init__.py

@@ -94,10 +94,15 @@ async def record(self, task_id: uuid.UUID) -> Optional[JobRecord]:
         """
         if not task_id:
             return None
+        if isinstance(task_id, uuid.UUID):
+            task_id = str(task_id.hex)
         if isinstance(task_id, str):
-            task_id = uuid.UUID(task_id)
-        if not isinstance(task_id, uuid.UUID):
-            raise ValueError("task_id must be a UUID or a string representation of a UUID")
-        if task_id not in self.tracker._jobs:
+            task_id = uuid.UUID(task_id).hex
+        if not isinstance(task_id, str):
+            raise ValueError(
+                "task_id must be a UUID, a hex-string, or None"
+            )
+        if not await self.tracker.exists(task_id):
             return None
+
         return await self.tracker.status(task_id)

navigator/background/tracker/memory.py

@@ -11,7 +11,7 @@ class JobTracker:
     Replace by a DB or Redis backend later.
     """
     def __init__(self) -> None:
-        self._jobs: Dict[uuid.UUID, JobRecord] = {}
+        self._jobs: Dict[str, JobRecord] = {}
         self._lock = asyncio.Lock()
 
     # -----------------------------------------------------------
@@ -28,30 +28,34 @@ async def create_job(self, **kwargs) -> JobRecord:
             self._jobs[record.task_id] = record
         return record
 
-    async def set_running(self, job_id: uuid.UUID) -> None:
+    async def set_running(self, job_id: str) -> None:
         async with self._lock:
             rec = self._jobs[job_id]
             rec.status = "running"
             rec.started_at = time_now()
 
-    async def set_done(self, job_id: uuid.UUID, result: Any = None) -> None:
+    async def set_done(self, job_id: str, result: Any = None) -> None:
         async with self._lock:
             rec = self._jobs[job_id]
             rec.status = "done"
             rec.finished_at = time_now()
             rec.result = result
 
-    async def set_failed(self, job_id: uuid.UUID, exc: Exception) -> None:
+    async def set_failed(self, job_id: str, exc: Exception) -> None:
         async with self._lock:
             rec = self._jobs[job_id]
             rec.status = "failed"
             rec.finished_at = time_now()
             rec.error = f"{type(exc).__name__}: {exc}"
 
-    async def status(self, job_id: uuid.UUID) -> Optional[JobRecord]:
+    async def status(self, job_id: str) -> Optional[JobRecord]:
         async with self._lock:
             return self._jobs.get(job_id)
 
-    async def list_jobs(self) -> Dict[uuid.UUID, JobRecord]:
+    async def list_jobs(self) -> Dict[str, JobRecord]:
         async with self._lock:
             return dict(self._jobs)
+
+    async def exists(self, job_id: str) -> bool:
+        async with self._lock:
+            return job_id in self._jobs

navigator/background/tracker/models.py

@@ -4,6 +4,9 @@
 from datetime import datetime
 from datamodel import BaseModel, Field
 
+def gen_uuid() -> str:
+    """Generate a new UUID."""
+    return str(uuid.uuid4().hex)
 
 def time_now() -> int:
     """Get the current time in milliseconds."""
@@ -14,7 +17,7 @@ class JobRecord(BaseModel):
 
     Job Record for Background Task Execution.
     """
-    task_id: str = Field(default=str(uuid.uuid4().hex))
+    task_id: str = Field(default=gen_uuid)
     name: str = None
     status: str = 'pending'
     attributes: Dict[str, Any] = Field(default_factory=dict)

navigator/background/tracker/redis.py

@@ -44,6 +44,9 @@ async def create_job(self, **kwargs) -> JobRecord:
             await self._redis.sadd(self._set_key, record.task_id)
         return record
 
+    async def exists(self, job_id: str) -> bool:
+        return await self._redis.exists(self._key(job_id)) == 1
+
     async def _update(self, job_id: str, **patch) -> None:
         """
         Update a job record with the given patch.

navigator/background/wrappers/__init__.py

@@ -9,24 +9,41 @@
 
 
 coroutine = Callable[[int], Coroutine[Any, Any, str]]
+OnCompleteFn = Callable[[Any, Optional[Exception]], Awaitable[None]]
 
 
-def coroutine_in_thread(coro: coroutine, callback: Optional[coroutine] = None):
+def coroutine_in_thread(
+    coro: coroutine,
+    callback: Optional[coroutine] = None,
+    on_complete: OnCompleteFn = None,
+) -> threading.Event:
     """Run a coroutine in a new thread with its own event loop."""
+    parent_loop = asyncio.get_running_loop()
     done_event = threading.Event()
 
-    def run():
+    def _runner():
         new_loop = asyncio.new_event_loop()
         asyncio.set_event_loop(new_loop)
-        result = new_loop.run_until_complete(coro)
-        # if callback exists:
-        if callback:
-            new_loop.run_until_complete(callback(result))
-        new_loop.close()
-        done_event.set()  # Signal that the coroutine has completed
-    thread = threading.Thread(target=run, daemon=True)
-    thread.start()
+        result, exc = None, None
+        try:
+            result = new_loop.run_until_complete(coro)
+        except Exception as e:         # noqa: BLE001
+            exc = e
+        finally:
+            if callback:
+                new_loop.run_until_complete(
+                    callback(result, exc, loop=new_loop)
+                )
+            new_loop.close()
+            done_event.set()  # Signal that the coroutine has completed
+            if on_complete is None:
+                return
+            fut = asyncio.run_coroutine_threadsafe(
+                on_complete(result, exc), parent_loop
+            )
+            fut.result()  # Wait for the completion of the callback
 
+    threading.Thread(target=_runner, daemon=True).start()
     return done_event
 
 
@@ -115,16 +132,33 @@ async def __call__(self):
             # Delay the execution by jitter seconds
             await asyncio.sleep(delay)
         try:
+            async def _finish(result: Any, exc: Exception):
+                """Callback to handle the completion of the coroutine."""
+                if exc:
+                    self.logger.error(
+                        f"TaskWrapper {self._name} failed with exception: {exc}"
+                    )
+                    result = {
+                        "status": "failed",
+                        "error": str(exc)
+                    }
+                    if self.tracker:
+                        await self.tracker.set_failed(self.task_uuid, exc)
+                else:
+                    self.logger.debug(
+                        f"TaskWrapper {self._name} completed successfully."
+                    )
+                    result = {
+                        "status": "done",
+                        "result": result
+                    }
+                    if self.tracker:
+                        await self.tracker.set_done(self.task_uuid, result)
+                return result
             with ThreadPoolExecutor(max_workers=1) as executor:
                 coro = self.fn(*self.args, **self.kwargs)
-                coroutine_in_thread(coro, self._callback_)
-                result = {
-                    "status": "done"
-                }
-                if self.tracker:
-                    # Set the job as done in the tracker
-                    await self.tracker.set_done(self.task_uuid, result)
-                return result
+                coroutine_in_thread(coro, self._callback_, on_complete=_finish)
+                return {"status": "running"}
         except asyncio.CancelledError:
             self.logger.warning(
                 f"TaskWrapper {self.fn.__name__} was cancelled."

navigator/version.py

@@ -4,7 +4,7 @@
 __description__ = (
     "Navigator Web Framework based on aiohttp, " "with batteries included."
 )
-__version__ = "2.12.29"
+__version__ = "2.12.30"
 __copyright__ = "Copyright (c) 2020-2024 Jesus Lara"
 __author__ = "Jesus Lara"
 __author_email__ = "jesuslarag@gmail.com"