broccoli-workers 0.1.0__py3-none-any.whl

broccoli/core/chain/chain_mixin.py

@@ -0,0 +1,49 @@
+# video_scheduler/core/chain_mixin.py
+from typing import Any
+
+from broccoli.core.task.task import Task
+
+
+class ChainWorkerMixin:
+    """Mixin to add chain support to a worker."""
+
+    def post_process(self, task: Task, success: bool) -> None:
+        """Override this in your worker and call super().post_process()."""
+        chain_id = task.payload.get("__chain_id")
+
+        if chain_id:
+            # Delete the chain-prefixed task key now that it's done
+            self._redis.delete(f"chain:{task.task_id}")
+
+            # Completion tasks have __chain_id set (so they get cleaned up above)
+            # but must not call continue_chain — they are not sequential chain steps
+            is_completion_task = task.payload.get("__chain_position") is None
+            if is_completion_task:
+                # Run completion handlers for the chain
+                if success:
+                    self._run_completion_handlers(task, task.result)
+                else:
+                    self._run_failure_handlers(
+                        task,
+                        Exception(task.error)
+                        if task.error
+                        else Exception("Chain task failed"),
+                    )
+                return
+
+            if success:
+                from broccoli.core.chain.task_chain import TaskChain
+
+                chain = TaskChain()
+                finished = chain.continue_chain(
+                    task, task.result, push_completion_task=True
+                )
+                if finished:
+                    self.on_finish(chain_id, task.result)
+
+                # Run completion handlers for chain steps
+                self._run_completion_handlers(task, task.result)
+
+    def on_finish(self, chain_id: str, final_result: Any) -> None:
+        """Hook called when the entire chain is finished. Override in your worker."""
+        pass  # Your existing on_finish logic here

broccoli/core/chain/chain_queue.py

@@ -0,0 +1,299 @@
+# broccoli/core/chain/chain_queue.py
+import logging
+import time
+from typing import Optional
+
+from broccoli.core.redis_controller import RedisController
+from broccoli.core.task.task import Task
+
+logger = logging.getLogger(__name__)
+
+# Max tasks per priority tier before FIFO scores would collide.
+# 10^12 tasks per tier is effectively unlimited in practice.
+_FIFO_TIER_SIZE = 1_000_000_000_000
+
+# Redis field used to carry a task's original priority through the waiting state
+# so it can be restored when the dependency is released.
+_PRIORITY_FIELD = "_broccoli_priority"
+
+
+class ChainQueue:
+    """
+    Priority queue backed by three Redis structures:
+
+    ``<base>:queue``       — sorted set of runnable task IDs
+                             score = priority_tier * _FIFO_TIER_SIZE + monotonic_seq
+    ``<base>:processing``  — sorted set of in-flight task IDs
+                             score = Unix timestamp of when the task was popped
+                             (used for crash recovery, not for ordering)
+    ``<base>:sequence``    — monotonic counter for FIFO ordering within a priority tier
+    ``dependency:<id>``    — set of task IDs waiting for task <id> to complete
+
+    Dependency resolution happens at *push time* and *completion time*, never
+    inside ``pop()``.  Workers only ever see runnable tasks.
+    """
+
+    def __init__(
+        self,
+        redis_url: str = "redis://localhost:6379",
+        queue_name="chain_tasks:queue",
+        task_prefix="chain",
+    ):
+        self.redis_url = redis_url
+        self._redis = RedisController(redis_url).get_client()
+        self.queue_key = queue_name
+        self.task_prefix = task_prefix
+
+        # Derive sibling keys from queue_name so all keys share a namespace.
+        # e.g. "tasks:queue" → "tasks:processing", "tasks:sequence"
+        base = (
+            queue_name[: -len(":queue")]
+            if queue_name.endswith(":queue")
+            else queue_name
+        )
+        self.processing_key = f"{base}:processing"
+        self.sequence_key = f"{base}:sequence"
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def push(self, task: Task, priority: int = 0) -> str:
+        """
+        Persist task metadata then either enqueue it or register it as waiting.
+
+        If the task has no dependency, or its dependency is already completed,
+        it goes straight into the runnable queue.  Otherwise it is stored as
+        ``waiting`` and added to ``dependency:<parent_id>`` so ``complete()``
+        can release it later.
+
+        Race condition handled with a register-then-verify pattern:
+          1. Register in the dependency set first.
+          2. Re-check the parent's status.
+          3. If parent is already completed (raced past us), release immediately.
+        This eliminates the TOCTOU window present in a check-then-register approach.
+        """
+        # Always persist task metadata first so workers can read it after pop().
+        self._redis.hset(f"{self.task_prefix}:{task.task_id}", mapping=task.to_dict())
+
+        if task.depends_on:
+            # Mark as waiting immediately and register in the dependency set.
+            # We do this BEFORE reading the parent's status to close the
+            # check-then-register race window.
+            task.status = "waiting"
+            pipe = self._redis.pipeline()
+            pipe.hset(f"{self.task_prefix}:{task.task_id}", "status", "waiting")
+            # Store the requested priority so complete() can restore it.
+            pipe.hset(
+                f"{self.task_prefix}:{task.task_id}", _PRIORITY_FIELD, str(priority)
+            )
+            pipe.sadd(f"dependency:{task.depends_on}", task.task_id)
+            pipe.execute()
+
+            # Now verify: if the parent already completed before our sadd, release now.
+            dep_status = self._redis.hget(
+                f"{self.task_prefix}:{task.depends_on}", "status"
+            )
+            parent_done = (
+                dep_status
+                and (
+                    dep_status.decode() if isinstance(dep_status, bytes) else dep_status
+                )
+                == "completed"
+            )
+
+            if parent_done:
+                # Parent completed before (or concurrently with) our registration.
+                # Release ourselves immediately and clean up.
+                self._redis.srem(f"dependency:{task.depends_on}", task.task_id)
+                self._enqueue(task.task_id, priority)
+                logger.debug(
+                    f"Task {task.task_id}: dependency {task.depends_on} already "
+                    "completed, enqueued immediately"
+                )
+            else:
+                logger.debug(
+                    f"Task {task.task_id} waiting on dependency {task.depends_on}"
+                )
+        else:
+            self._enqueue(task.task_id, priority)
+
+        return task.task_id
+
+    def pop(self) -> Optional[Task]:
+        """
+        Pop the highest-priority runnable task and move it to the processing set.
+
+        No dependency checking happens here — that was handled at push / complete
+        time.  Everything in the queue is ready to run.
+
+        Returns ``None`` if the queue is empty after the blocking timeout.
+        """
+        result = self._redis.bzpopmin(self.queue_key, timeout=1)
+        if result is None:
+            return None
+
+        _, task_id, _fifo_score = result
+        task_id = task_id.decode() if isinstance(task_id, bytes) else task_id
+
+        task_data = self._redis.hgetall(f"{self.task_prefix}:{task_id}")
+        if not task_data:
+            logger.error(
+                f"Task data missing for {task_id!r}; moving to dead-letter set"
+            )
+            self._redis.sadd(f"{self.task_prefix}:dead_letter", task_id)
+            return None
+
+        # Score in the processing set is a Unix timestamp so recover_stalled()
+        # can identify tasks that have been in-flight longer than expected.
+        now = time.time()
+        self._redis.zadd(self.processing_key, {task_id: now})
+        self._redis.hset(f"{self.task_prefix}:{task_id}", "status", "in_progress")
+
+        task = Task.from_dict(task_data)
+        task.status = "in_progress"
+        return task
+
+    def complete(self, task: Task) -> None:
+        """
+        Called after a task finishes successfully.
+
+        1. Removes the task from the processing set.
+        2. Reads any tasks that were waiting on this one and enqueues them,
+           restoring their original priority.
+        3. Deletes the dependency set.
+        """
+        task_id = task.task_id
+        self._redis.zrem(self.processing_key, task_id)
+
+        dep_key = f"dependency:{task_id}"
+        waiting_ids = self._redis.smembers(dep_key)
+        if waiting_ids:
+            for raw_id in waiting_ids:
+                wid = raw_id.decode() if isinstance(raw_id, bytes) else raw_id
+                # Restore the priority the caller supplied when pushing the waiting task.
+                raw_pri = self._redis.hget(f"{self.task_prefix}:{wid}", _PRIORITY_FIELD)
+                priority = int(raw_pri) if raw_pri else 0
+                self._enqueue(wid, priority)
+                logger.info(
+                    f"Released waiting task {wid} "
+                    f"(dependency {task_id} completed, priority={priority})"
+                )
+            self._redis.delete(dep_key)
+
+    def fail(self, task: Task) -> None:
+        """Remove a permanently failed (or skipped) task from the processing set."""
+        self._redis.zrem(self.processing_key, task.task_id)
+
+    def requeue(self, task_id: str, priority: Optional[int] = None) -> None:
+        """
+        Move a task from the processing set back to the runnable queue (retry).
+
+        A fresh FIFO sequence number is issued so the retry goes to the back of
+        its priority tier rather than jumping the queue.
+
+        If ``priority`` isn't given explicitly, the task's originally-requested
+        priority (stashed under ``_PRIORITY_FIELD`` at push time) is restored,
+        matching the behaviour of ``complete()`` when it releases dependents.
+        Without this, every retry would silently reset to priority 0 and could
+        jump ahead of lower-priority tasks still waiting in the queue.
+        """
+        self._redis.zrem(self.processing_key, task_id)
+        if priority is None:
+            raw_pri = self._redis.hget(f"{self.task_prefix}:{task_id}", _PRIORITY_FIELD)
+            priority = int(raw_pri) if raw_pri else 0
+        self._enqueue(task_id, priority)
+
+    def recover_stalled(self, timeout_seconds: int = 3600) -> int:
+        """
+        Crash-recovery helper: re-enqueue tasks that have been in the processing
+        set for longer than ``timeout_seconds``.
+
+        The processing set uses Unix timestamps as scores (set by ``pop()``),
+        so the threshold is simply ``time.time() - timeout_seconds``.
+
+        Returns the number of tasks recovered.
+        """
+        threshold = time.time() - timeout_seconds
+        stalled = self._redis.zrangebyscore(self.processing_key, "-inf", threshold)
+        for raw_id in stalled:
+            tid = raw_id.decode() if isinstance(raw_id, bytes) else raw_id
+            self.requeue(tid)
+            logger.warning(f"Crash-recovery: re-enqueued stalled task {tid}")
+        return len(stalled)
+
+    # ------------------------------------------------------------------
+    # Convenience helpers
+    # ------------------------------------------------------------------
+
+    def get_task(self, task_id: str) -> Optional[Task]:
+        """Fetch task metadata by ID, or ``None`` if it doesn't exist."""
+        task_data = self._redis.hgetall(f"{self.task_prefix}:{task_id}")
+        if not task_data:
+            return None
+        return Task.from_dict(task_data)
+
+    def get_queue_name(self) -> str:
+        return self.queue_key
+
+    def is_empty(self) -> bool:
+        return self._redis.zcard(self.queue_key) == 0
+
+    def pop_with_timeout(self, timeout: int = 1) -> Optional[Task]:
+        """Alias for ``pop()``; timeout is already handled internally."""
+        return self.pop()
+
+    # ------------------------------------------------------------------
+    # Private helpers
+    # ------------------------------------------------------------------
+
+    def _enqueue(self, task_id: str, priority: int = 0) -> None:
+        """
+        Add *task_id* to the runnable sorted set.
+
+        Score = priority * _FIFO_TIER_SIZE + monotonic_seq
+
+        This gives strict priority ordering across tiers *and* FIFO within a
+        tier, without requiring timestamps or Lua scripts.
+        """
+        seq = self._redis.incr(self.sequence_key)
+        score = priority * _FIFO_TIER_SIZE + seq
+        self._redis.zadd(self.queue_key, {task_id: score})
+        self._redis.hset(
+            f"{self.task_prefix}:{task_id}",
+            mapping={"status": "pending", _PRIORITY_FIELD: str(priority)},
+        )
+        logger.debug(f"Enqueued task {task_id} (priority={priority}, seq={seq})")
+
+    def is_fully_drained(self) -> bool:
+        """True only when both the runnable queue and processing set are empty."""
+        pipe = self._redis.pipeline()
+        pipe.zcard(self.queue_key)
+        pipe.zcard(self.processing_key)
+        runnable, processing = pipe.execute()
+        return runnable == 0 and processing == 0
+
+    def get_waiting_for(self, task_id: str) -> list:
+        """Return the IDs of tasks currently blocked on ``task_id``."""
+        return [
+            tid.decode() if isinstance(tid, bytes) else tid
+            for tid in self._redis.smembers(f"dependency:{task_id}")
+        ]
+
+    def stats(self) -> dict:
+        """
+        Snapshot of queue depth for monitoring/diagnostics.
+
+        ``dead_letter`` reflects tasks whose hash went missing on pop() —
+        see ``pop()`` for how they get there.
+        """
+        pipe = self._redis.pipeline()
+        pipe.zcard(self.queue_key)
+        pipe.zcard(self.processing_key)
+        pipe.scard(f"{self.task_prefix}:dead_letter")
+        runnable, processing, dead_letter = pipe.execute()
+        return {
+            "runnable": runnable,
+            "processing": processing,
+            "dead_letter": dead_letter,
+        }

broccoli/core/chain/task_chain.py

@@ -0,0 +1,201 @@
+# video_scheduler/core/task_chain.py
+import json
+import logging
+import uuid
+from typing import Any, Dict, List
+from xmlrpc.client import Boolean
+
+from broccoli.core.chain.chain import Chain
+from broccoli.core.redis_controller import RedisController
+from broccoli.core.task.task import Task
+from broccoli.core.task.task_queue import TaskQueue
+from broccoli.core.task.task_registry import TaskRegistry
+
+logger = logging.getLogger(__name__)
+
+
+class TaskChain:
+    """Chain multiple tasks together where each task passes its result to the next."""
+
+    def __init__(self, redis_url: str = "redis://localhost:6379", chain_id: str = None):
+        self.queue = TaskQueue(
+            queue_name="chain_tasks:queue", redis_url=redis_url, task_prefix="chain"
+        )
+        self._redis = RedisController(redis_url).get_client()
+        self.registry = TaskRegistry()
+        self.chain_id = chain_id or str(uuid.uuid4())
+
+    def chain(
+        self,
+        tasks: List[Dict[str, Any]],
+        shared_payload: Dict[str, Any] = None,
+        completion_task: str = None,
+    ) -> str:
+        """
+        Chain tasks together.
+
+        Args:
+            tasks: List of task configurations, each with 'task_type' and 'payload'
+            shared_payload: Data passed to all tasks in the chain
+            completion_task: Optional task to run after the chain completes
+
+        Returns:
+            chain_id: ID for tracking the entire chain
+
+        Example:
+            chain = TaskChain()
+            chain.chain([
+                {"task_type": "download_file", "payload": {"url": "https://..."}},
+                {"task_type": "process_video", "payload": {"quality": "high"}},
+                {"task_type": "upload_file", "payload": {"bucket": "my-bucket"}}
+            ], completion_task="on_chain_finished")
+        """
+        if not tasks:
+            raise ValueError("Cannot chain empty task list")
+
+        # Store chain metadata
+        chain = Chain(
+            chain_id=self.chain_id,
+            total_tasks=len(tasks),
+            completed_tasks=0,
+            failed=False,
+            status="pending",
+        )
+
+        if completion_task:
+            chain.completion_task = completion_task
+
+        # Save chain metadata to Redis
+        self._redis.hset(
+            f"chain:{self.chain_id}",
+            mapping={k: str(v) for k, v in chain.to_dict().items()},
+        )
+
+        # Assign task_ids to ALL tasks first, before creating any Task objects.
+        # This ensures the IDs stored in Redis match the Task objects pushed to the queue.
+        for t in tasks:
+            if "task_id" not in t:
+                t["task_id"] = str(uuid.uuid4())
+
+        # Store all task configs for reference
+        self._redis.set(
+            ex=86400,  # 24 hours TTL
+            name=f"chain:{self.chain_id}:tasks",
+            value=json.dumps(tasks),
+        )
+
+        # Create and push first task using the now-stable task_id
+        first_task = tasks[0]
+        task = Task(
+            task_type=first_task["task_type"],
+            task_id=first_task["task_id"],
+            payload={
+                **first_task.get("payload", {}),
+                **(shared_payload or {}),
+                "__chain_id": self.chain_id,
+                "__chain_position": 0,
+                "__total_tasks": len(tasks),
+                "__is_first": True,
+            },
+            max_retries=first_task.get("max_retries", 3),
+        )
+
+        # Push first task
+        self.queue.push(task)
+        logger.info(f"Chain {self.chain_id} started with task {task.task_id}")
+
+        return self.chain_id
+
+    def continue_chain(
+        self,
+        previous_task: Task,
+        previous_result: Any,
+        push_completion_task: bool = True,
+    ) -> Boolean | None:
+        """
+        Continue the chain after a task completes.
+        Called by the worker's post_process hook.
+
+        Args:
+            push_completion_task: If True, enqueue the completion task when the chain
+                finishes. Set to False when the caller (e.g. ChainWorkerMixin.on_finish)
+                handles chain completion directly, to avoid running it twice.
+        """
+        chain_id = previous_task.payload.get("__chain_id")
+        if not chain_id:
+            return None
+
+        position = previous_task.payload.get("__chain_position", 0)
+        total_tasks = previous_task.payload.get("__total_tasks", 0)
+
+        # Update chain progress
+        self._redis.hincrby(f"chain:{chain_id}", "completed_tasks", 1)
+        self._redis.hset(f"chain:{chain_id}", "current_task", position + 1)
+
+        # Check if chain is complete
+        if position + 1 >= total_tasks:
+            self._redis.hset(f"chain:{chain_id}", "status", "completed")
+            logger.info(f"Chain {chain_id} completed successfully")
+
+            if push_completion_task:
+                metadata = self._redis.hgetall(f"chain:{chain_id}")
+                completion_task = metadata.get("completion_task")
+                if completion_task:
+                    task = Task(
+                        task_type=completion_task,
+                        payload={
+                            "chain_id": chain_id,
+                            "result": previous_result,
+                            # Mark as chain task so ChainWorkerMixin.post_process
+                            # deletes chain:{task_id} after it runs
+                            "__chain_id": chain_id,
+                        },
+                    )
+                    self.queue.push(task)
+
+            return True
+
+        # Get next task configuration
+        tasks_json = self._redis.get(f"chain:{chain_id}:tasks")
+        if not tasks_json:
+            logger.error(f"Chain {chain_id} tasks not found")
+            return None
+
+        tasks = json.loads(tasks_json)
+        next_task_config = tasks[position + 1]
+
+        # Create next task using the pre-assigned task_id from the stored config
+        next_task = Task(
+            task_type=next_task_config["task_type"],
+            task_id=next_task_config["task_id"],
+            payload={
+                **next_task_config.get("payload", {}),
+                "__chain_id": chain_id,
+                "__chain_position": position + 1,
+                "__total_tasks": total_tasks,
+                "__previous_result": previous_result,
+                "__is_first": False,
+            },
+            max_retries=next_task_config.get("max_retries", 3),
+        )
+
+        # Push next task
+        self.queue.push(next_task)
+        logger.info(
+            f"Chain {chain_id} continuing with task {next_task.task_id} (position {position + 1})"
+        )
+
+    def get_chain_status(self, chain_id: str) -> Dict[str, Any]:
+        """Get the status of a chain."""
+        data = self._redis.hgetall(f"chain:{chain_id}")
+        if not data:
+            return {"status": "not_found"}
+
+        return {
+            "chain_id": chain_id,
+            "total_tasks": int(data.get("total_tasks", 0)),
+            "completed_tasks": int(data.get("completed_tasks", 0)),
+            "current_task": int(data.get("current_task", 0)),
+            "status": data.get("status", "unknown"),
+            "failed": data.get("failed", "False") == "True",
+        }

broccoli/core/health.py

@@ -0,0 +1,20 @@
+# video_scheduler/core/health.py
+class HealthCheck:
+    def __init__(self, worker):
+        self.worker = worker
+
+    def check(self) -> dict:
+        return {
+            "status": "healthy" if self.worker.running else "unhealthy",
+            "worker_id": self.worker.worker_id,
+            "redis": self._check_redis(),
+            "tasks_processed": self.worker.tasks_processed,
+            "active_task": self.worker.current_task_id,
+        }
+
+    def _check_redis(self):
+        try:
+            self.worker.queue.redis.ping()
+            return {"status": "connected"}
+        except:
+            return {"status": "disconnected"}

broccoli/core/redis_controller.py

@@ -0,0 +1,35 @@
+import redis
+
+
+class RedisController:
+    def __init__(self, redis_url: str = "redis://localhost:6379"):
+        self._redis = redis.from_url(redis_url, decode_responses=True)
+
+    def delete_all_keys(self):
+        """Delete all keys in Redis (use with caution)."""
+        return self._redis.flushdb()
+
+    def get_client(self):
+        """Get the Redis client instance."""
+        return self._redis
+
+    def view_all_keys(self):
+        """View all keys in Redis."""
+        return self._redis.keys("*")
+
+    def view_key(self, key: str):
+        """View value of a specific key."""
+        return self._redis.get(key)
+
+    def delete_key(self, key: str):
+        """Delete a specific key."""
+        return self._redis.delete(key)
+
+    def get_count(self, pattern="*"):
+        """Get count of keys matching a pattern."""
+        return len(self._redis.keys(pattern))
+
+    def view_all(self, pattern="*"):
+        """View all keys and their values matching a pattern."""
+        keys = self._redis.keys(pattern)
+        return {key.decode(): self._redis.get(key).decode() for key in keys}

broccoli/core/result.py

@@ -0,0 +1,78 @@
+# video_scheduler/core/result.py
+import json
+from dataclasses import dataclass
+
+import redis
+
+from broccoli.core.chain.chain import Chain
+from broccoli.core.task.task import Task
+
+
+class ResultBackend:
+    def __init__(self, redis_url: str):
+        self._redis = redis.from_url(redis_url)
+        self.ttl = 3600  # Default TTL for results in seconds
+
+    def store_task(self, task: Task) -> None:
+        """Store task result with TTL."""
+        key = f"result:{task.task_id}"
+        result_mapping = ResultMapping(
+            id=task.task_id,
+            result=task.result,
+            status=task.status,
+            chain=False,
+            error=task.error or "",
+        )
+
+        json_string = json.dumps(result_mapping.to_dict())
+        self._redis.set(name=key, ex=self.ttl, value=json_string)
+
+    def store_chain(self, chain: Chain) -> None:
+        """Store chain result with TTL."""
+        key = f"result:{chain.chain_id}"
+        result_mapping = ResultMapping(
+            id=chain.chain_id,
+            result=chain.result,
+            status=chain.status,
+            chain=True,
+            error="",
+        )
+
+        json_string = json.dumps(result_mapping.to_dict())
+        self._redis.set(name=key, ex=self.ttl, value=json_string)
+
+    def get(self, id: str) -> any:
+        """Retrieve task result."""
+        key = f"result:{id}"
+        data = self._redis.get(key)
+        if data:
+            return ResultMapping.from_dict(json.loads(data))
+        return None
+
+
+@dataclass
+class ResultMapping:
+    id: str
+    result: any
+    status: str
+    chain: bool
+    error: str = ""  # Optional error field for tasks
+
+    def to_dict(self) -> dict:
+        return {
+            "id": self.id,
+            "result": json.dumps(self.result),
+            "status": self.status,
+            "chain": self.chain,
+            "error": self.error,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ResultMapping":
+        return cls(
+            id=data.get("id"),
+            result=json.loads(data.get("result", "{}")),
+            status=data.get("status", "unknown"),
+            chain=data.get("chain", False),
+            error=data.get("error", ""),
+        )