opengris-scaler 1.12.7__cp311-cp311-manylinux_2_28_x86_64.whl

scaler/io/ymq/message_connection_tcp.h

@@ -0,0 +1,96 @@
+#pragma once
+
+#include <deque>
+#include <memory>
+#include <optional>
+#include <queue>
+
+#include "scaler/io/ymq/configuration.h"
+#include "scaler/io/ymq/io_socket.h"
+#include "scaler/io/ymq/message_connection.h"
+#include "scaler/io/ymq/tcp_operations.h"
+
+namespace scaler {
+namespace ymq {
+
+class EventLoopThread;
+class EventManager;
+
+class MessageConnectionTCP: public MessageConnection {
+public:
+    using SendMessageCallback = Configuration::SendMessageCallback;
+    using RecvMessageCallback = Configuration::RecvMessageCallback;
+
+    MessageConnectionTCP(
+        std::shared_ptr<EventLoopThread> eventLoopThread,
+        int connFd,
+        sockaddr localAddr,
+        sockaddr remoteAddr,
+        std::string localIOSocketIdentity,
+        bool responsibleForRetry,
+        std::shared_ptr<std::queue<RecvMessageCallback>> _pendingRecvMessageCallbacks) noexcept;
+
+    MessageConnectionTCP(
+        std::shared_ptr<EventLoopThread> eventLoopThread,
+        std::string localIOSocketIdentity,
+        std::string remoteIOSocketIdentity,
+        std::shared_ptr<std::queue<RecvMessageCallback>> _pendingRecvMessageCallbacks) noexcept;
+
+    ~MessageConnectionTCP() noexcept;
+
+    void onCreated();
+
+    void sendMessage(Message msg, SendMessageCallback onMessageSent);
+    bool recvMessage();
+    void disconnect();
+
+    std::shared_ptr<EventLoopThread> _eventLoopThread;
+    const sockaddr _remoteAddr;
+    const bool _responsibleForRetry;
+    std::optional<std::string> _remoteIOSocketIdentity;
+
+private:
+    enum class IOError {
+        Drained,
+        Aborted,
+        Disconnected,
+        MessageTooLarge,
+    };
+
+    void onRead();
+    void onWrite();
+    void onClose();
+    void onError()
+    {
+        onRead();
+        // onClose();
+    };
+
+    std::expected<void, IOError> tryReadOneMessage();
+    std::expected<void, IOError> tryReadMessages();
+    std::expected<size_t, IOError> trySendQueuedMessages();
+    void updateWriteOperations(size_t n);
+    void updateReadOperation();
+
+    void setRemoteIdentity() noexcept;
+
+    std::unique_ptr<EventManager> _eventManager;
+    int _connFd;
+    sockaddr _localAddr;
+    std::string _localIOSocketIdentity;
+
+    std::deque<TcpWriteOperation> _writeOperations;
+    size_t _sendCursor;
+
+    std::shared_ptr<std::queue<RecvMessageCallback>> _pendingRecvMessageCallbacks;
+    std::queue<TcpReadOperation> _receivedReadOperations;
+
+    bool _disconnect;
+
+    constexpr static bool isCompleteMessage(const TcpReadOperation& x);
+    friend void IOSocket::onConnectionIdentityReceived(MessageConnectionTCP* conn) noexcept;
+    friend void IOSocket::onConnectionDisconnected(MessageConnectionTCP* conn, bool keepInBook) noexcept;
+};
+
+}  // namespace ymq
+}  // namespace scaler

scaler/io/ymq/network_utils.h

@@ -0,0 +1,179 @@
+#pragma once
+
+#ifdef __linux__
+#include <arpa/inet.h>
+#include <netinet/in.h>
+#include <netinet/tcp.h>
+#include <unistd.h>
+#endif  // __linux__
+#ifdef _WIN32
+// clang-format off
+#include <windows.h>
+#include <winsock2.h>
+#include <ws2tcpip.h>
+// clang-format on
+#endif  // _WIN32
+
+#include <string.h>
+
+#include <cassert>
+#include <expected>
+
+#include "scaler/io/ymq/error.h"
+
+namespace scaler {
+namespace ymq {
+
+inline auto GetErrorCode()
+{
+#ifdef __linux__
+    return errno;
+#endif  // __linux__
+#ifdef _WIN32
+    return WSAGetLastError();
+#endif  // _WIN32
+}
+
+inline constexpr void CloseAndZeroSocket(auto& fd)
+{
+#ifdef __linux__
+    close(fd);
+#endif  // __linux__
+#ifdef _WIN32
+    closesocket(fd);
+#endif  // _WIN32
+    fd = 0;
+}
+
+inline std::expected<sockaddr, int> stringToSockaddr(const std::string& address)
+{
+    // Check and strip the "tcp://" prefix
+    static const std::string prefix = "tcp://";
+    if (address.substr(0, prefix.size()) != prefix) {
+        unrecoverableError({
+            Error::ErrorCode::InvalidAddressFormat,
+            "Originated from",
+            __PRETTY_FUNCTION__,
+            "Your input is",
+            address,
+        });
+    }
+
+    std::string addr_part = address.substr(prefix.size());
+    size_t colon_pos      = addr_part.find(':');
+    if (colon_pos == std::string::npos) {
+        unrecoverableError({
+            Error::ErrorCode::InvalidAddressFormat,
+            "Originated from",
+            __PRETTY_FUNCTION__,
+            "Your input is",
+            address,
+        });
+    }
+
+    std::string ip       = addr_part.substr(0, colon_pos);
+    std::string port_str = addr_part.substr(colon_pos + 1);
+
+    int port = 0;
+    try {
+        port = std::stoi(port_str);
+    } catch (...) {
+        unrecoverableError({
+            Error::ErrorCode::InvalidAddressFormat,
+            "Originated from",
+            __PRETTY_FUNCTION__,
+            "Your input is",
+            address,
+        });
+    }
+
+    sockaddr_in out_addr {};
+    out_addr.sin_family = AF_INET;
+    out_addr.sin_port   = htons(port);
+
+    int res = inet_pton(AF_INET, ip.c_str(), &out_addr.sin_addr);
+    if (res == 0) {
+        unrecoverableError({
+            Error::ErrorCode::InvalidAddressFormat,
+            "Originated from",
+            __PRETTY_FUNCTION__,
+            "Your input is",
+            address,
+        });
+    }
+
+    if (res == -1) {
+        unrecoverableError({
+            Error::ErrorCode::ConfigurationError,
+            "Originated from",
+            __PRETTY_FUNCTION__,
+            "Errno is",
+            strerror(errno),
+            "out_addr.sin_family",
+            out_addr.sin_family,
+            "out_addr.sin_port",
+            out_addr.sin_port,
+        });
+    }
+
+    return *(sockaddr*)&out_addr;
+}
+
+inline int setNoDelay(int fd)
+{
+    int optval = 1;
+    if (setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, (const char*)&optval, sizeof(optval)) == -1) {
+        unrecoverableError({
+            Error::ErrorCode::ConfigurationError,
+            "Originated from",
+            __PRETTY_FUNCTION__,
+            "Errno is",
+            strerror(errno),
+            "fd",
+            fd,
+        });
+    }
+
+    return fd;
+}
+
+inline sockaddr getLocalAddr(int fd)
+{
+    sockaddr localAddr     = {};
+    socklen_t localAddrLen = sizeof(localAddr);
+    if (getsockname(fd, &localAddr, &localAddrLen) == -1) {
+        unrecoverableError({
+            Error::ErrorCode::ConfigurationError,
+            "Originated from",
+            __PRETTY_FUNCTION__,
+            "Errno is",
+            strerror(errno),
+            "fd",
+            fd,
+        });
+    }
+    return localAddr;
+}
+
+inline sockaddr getRemoteAddr(int fd)
+{
+    sockaddr remoteAddr     = {};
+    socklen_t remoteAddrLen = sizeof(remoteAddr);
+
+    if (getpeername(fd, &remoteAddr, &remoteAddrLen) == -1) {
+        unrecoverableError({
+            Error::ErrorCode::ConfigurationError,
+            "Originated from",
+            __PRETTY_FUNCTION__,
+            "Errno is",
+            strerror(errno),
+            "fd",
+            fd,
+        });
+    }
+
+    return remoteAddr;
+}
+
+}  // namespace ymq
+}  // namespace scaler

scaler/io/ymq/pymod_ymq/bytes.h

@@ -0,0 +1,113 @@
+#pragma once
+
+// Python
+#include "scaler/io/ymq/pymod_ymq/python.h"
+
+// First-party
+#include "scaler/io/ymq/bytes.h"
+
+struct PyBytesYMQ {
+    PyObject_HEAD;
+    Bytes bytes;
+};
+
+extern "C" {
+
+static int PyBytesYMQ_init(PyBytesYMQ* self, PyObject* args, PyObject* kwds)
+{
+    Py_buffer view {.buf = nullptr};
+    const char* keywords[] = {"bytes", nullptr};
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|y*", (char**)keywords, &view))
+        return -1;  // Error parsing arguments
+
+    if (!view.buf) {
+        // If no bytes were provided, initialize with an empty Bytes object
+        self->bytes = Bytes();
+        return 0;
+    }
+
+    // copy the data into the Bytes object
+    // it might be possible to make this zero-copy in the future
+    self->bytes = Bytes((char*)view.buf, view.len);
+
+    PyBuffer_Release(&view);
+    return 0;
+}
+
+static void PyBytesYMQ_dealloc(PyBytesYMQ* self)
+{
+    try {
+        self->bytes.~Bytes();  // Call the destructor to free the Bytes object
+    } catch (...) {
+        PyErr_SetString(PyExc_RuntimeError, "Failed to deallocate Bytes");
+        PyErr_WriteUnraisable((PyObject*)self);
+    }
+
+    auto* tp = Py_TYPE(self);
+    tp->tp_free(self);
+    Py_DECREF(tp);
+}
+
+static PyObject* PyBytesYMQ_repr(PyBytesYMQ* self)
+{
+    if (self->bytes.is_null()) {
+        return PyUnicode_FromString("<Bytes: empty>");
+    } else {
+        return PyUnicode_FromFormat("<Bytes: %db>", self->bytes.len());
+    }
+}
+
+static PyObject* PyBytesYMQ_data_getter(PyBytesYMQ* self)
+{
+    if (self->bytes.is_null())
+        Py_RETURN_NONE;
+
+    return PyBytes_FromStringAndSize((const char*)self->bytes.data(), self->bytes.len());
+}
+
+static Py_ssize_t PyBytesYMQ_len(PyBytesYMQ* self)
+{
+    return self->bytes.len();
+}
+
+static PyObject* PyBytesYMQ_len_getter(PyBytesYMQ* self)
+{
+    return PyLong_FromSize_t(self->bytes.len());
+}
+
+static int PyBytesYMQ_getbuffer(PyBytesYMQ* self, Py_buffer* view, int flags)
+{
+    return PyBuffer_FillInfo(view, (PyObject*)self, (void*)self->bytes.data(), self->bytes.len(), true, flags);
+}
+
+static void PyBytesYMQ_releasebuffer(PyBytesYMQ* self, Py_buffer* view)
+{
+}
+}
+
+static PyGetSetDef PyBytesYMQ_properties[] = {
+    {"data", (getter)PyBytesYMQ_data_getter, nullptr, PyDoc_STR("Data of the Bytes object"), nullptr},
+    {"len", (getter)PyBytesYMQ_len_getter, nullptr, PyDoc_STR("Length of the Bytes object"), nullptr},
+    {nullptr, nullptr, nullptr, nullptr, nullptr},  // Sentinel
+};
+
+static PyType_Slot PyBytesYMQ_slots[] = {
+    {Py_tp_init, (void*)PyBytesYMQ_init},
+    {Py_tp_dealloc, (void*)PyBytesYMQ_dealloc},
+    {Py_tp_repr, (void*)PyBytesYMQ_repr},
+    {Py_mp_length, (void*)PyBytesYMQ_len},
+    {Py_tp_getset, (void*)PyBytesYMQ_properties},
+#if PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION > 8
+    {Py_bf_getbuffer, (void*)PyBytesYMQ_getbuffer},
+    {Py_bf_releasebuffer, (void*)PyBytesYMQ_releasebuffer},
+#endif
+    {0, nullptr},
+};
+
+static PyType_Spec PyBytesYMQ_spec = {
+    .name      = "_ymq.Bytes",
+    .basicsize = sizeof(PyBytesYMQ),
+    .itemsize  = 0,
+    .flags     = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE,
+    .slots     = PyBytesYMQ_slots,
+};

scaler/io/ymq/pymod_ymq/exception.h

@@ -0,0 +1,124 @@
+#pragma once
+
+// Python
+#include "scaler/io/ymq/pymod_ymq/python.h"
+
+// First-party
+#include "scaler/io/ymq/error.h"
+#include "scaler/io/ymq/pymod_ymq/ymq.h"
+
+using namespace scaler::ymq;
+
+// the order of the members in the exception args tuple
+const Py_ssize_t YMQException_errorCodeIndex = 0;
+const Py_ssize_t YMQException_messageIndex   = 1;
+
+struct YMQException {
+    PyException_HEAD;
+};
+
+extern "C" {
+
+static int YMQException_init(YMQException* self, PyObject* args, PyObject* kwds)
+{
+    auto state = YMQStateFromSelf((PyObject*)self);
+    if (!state)
+        return -1;
+
+    // no need to incref these because we don't store them
+    // Furthermore, this fn does not create a strong reference to the args
+    PyObject* errorCode    = nullptr;
+    PyObject* errorMessage = nullptr;
+    if (!PyArg_ParseTuple(args, "OO", &errorCode, &errorMessage))
+        return -1;
+
+    if (!PyObject_IsInstance(errorCode, *state->PyErrorCodeType)) {
+        PyErr_SetString(PyExc_TypeError, "expected code to be of type ErrorCode");
+        return -1;
+    }
+
+    if (!PyUnicode_Check(errorMessage)) {
+        PyErr_SetString(PyExc_TypeError, "expected message to be a string");
+        return -1;
+    }
+
+    // delegate to the base class init
+    return self->ob_base.ob_type->tp_base->tp_init((PyObject*)self, args, kwds);
+}
+
+static void YMQException_dealloc(YMQException* self)
+{
+    self->ob_base.ob_type->tp_base->tp_dealloc((PyObject*)self);
+
+    // we still need to release the reference to the heap type
+    auto* tp = Py_TYPE(self);
+    Py_DECREF(tp);
+}
+
+static PyObject* YMQException_code_getter(YMQException* self, void* Py_UNUSED(closure))
+{
+    return PySequence_GetItem(self->args, YMQException_errorCodeIndex);
+}
+
+static PyObject* YMQException_message_getter(YMQException* self, void* Py_UNUSED(closure))
+{
+    return PySequence_GetItem(self->args, YMQException_messageIndex);
+}
+}
+
+static PyGetSetDef YMQException_getset[] = {
+    {"code", (getter)YMQException_code_getter, nullptr, PyDoc_STR("error code"), nullptr},
+    {"message", (getter)YMQException_message_getter, nullptr, PyDoc_STR("error message"), nullptr},
+    {nullptr}  // Sentinel
+};
+
+static PyType_Slot YMQException_slots[] = {
+    {Py_tp_init, (void*)YMQException_init},
+    {Py_tp_dealloc, (void*)YMQException_dealloc},
+    {Py_tp_getset, (void*)YMQException_getset},
+    {0, 0},
+};
+
+static PyType_Spec YMQException_spec = {
+    "_ymq.YMQException", sizeof(YMQException), 0, Py_TPFLAGS_DEFAULT, YMQException_slots};
+
+inline OwnedPyObject<> YMQException_argtupleFromCoreError(YMQState* state, const Error* error)
+{
+    OwnedPyObject code = PyLong_FromLong(static_cast<long>(error->_errorCode));
+
+    if (!code)
+        return nullptr;
+
+    OwnedPyObject pyCode = PyObject_CallFunction(*state->PyErrorCodeType, "O", *code);
+
+    if (!pyCode)
+        return nullptr;
+
+    OwnedPyObject message = PyUnicode_FromString(error->what());
+
+    if (!message)
+        return nullptr;
+
+    return PyTuple_Pack(2, *pyCode, *message);
+}
+
+inline void YMQException_setFromCoreError(YMQState* state, const Error* error)
+{
+    auto tuple = YMQException_argtupleFromCoreError(state, error);
+    if (!tuple)
+        return;
+
+    PyErr_SetObject(*state->PyExceptionType, *tuple);
+}
+
+inline PyObject* YMQException_createFromCoreError(YMQState* state, const Error* error)
+{
+    auto tuple = YMQException_argtupleFromCoreError(state, error);
+    if (!tuple)
+        return nullptr;
+
+    OwnedPyObject exc = PyObject_CallObject(*state->PyExceptionType, *tuple);
+
+    // transfer ownership to caller
+    return exc.take();
+}

scaler/io/ymq/pymod_ymq/gil.h

@@ -0,0 +1,15 @@
+#include "scaler/io/ymq/pymod_ymq/python.h"
+
+class AcquireGIL {
+public:
+    AcquireGIL() : _state(PyGILState_Ensure()) {}
+    ~AcquireGIL() { PyGILState_Release(_state); }
+
+    AcquireGIL(const AcquireGIL&) = delete;
+    AcquireGIL& operator=(const AcquireGIL&) = delete;
+    AcquireGIL(AcquireGIL&&) = delete;
+    AcquireGIL& operator=(AcquireGIL&&) = delete;
+
+private:
+    PyGILState_STATE _state;
+};

scaler/io/ymq/pymod_ymq/io_context.h

@@ -0,0 +1,166 @@
+#pragma once
+
+// Python
+#include "scaler/io/ymq/pymod_ymq/python.h"
+
+// C++
+#include <memory>
+
+// First-party
+#include "scaler/io/ymq/configuration.h"
+#include "scaler/io/ymq/io_context.h"
+#include "scaler/io/ymq/io_socket.h"
+#include "scaler/io/ymq/pymod_ymq/io_socket.h"
+#include "scaler/io/ymq/pymod_ymq/ymq.h"
+
+// TODO: move ymq's python module into this namespace
+using namespace scaler::ymq;
+using Identity = Configuration::IOSocketIdentity;
+
+struct PyIOContext {
+    PyObject_HEAD;
+    std::shared_ptr<IOContext> ioContext;
+};
+
+extern "C" {
+
+static int PyIOContext_init(PyIOContext* self, PyObject* args, PyObject* kwds)
+{
+    // default to 1 thread if not specified
+    Py_ssize_t numThreads = 1;
+    const char* kwlist[]  = {"num_threads", nullptr};
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|n", (char**)kwlist, &numThreads))
+        return -1;
+
+    try {
+        new (&self->ioContext) std::shared_ptr<IOContext>();
+        self->ioContext = std::make_shared<IOContext>(numThreads);
+    } catch (...) {
+        PyErr_SetString(PyExc_RuntimeError, "Failed to create IOContext");
+        return -1;
+    }
+
+    return 0;
+}
+
+static void PyIOContext_dealloc(PyIOContext* self)
+{
+    try {
+        self->ioContext.~shared_ptr();
+    } catch (...) {
+        PyErr_SetString(PyExc_RuntimeError, "Failed to deallocate IOContext");
+        PyErr_WriteUnraisable((PyObject*)self);
+    }
+
+    auto* tp = Py_TYPE(self);
+    tp->tp_free(self);
+    Py_DECREF(tp);
+}
+
+static PyObject* PyIOContext_repr(PyIOContext* self)
+{
+    return PyUnicode_FromFormat("<IOContext at %p>", (void*)self->ioContext.get());
+}
+
+static PyObject* PyIOContext_numThreads_getter(PyIOContext* self, void* Py_UNUSED(closure))
+{
+    return PyLong_FromSize_t(self->ioContext->numThreads());
+}
+
+static PyObject* PyIOContext_createIOSocket(PyIOContext* self, PyObject* args, PyObject* kwargs)
+{
+    YMQState* state = YMQStateFromSelf((PyObject*)self);
+    if (!state)
+        return nullptr;
+
+    PyObject* callback     = nullptr;
+    const char* identity   = nullptr;
+    Py_ssize_t identityLen = 0;
+    PyObject* pySocketType = nullptr;
+    const char* kwlist[]   = {"", "identity", "socket_type", nullptr};
+    if (!PyArg_ParseTupleAndKeywords(
+            args, kwargs, "Os#O", (char**)kwlist, &callback, &identity, &identityLen, &pySocketType))
+        return nullptr;
+
+    if (!PyObject_IsInstance(pySocketType, *state->PyIOSocketEnumType)) {
+        PyErr_SetString(PyExc_TypeError, "Expected socket_type to be an instance of IOSocketType");
+        return nullptr;
+    }
+
+    OwnedPyObject value = PyObject_GetAttrString(pySocketType, "value");
+    if (!value)
+        return nullptr;
+
+    if (!PyLong_Check(*value)) {
+        PyErr_SetString(PyExc_TypeError, "Expected socket_type to be an integer");
+        return nullptr;
+    }
+
+    long socketTypeValue = PyLong_AsLong(*value);
+    if (socketTypeValue < 0 && PyErr_Occurred())
+        return nullptr;
+
+    IOSocketType socketType            = static_cast<IOSocketType>(socketTypeValue);
+    OwnedPyObject<PyIOSocket> ioSocket = PyObject_New(PyIOSocket, (PyTypeObject*)*state->PyIOSocketType);
+    if (!ioSocket)
+        return nullptr;
+
+    Py_INCREF(callback);
+
+    try {
+        // ensure the fields are init
+        new (&ioSocket->socket) std::shared_ptr<IOSocket>();
+        new (&ioSocket->ioContext) std::shared_ptr<IOContext>();
+        ioSocket->ioContext = self->ioContext;
+
+        self->ioContext->createIOSocket(
+            std::string(identity, identityLen), socketType, [callback, ioSocket](auto socket) {
+                AcquireGIL _;
+
+                ioSocket->socket = socket;
+                OwnedPyObject _result  = PyObject_CallFunctionObjArgs(callback, *ioSocket, nullptr);
+                Py_DECREF(callback);
+            });
+    } catch (...) {
+        PyErr_SetString(PyExc_RuntimeError, "Failed to create IOSocket");
+        return nullptr;
+    }
+
+    Py_RETURN_NONE;
+}
+
+}  // extern "C"
+
+static PyMethodDef PyIOContext_methods[] = {
+    {"createIOSocket",
+     (PyCFunction)PyIOContext_createIOSocket,
+     METH_VARARGS | METH_KEYWORDS,
+     PyDoc_STR("Create a new IOSocket")},
+    {nullptr, nullptr, 0, nullptr},
+};
+
+static PyGetSetDef PyIOContext_properties[] = {
+    {"num_threads",
+     (getter)PyIOContext_numThreads_getter,
+     nullptr,
+     PyDoc_STR("Get the number of threads in the IOContext"),
+     nullptr},
+    {nullptr, nullptr, nullptr, nullptr, nullptr},
+};
+
+static PyType_Slot PyIOContext_slots[] = {
+    {Py_tp_init, (void*)PyIOContext_init},
+    {Py_tp_dealloc, (void*)PyIOContext_dealloc},
+    {Py_tp_repr, (void*)PyIOContext_repr},
+    {Py_tp_methods, (void*)PyIOContext_methods},
+    {Py_tp_getset, (void*)PyIOContext_properties},
+    {0, nullptr},
+};
+
+static PyType_Spec PyIOContext_spec = {
+    .name      = "_ymq.BaseIOContext",
+    .basicsize = sizeof(PyIOContext),
+    .itemsize  = 0,
+    .flags     = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE | Py_TPFLAGS_BASETYPE,
+    .slots     = PyIOContext_slots,
+};