tigerbeetle 0.16.16__py3-none-any.whl

tigerbeetle/client.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+
+import asyncio
+import ctypes
+import logging
+import os
+import threading
+import time
+from collections.abc import Callable  # noqa: TCH003
+from dataclasses import dataclass
+from typing import Any
+
+from . import bindings
+from .lib import tb_assert
+
+logger = logging.getLogger("tigerbeetle")
+
+
+class AtomicInteger:
+    def __init__(self, value=0):
+        self._value = value
+        self._lock = threading.Lock()
+
+    def increment(self):
+        with self._lock:
+            self._value += 1
+            return self._value
+
+
+@dataclass
+class CompletionContextSync:
+    event: threading.Event
+
+
+@dataclass
+class CompletionContextAsync:
+    loop: asyncio.AbstractEventLoop
+    event: asyncio.Event
+
+
+@dataclass
+class InflightPacket:
+    packet: bindings.CPacket
+    response: Any
+    operation: bindings.Operation
+    c_event_type: Any
+    c_result_type: Any
+    on_completion: Callable | None
+    on_completion_context: CompletionContextSync | CompletionContextAsync | None
+
+
+
+def id() -> int:
+    """
+    Generates a Universally Unique and Sortable Identifier as a 128-bit integer. Based on ULIDs.
+    """
+    time_ms = time.time_ns() // (1000 * 1000)
+
+    # Ensure time_ms monotonically increases.
+    time_ms_last = getattr(id, "_time_ms_last", 0)
+    if time_ms <= time_ms_last:
+        time_ms = time_ms_last
+    else:
+        id._time_ms_last = time_ms
+
+    randomness = os.urandom(10)
+
+    return int.from_bytes(
+        time_ms.to_bytes(6, "big") + randomness,
+        "big",
+    )
+
+
+amount_max = (2 ** 128) - 1
+
+
+class InitError(Exception):
+    pass
+
+class PacketError(Exception):
+    pass
+
+
+class Client:
+    _clients: dict[int, Any] = {}
+    _counter = AtomicInteger()
+
+    def __init__(self, cluster_id: int, replica_addresses: str):
+        self._client_key = Client._counter.increment()
+        self._client = bindings.Client()
+
+        self._inflight_packets: dict[int, InflightPacket] = {}
+
+        init_status = bindings.tb_client_init(
+            ctypes.byref(self._client),
+            cluster_id,
+            replica_addresses.encode("ascii"),
+            len(replica_addresses),
+            self._client_key,
+            self._c_on_completion
+        )
+        if init_status != bindings.Status.SUCCESS:
+            raise InitError(init_status)
+
+        Client._clients[self._client_key] = self
+
+
+    def _acquire_packet(self, operation: bindings.Operation, operations: Any,
+                        c_event_type: Any, c_result_type: Any) -> InflightPacket:
+        packet = bindings.CPacket()
+        packet.next = None
+        packet.user_data = Client._counter.increment()
+        packet.operation = operation
+        packet.status = bindings.PacketStatus.OK
+
+        operations_array_type = c_event_type * len(operations)
+        operations_array = operations_array_type(*map(c_event_type.from_param, operations))
+
+        packet.data_size = ctypes.sizeof(operations_array)
+        packet.data = ctypes.cast(operations_array, ctypes.c_void_p)
+
+        return InflightPacket(
+            packet=packet,
+            response=None,
+            on_completion=None,
+            on_completion_context=None,
+            operation=operation,
+            c_event_type=c_event_type,
+            c_result_type=c_result_type)
+
+    def close(self):
+        bindings.tb_client_deinit(self._client)
+        tb_assert(self._client is not None)
+        tb_assert(len(self._inflight_packets) == 0)
+        del Client._clients[self._client_key]
+
+    @staticmethod
+    @bindings.OnCompletion
+    def _c_on_completion(completion_ctx, tb_client, packet, timestamp, bytes_ptr, len_):
+        """
+        Invoked in a separate thread
+        """
+        self: Client = Client._clients[completion_ctx]
+        tb_assert(self._client.value == tb_client)
+
+        packet = ctypes.cast(packet, ctypes.POINTER(bindings.CPacket))
+        inflight_packet = self._inflight_packets[packet[0].user_data]
+
+        if packet[0].status != bindings.PacketStatus.OK.value:
+            inflight_packet.response = PacketError(repr(bindings.PacketStatus(packet[0].status)))
+            tb_assert(inflight_packet.on_completion is not None)
+            inflight_packet.on_completion(inflight_packet)
+            return
+
+        c_result_type = inflight_packet.c_result_type
+        tb_assert(len_ % ctypes.sizeof(c_result_type) == 0)
+
+        # The memory referenced in bytes_ptr is only valid for the duration of this callback. Copy
+        # it to a fresh, Python owned buffer and do the conversion from the raw C type to the Python
+        # dataclass.
+        results_slice = ctypes.cast(
+            bytes_ptr,
+            ctypes.POINTER(c_result_type)
+        )[0:(len_ // ctypes.sizeof(c_result_type))]
+        results = [result.to_python() for result in results_slice]
+
+        inflight_packet.response = results
+
+        tb_assert(inflight_packet.on_completion is not None)
+        inflight_packet.on_completion(inflight_packet)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+
+class ClientSync(Client, bindings.StateMachineMixin):
+    def _on_completion(self, inflight_packet):
+        inflight_packet.on_completion_context.event.set()
+
+    def _submit(self, operation: bindings.Operation, operations: list[Any],
+                c_event_type: Any, c_result_type: Any):
+        inflight_packet = self._acquire_packet(operation, operations, c_event_type, c_result_type)
+        self._inflight_packets[inflight_packet.packet.user_data] = inflight_packet
+
+        inflight_packet.on_completion = self._on_completion
+        inflight_packet.on_completion_context = CompletionContextSync(event=threading.Event())
+
+        bindings.tb_client_submit(self._client, ctypes.byref(inflight_packet.packet))
+        inflight_packet.on_completion_context.event.wait()
+
+        del self._inflight_packets[inflight_packet.packet.user_data]
+
+        if isinstance(inflight_packet.response, Exception):
+            raise inflight_packet.response
+
+        return inflight_packet.response
+
+
+class ClientAsync(Client, bindings.AsyncStateMachineMixin):
+    def _on_completion(self, inflight_packet):
+        """
+        Called by Client._c_on_completion, which itself is called from a different thread. Use
+        `call_soon_threadsafe` to return to the thread of the event loop the request was invoked
+        from, so _trigger_event() can trigger the async event and allow the client to progress.
+        """
+        inflight_packet.on_completion_context.loop.call_soon_threadsafe(
+            self._trigger_event,
+            inflight_packet
+        )
+
+    def _trigger_event(self, inflight_packet):
+        inflight_packet.on_completion_context.event.set()
+
+    async def _submit(self, operation: bindings.Operation, operations: Any,
+                      c_event_type: Any, c_result_type: Any):
+        inflight_packet = self._acquire_packet(operation, operations, c_event_type, c_result_type)
+        self._inflight_packets[inflight_packet.packet.user_data] = inflight_packet
+
+        inflight_packet.on_completion = self._on_completion
+        inflight_packet.on_completion_context = CompletionContextAsync(
+            loop=asyncio.get_event_loop(),
+            event=asyncio.Event()
+        )
+
+        bindings.tb_client_submit(self._client, ctypes.byref(inflight_packet.packet))
+        await inflight_packet.on_completion_context.event.wait()
+
+        del self._inflight_packets[inflight_packet.packet.user_data]
+
+        if isinstance(inflight_packet.response, Exception):
+            raise inflight_packet.response
+
+        return inflight_packet.response

tigerbeetle/lib.py

@@ -0,0 +1,88 @@
+import ctypes
+import dataclasses
+import platform
+from pathlib import Path
+
+
+class NativeError(Exception):
+    pass
+
+class IntegerOverflowError(ValueError):
+    pass
+
+
+def _load_tbclient():
+    prefix = ""
+    arch = ""
+    system = ""
+    linux_libc = ""
+    suffix = ""
+
+    platform_machine = platform.machine().lower()
+
+    if platform_machine == "x86_64" or platform_machine == "amd64":
+        arch = "x86_64"
+    elif platform_machine == "aarch64" or platform_machine == "arm64":
+        arch = "aarch64"
+    else:
+        raise NativeError("Unsupported machine: " + platform.machine())
+
+    if platform.system() == "Linux":
+        prefix = "lib"
+        system = "linux"
+        suffix = ".so"
+        libc = platform.libc_ver()[0]
+        if libc == "glibc":
+            linux_libc = "-gnu.2.27"
+        elif libc == "musl":
+            linux_libc = "-musl"
+        else:
+            raise NativeError("Unsupported libc: " + libc)
+    elif platform.system() == "Darwin":
+        prefix = "lib"
+        system = "macos"
+        suffix = ".dylib"
+    elif platform.system() == "Windows":
+        system = "windows"
+        suffix = ".dll"
+    else:
+        raise NativeError("Unsupported system: " + platform.system())
+
+    source_path = Path(__file__)
+    source_dir = source_path.parent
+    library_path = source_dir / "lib" / f"{arch}-{system}{linux_libc}" / f"{prefix}tb_client{suffix}"
+    return ctypes.CDLL(str(library_path))
+
+
+def validate_uint(*, bits: int, name: str, number: int):
+    if number > 2**bits - 1:
+        raise IntegerOverflowError(f"{name}=={number} is too large to fit in {bits} bits")
+    if number < 0:
+        raise IntegerOverflowError(f"{name}=={number} cannot be negative")
+
+
+class c_uint128(ctypes.Structure): # noqa: N801
+    _fields_ = [("_low", ctypes.c_uint64), ("_high", ctypes.c_uint64)] # noqa: RUF012
+
+    @classmethod
+    def from_param(cls, obj):
+        return cls(_high=obj >> 64, _low=obj & 0xffffffffffffffff)
+
+    def to_python(self):
+        return self._high << 64 | self._low
+
+
+# Use slots=True if the version of Python is new enough (3.10+) to support it.
+try:
+    dataclass = dataclasses.dataclass(slots=True)
+except TypeError:
+    dataclass = dataclasses.dataclass()
+
+def tb_assert(value):
+    """
+    Python's built-in assert can be silently disabled if Python is run with -O.
+    """
+    if not value:
+        raise AssertionError()
+
+tbclient = _load_tbclient()