99 lines
3.6 KiB
Python
99 lines
3.6 KiB
Python
"""
|
|
ThinkStorm Job Queue & Background Worker
|
|
Manages priority queues (foreground intake vs background enrichment) with bounded execution limits.
|
|
"""
|
|
|
|
import asyncio
|
|
from typing import Dict, Any, Optional, List
|
|
from ..config import config
|
|
from ..processors.pipeline import execute_intake_pipeline, execute_work_track_workflow
|
|
|
|
class ThinkStormQueue:
|
|
def __init__(self):
|
|
self._foreground_queue: Optional[asyncio.Queue] = None
|
|
self._background_queue: Optional[asyncio.Queue] = None
|
|
self._semaphore: Optional[asyncio.Semaphore] = None
|
|
self.running_jobs: Dict[str, Dict[str, Any]] = {}
|
|
self.worker_task: Optional[asyncio.Task] = None
|
|
|
|
@property
|
|
def foreground_queue(self) -> asyncio.Queue:
|
|
if self._foreground_queue is None:
|
|
self._foreground_queue = asyncio.Queue()
|
|
return self._foreground_queue
|
|
|
|
@property
|
|
def background_queue(self) -> asyncio.Queue:
|
|
if self._background_queue is None:
|
|
self._background_queue = asyncio.Queue()
|
|
return self._background_queue
|
|
|
|
@property
|
|
def semaphore(self) -> asyncio.Semaphore:
|
|
if self._semaphore is None:
|
|
self._semaphore = asyncio.Semaphore(config.max_concurrent_background_jobs)
|
|
return self._semaphore
|
|
|
|
async def enqueue_foreground(self, job_type: str, job_id: str, data: Dict[str, Any] = None):
|
|
"""Enqueues high priority job (new submissions, user activations)."""
|
|
await self.foreground_queue.put({"type": job_type, "id": job_id, "data": data or {}})
|
|
|
|
async def enqueue_background(self, job_type: str, job_id: str, data: Dict[str, Any] = None):
|
|
"""Enqueues lower priority background enrichment job."""
|
|
await self.background_queue.put({"type": job_type, "id": job_id, "data": data or {}})
|
|
|
|
async def start(self):
|
|
"""Starts worker loop."""
|
|
if self.worker_task is None:
|
|
self.worker_task = asyncio.create_task(self._worker_loop())
|
|
|
|
async def _worker_loop(self):
|
|
while True:
|
|
job = None
|
|
try:
|
|
# 1. Check foreground queue first
|
|
if not self.foreground_queue.empty():
|
|
job = await self.foreground_queue.get()
|
|
elif not self.background_queue.empty():
|
|
job = await self.background_queue.get()
|
|
else:
|
|
# Wait for any job
|
|
job = await self.foreground_queue.get()
|
|
|
|
job_key = f"{job['type']}:{job['id']}"
|
|
self.running_jobs[job_key] = {
|
|
"type": job["type"],
|
|
"id": job["id"],
|
|
"started_at": asyncio.get_event_loop().time()
|
|
}
|
|
|
|
async with self.semaphore:
|
|
await self._process_job(job)
|
|
|
|
self.running_jobs.pop(job_key, None)
|
|
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as e:
|
|
print(f"[Queue Worker Exception] {e}")
|
|
await asyncio.sleep(1.0)
|
|
|
|
async def _process_job(self, job: Dict[str, Any]):
|
|
job_type = job["type"]
|
|
job_id = job["id"]
|
|
|
|
if job_type == "intake":
|
|
await execute_intake_pipeline(job_id)
|
|
elif job_type == "work_track":
|
|
await execute_work_track_workflow(job_id)
|
|
|
|
def get_status(self) -> Dict[str, Any]:
|
|
return {
|
|
"foreground_queued": self.foreground_queue.qsize(),
|
|
"background_queued": self.background_queue.qsize(),
|
|
"active_jobs_count": len(self.running_jobs),
|
|
"running_jobs": list(self.running_jobs.values())
|
|
}
|
|
|
|
job_queue = ThinkStormQueue()
|