Add YMQ Tests

Author: magniloquency

.github/actions/compile-library/action.yml

@@ -14,6 +14,7 @@ runs:
       shell: bash
       run: |
         CXX=$(which g++-14) ./scripts/build.sh
+        sudo ./scripts/test.sh
 
     - name: Build and test C++ Components (Windows)
       if: inputs.os == 'Windows'

.github/actions/setup-env/action.yml

@@ -23,11 +23,12 @@ runs:
         sudo chmod 755 ./scripts/download_install_libraries.sh
         sudo chmod 755 ./scripts/build.sh
 
-    - name: Install Python Base Packages
+    - name: Install Python Packages
       shell: bash
       run: |
         pip install uv
         uv pip install --system --upgrade pip
+        uv pip install --system 'scapy==2.*'
 
     - name: Cache Library Install
       if: inputs.os == 'Linux'

examples/cpp/ymq/automated_echo_client.cpp

@@ -70,7 +70,7 @@ int main()
         auto future = x.get_future();
         Message msg = future.get().first;
         if (msg.payload.as_string() != longStr) {
-            printf("Checksum failed, %s\n", msg.payload.as_string().c_str());
+            printf("Checksum failed, %s\n", msg.payload.as_string()->c_str());
             exit(1);
         }
     }

examples/task_capabilities.py

@@ -56,18 +56,12 @@ def main():
 
         # Submit a task that requires GPU capabilities, this will be redirected to the GPU worker.
         gpu_future = client.submit_verbose(
-            gpu_task,
-            args=(16.0,),
-            kwargs={},
-            capabilities={"gpu": 1}  # Requires a GPU capability
+            gpu_task, args=(16.0,), kwargs={}, capabilities={"gpu": 1}  # Requires a GPU capability
         )
 
         # Submit a task that does not require GPU capabilities, this will be routed to any available worker.
         cpu_future = client.submit_verbose(
-            cpu_task,
-            args=(16.0,),
-            kwargs={},
-            capabilities={}  # No GPU capability required
+            cpu_task, args=(16.0,), kwargs={}, capabilities={}  # No GPU capability required
         )
 
         # Waits for the tasks for finish

pyproject.toml

@@ -51,10 +51,12 @@ all = [
 
 [dependency-groups]
 dev = [
-    "black>=25.1.0",
+    "black>=24.8.0; python_version == '3.8'",
+    "black>=25.1.0; python_version >= '3.8'",
     "flake8>=7.3.0",
     "flake8-pyproject>=1.2.3",
     "mypy>=1.17.1",
+    "scapy==2.*",
 ]
 
 [tool.scikit-build.metadata.version]

scaler/io/ymq/bytes.h

@@ -10,6 +10,7 @@
 #include <cstring>
 
 // C++
+#include <optional>
 #include <string>
 
 // First-party
@@ -33,10 +34,9 @@ class Bytes {
 public:
     Bytes(char* data, size_t len): _data(datadup((uint8_t*)data, len)), _len(len) {}
 
-    Bytes(): _data {}, _len {} {}
+    Bytes(const std::string& s): _data(datadup((uint8_t*)s.data(), s.length())), _len(s.length()) {}
 
-    // For debug and convenience only
-    explicit Bytes(const std::string& str): Bytes((char*)str.c_str(), str.size()) {}
+    Bytes(): _data {}, _len {} {}
 
     Bytes(const Bytes& other) noexcept
     {
@@ -92,11 +92,10 @@ class Bytes {
 
     [[nodiscard]] constexpr bool is_null() const noexcept { return !this->_data; }
 
-    // debugging utility
-    std::string as_string() const
+    std::optional<std::string> as_string() const
     {
         if (is_null())
-            return "[EMPTY]";
+            return std::nullopt;
 
         return std::string((char*)_data, _len);
     }

scaler/io/ymq/message_connection_tcp.cpp

@@ -2,6 +2,7 @@
 #include "scaler/io/ymq/message_connection_tcp.h"
 
 #include <future>
+#include <print>
 
 #include "scaler/io/ymq/configuration.h"
 
@@ -269,7 +270,6 @@ void MessageConnectionTCP::updateReadOperation()
             _pendingRecvMessageCallbacks->pop();
 
             recvMessageCallback({Message(std::move(address), std::move(payload)), {}});
-
         } else {
             assert(_pendingRecvMessageCallbacks->size());
             break;

scaler/io/ymq/network_utils.h

@@ -4,6 +4,7 @@
 #include <arpa/inet.h>
 #include <netinet/in.h>
 #include <netinet/tcp.h>
+#include <unistd.h>
 #endif  // __linux__
 #ifdef _WIN32
 // clang-format off

scaler/io/ymq/pymod_ymq/async.h

@@ -10,7 +10,7 @@
 #include "scaler/io/ymq/pymod_ymq/ymq.h"
 
 // wraps an async callback that accepts a Python asyncio future
-static PyObject* async_wrapper(PyObject* self, const std::function<void(YMQState* state, PyObject* future)>& callback)
+static PyObject* async_wrapper(PyObject* self, const std::function<void(YMQState* state, PyObject* future)>&& callback)
 {
     auto state = YMQStateFromSelf(self);
     if (!state)
@@ -23,7 +23,6 @@ static PyObject* async_wrapper(PyObject* self, const std::function<void(YMQState
     }
 
     OwnedPyObject future = PyObject_CallMethod(*loop, "create_future", nullptr);
-
     if (!future) {
         PyErr_SetString(PyExc_RuntimeError, "Failed to create future");
         return nullptr;

scaler/io/ymq/pymod_ymq/io_socket.h

@@ -154,10 +154,6 @@ static PyObject* PyIOSocket_recv(PyIOSocket* self, PyObject* args)
                     if (!pyMessage)
                         return YMQ_GetRaisedException();
 
-                    // TODO: why is leaking necessary?
-                    address.forget();
-                    payload.forget();
-
                     return (PyObject*)pyMessage.take();
                 });
             } catch (...) {
@@ -219,10 +215,6 @@ static PyObject* PyIOSocket_recv_sync(PyIOSocket* self, PyObject* args)
     if (!pyMessage)
         return nullptr;
 
-    // TODO: why is leaking necessary?
-    address.forget();
-    payload.forget();
-
     return (PyObject*)pyMessage.take();
 }