orchestrata_ultimate.postman.priority

Reference › orchestrata_ultimate › postman › priority

Priority ordering for queued messages, SPEC-002 story FR-004 (REQ-002-FR4-003, -005 to -013) and story US-004 (REQ-002-US-004-006 to -010). PriorityQueueImpl implements IPriorityQueue from interfaces with a heap, and PriorityScheduler wraps a queue to stop low-priority messages from waiting forever. Priority values (0 critical to 4 background) and the age threshold come from constants. No other module imports it.

PriorityQueueImpl stores entries in a heap ordered by priority number and then by a sequence counter, so equal priorities leave in arrival order. enqueue_with_priority rejects a priority outside 0 to 4 with QueueError; enqueue uses normal priority. dequeue pops the lowest priority number, skipping heap items whose id was already removed; peek, size, is_empty, clear and get_by_priority read the state. boost_priority moves an entry to a higher priority. With a checkpoint path, checkpoint writes the entries, priorities, enqueue times and sequence to a JSON file and recover reads it back, returning the number restored.

PriorityScheduler.get_next first calls check_starvation, which lists entries other than critical ones that are older than the age threshold (60 seconds), and _apply_starvation_boost raises each one by one level per elapsed threshold, down to critical, before dequeuing. get_priority_stats counts entries per level, and detect_priority_inversion reports pairs where a lower-priority entry has waited more than ten seconds longer than a higher-priority one; its docstring calls it a heuristic.

At a glance

Class diagram

Dependency diagram

API

Classes

PriorityQueueImpl

Bases: IPriorityQueue · Decorators: traces(spec='002', story='FR-004', reqs=['REQ-002-FR4-003', 'REQ-002-FR4-005']) · Line 46

Priority queue with age-based starvation prevention.

PriorityScheduler

Decorators: traces(spec='002', story='FR-004', reqs=['REQ-002-US-004-008']) · Line 304

Fair scheduler with starvation prevention.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/postman/priority.py

# mypy: ignore-errors
"""Priority queue implementation with starvation prevention.

Provides priority-based message processing with:
- Multiple priority levels (CRITICAL to BACKGROUND)
- Preemption for CRITICAL messages
- Age-based starvation prevention
- Priority inversion detection
- FIFO ordering within same priority level
"""

import asyncio
import heapq
import json
import logging
import uuid
from collections import defaultdict
from dataclasses import asdict
from datetime import UTC, datetime
from pathlib import Path

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.postman.constants import (
    AGE_BOOST_THRESHOLD_SECONDS,
    PRIORITY_BACKGROUND,
    PRIORITY_CRITICAL,
    PRIORITY_HIGH,
    PRIORITY_LOW,
    PRIORITY_NORMAL,
)
from orchestrata_ultimate.postman.exceptions import QueueError
from orchestrata_ultimate.postman.interfaces import IPriorityQueue
from orchestrata_ultimate.postman.types import QueueEntry

logger = logging.getLogger(__name__)


@traces(
    spec="002",
    story="FR-004",
    reqs=[
        "REQ-002-FR4-003",
        "REQ-002-FR4-005",
    ],
)
class PriorityQueueImpl(IPriorityQueue):
    """Priority queue with age-based starvation prevention.

    Uses heapq for efficient priority ordering with FIFO within same priority.
    Thread-safe with asyncio.Lock.
    """

    def __init__(self, checkpoint_path: Path | None = None) -> None:
        self._lock = asyncio.Lock()
        # Heap of (priority, sequence, entry_id, entry)
        self._heap: list[tuple[int, int, str, QueueEntry]] = []
        # Entry lookup by ID
        self._entries: dict[str, QueueEntry] = {}
        # Priority lookup by entry ID
        self._priorities: dict[str, int] = {}
        # Enqueue time by entry ID for age tracking
        self._enqueue_times: dict[str, datetime] = {}
        # Sequence counter for FIFO within same priority
        self._sequence = 0
        # Checkpoint path for persistence
        self._checkpoint_path = checkpoint_path

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-FR4-006"])
    async def enqueue(self, entry: QueueEntry) -> str:
        """Enqueue with default NORMAL priority."""
        return await self.enqueue_with_priority(entry, PRIORITY_NORMAL)

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-FR4-006"])
    async def enqueue_with_priority(self, entry: QueueEntry, priority: int) -> str:
        """Add entry with explicit priority level."""
        async with self._lock:
            entry_id = entry.entry_id or str(uuid.uuid4())
            entry.entry_id = entry_id

            # Validate priority range
            if priority < PRIORITY_CRITICAL or priority > PRIORITY_BACKGROUND:
                raise QueueError(f"Invalid priority {priority}, must be 0-4")

            # Store entry
            self._entries[entry_id] = entry
            self._priorities[entry_id] = priority
            self._enqueue_times[entry_id] = datetime.now(UTC)

            # Add to heap with sequence for FIFO within priority
            heapq.heappush(self._heap, (priority, self._sequence, entry_id, entry))
            self._sequence += 1

            return entry_id

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-FR4-007"])
    async def dequeue(self) -> QueueEntry | None:
        """Get highest priority entry (lowest priority number)."""
        async with self._lock:
            while self._heap:
                priority, sequence, entry_id, entry = heapq.heappop(self._heap)

                # Check if entry still valid (not removed)
                if entry_id in self._entries:
                    # Clean up tracking
                    del self._entries[entry_id]
                    del self._priorities[entry_id]
                    del self._enqueue_times[entry_id]
                    return entry

            return None

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-FR4-008"])
    async def peek(self) -> QueueEntry | None:
        """Get highest priority entry without removing."""
        async with self._lock:
            while self._heap:
                priority, sequence, entry_id, entry = self._heap[0]

                # Check if entry still valid
                if entry_id in self._entries:
                    return entry
                else:
                    # Remove stale entry
                    heapq.heappop(self._heap)

            return None

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-FR4-009"])
    async def size(self) -> int:
        """Get number of entries in queue."""
        async with self._lock:
            return len(self._entries)

    async def is_empty(self) -> bool:
        """Check if queue is empty."""
        async with self._lock:
            return len(self._entries) == 0

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-FR4-010"])
    async def clear(self) -> int:
        """Remove all entries from queue.

        Returns:
            Number of entries cleared (matches IQueue.clear's contract —
            the other IQueue implementations in postman/queue.py already
            return this; this one silently discarded it).
        """
        async with self._lock:
            count = len(self._entries)
            self._heap.clear()
            self._entries.clear()
            self._priorities.clear()
            self._enqueue_times.clear()
            self._sequence = 0
            return count

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-US-004-012"])
    async def boost_priority(self, entry_id: str, new_priority: int) -> bool:
        """Boost entry to higher priority (lower number).

        Returns True if priority was boosted, False if entry not found.
        """
        async with self._lock:
            if entry_id not in self._entries:
                return False

            current_priority = self._priorities[entry_id]

            # Only boost if new priority is higher (lower number)
            if new_priority >= current_priority:
                return False

            # Validate new priority
            if new_priority < PRIORITY_CRITICAL or new_priority > PRIORITY_BACKGROUND:
                raise QueueError(f"Invalid priority {new_priority}, must be 0-4")

            # Update priority
            self._priorities[entry_id] = new_priority
            entry = self._entries[entry_id]

            # Re-insert with new priority
            heapq.heappush(self._heap, (new_priority, self._sequence, entry_id, entry))
            self._sequence += 1

            return True

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-FR4-011"])
    async def get_by_priority(self, priority: int) -> list[QueueEntry]:
        """Get all entries at specific priority level."""
        async with self._lock:
            return [
                entry
                for eid, entry in self._entries.items()
                if self._priorities[eid] == priority
            ]

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-FR4-013"])
    async def checkpoint(self) -> None:
        """Create persistence checkpoint."""
        if not self._checkpoint_path:
            return

        async with self._lock:
            # Convert entries with datetime handling
            entries_data = []
            for entry_id, entry in self._entries.items():
                entry_dict = asdict(entry)
                # Convert datetime objects to ISO format
                if "created_at" in entry_dict and entry_dict["created_at"]:
                    entry_dict["created_at"] = entry_dict["created_at"].isoformat()
                if "ttl_expires_at" in entry_dict and entry_dict["ttl_expires_at"]:
                    entry_dict["ttl_expires_at"] = entry_dict[
                        "ttl_expires_at"
                    ].isoformat()

                entries_data.append(
                    {
                        "id": entry_id,
                        "entry": entry_dict,
                        "priority": self._priorities[entry_id],
                        "enqueue_time": self._enqueue_times[entry_id].isoformat(),
                    }
                )

            checkpoint_data = {
                "sequence": self._sequence,
                "entries": entries_data,
            }

            # Write to checkpoint file
            self._checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
            self._checkpoint_path.write_text(json.dumps(checkpoint_data, indent=2))

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-FR4-014"])
    async def recover(self) -> int:
        """Recover from last checkpoint.

        Returns:
            Number of entries recovered
        """
        if not self._checkpoint_path or not self._checkpoint_path.exists():
            return 0

        async with self._lock:
            # Clear current state
            self._heap.clear()
            self._entries.clear()
            self._priorities.clear()
            self._enqueue_times.clear()

            # Load checkpoint
            checkpoint_data = json.loads(self._checkpoint_path.read_text())
            self._sequence = checkpoint_data.get("sequence", 0)

            recovered_count = 0
            for item in checkpoint_data.get("entries", []):
                entry_id = item["id"]
                entry_dict = item["entry"]

                # Convert ISO strings back to datetime
                if "created_at" in entry_dict and entry_dict["created_at"]:
                    entry_dict["created_at"] = datetime.fromisoformat(
                        entry_dict["created_at"]
                    )
                if "ttl_expires_at" in entry_dict and entry_dict["ttl_expires_at"]:
                    entry_dict["ttl_expires_at"] = datetime.fromisoformat(
                        entry_dict["ttl_expires_at"]
                    )

                entry = QueueEntry(**entry_dict)
                priority = item["priority"]
                enqueue_time = datetime.fromisoformat(item["enqueue_time"])

                # Restore entry
                self._entries[entry_id] = entry
                self._priorities[entry_id] = priority
                self._enqueue_times[entry_id] = enqueue_time

                # Add to heap
                heapq.heappush(self._heap, (priority, self._sequence, entry_id, entry))
                self._sequence += 1
                recovered_count += 1

            return recovered_count

    @traces(spec="002", story="US-004", reqs=["REQ-002-US-004-009"])
    def _get_entry_age(self, entry_id: str) -> float:
        """Get entry age in seconds."""
        if entry_id not in self._enqueue_times:
            return 0.0

        enqueue_time = self._enqueue_times[entry_id]
        now = datetime.now(UTC)
        return (now - enqueue_time).total_seconds()


@traces(
    spec="002",
    story="FR-004",
    reqs=[
        "REQ-002-US-004-008",
    ],
)
class PriorityScheduler:
    """Fair scheduler with starvation prevention.

    Monitors entry age and boosts priority to prevent starvation:
    - After AGE_BOOST_THRESHOLD_SECONDS: boost by 1 level
    - After 2x threshold: boost by 2 levels
    - Never boost above CRITICAL
    """

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-US-004-014"])
    def __init__(
        self,
        queue: IPriorityQueue,
        age_boost_threshold: int = AGE_BOOST_THRESHOLD_SECONDS,
    ) -> None:
        """Initialize scheduler.

        Args:
            queue: Priority queue to manage
            age_boost_threshold: Seconds before boosting priority
        """
        self.queue = queue
        self.age_boost_threshold = age_boost_threshold

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-US-004-006"])
    async def get_next(self) -> QueueEntry | None:
        """Get next entry with starvation prevention applied.

        Checks for starving entries and boosts their priority before dequeue.
        """
        # Check for starvation and boost if needed
        starving_ids = await self.check_starvation()
        for entry_id in starving_ids:
            await self._apply_starvation_boost(entry_id)

        # Dequeue highest priority
        return await self.queue.dequeue()

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-US-004-008"])
    async def check_starvation(self) -> list[str]:
        """Return entry_ids that need priority boost due to age."""
        starving: list[str] = []

        # Access private method to get entry ages
        if not isinstance(self.queue, PriorityQueueImpl):
            return starving

        async with self.queue._lock:
            for entry_id in self.queue._entries.keys():
                age = self.queue._get_entry_age(entry_id)
                priority = self.queue._priorities[entry_id]

                # Skip CRITICAL entries (already highest priority)
                if priority == PRIORITY_CRITICAL:
                    continue

                # Check if age exceeds threshold
                if age >= self.age_boost_threshold:
                    starving.append(entry_id)

        return starving

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-US-004-009"])
    async def _apply_starvation_boost(self, entry_id: str) -> None:
        """Apply age-based priority boost to entry."""
        if not isinstance(self.queue, PriorityQueueImpl):
            return

        async with self.queue._lock:
            if entry_id not in self.queue._entries:
                return

            age = self.queue._get_entry_age(entry_id)
            current_priority = self.queue._priorities[entry_id]

            # Calculate boost level based on age (guard against division by zero)
            if self.age_boost_threshold > 0:
                boost_level = int(age / self.age_boost_threshold)
            else:
                # If threshold is 0 or negative, boost immediately to max
                boost_level = current_priority
            new_priority = max(PRIORITY_CRITICAL, current_priority - boost_level)

            # Check if priority boost is needed
            should_boost = new_priority < current_priority

        # Call boost_priority outside the lock only if needed
        if should_boost:
            logger.debug(
                f"Boosting entry {entry_id} from priority {current_priority} to {new_priority}"
            )
            await self.queue.boost_priority(entry_id, new_priority)

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-US-004-016"])
    async def get_priority_stats(self) -> dict[int, int]:
        """Get count of entries at each priority level."""
        stats = defaultdict(int)

        for priority in [
            PRIORITY_CRITICAL,
            PRIORITY_HIGH,
            PRIORITY_NORMAL,
            PRIORITY_LOW,
            PRIORITY_BACKGROUND,
        ]:
            entries = await self.queue.get_by_priority(priority)
            stats[priority] = len(entries)

        return dict(stats)

    @traces(spec="002", story="US-002-001", reqs=["REQ-002-US-004-010"])
    async def detect_priority_inversion(self) -> list[tuple[str, str]]:
        """Detect potential priority inversions.

        Returns list of (low_priority_id, high_priority_id) tuples where
        low priority entry is blocking high priority entry.

        Note: This is a simple heuristic based on age and priority.
        """
        inversions: list[tuple[str, str]] = []

        if not isinstance(self.queue, PriorityQueueImpl):
            return inversions

        async with self.queue._lock:
            entries = list(self.queue._entries.items())

            for i, (id1, entry1) in enumerate(entries):
                for id2, entry2 in entries[i + 1 :]:
                    prio1 = self.queue._priorities[id1]
                    prio2 = self.queue._priorities[id2]
                    age1 = self.queue._get_entry_age(id1)
                    age2 = self.queue._get_entry_age(id2)

                    # Inversion: lower priority is older (blocking)
                    if prio1 > prio2 and age1 > age2 + 10:
                        inversions.append((id1, id2))
                    elif prio2 > prio1 and age2 > age1 + 10:
                        inversions.append((id2, id1))

        return inversions