opengris-scaler 1.12.7__cp310-cp310-manylinux_2_28_x86_64.whl

scaler/io/async_connector.py

@@ -0,0 +1,95 @@
+import logging
+import os
+import uuid
+from typing import Awaitable, Callable, Literal, Optional
+
+import zmq.asyncio
+
+from scaler.io.mixins import AsyncConnector
+from scaler.io.utility import deserialize, serialize
+from scaler.protocol.python.mixins import Message
+from scaler.config.types.zmq import ZMQConfig
+
+
+class ZMQAsyncConnector(AsyncConnector):
+    def __init__(
+        self,
+        context: zmq.asyncio.Context,
+        name: str,
+        socket_type: int,
+        address: ZMQConfig,
+        bind_or_connect: Literal["bind", "connect"],
+        callback: Optional[Callable[[Message], Awaitable[None]]],
+        identity: Optional[bytes],
+    ):
+        self._address = address
+
+        self._context = context
+        self._socket = self._context.socket(socket_type)
+
+        if identity is None:
+            identity = f"{os.getpid()}|{name}|{uuid.uuid4().bytes.hex()}".encode()
+        self._identity = identity
+
+        # set socket option
+        self._socket.setsockopt(zmq.IDENTITY, self._identity)
+        self._socket.setsockopt(zmq.SNDHWM, 0)
+        self._socket.setsockopt(zmq.RCVHWM, 0)
+
+        if bind_or_connect == "bind":
+            self._socket.bind(self.address)
+        elif bind_or_connect == "connect":
+            self._socket.connect(self.address)
+        else:
+            raise TypeError("bind_or_connect has to be 'bind' or 'connect'")
+
+        self._callback: Optional[Callable[[Message], Awaitable[None]]] = callback
+
+    def __del__(self):
+        self.destroy()
+
+    def destroy(self):
+        if self._socket.closed:
+            return
+
+        self._socket.close(linger=1)
+
+    @property
+    def identity(self) -> bytes:
+        return self._identity
+
+    @property
+    def socket(self) -> zmq.asyncio.Socket:
+        return self._socket
+
+    @property
+    def address(self) -> str:
+        return self._address.to_address()
+
+    async def routine(self):
+        if self._callback is None:
+            return
+
+        message: Optional[Message] = await self.receive()
+        if message is None:
+            return
+
+        await self._callback(message)
+
+    async def receive(self) -> Optional[Message]:
+        if self._context.closed:
+            return None
+
+        if self._socket.closed:
+            return None
+
+        payload = await self._socket.recv(copy=False)
+        result: Optional[Message] = deserialize(payload.bytes)
+        if result is None:
+            logging.error(f"received unknown message: {payload.bytes!r}")
+            return None
+
+        return result
+
+    async def send(self, message: Message):
+        await self._socket.send(serialize(message), copy=False)

scaler/io/async_object_storage_connector.py

@@ -0,0 +1,185 @@
+import asyncio
+import logging
+import os
+import socket
+import uuid
+from typing import Dict, Optional, Tuple
+
+from scaler.io.mixins import AsyncObjectStorageConnector
+from scaler.io.ymq.ymq import IOSocketType, IOContext, Message, YMQException
+from scaler.protocol.capnp._python import _object_storage  # noqa
+from scaler.protocol.python.object_storage import ObjectRequestHeader, ObjectResponseHeader, to_capnp_object_id
+from scaler.utility.exceptions import ObjectStorageException
+from scaler.utility.identifiers import ObjectID
+
+
+class PyAsyncObjectStorageConnector(AsyncObjectStorageConnector):
+    """An asyncio connector that uses an raw TCP socket to connect to a Scaler's object storage instance."""
+
+    def __init__(self):
+        self._host: Optional[str] = None
+        self._port: Optional[int] = None
+
+        self._connected_event = asyncio.Event()
+
+        self._next_request_id = 0
+        self._pending_get_requests: Dict[ObjectID, asyncio.Future] = {}
+
+        self._lock = asyncio.Lock()
+        self._identity: str = f"{os.getpid()}|{socket.gethostname().split('.')[0]}|{uuid.uuid4()}"
+        self._io_context: IOContext = IOContext()
+        self._io_socket = self._io_context.createIOSocket_sync(self._identity, IOSocketType.Connector)
+
+    def __del__(self):
+        if not self.is_connected():
+            return
+        self._io_socket = None
+
+    async def connect(self, host: str, port: int):
+        self._host = host
+        self._port = port
+
+        if self.is_connected():
+            raise ObjectStorageException("connector is already connected.")
+        await self._io_socket.connect(self.address)
+        self._connected_event.set()
+
+    async def wait_until_connected(self):
+        await self._connected_event.wait()
+
+    def is_connected(self) -> bool:
+        return self._connected_event.is_set()
+
+    async def destroy(self):
+        if not self.is_connected():
+            return
+        self._io_socket = None
+
+    @property
+    def address(self) -> str:
+        return f"tcp://{self._host}:{self._port}"
+
+    async def routine(self):
+        await self.wait_until_connected()
+
+        response = await self.__receive_response()
+        if response is None:
+            return
+
+        header, payload = response
+
+        if header.response_type != ObjectResponseHeader.ObjectResponseType.GetOK:
+            return
+
+        pending_get_future = self._pending_get_requests.pop(header.object_id, None)
+
+        if pending_get_future is None:
+            logging.warning(f"unknown get-ok response for unrequested object_id={repr(header.object_id)}.")
+            return
+
+        pending_get_future.set_result(payload)
+
+    async def set_object(self, object_id: ObjectID, payload: bytes) -> None:
+        await self.__send_request(object_id, len(payload), ObjectRequestHeader.ObjectRequestType.SetObject, payload)
+
+    async def get_object(self, object_id: ObjectID, max_payload_length: int = 2**64 - 1) -> bytes:
+        pending_get_future = self._pending_get_requests.get(object_id)
+
+        if pending_get_future is None:
+            pending_get_future = asyncio.Future()
+            self._pending_get_requests[object_id] = pending_get_future
+
+            await self.__send_request(
+                object_id, max_payload_length, ObjectRequestHeader.ObjectRequestType.GetObject, None
+            )
+
+        return await pending_get_future
+
+    async def delete_object(self, object_id: ObjectID) -> None:
+        await self.__send_request(object_id, 0, ObjectRequestHeader.ObjectRequestType.DeleteObject, None)
+
+    async def duplicate_object_id(self, object_id: ObjectID, new_object_id: ObjectID) -> None:
+        object_id_payload = to_capnp_object_id(object_id).to_bytes()
+
+        await self.__send_request(
+            new_object_id,
+            len(object_id_payload),
+            ObjectRequestHeader.ObjectRequestType.DuplicateObjectID,
+            object_id_payload,
+        )
+
+    def __ensure_is_connected(self):
+        if self._io_socket is None:
+            raise ObjectStorageException("connector is not connected.")
+
+    async def __send_request(
+        self,
+        object_id: ObjectID,
+        payload_length: int,
+        request_type: ObjectRequestHeader.ObjectRequestType,
+        payload: Optional[bytes],
+    ):
+        self.__ensure_is_connected()
+
+        request_id = self._next_request_id
+        self._next_request_id += 1
+        self._next_request_id %= 2**64 - 1  # UINT64_MAX
+
+        header = ObjectRequestHeader.new_msg(object_id, payload_length, request_id, request_type)
+
+        try:
+            async with self._lock:
+                await self.__write_request_header(header)
+
+                if payload is not None:
+                    await self.__write_request_payload(payload)
+
+        except YMQException:
+            self._io_socket = None
+            self.__raise_connection_failure()
+
+    async def __write_request_header(self, header: ObjectRequestHeader):
+        assert self._io_socket is not None
+        await self._io_socket.send(Message(address=None, payload=header.get_message().to_bytes()))
+
+    async def __write_request_payload(self, payload: bytes):
+        assert self._io_socket is not None
+        await self._io_socket.send(Message(address=None, payload=payload))
+
+    async def __receive_response(self) -> Optional[Tuple[ObjectResponseHeader, bytes]]:
+        if self._io_socket is None:
+            return None
+
+        try:
+            header = await self.__read_response_header()
+            payload = await self.__read_response_payload(header)
+        except YMQException:
+            self._io_socket = None
+            self.__raise_connection_failure()
+
+        return header, payload
+
+    async def __read_response_header(self) -> ObjectResponseHeader:
+        assert self._io_socket is not None
+
+        msg = await self._io_socket.recv()
+        header_data = msg.payload.data
+        assert len(header_data) == ObjectResponseHeader.MESSAGE_LENGTH
+
+        with _object_storage.ObjectResponseHeader.from_bytes(header_data) as header_message:
+            return ObjectResponseHeader(header_message)
+
+    async def __read_response_payload(self, header: ObjectResponseHeader) -> bytes:
+        assert self._io_socket is not None
+        # assert self._reader is not None
+
+        if header.payload_length > 0:
+            res = await self._io_socket.recv()
+            assert len(res.payload) == header.payload_length
+            return res.payload.data
+        else:
+            return b""
+
+    @staticmethod
+    def __raise_connection_failure():
+        raise ObjectStorageException("connection failure to object storage server.")

scaler/io/mixins.py

@@ -0,0 +1,154 @@
+import abc
+from typing import Awaitable, Callable, Optional
+
+from scaler.protocol.python.mixins import Message
+from scaler.protocol.python.status import BinderStatus
+from scaler.utility.identifiers import ObjectID
+from scaler.utility.mixins import Looper, Reporter
+
+
+class AsyncBinder(Looper, Reporter, metaclass=abc.ABCMeta):
+    @property
+    @abc.abstractmethod
+    def identity(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def destroy(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def register(self, callback: Callable[[bytes, Message], Awaitable[None]]):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def send(self, to: bytes, message: Message):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_status(self) -> BinderStatus:
+        raise NotImplementedError()
+
+
+class AsyncConnector(Looper, metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def destroy(self):
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def identity(self) -> bytes:
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def address(self) -> str:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def send(self, message: Message):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def receive(self) -> Optional[Message]:
+        raise NotImplementedError()
+
+
+class SyncConnector(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def destroy(self):
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def identity(self) -> bytes:
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def address(self) -> str:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def send(self, message: Message):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def receive(self) -> Optional[Message]:
+        raise NotImplementedError()
+
+
+class AsyncObjectStorageConnector(Looper, metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def connect(self, host: str, port: int):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def wait_until_connected(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def is_connected(self) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def destroy(self):
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def address(self) -> str:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def set_object(self, object_id: ObjectID, payload: bytes) -> None:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def get_object(self, object_id: ObjectID, max_payload_length: int = 2**64 - 1) -> bytes:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def delete_object(self, object_id: ObjectID) -> None:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def duplicate_object_id(self, object_id: ObjectID, new_object_id: ObjectID) -> None:
+        raise NotImplementedError()
+
+
+class SyncObjectStorageConnector(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def destroy(self):
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def address(self) -> str:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def set_object(self, object_id: ObjectID, payload: bytes):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_object(self, object_id: ObjectID, max_payload_length: int = 2**64 - 1) -> bytearray:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def delete_object(self, object_id: ObjectID) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def duplicate_object_id(self, object_id: ObjectID, new_object_id: ObjectID) -> None:
+        raise NotImplementedError()
+
+
+class SyncSubscriber(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def destroy(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def run(self) -> None:
+        raise NotImplementedError()

scaler/io/sync_connector.py

@@ -0,0 +1,68 @@
+import logging
+import os
+import socket
+import threading
+import uuid
+from typing import Optional
+
+import zmq
+
+from scaler.io.mixins import SyncConnector
+from scaler.io.utility import deserialize, serialize
+from scaler.protocol.python.mixins import Message
+from scaler.config.types.zmq import ZMQConfig
+
+
+class ZMQSyncConnector(SyncConnector):
+    def __init__(self, context: zmq.Context, socket_type: int, address: ZMQConfig, identity: Optional[bytes]):
+        self._address = address
+
+        self._context = context
+        self._socket = self._context.socket(socket_type)
+
+        self._identity: bytes = (
+            f"{os.getpid()}|{socket.gethostname().split('.')[0]}|{uuid.uuid4()}".encode()
+            if identity is None
+            else identity
+        )
+
+        # set socket option
+        self._socket.setsockopt(zmq.IDENTITY, self._identity)
+        self._socket.setsockopt(zmq.SNDHWM, 0)
+        self._socket.setsockopt(zmq.RCVHWM, 0)
+
+        self._socket.connect(self._address.to_address())
+
+        self._lock = threading.Lock()
+
+    def destroy(self):
+        self._socket.close()
+
+    @property
+    def address(self) -> str:
+        return self._address.to_address()
+
+    @property
+    def identity(self) -> bytes:
+        return self._identity
+
+    def send(self, message: Message):
+        with self._lock:
+            self._socket.send(serialize(message), copy=False)
+
+    def receive(self) -> Optional[Message]:
+        with self._lock:
+            payload = self._socket.recv(copy=False)
+
+        return self.__compose_message(payload.bytes)
+
+    def __compose_message(self, payload: bytes) -> Optional[Message]:
+        result: Optional[Message] = deserialize(payload)
+        if result is None:
+            logging.error(f"{self.__get_prefix()}: received unknown message: {payload!r}")
+            return None
+
+        return result
+
+    def __get_prefix(self):
+        return f"{self.__class__.__name__}[{self._identity.decode()}]:"

scaler/io/sync_object_storage_connector.py

@@ -0,0 +1,185 @@
+import os
+import socket
+import uuid
+from threading import Lock
+from typing import Iterable, Optional
+
+from scaler.io.mixins import SyncObjectStorageConnector
+from scaler.io.ymq.ymq import IOContext, IOSocket, IOSocketType, Message, YMQException
+from scaler.protocol.capnp._python import _object_storage  # noqa
+from scaler.protocol.python.object_storage import ObjectRequestHeader, ObjectResponseHeader, to_capnp_object_id
+from scaler.utility.exceptions import ObjectStorageException
+from scaler.utility.identifiers import ObjectID
+
+# Some OSes raise an OSError when sending buffers too large with send() or sendmsg().
+MAX_CHUNK_SIZE = 128 * 1024 * 1024
+
+
+class PySyncObjectStorageConnector(SyncObjectStorageConnector):
+    """An synchronous connector that uses an raw TCP socket to connect to a Scaler's object storage instance."""
+
+    def __init__(self, host: str, port: int):
+        self._host = host
+        self._port = port
+
+        self._next_request_id = 0
+
+        self._socket_lock = Lock()
+
+        self._identity: str = f"{os.getpid()}|{socket.gethostname().split('.')[0]}|{uuid.uuid4()}"
+        self._io_context = IOContext()
+        self._io_socket: IOSocket = self._io_context.createIOSocket_sync(self._identity, IOSocketType.Connector)
+        self._io_socket.connect_sync(self.address)
+
+    def __del__(self):
+        self.destroy()
+
+    def destroy(self):
+        with self._socket_lock:
+            if self._io_socket is not None:
+                self._io_socket = None
+
+    @property
+    def address(self) -> str:
+        return f"tcp://{self._host}:{self._port}"
+
+    def set_object(self, object_id: ObjectID, payload: bytes):
+        """
+        Sets the object's payload on the object storage server.
+        """
+
+        with self._socket_lock:
+            self.__send_request(object_id, len(payload), ObjectRequestHeader.ObjectRequestType.SetObject, payload)
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(response_header, [ObjectResponseHeader.ObjectResponseType.SetOK])
+        self.__ensure_empty_payload(response_payload)
+
+    def get_object(self, object_id: ObjectID, max_payload_length: int = 2**64 - 1) -> bytearray:
+        """
+        Returns the object's payload from the object storage server.
+
+        Will block until the object is available.
+        """
+
+        with self._socket_lock:
+            self.__send_request(object_id, max_payload_length, ObjectRequestHeader.ObjectRequestType.GetObject)
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(response_header, [ObjectResponseHeader.ObjectResponseType.GetOK])
+
+        return response_payload
+
+    def delete_object(self, object_id: ObjectID) -> bool:
+        """
+        Removes the object from the object storage server.
+
+        Returns `False` if the object wasn't found in the server. Otherwise returns `True`.
+        """
+
+        with self._socket_lock:
+            self.__send_request(object_id, 0, ObjectRequestHeader.ObjectRequestType.DeleteObject)
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(
+            response_header,
+            [ObjectResponseHeader.ObjectResponseType.DelOK, ObjectResponseHeader.ObjectResponseType.DelNotExists],
+        )
+        self.__ensure_empty_payload(response_payload)
+
+        return response_header.response_type == ObjectResponseHeader.ObjectResponseType.DelOK
+
+    def duplicate_object_id(self, object_id: ObjectID, new_object_id: ObjectID) -> None:
+        """
+        Link an object's content to a new object ID on the object storage server.
+        """
+
+        object_id_payload = to_capnp_object_id(object_id).to_bytes()
+
+        with self._socket_lock:
+            self.__send_request(
+                new_object_id,
+                len(object_id_payload),
+                ObjectRequestHeader.ObjectRequestType.DuplicateObjectID,
+                object_id_payload,
+            )
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(response_header, [ObjectResponseHeader.ObjectResponseType.DuplicateOK])
+        self.__ensure_empty_payload(response_payload)
+
+    def __ensure_is_connected(self):
+        if self._io_socket is None:
+            raise ObjectStorageException("connector is closed.")
+
+    def __ensure_response_type(
+        self, header: ObjectResponseHeader, valid_response_types: Iterable[ObjectResponseHeader.ObjectResponseType]
+    ):
+        if header.response_type not in valid_response_types:
+            raise RuntimeError(f"unexpected object storage response_type={header.response_type}.")
+
+    def __ensure_empty_payload(self, payload: bytearray):
+        if len(payload) != 0:
+            raise RuntimeError(f"unexpected response payload_length={len(payload)}, expected 0.")
+
+    def __send_request(
+        self,
+        object_id: ObjectID,
+        payload_length: int,
+        request_type: ObjectRequestHeader.ObjectRequestType,
+        payload: Optional[bytes] = None,
+    ):
+        self.__ensure_is_connected()
+        assert self._io_socket is not None
+
+        request_id = self._next_request_id
+        self._next_request_id += 1
+        self._next_request_id %= 2**64 - 1  # UINT64_MAX
+
+        header = ObjectRequestHeader.new_msg(object_id, payload_length, request_id, request_type)
+        header_bytes = header.get_message().to_bytes()
+
+        if payload is not None:
+            self._io_socket.send_sync(Message(address=None, payload=header_bytes))
+            self._io_socket.send_sync(Message(address=None, payload=payload))
+        else:
+            self._io_socket.send_sync(Message(address=None, payload=header_bytes))
+
+    def __receive_response(self):
+        assert self._io_socket is not None
+
+        try:
+            header = self.__read_response_header()
+            payload = self.__read_response_payload(header)
+            return header, payload
+        except YMQException:
+            self.__raise_connection_failure()
+
+    def __read_response_header(self) -> ObjectResponseHeader:
+        assert self._io_socket is not None
+
+        header_bytes = self._io_socket.recv_sync().payload.data
+        if header_bytes is None:
+            self.__raise_connection_failure()
+
+        # pycapnp does not like to read from a bytearray object. This look like an not-yet-resolved issue.
+        # That's is annoying because it leads to an unnecessary copy of the header's buffer.
+        # See https://github.com/capnproto/pycapnp/issues/153
+        # header_bytes = bytes(header_bytearray)
+
+        with _object_storage.ObjectResponseHeader.from_bytes(header_bytes) as header_message:
+            return ObjectResponseHeader(header_message)
+
+    def __read_response_payload(self, header: ObjectResponseHeader) -> bytearray:
+        if header.payload_length > 0:
+            res = self._io_socket.recv_sync().payload.data
+            if res is None:
+                self.__raise_connection_failure()
+            assert len(res) == header.payload_length
+            return bytearray(res)
+        else:
+            return bytearray()
+
+    @staticmethod
+    def __raise_connection_failure():
+        raise ObjectStorageException("connection failure to object storage server.")