Merge pull request #357 from phenobarbital/dev

redis support on background service

Author: phenobarbital

navigator/background/__init__.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-from .tracker import JobTracker, JobRecord
+from .tracker import JobTracker, RedisJobTracker, JobRecord
 from .wrappers import TaskWrapper
 from .service import BackgroundService
 from .queue import BackgroundQueue, BackgroundTask, SERVICE_NAME

navigator/background/queue/__init__.py

@@ -1,3 +1,4 @@
+import random
 from typing import (
     Union,
     Optional,
@@ -38,8 +39,8 @@ class BackgroundQueue:
 
     TODO:
     - Add Task Timeout
-    - Add Task Retry
-    - Added Wrapper Support
+    - Add Task Retry (done)
+    - Added Wrapper Support (done)
     """
     service_name: str = SERVICE_NAME
 
@@ -219,7 +220,6 @@ async def _execute_callable(self, func: Callable, *args, **kwargs):
         result = None
         try:
             loop = asyncio.get_running_loop()
-            # with ThreadPoolExecutor(max_workers=1) as executor:
             result = await loop.run_in_executor(
                 self.executor,
                 func,
@@ -281,11 +281,8 @@ async def process_queue(self):
                             f"Invalid Function {func} in Queue"
                         )
                         continue
-            except Exception as e:  # Catch all exceptions
-                print('ERROR > ', e)
-                self.logger.error(
-                    f"Error executing task {func.__name__}: {e}"
-                )
+            except Exception as exc:  # Catch all exceptions
+                await self._handle_failure(task, exc)
                 continue
             finally:
                 if self._enable_profiling is True:
@@ -305,7 +302,7 @@ async def process_queue(self):
                         Peak Memory Usage: {peak_memory / (1024 ** 2):.2f} MB
                     """)
                 except Exception as e:
-                    print('LOG ERROR > ', e)
+                    print('TASK LOG ERROR > ', e)
                 # Call your task completion callback (if any)
                 try:
                     await self._callback(task, result=result)
@@ -330,6 +327,41 @@ async def fire_consumers(self):
             )
             self.consumers.append(task)
 
+    async def _requeue(self, task: TaskWrapper, exc: Exception) -> None:
+        """Internal: re-enqueues `task` after updating retry-counters."""
+        task.retries_done += 1
+        # optional exponential back-off
+        if task.retry_delay:
+            await asyncio.sleep(
+                random.uniform(0.8, 1.2) * task.retry_delay * task.retries_done
+            )
+
+        self.logger.warning(
+            f"Retry {task.retries_done}/{task.max_retries} for {task!r} "
+            f"after error: {exc}"
+        )
+        if hasattr(task, "tracker"):
+            # set status = "retrying"
+            await task.tracker.set_running(task.task_uuid)
+        await self.queue.put(task)
+
+    async def _handle_failure(
+        self,
+        task: Any,
+        exc: Exception
+    ) -> None:
+        """Central place that decides whether we retry or finally give up."""
+        if (
+            isinstance(task, TaskWrapper) and task.retries_done < task.max_retries
+        ):
+            await self._requeue(task, exc)
+        else:
+            self.logger.error(
+                f"Task {task!r} failed permanently after "
+                f"{getattr(task, 'retries_done', 0)} attempt(s)."
+            )
+            await self._callback(task, result=dict(status="failed", error=exc))
+
 
 class BackgroundTask:
     """BackgroundTask.

navigator/background/service/__init__.py

@@ -2,8 +2,9 @@
 import uuid
 from aiohttp import web
 from ..queue import BackgroundQueue
-from ..tracker import JobTracker, JobRecord
+from ..tracker import JobTracker, RedisJobTracker, JobRecord
 from ..wrappers import TaskWrapper
+from ...conf import CACHE_URL
 
 
 class BackgroundService:
@@ -16,11 +17,20 @@ def __init__(
         app: web.Application,
         queue: Optional[BackgroundQueue] = None,
         tracker: Optional[JobTracker] = None,
+        tracker_type: str = 'memory',
         **kwargs
     ) -> None:
         self.queue = queue or BackgroundQueue(app, **kwargs)
         # Create a new JobTracker if not provided
-        self.tracker = tracker or JobTracker()
+        self.tracker = tracker
+        if not tracker:
+            if tracker_type == 'redis':
+                self.tracker = RedisJobTracker(
+                    url=kwargs.get('redis_url', CACHE_URL),
+                    prefix=kwargs.get('tracker_prefix', 'job:')
+                )
+            else:
+                self.tracker = JobTracker()
         # Register the queue and tracker in the application
         app['background_service'] = self
         app['service_tracker'] = self.tracker

navigator/background/tracker/__init__.py

@@ -1,82 +1,11 @@
-from typing import Dict, Any, Optional
-import asyncio
-import uuid
-import time
-from datetime import datetime
-from datamodel import BaseModel, Field
+from .models import JobRecord, time_now
+from .memory import JobTracker
+from .redis import RedisJobTracker
 
 
-def time_now() -> int:
-    """Get the current time in milliseconds."""
-    return int(time.time() * 1000)
-
-class JobRecord(BaseModel):
-    """JobRecord.
-
-    Job Record for Background Task Execution.
-    """
-    task_id: str = Field(default=uuid.uuid4().hex)
-    name: str = None
-    status: str = 'pending'
-    attributes: Dict[str, Any] = Field(default_factory=dict)
-    result: Optional[Any] = None
-    error: Optional[str] = None
-    created_at: datetime = Field(default=datetime.now())
-    started_at: Optional[int]
-    finished_at: Optional[int]
-    result: Any = None
-    error: Optional[str] = None
-    stacktrace: Optional[str] = None
-
-    class Meta:
-        strict = True
-
-    def __repr__(self):
-        return f"<JobRecord {self.name} ({self.id})>"
-
-
-class JobTracker:
-    """
-    A very small, coroutine-safe in-memory job store.
-    Replace by a DB or Redis backend later.
-    """
-    def __init__(self) -> None:
-        self._jobs: Dict[uuid.UUID, JobRecord] = {}
-        self._lock = asyncio.Lock()
-
-    # -----------------------------------------------------------
-    # Public helpers
-    # -----------------------------------------------------------
-    async def create_job(self, **kwargs) -> JobRecord:
-        record = JobRecord(**kwargs)
-        async with self._lock:
-            self._jobs[record.task_id] = record
-        return record
-
-    async def set_running(self, job_id: uuid.UUID) -> 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 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 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 with self._lock:
-            return self._jobs.get(job_id)
-
-    async def list_jobs(self) -> Dict[uuid.UUID, JobRecord]:
-        async with self._lock:
-            return dict(self._jobs)
+__all__ = (
+    'JobTracker',
+    'RedisJobTracker',
+    'JobRecord',
+    'time_now',
+)

navigator/background/tracker/memory.py

@@ -0,0 +1,57 @@
+from typing import Dict, Any, Optional
+import asyncio
+import uuid
+from datamodel.exceptions import ValidationError
+from .models import JobRecord, time_now
+
+
+class JobTracker:
+    """
+    A very small, coroutine-safe in-memory job store.
+    Replace by a DB or Redis backend later.
+    """
+    def __init__(self) -> None:
+        self._jobs: Dict[uuid.UUID, JobRecord] = {}
+        self._lock = asyncio.Lock()
+
+    # -----------------------------------------------------------
+    # Public helpers
+    # -----------------------------------------------------------
+    async def create_job(self, **kwargs) -> JobRecord:
+        try:
+            record = JobRecord(**kwargs)
+        except ValidationError as exc:
+            raise ValueError(
+                f"Invalid job record data: {exc}, payload: {exc.payload}"
+            ) from exc
+        async with self._lock:
+            self._jobs[record.task_id] = record
+        return record
+
+    async def set_running(self, job_id: uuid.UUID) -> 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 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 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 with self._lock:
+            return self._jobs.get(job_id)
+
+    async def list_jobs(self) -> Dict[uuid.UUID, JobRecord]:
+        async with self._lock:
+            return dict(self._jobs)

navigator/background/tracker/models.py

@@ -0,0 +1,34 @@
+from typing import Dict, Any, Optional
+import uuid
+import time
+from datetime import datetime
+from datamodel import BaseModel, Field
+
+
+def time_now() -> int:
+    """Get the current time in milliseconds."""
+    return int(time.time() * 1000)
+
+class JobRecord(BaseModel):
+    """JobRecord.
+
+    Job Record for Background Task Execution.
+    """
+    task_id: str = Field(default=str(uuid.uuid4().hex))
+    name: str = None
+    status: str = 'pending'
+    attributes: Dict[str, Any] = Field(default_factory=dict)
+    result: Optional[Any] = None
+    error: Optional[str] = None
+    created_at: datetime = Field(default=datetime.now())
+    started_at: Optional[int]
+    finished_at: Optional[int]
+    result: Any = None
+    error: Optional[str] = None
+    stacktrace: Optional[str] = None
+
+    class Meta:
+        strict = True
+
+    def __repr__(self):
+        return f"<JobRecord {self.name} ({self.id})>"