Class PartitionedMemoryQueue

methods/MIRIX/mirix/queue/memory_queue.py:52–183

Partitioned async in-memory queue that simulates Kafka's partitioning. Messages are routed to partitions based on user_id hash, ensuring: - All messages for the same user go to the same partition - Each partition is consumed by exactly one worker task - Per-user message orderin

class PartitionedMemoryQueue(QueueInterface):
    """
    Partitioned async in-memory queue that simulates Kafka's partitioning.

    Messages are routed to partitions based on user_id hash, ensuring:
    - All messages for the same user go to the same partition
    - Each partition is consumed by exactly one worker task
    - Per-user message ordering is preserved (same as Kafka behavior)

    Partitioning modes:
    - Hash (default): Uses hash(key) % num_partitions (Kafka-like)
    - Round-robin: Sequential assignment (user1->p0, user2->p1, ...)
    """

    def __init__(self, num_partitions: int = 1, round_robin: bool = False):
        self._num_partitions = max(1, num_partitions)
        self._round_robin = round_robin
        self._partitions: List[asyncio.Queue[QueueMessage]] = [
            asyncio.Queue() for _ in range(self._num_partitions)
        ]

        self._user_partition_map: Dict[str, int] = {}
        self._next_partition: int = 0
        self._partition_lock = asyncio.Lock()

        mode = "round-robin" if round_robin else "hash"
        logger.info(
            "Initialized PartitionedMemoryQueue with %d partitions, mode=%s",
            self._num_partitions,
            mode,
        )

    @property
    def num_partitions(self) -> int:
        return self._num_partitions

    @property
    def round_robin(self) -> bool:
        return self._round_robin

    async def get_partition_stats(self) -> Dict[str, any]:
        """Get statistics about partition distribution."""
        async with self._partition_lock:
            partition_counts = [0] * self._num_partitions
            for partition_id in self._user_partition_map.values():
                partition_counts[partition_id] += 1

            return {
                "mode": "round-robin" if self._round_robin else "hash",
                "num_partitions": self._num_partitions,
                "total_users": len(self._user_partition_map),
                "users_per_partition": partition_counts,
                "queue_sizes": [p.qsize() for p in self._partitions],
            }

    def _get_partition_key(self, message: QueueMessage) -> str:
        if message.HasField("user_id") and message.user_id:
            return message.user_id