From 854f346fd11f8f080b7534870918288c0aa50375 Mon Sep 17 00:00:00 2001 From: Joao Mario Lago Date: Wed, 22 Jul 2026 11:03:07 -0300 Subject: [PATCH 1/2] core:services:kraken: Allow consuming jobs results as stream * Allow users to consume jobs result as streams once they are executing --- core/services/kraken/api/v2/routers/jobs.py | 11 ++++- core/services/kraken/jobs/exceptions.py | 4 ++ core/services/kraken/jobs/jobs.py | 47 +++++++++++++++---- core/services/kraken/jobs/models.py | 8 ++++ core/services/kraken/jobs/stream.py | 52 +++++++++++++++++++++ 5 files changed, 111 insertions(+), 11 deletions(-) create mode 100644 core/services/kraken/jobs/stream.py diff --git a/core/services/kraken/api/v2/routers/jobs.py b/core/services/kraken/api/v2/routers/jobs.py index af20a2b329..7a760fbfef 100644 --- a/core/services/kraken/api/v2/routers/jobs.py +++ b/core/services/kraken/api/v2/routers/jobs.py @@ -3,9 +3,10 @@ from typing import Any, Callable, List, Tuple from fastapi import APIRouter, Body, HTTPException, status +from fastapi.responses import StreamingResponse from fastapi_versioning import versioned_api_route from jobs import JobsManager -from jobs.exceptions import JobNotFound +from jobs.exceptions import JobIsRunning, JobNotFound from jobs.models import Job, JobMethod jobs_router_v2 = APIRouter( @@ -21,6 +22,8 @@ def jobs_to_http_exception(endpoint: Callable[..., Any]) -> Callable[..., Any]: async def wrapper(*args: Tuple[Any], **kwargs: dict[str, Any]) -> Any: try: return await endpoint(*args, **kwargs) + except JobIsRunning as error: + raise HTTPException(status_code=status.HTTP_423_LOCKED, detail=str(error)) from error except JobNotFound as error: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(error)) from error except Exception as error: @@ -45,6 +48,12 @@ async def fetch() -> List[Job]: return JobsManager.get() +@jobs_router_v2.get("/{identifier}/stream", status_code=status.HTTP_200_OK) +@jobs_to_http_exception +async def stream(identifier: str) -> StreamingResponse: + return StreamingResponse(JobsManager.stream(identifier), media_type="text/plain") + + @jobs_router_v2.get("/{identifier}", status_code=status.HTTP_200_OK) @jobs_to_http_exception async def fetch_by_identifier(identifier: str) -> Job: diff --git a/core/services/kraken/jobs/exceptions.py b/core/services/kraken/jobs/exceptions.py index 32fe30fd80..a174a0a799 100644 --- a/core/services/kraken/jobs/exceptions.py +++ b/core/services/kraken/jobs/exceptions.py @@ -1,2 +1,6 @@ class JobNotFound(Exception): pass + + +class JobIsRunning(Exception): + pass diff --git a/core/services/kraken/jobs/jobs.py b/core/services/kraken/jobs/jobs.py index d3ff8d11eb..ca2b4d69ca 100644 --- a/core/services/kraken/jobs/jobs.py +++ b/core/services/kraken/jobs/jobs.py @@ -1,15 +1,17 @@ import asyncio -from typing import List, Optional +from typing import AsyncGenerator, Dict, List, Optional import aiohttp -from jobs.exceptions import JobNotFound -from jobs.models import Job +from jobs.exceptions import JobIsRunning, JobNotFound +from jobs.models import Job, JobStatus +from jobs.stream import JobStream from loguru import logger class JobsManager: _jobs: List[Job] = [] _executing_job: Optional[Job] = None + _streams: Dict[str, JobStream] = {} def __init__(self) -> None: self.is_running = True @@ -18,6 +20,7 @@ def __init__(self) -> None: async def execute_job(self, job: Job) -> None: job_name = f"{job.method.value} - {job.route}" logger.info(f"Executing job {job_name}") + stream = self._streams[job.id] async with aiohttp.ClientSession() as session: for i in range(job.retries): try: @@ -25,20 +28,31 @@ async def execute_job(self, job: Job) -> None: method=job.method, url=f"{self.base_host}/{job.route}", json=job.body ) as response: response.raise_for_status() - await response.read() + async for chunk in response.content.iter_any(): + stream.publish(chunk) + job.status = JobStatus.SUCCESS return except Exception: logger.warning(f"Failed job {job_name} attempt {i + 1}/{job.retries}") + stream.reset() await asyncio.sleep(5) + job.status = JobStatus.ERROR logger.error(f"Job {job_name} failed to be executed") async def start(self) -> None: while self.is_running: await asyncio.sleep(1) - if self._jobs: - self._executing_job = self._jobs.pop(0) - await self.execute_job(self._executing_job) - self._executing_job = None + if JobsManager._jobs: + JobsManager._executing_job = JobsManager._jobs.pop(0) + JobsManager._executing_job.status = JobStatus.RUNNING + stream = JobsManager._streams[JobsManager._executing_job.id] + try: + await self.execute_job(JobsManager._executing_job) + finally: + stream.close() + await stream.wait_drained() + JobsManager._streams.pop(JobsManager._executing_job.id, None) + JobsManager._executing_job = None async def stop(self) -> None: self.is_running = False @@ -48,6 +62,7 @@ def set_base_host(self, host: str) -> None: @classmethod def add(cls, job: Job) -> None: + cls._streams[job.id] = JobStream() cls._jobs.append(job) @classmethod @@ -61,7 +76,19 @@ def get_by_identifier(cls, identifier: str) -> Job: raise JobNotFound(f"Job with id {identifier} not found") return job + @classmethod + def stream(cls, identifier: str) -> AsyncGenerator[bytes, None]: + cls.get_by_identifier(identifier) + return cls._streams[identifier].subscribe() + @classmethod def delete(cls, identifier: str) -> None: - job = cls.get_by_identifier(identifier) - cls._jobs.remove(job) + cls.get_by_identifier(identifier) + job = next((job for job in cls._jobs if job.id == identifier), None) + if job: + cls._jobs.remove(job) + stream = cls._streams.pop(identifier, None) + if stream: + stream.close() + elif cls._executing_job and cls._executing_job.id == identifier: + raise JobIsRunning(f"Cannot delete job {identifier} while it is running") diff --git a/core/services/kraken/jobs/models.py b/core/services/kraken/jobs/models.py index f5d72d46bf..d655f52b55 100644 --- a/core/services/kraken/jobs/models.py +++ b/core/services/kraken/jobs/models.py @@ -13,9 +13,17 @@ class JobMethod(str, Enum): DELETE = "DELETE" +class JobStatus(str, Enum): + PENDING = "PENDING" + RUNNING = "RUNNING" + SUCCESS = "SUCCESS" + ERROR = "ERROR" + + class Job(BaseModel): id: str route: str method: JobMethod body: Any retries: int = 5 + status: JobStatus = JobStatus.PENDING diff --git a/core/services/kraken/jobs/stream.py b/core/services/kraken/jobs/stream.py new file mode 100644 index 0000000000..47e8ed4ef4 --- /dev/null +++ b/core/services/kraken/jobs/stream.py @@ -0,0 +1,52 @@ +import asyncio +from typing import AsyncGenerator, List + + +class JobStream: + def __init__(self) -> None: + self._chunks: List[bytes] = [] + self._updated = asyncio.Event() + self._detached = asyncio.Event() + self._done = False + self._subscribers = 0 + self._retries = 0 + + def publish(self, chunk: bytes) -> None: + self._chunks.append(chunk) + self._updated.set() + + def reset(self) -> None: + self._chunks = [] + self._retries += 1 + self._updated.set() + + def close(self) -> None: + self._done = True + self._updated.set() + + async def wait_drained(self) -> None: + while True: + self._detached.clear() + if self._subscribers == 0: + return + await self._detached.wait() + + async def subscribe(self) -> AsyncGenerator[bytes, None]: + self._subscribers += 1 + try: + index = 0 + retries = self._retries + while True: + self._updated.clear() + if retries != self._retries: + retries = self._retries + index = 0 + while index < len(self._chunks): + yield self._chunks[index] + index += 1 + if self._done: + return + await self._updated.wait() + finally: + self._subscribers -= 1 + self._detached.set() From c794e691cb4b90b1c53d4e4289b1a64e68557c3c Mon Sep 17 00:00:00 2001 From: Joao Mario Lago Date: Wed, 22 Jul 2026 12:04:36 -0300 Subject: [PATCH 2/2] WIP: Testing not tracking references --- core/services/kraken/jobs/jobs.py | 1 - core/services/kraken/jobs/stream.py | 38 +++++++++-------------------- 2 files changed, 12 insertions(+), 27 deletions(-) diff --git a/core/services/kraken/jobs/jobs.py b/core/services/kraken/jobs/jobs.py index ca2b4d69ca..268e68d36c 100644 --- a/core/services/kraken/jobs/jobs.py +++ b/core/services/kraken/jobs/jobs.py @@ -50,7 +50,6 @@ async def start(self) -> None: await self.execute_job(JobsManager._executing_job) finally: stream.close() - await stream.wait_drained() JobsManager._streams.pop(JobsManager._executing_job.id, None) JobsManager._executing_job = None diff --git a/core/services/kraken/jobs/stream.py b/core/services/kraken/jobs/stream.py index 47e8ed4ef4..fc7c688c9f 100644 --- a/core/services/kraken/jobs/stream.py +++ b/core/services/kraken/jobs/stream.py @@ -6,9 +6,7 @@ class JobStream: def __init__(self) -> None: self._chunks: List[bytes] = [] self._updated = asyncio.Event() - self._detached = asyncio.Event() self._done = False - self._subscribers = 0 self._retries = 0 def publish(self, chunk: bytes) -> None: @@ -24,29 +22,17 @@ def close(self) -> None: self._done = True self._updated.set() - async def wait_drained(self) -> None: + async def subscribe(self) -> AsyncGenerator[bytes, None]: + index = 0 + retries = self._retries while True: - self._detached.clear() - if self._subscribers == 0: + self._updated.clear() + if retries != self._retries: + retries = self._retries + index = 0 + while index < len(self._chunks): + yield self._chunks[index] + index += 1 + if self._done: return - await self._detached.wait() - - async def subscribe(self) -> AsyncGenerator[bytes, None]: - self._subscribers += 1 - try: - index = 0 - retries = self._retries - while True: - self._updated.clear() - if retries != self._retries: - retries = self._retries - index = 0 - while index < len(self._chunks): - yield self._chunks[index] - index += 1 - if self._done: - return - await self._updated.wait() - finally: - self._subscribers -= 1 - self._detached.set() + await self._updated.wait()