opengris-scaler 1.12.37__cp38-cp38-musllinux_1_2_x86_64.whl

scaler/io/async_object_storage_connector.py

@@ -0,0 +1,225 @@
+import asyncio
+import logging
+import socket
+import struct
+import uuid
+from typing import Dict, Optional, Tuple
+
+from scaler.io.mixins import AsyncObjectStorageConnector
+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._reader: Optional[asyncio.StreamReader] = None
+        self._writer: Optional[asyncio.StreamWriter] = None
+
+        self._next_request_id = 0
+        self._pending_get_requests: Dict[ObjectID, asyncio.Future] = {}
+
+        self._identity: bytes = (
+            f"{self.__class__.__name__}|{socket.gethostname().split('.')[0]}|{uuid.uuid4()}".encode()
+        )
+
+    def __del__(self):
+        if not self.is_connected():
+            return
+
+        self._writer.close()
+
+    async def connect(self, host: str, port: int):
+        self._host = host
+        self._port = port
+
+        if self.is_connected():
+            raise ObjectStorageException("connector is already connected.")
+
+        self._reader, self._writer = await asyncio.open_connection(self._host, self._port)
+        await self.__read_framed_message()
+        self.__write_framed(self._identity)
+
+        try:
+            await self._writer.drain()
+        except ConnectionResetError:
+            self.__raise_connection_failure()
+
+        # Makes sure the socket is TCP_NODELAY. It seems to be the case by default, but that's not specified in the
+        # asyncio's documentation and might change in the future.
+        self._writer.get_extra_info("socket").setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+
+        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
+
+        if not self._writer.is_closing:
+            self._writer.close()
+
+        await self._writer.wait_closed()
+
+    @property
+    def reader(self) -> Optional[asyncio.StreamReader]:
+        return self._reader
+
+    @property
+    def writer(self) -> Optional[asyncio.StreamWriter]:
+        return self._writer
+
+    @property
+    def address(self) -> str:
+        self.__ensure_is_connected()
+        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._writer is None:
+            raise ObjectStorageException("connector is not connected.")
+
+        if self._writer.is_closing():
+            raise ObjectStorageException("connection is closed.")
+
+    async def __send_request(
+        self,
+        object_id: ObjectID,
+        payload_length: int,
+        request_type: ObjectRequestHeader.ObjectRequestType,
+        payload: Optional[bytes],
+    ):
+        self.__ensure_is_connected()
+        assert self._writer 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)
+
+        self.__write_request_header(header)
+
+        if payload is not None:
+            self.__write_request_payload(payload)
+
+        try:
+            await self._writer.drain()
+        except ConnectionResetError:
+            self.__raise_connection_failure()
+
+    def __write_request_header(self, header: ObjectRequestHeader):
+        assert self._writer is not None
+        self.__write_framed(header.get_message().to_bytes())
+
+    def __write_request_payload(self, payload: bytes):
+        assert self._writer is not None
+        self.__write_framed(payload)
+
+    async def __receive_response(self) -> Optional[Tuple[ObjectResponseHeader, bytes]]:
+        assert self._reader is not None
+
+        if self._writer.is_closing():
+            return None
+
+        try:
+            header = await self.__read_response_header()
+            payload = await self.__read_response_payload(header)
+        except asyncio.IncompleteReadError:
+            self.__raise_connection_failure()
+
+        return header, payload
+
+    async def __read_response_header(self) -> ObjectResponseHeader:
+        assert self._reader is not None
+
+        header_data = await self.__read_framed_message()
+        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._reader is not None
+
+        if header.payload_length > 0:
+            res = await self.__read_framed_message()
+            assert len(res) == header.payload_length
+            return res
+        else:
+            return b""
+
+    async def __read_framed_message(self) -> bytes:
+        length_bytes = await self._reader.readexactly(8)
+        (payload_length,) = struct.unpack("<Q", length_bytes)
+        return await self._reader.readexactly(payload_length) if payload_length > 0 else bytes()
+
+    def __write_framed(self, payload: bytes):
+        self._writer.write(struct.pack("<Q", len(payload)))
+        self._writer.write(payload)
+        return
+
+    @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.config.types.zmq import ZMQConfig
+from scaler.io.mixins import SyncConnector
+from scaler.io.utility import deserialize, serialize
+from scaler.protocol.python.mixins import Message
+
+
+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,249 @@
+import collections
+import socket
+import struct
+import uuid
+from threading import Lock
+from typing import Iterable, List, Optional, Tuple
+
+from scaler.io.mixins import SyncObjectStorageConnector
+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._identity: bytes = (
+            f"{self.__class__.__name__}|{socket.gethostname().split('.')[0]}|{uuid.uuid4()}".encode()
+        )
+
+        self._socket: Optional[socket.socket] = socket.create_connection((self._host, self._port))
+        self._socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+
+        self._next_request_id = 0
+
+        self._socket_lock = Lock()
+
+        self.__send_buffers([struct.pack("<Q", len(self._identity)), self._identity])
+        self.__read_framed_message()  # receive server identity
+
+    def __del__(self):
+        self.destroy()
+
+    def destroy(self):
+        with self._socket_lock:
+            if self._socket is not None:
+                self._socket.close()
+                self._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._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._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.__send_buffers(
+                [struct.pack("<Q", len(header_bytes)), header_bytes, struct.pack("<Q", len(payload)), payload]
+            )
+        else:
+            self.__send_buffers([struct.pack("<Q", len(header_bytes)), header_bytes])
+
+    def __send_buffers(self, buffers: List[bytes]) -> None:
+        if len(buffers) < 1:
+            return
+
+        assert self._socket is not None
+
+        total_size = sum(len(buffer) for buffer in buffers)
+
+        # If the message is small enough, first try to send it at once with sendmsg(). This would ensure the message can
+        # be transmitted within a single TCP segment.
+        if total_size < MAX_CHUNK_SIZE:
+            sent = self._socket.sendmsg(buffers)  # type: ignore[attr-defined]
+
+            if sent <= 0:
+                self.__raise_connection_failure()
+
+            remaining_buffers = collections.deque(buffers)
+            while sent > len(remaining_buffers[0]):
+                removed_buffer = remaining_buffers.popleft()
+                sent -= len(removed_buffer)
+
+            if sent > 0:
+                # Truncate the first partially sent buffer
+                remaining_buffers[0] = memoryview(remaining_buffers[0])[sent:]
+
+            buffers = list(remaining_buffers)
+
+        # Send the remaining buffers sequentially
+        for buffer in buffers:
+            self.__send_buffer(buffer)
+
+    def __send_buffer(self, buffer: bytes) -> None:
+        buffer_view = memoryview(buffer)
+
+        total_sent = 0
+        while total_sent < len(buffer):
+            sent = self._socket.send(buffer_view[total_sent : MAX_CHUNK_SIZE + total_sent])
+
+            if sent <= 0:
+                self.__raise_connection_failure()
+
+            total_sent += sent
+
+    def __receive_response(self) -> Tuple[ObjectResponseHeader, bytearray]:
+        assert self._socket is not None
+
+        header = self.__read_response_header()
+        payload = self.__read_response_payload(header)
+
+        return header, payload
+
+    def __read_response_header(self) -> ObjectResponseHeader:
+        assert self._socket is not None
+
+        header_bytearray = self.__read_framed_message()
+
+        # 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.__read_framed_message()
+            assert len(res) == header.payload_length
+            return res
+        else:
+            return bytearray()
+
+    def __read_exactly(self, length: int) -> bytearray:
+        buffer = bytearray(length)
+
+        total_received = 0
+        while total_received < length:
+            chunk_size = min(MAX_CHUNK_SIZE, length - total_received)
+            received = self._socket.recv_into(memoryview(buffer)[total_received:], chunk_size)
+
+            if received <= 0:
+                self.__raise_connection_failure()
+
+            total_received += received
+
+        return buffer
+
+    def __read_framed_message(self) -> bytearray:
+        length_bytes = self.__read_exactly(8)
+        (payload_length,) = struct.unpack("<Q", length_bytes)
+        return self.__read_exactly(payload_length) if payload_length > 0 else bytearray()
+
+    @staticmethod
+    def __raise_connection_failure():
+        raise ObjectStorageException("connection failure to object storage server.")