twainmq 0.2.3__py3-none-any.whl

twainmq/__init__.py

@@ -0,0 +1,24 @@
+"""
+
+TwainMQ is a broker free messaging system.  Unlike similar tools like Apache Kafka and RabbitMQ, TwainMQ does not have any running services or brokers to coordinate messages, being run entirely from files.  Think of it like SQLite is to traditional database systems.
+
+To set up a TwainMQ message queue you simply need to specify a directory for it to use, either locally or on a network share.
+
+Messages are limited in size to ensure that they can be written quickly.  The size limit is currently set at slightly less than 4kb.
+
+Messages can be either plain UTF-8 or raw bytes.
+
+Messages are organised by key, partition and topic.
+
+## Keys
+
+u8
+u16
+u32
+u64
+str
+
+"""
+
+from .core import Twain, TwainMQConsumer, TwainMQProducer
+from .utils import multi_poll

twainmq/agent_base.py

@@ -0,0 +1,44 @@
+from abc import ABC
+from datetime import datetime, timezone
+import json
+
+from twainmq.encoding import find_key_char_width
+from twainmq.errors import ConfigNotFoundError
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .core import Twain
+
+class TwainMQBase(ABC):
+    def __init__(self, twain: "Twain", topic: str):
+        self._topic = topic
+        self._twain = twain
+        config_path = self._twain._config_path(topic)
+        if not config_path.is_file():
+            raise ConfigNotFoundError(f"No config for {topic}: cannot find file {config_path}")
+        with config_path.open("r") as f:
+            config = json.load(f)
+        self._key_width = config["key_width"]
+        self._message_types = config["message_types"]
+        self._message_types_rev = {i: msg_type for msg_type, i in self._message_types.items()}
+        self._n_partitions = int(config["partitions"])
+        self._key_chars = find_key_char_width(self._key_width)
+    
+    @property
+    def topic(self):
+        return self._topic
+    
+    @property
+    def _topic_dir(self):
+        return self._twain.root_dir / self._topic
+    
+    @property
+    def key_width(self):
+        """The width of the key in bytes.  0 indicates a string key.
+        """
+        return self._key_width
+
+    @property
+    def chunk_str_now(self):
+        return f"{datetime.now(timezone.utc):%Y%m%d}"

twainmq/atomic_append.py

@@ -0,0 +1,100 @@
+import os
+from pathlib import Path
+import errno
+import ctypes
+from ctypes import wintypes
+from pathlib import Path
+import os
+import errno
+from pathlib import Path
+
+# Win32 constants
+FILE_APPEND_DATA        = 0x0004
+FILE_SHARE_READ         = 0x00000001
+FILE_SHARE_WRITE        = 0x00000002
+OPEN_ALWAYS             = 4
+FILE_ATTRIBUTE_NORMAL   = 0x80
+INVALID_HANDLE_VALUE    = wintypes.HANDLE(-1).value
+
+if os.name == "nt":
+    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
+
+def atomic_append_win(path, data: bytes) -> None:
+    """Atomic appending to a file under Windows.
+    
+    This makes use of the windows FILE_APPEND_DATA mode which doesn't check the file length first or use a cursor
+    it simply dumps the data on the end of the file.
+    """
+    # Accept Path or str
+    if isinstance(path, Path):
+        path = str(path)
+
+    # 1. Open with FILE_APPEND_DATA only (critical)
+    handle = kernel32.CreateFileW(
+        wintypes.LPCWSTR(path),
+        FILE_APPEND_DATA,                         # <-- atomic append mode
+        FILE_SHARE_READ | FILE_SHARE_WRITE,
+        None,
+        OPEN_ALWAYS,
+        FILE_ATTRIBUTE_NORMAL,
+        None,
+    )
+
+    if handle == INVALID_HANDLE_VALUE:
+        raise OSError(ctypes.get_last_error(), "CreateFileW failed")
+
+    try:
+        # 2. WriteFile (one syscall, atomic append)
+        written = wintypes.DWORD(0)
+        ok = kernel32.WriteFile(
+            handle,
+            data,
+            len(data),
+            ctypes.byref(written),
+            None,
+        )
+
+        if not ok or written.value != len(data):
+            raise OSError(ctypes.get_last_error(), "WriteFile failed")
+
+    finally:
+        kernel32.CloseHandle(handle)
+        
+def atomic_append_posix(path, data: bytes, mode: int = 0o644) -> None:
+    """Atomic append function for POSIX systems.
+
+    This should work by using O_APPEND but this is not properly tested yet!
+
+    Short write might in theory happen, but shouldn't, when we can test this better
+    we will know if this is a problem and come up with a solution.
+    """
+    # Accept Path or str
+    if isinstance(path, Path):
+        path = os.fspath(path)
+
+    # Open with O_APPEND so the kernel moves the file offset atomically
+    flags = os.O_WRONLY | os.O_APPEND | os.O_CREAT
+
+    # Add close-on-exec if available (good hygiene)
+    flags |= getattr(os, "O_CLOEXEC", 0)
+
+    fd = os.open(path, flags, mode)
+    try:
+        # Exactly one write() syscall
+        n = os.write(fd, data)
+
+        # POSIX guarantees full writes to regular files unless an error occurs.
+        # If we ever see a short write, treat it as fatal.
+        if n != len(data):
+            raise OSError(errno.EIO, f"short write: wrote {n} of {len(data)} bytes")
+
+    finally:
+        os.close(fd)
+
+def atomic_append(path, data: bytes, mode: int = 0o644) -> None:
+    """Atomic writing used by producers to allow multiple producers to write to the same 
+    log file at the same time without locks"""
+    if os.name == "nt":
+        atomic_append_win(path, data)
+    else:
+        atomic_append_posix(path, data)

twainmq/consumer.py

@@ -0,0 +1,357 @@
+
+import base64
+import bisect
+from datetime import datetime, timezone
+import json
+import random
+import time
+from typing import TYPE_CHECKING
+import zlib
+
+from .agent_base import TwainMQBase
+from .encoding import _DATACLASS_MAGIC, _GZIP_MAGIC, Message, base85_to_key, dataclass_from_dict, decode_datetime
+from .errors import TopicCorruptError
+from .consumer_groups import CONSUMER_GROUP_MESSAGE_SET, REBAL_LENGTH, AbortJoin, BeginRebal, Commit, EndRebal, Joined, RebalConfirm, RebalInProgress, RebalOffer, _group_topic_name
+
+if TYPE_CHECKING:
+    from .core import Twain
+
+class TwainMQConsumer(TwainMQBase):
+    """
+    A TwainMQ consumer.
+
+    Consumers may operate independently or as part of a consumer group. A
+    standalone consumer can read messages but cannot commit progress. A consumer
+    that belongs to a group (even a group of one) can commit offsets, allowing
+    any consumer joining that group to resume from the last committed position.
+
+    The `start_from` parameter controls the initial read position only when:
+    - the consumer is not part of a group, or
+    - the group has no existing commit record.
+    """
+    def __init__(self, twain: "Twain", topic: str, start_from = "start", group = None):
+        """
+        Create a consumer for a given topic.
+
+        Parameters
+        ----------
+        twain : Twain
+            The TwainMQ instance backing this consumer.
+        topic : str
+            The topic to consume from.
+        start_from : {"start", "now"}, optional
+            Initial read position when no committed offset exists. `"start"`
+            begins at offset 0; `"now"` begins at the current end of each
+            partition.
+        group : str or None, optional
+            The consumer-group identifier. If `None`, the consumer operates
+            independently and cannot commit offsets. If provided, the consumer
+            joins the group and participates in consensus-based partition
+            assignment and offset commits.
+        """
+        super().__init__(twain, topic)
+        self._twain = twain
+        self._topic = topic
+        self._group = group
+        self._rebal_ending = None
+        self._consumer_id = None
+        self._rebal_in_progress = None
+        self._commit_record = dict()
+        if group is None:
+            self._partitions = [i for i in range(self._n_partitions)]
+            if start_from == "start":
+                self._consumerlets = {p: TwainMQConsumerlet(twain, topic, p, offset=0) for p in self._partitions}
+            elif start_from == "now":
+                self._consumerlets = {p: TwainMQConsumerlet(twain, topic, p, offset=None) for p in self._partitions}
+            else:
+                raise ValueError(f'start_from should be either "start" of "now", received: "{start_from}"')
+        else:
+            self._partitions = []      
+            self._consumerlets = dict()      
+            self._join_group(group, start_from)
+        self._last_polled = 0
+        self.heartbeat()
+
+    def _create_consumer_group(self, group_topic, start_from):
+        self._twain.create_topic(group_topic, key_type = "u32", partitions = 1, message_types = CONSUMER_GROUP_MESSAGE_SET)
+        with self._twain.producer(group_topic) as gprod:
+            if start_from == "start":
+                init_commits = [(i, 0) for i in range(self._n_partitions)]
+            elif start_from == "now":
+                init_commits = [(i, -1) for i in range(self._n_partitions)]
+            else:
+                raise ValueError(f'start_from should be either "start" of "now", received: "{start_from}"')
+            gprod.write_message(0, Commit(init_commits))
+
+    def _join_group(self, group, start_from, retry=3):
+        """Called when a consumer joins a group to handle id registration with the group"""
+        group_topic = _group_topic_name(group)
+        if not self._twain.topic_exists(group_topic):
+            self._create_consumer_group(group_topic, start_from)
+        gprod = self._twain.producer(group_topic)
+        gcon = self._twain.consumer(group_topic, start_from = "start", group = None)
+        self._group_producer = gprod
+        self._group_consumer = gcon
+        
+        keys_taken = set()
+        while (msg := gcon.poll()) is not None:
+            keys_taken.add(msg.key)
+            self._process_group_msg(msg)
+            
+        consumer_id = max(keys_taken) + 1
+        join_key = random.getrandbits(64)
+        gprod.write_message(consumer_id, Joined(join_key))
+        
+        while True:
+            msg = gcon.poll()
+            if msg is None:
+                time.sleep(0.5)
+            elif isinstance(msg.message, Joined):
+                key = msg.message.tag
+                if join_key == key:
+                    self._consumer_id = consumer_id
+                    self._begin_rebal()
+                    return
+                else:
+                    gprod.write_message(consumer_id, AbortJoin(key))
+                    if retry > 0:
+                        time.sleep(random.random() * 0.1)
+                        return self._join_group(group, start_from, retry = retry - 1)
+                    else:
+                        return
+            else:
+                self._process_group_msg(msg)
+    
+    def trigger_rebal(self):
+        """
+        Initiate a consumer-group rebalance.
+
+        In TwainMQ, consumer-group membership is maintained through periodic
+        consensus. Triggering a rebalance is lightweight and safe, and is typically
+        done periodically to confirm that all consumers in the group are still
+        active.
+
+        If all consumers respond within the rebalance window, the rebalance becomes
+        a no-op and existing partition assignments remain unchanged.
+        """
+        self._begin_rebal()
+
+    def _begin_rebal(self):
+        gprod = self._group_producer
+        rebal_end_ts = (datetime.now(timezone.utc) + REBAL_LENGTH).timestamp()
+        gprod.write_message(self._consumer_id, BeginRebal(rebal_end_ts))
+        self.heartbeat()
+    
+    def _rebal_participate(self):
+        gprod = self._group_producer
+        gprod.write_message(self._consumer_id, RebalOffer(self._partitions, [i for i in range(self._n_partitions)]))
+
+    def _end_rebal(self):
+        gprod = self._group_producer
+        gprod.write_message(self._consumer_id, EndRebal())
+        self._confirm_rebal()
+
+    def _confirm_rebal(self):
+        gprod = self._group_producer
+        assert self._rebal_in_progress is not None
+        assignements = self._rebal_in_progress.get_assignments()
+        self._partitions = assignements.get(self._consumer_id, [])
+        self._consumerlets = {p: self._consumerlets.get(p, TwainMQConsumerlet(self._twain, self._topic, p, offset=self._commit_record.get(p,0))) for p in self._partitions}
+        gprod.write_message(self._consumer_id, RebalConfirm(sorted(list(self._partitions))))
+        self._rebal_in_progress = None
+
+    def commit(self, partitions = None):
+        if self._group is not None:
+            if partitions is None:
+                partitions = self._partitions
+            commit_list = [(p, self._consumerlets[p]._offset) for p in partitions if p in self._consumerlets]
+            print(commit_list)
+            gprod = self._group_producer
+            gprod.write_message(self._consumer_id, Commit(commit_list))
+
+    def _handle_commit_msg(self, commit_msg):
+        for p,o in commit_msg.partition_offsets:
+            self._commit_record[p] = o
+
+    def _process_group_msg(self, msg):
+        if isinstance(msg.message, BeginRebal):
+            if self._consumer_id:
+                self._rebal_in_progress = RebalInProgress(msg.message.timestamp)
+                self._rebal_participate()
+        elif isinstance(msg.message, RebalOffer):
+            if self._rebal_in_progress:
+                self._rebal_in_progress.add_offer(msg.key, msg.message)
+        elif isinstance(msg.message, EndRebal):
+            if self._rebal_in_progress:
+                self._confirm_rebal()
+        elif isinstance(msg.message, Commit):
+            self._handle_commit_msg(msg.message)
+        else:
+            pass
+
+    def heartbeat(self):
+        """
+        Maintain membership in a consumer group and process group-management events.
+
+        This method must be called periodically when using consumer groups; otherwise
+        the consumer may be considered inactive and be removed from the group.
+
+        Behaviour
+        ---------
+        - Processes any pending group-management messages (rebalance events, etc.).
+        - Automatically invoked at the start of `poll()`.
+        """
+        if self._group is not None:
+            while msg := self._group_consumer.poll():
+                self._process_group_msg(msg)
+            if self._rebal_in_progress and self._rebal_in_progress.rebal_ending < datetime.now(timezone.utc).timestamp():
+                self._end_rebal()
+        
+    def poll(self):
+        """
+        Poll partitions in round-robin order and return the next available message.
+
+        The consumer maintains a list of assigned partitions. Each call polls them
+        starting from the partition after the one last successfully polled. The first
+        non-`None` message returned by a partition is yielded.
+
+        This call does not block.
+
+        Returns
+        -------
+        message or None
+            The next available message, or `None` if all partitions return no data.
+        """
+        self.heartbeat()
+        n = len(self._partitions)
+        for p_offset in range(n):
+            partition_to_poll = self._partitions[(self._last_polled + p_offset + 1) % n]
+            this_consumerlet = self._consumerlets[partition_to_poll]
+            msg = this_consumerlet.poll()
+            if msg is not None:
+                self._last_polled = partition_to_poll
+                return msg
+    
+    def poll_many(self, n=10):
+        """
+        Poll up to `n` messages and return them as a list.
+
+        This is a convenience wrapper around `poll()`. Polling stops early if
+        `poll()` returns `None`.
+
+        Parameters
+        ----------
+        n : int, optional
+            Maximum number of messages to retrieve.
+
+        Returns
+        -------
+        list
+            A list of messages, possibly shorter than `n`.
+        """
+        msgs = []
+        for i in range(n):
+            msg = self.poll()
+            if msg is None:
+                break
+            msgs.append(msg)
+        return msgs
+
+    def __str__(self):
+        return f"TwainMQConsumer(topic={self._topic}, group={self._group})"
+
+class TwainMQConsumerlet(TwainMQBase):
+    """A consumerlet is a simple single partition consumer.  Usually you would not use a Consumerlet directly, rather use the Consumer container"""
+    def __init__(self, twain, topic, partition, offset=None):
+        super().__init__(twain, topic)
+        self._partition = partition
+        self._seek_active_file(offset)
+
+    def decode_message(self, message):
+        compressed = base64.b85decode(message.encode("utf-8"))
+        decoded = zlib.decompress(compressed, wbits=-15)
+        if decoded.startswith(_GZIP_MAGIC):
+            return decoded[1:]
+        elif decoded.startswith(_DATACLASS_MAGIC):
+            class_id = int.from_bytes(decoded[1:3], "big", signed=False)
+            message_type = self._twain._msg_cls_registry[self._message_types_rev[class_id]]
+            data = json.loads(decoded[3:])
+            return dataclass_from_dict(message_type, data)
+        else:
+            return decoded.decode("utf-8")
+    
+    def _seek_active_file(self, offset = None):
+        """Sets the file handle to the active file to read from and seeks to the end.
+        
+        If offset is None this is read from latest
+        
+        If the latest file is missing then it will return None.
+        """
+        part_files = self._list_partition_files()
+        if offset is None or len(part_files) == 0:
+            chunk_str = self.chunk_str_now
+            head_file = [f for f in part_files if f.stem.split("-")[1].split("_")[0] == chunk_str]
+            if len(head_file) == 0:
+                self._current_file_handle = None
+            elif len(head_file) == 1:
+                offset = int(head_file[0].stem.split("_")[1])
+                self._current_file_handle = head_file[0].open("r", encoding = "utf-8")
+                while self._current_file_handle.readline():
+                    offset += 1
+                self._offset = offset
+                self._chunk_str = chunk_str
+            else:
+                raise TopicCorruptError(f"Multiple message files for the same chunk partition {self._partition}-{chunk_str}")
+        elif offset >= 0:
+            offsets_files_chunks = sorted([(int(f.stem.split("_")[1]), f, f.stem.split("_")[0].split("-")[1]) for f in part_files])
+                        
+            offsets = [o for o, _, _ in offsets_files_chunks]
+            i = bisect.bisect_right(offsets, offset) - 1
+            this_offset, current_file, chunk_str = offsets_files_chunks[max(i, 0)]
+
+            self._current_file_handle = current_file.open("r", encoding = "utf-8")
+            self._chunk_str = chunk_str
+            self._offset = this_offset
+            while self._offset < offset:
+                line = self._current_file_handle.readline()
+                if line != "":
+                    self._offset += 1
+                else:
+                    break
+        else:
+            raise NotImplementedError("Seek back from end not yet implemented")
+    
+    def _list_partition_files(self):
+        return [f for f in self._topic_dir.iterdir() if f.stem.split("-")[0] == str(self._partition)]
+
+    def poll(self):
+        """Polls for the next message without blocking.  Returns None if no new message available."""
+        if self._current_file_handle is None:
+            self._seek_active_file()
+            if self._current_file_handle is None:
+                return None
+        msg_line = self._current_file_handle.readline()[:-1]
+        if msg_line == "":
+            if self.chunk_str_now == self._chunk_str:
+                return None
+            else:
+                self._seek_active_file(self._offset)
+                msg_line = self._current_file_handle.readline()[:-1]
+                if msg_line == "":
+                    return None
+        key = base85_to_key(msg_line[:self._key_chars], self._key_width)
+        timestamp = decode_datetime(msg_line[self._key_chars:self._key_chars+10])
+        message = self.decode_message(msg_line[self._key_chars+10:])
+        msg_tuple = Message(
+            offset=self._offset,
+            key=key,
+            timestamp=timestamp,
+            message=message
+        )
+        self._offset += 1
+        return msg_tuple
+
+    @property
+    def partition(self):
+        return self._partition
+

twainmq/consumer_groups.py

@@ -0,0 +1,143 @@
+from dataclasses import dataclass
+from datetime import timedelta
+from dataclasses_jsonschema import JsonSchemaMixin
+
+from .encoding import _is_safe
+from .errors import InvalidGroupNameError
+
+
+REBAL_LENGTH = timedelta(seconds = 60)
+
+## CONSUMER GROUP MESSAGES
+
+@dataclass
+class Joined(JsonSchemaMixin):
+    tag: int
+    __message_type__ = "~~TWAIN~~JOINED"
+    
+    def __init__(self, tag):
+        self.tag = tag
+
+@dataclass
+class AbortJoin(JsonSchemaMixin):
+    tag: int
+    __message_type__ = "~~TWAIN~~ABORTJOIN"
+    
+    def __init__(self, tag):
+        self.tag = tag
+
+@dataclass
+class RebalOffer(JsonSchemaMixin):
+    own: list[int]
+    partitions: list[int]
+    __message_type__ = "~~TWAIN~~REBALOFFER"
+
+    def __init__(self, own, partitions):
+        self.own = own
+        self.partitions = partitions
+
+@dataclass
+class RebalConfirm(JsonSchemaMixin):
+    partitions: list[int]
+    __message_type__ = "~~TWAIN~~REBALCONFIRM"
+
+    def __init__(self, partitions):
+        self.partitions = partitions
+
+@dataclass
+class BeginRebal(JsonSchemaMixin):
+    timestamp: float
+    __message_type__ = "~~TWAIN~~BEGINREBAL"
+
+    def __init__(self, timestamp):
+        self.timestamp = timestamp
+
+@dataclass
+class EndRebal(JsonSchemaMixin):
+    __message_type__ = "~~TWAIN~~ENDREBAL"
+
+    def __init__(self):
+        pass
+
+@dataclass
+class Commit(JsonSchemaMixin):
+    partition_offsets: list[tuple[int, int]]
+    __message_type__ = "~~TWAIN~~COMMIT"
+
+    def __init__(self, partition_offsets):
+        self.partition_offsets = partition_offsets
+
+CONSUMER_GROUP_MESSAGE_CLASSES = [Joined, AbortJoin, RebalOffer, RebalConfirm, BeginRebal, EndRebal, Commit]
+CONSUMER_GROUP_MESSAGE_SET = [c.__message_type__ for c in CONSUMER_GROUP_MESSAGE_CLASSES]
+
+def _group_topic_name(group):
+    if not _is_safe(group):
+        raise InvalidGroupNameError(f"Consumer group names must be [a-zA-Z0-9_.], {group} is not valid")
+    return f"--group--{group}"
+
+class RebalInProgress:
+    def __init__(self, rebal_ending):
+        self.rebal_ending = rebal_ending
+        self._offers = dict()
+
+    def add_offer(self, cons_id: int, rebal_offer: RebalOffer):
+        self._offers[cons_id] = rebal_offer
+
+    def get_assignments(self):
+        # --- Step 0: Clean ownership lists ---
+        cleaned = []
+        for cid, offer in self._offers.items():
+            allowed = set(offer.partitions)
+            owned = [p for p in offer.own if p in allowed]
+            cleaned.append((cid, allowed, owned))
+
+        # --- Step 1: Sort each consumer's partition list ---
+        # Owned first, then by partition ID
+        consumer_lists = {}
+        for cid, allowed, owned in cleaned:
+            owned_sorted = sorted(owned)
+            new_sorted = sorted(p for p in allowed if p not in owned)
+            consumer_lists[cid] = owned_sorted + new_sorted
+
+        # --- Step 2: Sort consumers ---
+        # a) number of partitions offered (ascending)
+        # b) number of partitions owned (descending)
+        # c) consumer id (ascending)
+        sorted_consumers = sorted(
+            cleaned,
+            key=lambda x: (len(x[1]), -len(x[2]), x[0])
+        )
+
+        # --- Step 3: Round-robin assignment ---
+        assigned = {cid: [] for cid, _, _ in cleaned}
+        remaining_lists = {cid: list(consumer_lists[cid]) for cid, _, _ in cleaned}
+
+        # Set of all partitions that appear anywhere
+        all_partitions = set()
+        for _, allowed, _ in cleaned:
+            all_partitions.update(allowed)
+
+        assigned_partitions = set()
+
+        # Continue until all partitions assigned
+        while len(assigned_partitions) < len(all_partitions):
+            for cid, allowed, owned in sorted_consumers:
+                lst = remaining_lists[cid]
+
+                # Skip already-taken partitions
+                while lst and lst[0] in assigned_partitions:
+                    lst.pop(0)
+
+                if not lst:
+                    continue
+
+                # Assign the next partition
+                p = lst.pop(0)
+                assigned[cid].append(p)
+                assigned_partitions.add(p)
+
+        # Sort results for determinism
+        for cid in assigned:
+            assigned[cid].sort()
+
+        return assigned