opengris-scaler 1.12.7__cp38-cp38-manylinux_2_28_x86_64.whl

scaler/io/ymq/pymod_ymq/io_socket.h

@@ -0,0 +1,285 @@
+#pragma once
+
+// Python
+#include "scaler/io/ymq/pymod_ymq/python.h"
+
+// C++
+#include <memory>
+#include <utility>
+
+// C
+#include <object.h>
+#include <semaphore.h>
+#include <sys/eventfd.h>
+#include <sys/poll.h>
+
+// First-party
+#include "scaler/io/ymq/bytes.h"
+#include "scaler/io/ymq/error.h"
+#include "scaler/io/ymq/io_context.h"
+#include "scaler/io/ymq/io_socket.h"
+#include "scaler/io/ymq/message.h"
+#include "scaler/io/ymq/pymod_ymq/bytes.h"
+#include "scaler/io/ymq/pymod_ymq/exception.h"
+#include "scaler/io/ymq/pymod_ymq/gil.h"
+#include "scaler/io/ymq/pymod_ymq/message.h"
+#include "scaler/io/ymq/pymod_ymq/ymq.h"
+
+using namespace scaler::ymq;
+
+struct PyIOSocket {
+    PyObject_HEAD;
+    std::shared_ptr<IOSocket> socket;
+    std::shared_ptr<IOContext> ioContext;
+};
+
+extern "C" {
+
+static void PyIOSocket_dealloc(PyIOSocket* self)
+{
+    try {
+        self->ioContext->removeIOSocket(self->socket);
+        self->ioContext.~shared_ptr();
+        self->socket.~shared_ptr();
+    } catch (...) {
+        PyErr_SetString(PyExc_RuntimeError, "Failed to deallocate IOSocket");
+        PyErr_WriteUnraisable((PyObject*)self);
+    }
+
+    auto* tp = Py_TYPE(self);
+    tp->tp_free(self);
+    Py_DECREF(tp);
+}
+
+static PyObject* PyIOSocket_send(PyIOSocket* self, PyObject* args, PyObject* kwargs)
+{
+    auto state = YMQStateFromSelf((PyObject*)self);
+    if (!state)
+        return nullptr;
+
+    PyObject* callback   = nullptr;
+    PyMessage* message   = nullptr;
+
+    // empty str -> positional only
+    const char* kwlist[] = {"", "message", nullptr};
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO", (char**)kwlist, &callback, &message))
+        return nullptr;
+
+    if (!PyObject_TypeCheck(message, (PyTypeObject*)*state->PyMessageType)) {
+        PyErr_SetString(PyExc_TypeError, "message must be a Message");
+        return nullptr;
+    }
+
+    auto address = message->address.is_none() ? Bytes() : std::move(message->address->bytes);
+    auto payload = std::move(message->payload->bytes);
+
+    Py_INCREF(callback);
+
+    try {
+        self->socket->sendMessage(
+            {.address = std::move(address), .payload = std::move(payload)}, [callback, state](auto result) {
+                AcquireGIL _;
+
+                if (result) {
+                    OwnedPyObject result = PyObject_CallFunctionObjArgs(callback, Py_None, nullptr);
+                } else {
+                    OwnedPyObject obj = YMQException_createFromCoreError(state, &result.error());
+                    OwnedPyObject _   = PyObject_CallFunctionObjArgs(callback, *obj, nullptr);
+                }
+
+                Py_DECREF(callback);
+            });
+    } catch (...) {
+        PyErr_SetString(PyExc_RuntimeError, "Failed to send message");
+        return nullptr;
+    }
+
+    Py_RETURN_NONE;
+}
+
+static PyObject* PyIOSocket_recv(PyIOSocket* self, PyObject* args, PyObject* kwargs)
+{
+    auto state = YMQStateFromSelf((PyObject*)self);
+    if (!state)
+        return nullptr;
+
+    PyObject* callback   = nullptr;
+    const char* kwlist[] = {"", nullptr};
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O", (char**)kwlist, &callback))
+        return nullptr;
+
+    Py_INCREF(callback);
+
+    try {
+        self->socket->recvMessage([callback, state](std::pair<Message, Error> result) {
+            AcquireGIL _;
+
+            if (result.second._errorCode != Error::ErrorCode::Uninit) {
+                OwnedPyObject obj = YMQException_createFromCoreError(state, &result.second);
+                OwnedPyObject _   = PyObject_CallFunctionObjArgs(callback, *obj, nullptr);
+                return;
+            }
+
+            auto message                      = result.first;
+            OwnedPyObject<PyBytesYMQ> address = (PyBytesYMQ*)PyObject_CallNoArgs(*state->PyBytesYMQType);
+            if (!address)
+                return completeCallbackWithRaisedException(callback);
+
+            address->bytes = std::move(message.address);
+
+            OwnedPyObject<PyBytesYMQ> payload = (PyBytesYMQ*)PyObject_CallNoArgs(*state->PyBytesYMQType);
+            if (!payload)
+                return completeCallbackWithRaisedException(callback);
+
+            payload->bytes = std::move(message.payload);
+
+            OwnedPyObject<PyMessage> pyMessage =
+                (PyMessage*)PyObject_CallFunction(*state->PyMessageType, "OO", *address, *payload);
+            if (!pyMessage)
+                return completeCallbackWithRaisedException(callback);
+
+            OwnedPyObject _result = PyObject_CallFunctionObjArgs(callback, *pyMessage, nullptr);
+            Py_DECREF(callback);
+        });
+    } catch (...) {
+        PyErr_SetString(PyExc_RuntimeError, "Failed to receive message");
+        return nullptr;
+    }
+
+    Py_RETURN_NONE;
+}
+
+static PyObject* PyIOSocket_bind(PyIOSocket* self, PyObject* args, PyObject* kwargs)
+{
+    auto state = YMQStateFromSelf((PyObject*)self);
+    if (!state)
+        return nullptr;
+
+    PyObject* callback    = nullptr;
+    const char* address   = nullptr;
+    Py_ssize_t addressLen = 0;
+    const char* kwlist[]  = {"", "address", nullptr};
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "Os#", (char**)kwlist, &callback, &address, &addressLen))
+        return nullptr;
+
+    Py_INCREF(callback);
+
+    try {
+        self->socket->bindTo(std::string(address, addressLen), [callback, state](auto result) {
+            AcquireGIL _;
+
+            if (!result) {
+                OwnedPyObject exc = YMQException_createFromCoreError(state, &result.error());
+                OwnedPyObject _   = PyObject_CallFunctionObjArgs(callback, *exc, nullptr);
+            } else {
+                OwnedPyObject _ = PyObject_CallFunctionObjArgs(callback, Py_None, nullptr);
+            }
+
+            Py_DECREF(callback);
+        });
+    } catch (...) {
+        PyErr_SetString(PyExc_RuntimeError, "Failed to bind to address");
+        return nullptr;
+    }
+
+    Py_RETURN_NONE;
+}
+
+static PyObject* PyIOSocket_connect(PyIOSocket* self, PyObject* args, PyObject* kwargs)
+{
+    auto state = YMQStateFromSelf((PyObject*)self);
+    if (!state)
+        return nullptr;
+
+    PyObject* callback    = nullptr;
+    const char* address   = nullptr;
+    Py_ssize_t addressLen = 0;
+    const char* kwlist[]  = {"", "address", nullptr};
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "Os#", (char**)kwlist, &callback, &address, &addressLen))
+        return nullptr;
+
+    Py_INCREF(callback);
+
+    try {
+        self->socket->connectTo(std::string(address, addressLen), [callback, state](auto result) {
+            AcquireGIL _;
+
+            if (result || result.error()._errorCode == Error::ErrorCode::InitialConnectFailedWithInProgress) {
+                OwnedPyObject _ = PyObject_CallFunctionObjArgs(callback, Py_None, nullptr);
+            } else {
+                OwnedPyObject exc = YMQException_createFromCoreError(state, &result.error());
+                OwnedPyObject _   = PyObject_CallFunctionObjArgs(callback, *exc, nullptr);
+            }
+
+            Py_DECREF(callback);
+        });
+    } catch (...) {
+        PyErr_SetString(PyExc_RuntimeError, "Failed to connect to address");
+        return nullptr;
+    }
+
+    Py_RETURN_NONE;
+}
+
+static PyObject* PyIOSocket_repr(PyIOSocket* self)
+{
+    return PyUnicode_FromFormat("<IOSocket at %p>", (void*)self->socket.get());
+}
+
+static PyObject* PyIOSocket_identity_getter(PyIOSocket* self, void* closure)
+{
+    return PyUnicode_FromStringAndSize(self->socket->identity().data(), self->socket->identity().size());
+}
+
+static PyObject* PyIOSocket_socket_type_getter(PyIOSocket* self, void* closure)
+{
+    auto state = YMQStateFromSelf((PyObject*)self);
+    if (!state)
+        return nullptr;
+
+    const IOSocketType socketType  = self->socket->socketType();
+    OwnedPyObject socketTypeIntObj = PyLong_FromLong((long)socketType);
+
+    if (!socketTypeIntObj)
+        return nullptr;
+
+    return PyObject_CallOneArg(*state->PyIOSocketEnumType, *socketTypeIntObj);
+}
+}
+
+static PyGetSetDef PyIOSocket_properties[] = {
+    {"identity", (getter)PyIOSocket_identity_getter, nullptr, PyDoc_STR("Get the identity of the IOSocket"), nullptr},
+    {"socket_type", (getter)PyIOSocket_socket_type_getter, nullptr, PyDoc_STR("Get the type of the IOSocket"), nullptr},
+    {nullptr, nullptr, nullptr, nullptr, nullptr},
+};
+
+static PyMethodDef PyIOSocket_methods[] = {
+    {"send", (PyCFunction)PyIOSocket_send, METH_VARARGS | METH_KEYWORDS, PyDoc_STR("Send data through the IOSocket")},
+    {"recv", (PyCFunction)PyIOSocket_recv, METH_VARARGS | METH_KEYWORDS, PyDoc_STR("Receive data from the IOSocket")},
+    {"bind",
+     (PyCFunction)PyIOSocket_bind,
+     METH_VARARGS | METH_KEYWORDS,
+     PyDoc_STR("Bind to an address and listen for incoming connections")},
+    {"connect",
+     (PyCFunction)PyIOSocket_connect,
+     METH_VARARGS | METH_KEYWORDS,
+     PyDoc_STR("Connect to a remote IOSocket")},
+    {nullptr, nullptr, 0, nullptr},
+};
+
+static PyType_Slot PyIOSocket_slots[] = {
+    {Py_tp_dealloc, (void*)PyIOSocket_dealloc},
+    {Py_tp_repr, (void*)PyIOSocket_repr},
+    {Py_tp_getset, (void*)PyIOSocket_properties},
+    {Py_tp_methods, (void*)PyIOSocket_methods},
+    {Py_tp_new, nullptr},
+    {0, nullptr},
+};
+
+static PyType_Spec PyIOSocket_spec = {
+    .name      = "_ymq.BaseIOSocket",
+    .basicsize = sizeof(PyIOSocket),
+    .itemsize  = 0,
+    .flags     = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE | Py_TPFLAGS_DISALLOW_INSTANTIATION,
+    .slots     = PyIOSocket_slots,
+};

scaler/io/ymq/pymod_ymq/message.h

@@ -0,0 +1,99 @@
+#pragma once
+
+// Python
+#include "scaler/io/ymq/pymod_ymq/python.h"
+
+// First-party
+#include "scaler/io/ymq/pymod_ymq/bytes.h"
+#include "scaler/io/ymq/pymod_ymq/ymq.h"
+
+struct PyMessage {
+    PyObject_HEAD;
+    OwnedPyObject<PyBytesYMQ> address;  // Address of the message; can be None
+    OwnedPyObject<PyBytesYMQ> payload;  // Payload of the message
+};
+
+extern "C" {
+
+static int PyMessage_init(PyMessage* self, PyObject* args, PyObject* kwds)
+{
+    auto state = YMQStateFromSelf((PyObject*)self);
+    if (!state)
+        return -1;
+
+    PyObject* address      = nullptr;
+    PyObject* payload      = nullptr;
+    const char* keywords[] = {"address", "payload", nullptr};
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO", (char**)keywords, &address, &payload))
+        return -1;
+
+    // address can be None, which means the message has no address
+    // check if the address and payload are of type PyBytesYMQ
+    if (PyObject_IsInstance(address, *state->PyBytesYMQType)) {
+        self->address = OwnedPyObject<PyBytesYMQ>::fromBorrowed((PyBytesYMQ*)address);
+    } else if (address == Py_None) {
+        self->address = OwnedPyObject<PyBytesYMQ>::none();
+    } else {
+        OwnedPyObject args = PyTuple_Pack(1, address);
+        self->address      = (PyBytesYMQ*)PyObject_CallObject(*state->PyBytesYMQType, *args);
+
+        if (!self->address)
+            return -1;
+    }
+
+    if (PyObject_IsInstance(payload, *state->PyBytesYMQType)) {
+        self->payload = OwnedPyObject<PyBytesYMQ>::fromBorrowed((PyBytesYMQ*)payload);
+    } else {
+        OwnedPyObject args = PyTuple_Pack(1, payload);
+        self->payload      = (PyBytesYMQ*)PyObject_CallObject(*state->PyBytesYMQType, *args);
+
+        if (!self->payload) {
+            return -1;
+        }
+    }
+
+    return 0;
+}
+
+static void PyMessage_dealloc(PyMessage* self)
+{
+    try {
+        self->address.~OwnedPyObject();
+        self->payload.~OwnedPyObject();
+    } catch (...) {
+        PyErr_SetString(PyExc_RuntimeError, "Failed to deallocate Message");
+        PyErr_WriteUnraisable((PyObject*)self);
+    }
+
+    auto* tp = Py_TYPE(self);
+    tp->tp_free(self);
+    Py_DECREF(tp);
+}
+
+static PyObject* PyMessage_repr(PyMessage* self)
+{
+    return PyUnicode_FromFormat("<Message address=%R payload=%R>", *self->address, *self->payload);
+}
+}
+
+static PyMemberDef PyMessage_members[] = {
+    {"address", T_OBJECT, offsetof(PyMessage, address), 0, PyDoc_STR("the address of the message")},
+    {"payload", T_OBJECT, offsetof(PyMessage, payload), 0, PyDoc_STR("the payload of the message")},
+    {nullptr},
+};
+
+static PyType_Slot PyMessage_slots[] = {
+    {Py_tp_init, (void*)PyMessage_init},
+    {Py_tp_dealloc, (void*)PyMessage_dealloc},
+    {Py_tp_repr, (void*)PyMessage_repr},
+    {Py_tp_members, (void*)PyMessage_members},
+    {0, nullptr},
+};
+
+static PyType_Spec PyMessage_spec = {
+    .name      = "_ymq.Message",
+    .basicsize = sizeof(PyMessage),
+    .itemsize  = 0,
+    .flags     = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE,
+    .slots     = PyMessage_slots,
+};

scaler/io/ymq/pymod_ymq/python.h

@@ -0,0 +1,153 @@
+// Python
+#pragma once
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+#include <structmember.h>
+
+#if PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION == 8
+static inline PyObject* Py_NewRef(PyObject* obj)
+{
+    Py_INCREF(obj);
+    return obj;
+}
+
+static inline PyObject* Py_XNewRef(PyObject* obj)
+{
+    Py_XINCREF(obj);
+    return obj;
+}
+
+// This is a very dirty hack, we basically place the raw pointer to this dict when init,
+// see scaler/io/ymq/pymod_ymq/ymq.h for more detail
+PyObject* PyType_GetModule(PyTypeObject* type)
+{
+    return PyObject_GetAttrString((PyObject*)(type), "__module_object__");
+}
+
+PyObject* PyType_GetModuleState(PyTypeObject* type)
+{
+    PyObject* module = PyObject_GetAttrString((PyObject*)(type), "__module_object__");
+    if (!module)
+        return nullptr;
+    void* state = PyModule_GetState(module);
+    Py_DECREF(module);
+    return (PyObject*)state;
+}
+
+static inline PyObject* PyObject_CallNoArgs(PyObject* callable)
+{
+    return PyObject_Call(callable, PyTuple_New(0), nullptr);
+}
+
+static inline int PyModule_AddObjectRef(PyObject* mod, const char* name, PyObject* value)
+{
+    Py_INCREF(value);  // Since PyModule_AddObject steals a ref, we balance it
+    return PyModule_AddObject(mod, name, value);
+}
+
+static inline PyObject* PyType_FromModuleAndSpec(PyObject* pymodule, PyType_Spec* spec, PyObject* bases)
+{
+    (void)pymodule;
+    if (!bases) {
+        bases = PyTuple_Pack(1, (PyObject*)&PyBaseObject_Type);
+        if (bases) {
+            PyObject* res = PyType_FromSpecWithBases(spec, bases);
+            Py_DECREF(bases);  // avoid leak
+            return res;
+        }
+        PyObject* res = PyType_FromSpecWithBases(spec, bases);
+        return res;
+    }
+    return PyType_FromSpecWithBases(spec, bases);
+}
+
+#define Py_TPFLAGS_IMMUTABLETYPE          (0)
+#define Py_TPFLAGS_DISALLOW_INSTANTIATION (0)
+#endif
+
+// NOTE: We define this no matter what version of Python we use.
+// an owned handle to a PyObject with automatic reference counting via RAII
+template <typename T = PyObject>
+class OwnedPyObject {
+public:
+    OwnedPyObject(): _ptr(nullptr) {}
+
+    // steals a reference
+    OwnedPyObject(T* ptr): _ptr(ptr) {}
+
+    OwnedPyObject(const OwnedPyObject& other) { this->_ptr = (T*)Py_XNewRef((PyObject*)other._ptr); }
+    OwnedPyObject(OwnedPyObject&& other) noexcept: _ptr(other._ptr) { other._ptr = nullptr; }
+    OwnedPyObject& operator=(const OwnedPyObject& other)
+    {
+        if (this == &other)
+            return *this;
+
+        this->free();
+        this->_ptr = (T*)Py_XNewRef((PyObject*)other._ptr);
+        return *this;
+    }
+    OwnedPyObject& operator=(OwnedPyObject&& other) noexcept
+    {
+        if (this == &other)
+            return *this;
+
+        this->free();
+        this->_ptr = other._ptr;
+        other._ptr = nullptr;
+        return *this;
+    }
+
+    ~OwnedPyObject() { this->free(); }
+
+    // creates a new OwnedPyObject from a borrowed reference
+    static OwnedPyObject fromBorrowed(T* ptr) { return OwnedPyObject((T*)Py_XNewRef((PyObject*)ptr)); }
+
+    // convenience method for creating an OwnedPyObject that holds Py_None
+    static OwnedPyObject none() { return OwnedPyObject((T*)Py_NewRef(Py_None)); }
+
+    bool is_none() const { return (PyObject*)_ptr == Py_None; }
+
+    // takes the pointer out of the OwnedPyObject
+    // without decrementing the reference count
+    // use this to transfer ownership to C code
+    T* take()
+    {
+        T* ptr     = this->_ptr;
+        this->_ptr = nullptr;
+        return ptr;
+    }
+
+    void forget() { this->_ptr = nullptr; }
+
+    // operator T*() const { return _ptr; }
+    explicit operator bool() const { return _ptr != nullptr; }
+    bool operator!() const { return _ptr == nullptr; }
+
+    T* operator->() const { return _ptr; }
+    T* operator*() const { return _ptr; }
+
+private:
+    T* _ptr;
+
+    void free()
+    {
+        if (!_ptr)
+            return;
+
+        if (!PyGILState_Check())
+            return;
+
+        Py_CLEAR(_ptr);
+    }
+};
+
+#if PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION == 8
+static inline PyObject* PyObject_CallOneArg(PyObject* callable, PyObject* arg)
+{
+    OwnedPyObject<> args = PyTuple_Pack(1, arg);
+    if (!args)
+        return nullptr;
+    PyObject* result = PyObject_Call(callable, *args, nullptr);
+    return result;
+}
+#endif

scaler/io/ymq/pymod_ymq/ymq.cpp

@@ -0,0 +1,23 @@
+#include "scaler/io/ymq/pymod_ymq/ymq.h"
+
+#include <cstdlib>
+
+#include "scaler/io/ymq/error.h"
+
+inline void ymqUnrecoverableError(scaler::ymq::Error e)
+{
+    PyGILState_STATE gstate = PyGILState_Ensure();
+    (void)gstate;  // silent the warning
+    PyErr_SetString(PyExc_SystemExit, e.what());
+    PyErr_WriteUnraisable(nullptr);
+    Py_Finalize();
+
+    std::exit(EXIT_FAILURE);
+}
+
+PyMODINIT_FUNC PyInit__ymq(void)
+{
+    unrecoverableErrorFunctionHookPtr = ymqUnrecoverableError;
+
+    return PyModuleDef_Init(&YMQ_module);
+}