Implement raw TCP/IP in jfjoch_broker and jfjoch_writer

image_puller/TCPImagePuller.cpp

@@ -0,0 +1,232 @@
+// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#include "TCPImagePuller.h"
+
+#include <sys/socket.h>
+#include <arpa/inet.h>
+#include <unistd.h>
+#include <cerrno>
+#include <limits>
+#include <netdb.h>
+
+static std::pair<std::string, uint16_t> ParseTcpAddressPull(const std::string& addr) {
+    const std::string prefix = "tcp://";
+    if (addr.rfind(prefix, 0) != 0)
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid TCP address: " + addr);
+
+    auto hp = addr.substr(prefix.size());
+    auto p = hp.find_last_of(':');
+    if (p == std::string::npos)
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid TCP address: " + addr);
+
+    const auto host = hp.substr(0, p);
+    const auto port_str = hp.substr(p + 1);
+
+    int port_i = 0;
+    try {
+        size_t parsed = 0;
+        port_i = std::stoi(port_str, &parsed);
+        if (parsed != port_str.size())
+            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid TCP port in address: " + addr);
+    } catch (...) {
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid TCP port in address: " + addr);
+    }
+
+    if (port_i < 1 || port_i > static_cast<int>(std::numeric_limits<uint16_t>::max()))
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "TCP port out of range in address: " + addr);
+
+    return {host, static_cast<uint16_t>(port_i)};
+}
+
+TCPImagePuller::TCPImagePuller(const std::string &tcp_addr,
+                               std::optional<int32_t> rcv_buffer_size)
+        : addr(tcp_addr),
+          receive_buffer_size(rcv_buffer_size) {
+    auto parsed = ParseTcpAddressPull(tcp_addr);
+    host = parsed.first;
+    port = parsed.second;
+
+    receiver_thread = std::thread(&TCPImagePuller::ReceiverThread, this);
+    cbor_thread = std::thread(&TCPImagePuller::CBORThread, this);
+}
+
+TCPImagePuller::~TCPImagePuller() {
+    TCPImagePuller::Disconnect();
+}
+
+void TCPImagePuller::CloseSocket() {
+    int old_fd = -1;
+    {
+        std::unique_lock ul(fd_mutex);
+        if (fd >= 0) {
+            old_fd = fd;
+            fd = -1;
+        }
+    }
+
+    if (old_fd >= 0) {
+        shutdown(old_fd, SHUT_RDWR);
+        close(old_fd);
+    }
+}
+
+bool TCPImagePuller::EnsureConnected() {
+    {
+        std::unique_lock ul(fd_mutex);
+        if (fd >= 0)
+            return true;
+    }
+
+    addrinfo hints{};
+    hints.ai_family = AF_UNSPEC;      // Allow IPv4 or IPv6
+    hints.ai_socktype = SOCK_STREAM;
+    hints.ai_protocol = IPPROTO_TCP;
+
+    addrinfo *res = nullptr;
+    const std::string port_str = std::to_string(port);
+    int gai_rc = getaddrinfo(host.c_str(), port_str.c_str(), &hints, &res);
+    if (gai_rc != 0) {
+        logger.Error(std::string("getaddrinfo failed for ") + host + ":" + port_str + " - " + gai_strerror(gai_rc));
+        return false;
+    }
+
+    int new_fd = -1;
+    for (addrinfo *ai = res; ai != nullptr; ai = ai->ai_next) {
+        new_fd = socket(ai->ai_family, ai->ai_socktype, ai->ai_protocol);
+        if (new_fd < 0)
+            continue;
+
+        if (receive_buffer_size)
+            setsockopt(new_fd, SOL_SOCKET, SO_RCVBUF, &receive_buffer_size.value(), sizeof(int32_t));
+
+        if (connect(new_fd, ai->ai_addr, ai->ai_addrlen) == 0)
+            break;
+
+        close(new_fd);
+        new_fd = -1;
+    }
+
+    freeaddrinfo(res);
+
+    if (new_fd < 0)
+        return false;
+
+    {
+        std::unique_lock ul(fd_mutex);
+        if (fd >= 0) {
+            close(new_fd);
+            return true;
+        }
+        fd = new_fd;
+    }
+
+    logger.Info("TCP connected to " + addr);
+    return true;
+}
+
+bool TCPImagePuller::ReadExact(void *buf, size_t size) {
+    auto p = static_cast<uint8_t *>(buf);
+    size_t got = 0;
+
+    while (got < size) {
+        if (disconnect)
+            return false;
+
+        int local_fd = -1;
+        {
+            std::unique_lock ul(fd_mutex);
+            local_fd = fd;
+        }
+        if (local_fd < 0)
+            return false;
+
+        ssize_t rc = recv(local_fd, p + got, size - got, 0);
+        if (rc == 0)
+            return false;
+        if (rc < 0) {
+            if (errno == EINTR)
+                continue;
+            return false;
+        }
+
+        got += static_cast<size_t>(rc);
+    }
+
+    return true;
+}
+
+void TCPImagePuller::ReceiverThread() {
+    try {
+        while (!disconnect) {
+            if (!EnsureConnected()) {
+                std::this_thread::sleep_for(std::chrono::milliseconds(20));
+                continue;
+            }
+
+            RawFrame frame{};
+            if (!ReadExact(&frame.header, sizeof(frame.header))) {
+                logger.Info("TCP receive failed, reconnecting to " + addr);
+                CloseSocket();
+                std::this_thread::sleep_for(std::chrono::milliseconds(20));
+                continue;
+            }
+
+            if (frame.header.magic != JFJOCH_TCP_MAGIC || frame.header.version != JFJOCH_TCP_VERSION) {
+                logger.Error("Invalid TCP frame header, reconnecting to " + addr);
+                CloseSocket();
+                std::this_thread::sleep_for(std::chrono::milliseconds(20));
+                continue;
+            }
+
+            ImagePullerOutput out;
+            out.tcp_msg = std::make_shared<RawFrame>();
+            out.tcp_msg->header = frame.header;
+            out.tcp_msg->payload.resize(frame.header.payload_size);
+
+            if (frame.header.payload_size > 0
+                && !ReadExact(out.tcp_msg->payload.data(), out.tcp_msg->payload.size())) {
+                logger.Info("TCP payload read failed, reconnecting to " + addr);
+                CloseSocket();
+                std::this_thread::sleep_for(std::chrono::milliseconds(20));
+                continue;
+            }
+
+            cbor_fifo.PutBlocking(out);
+        }
+    } catch (const JFJochException &e) {
+        logger.ErrorException(e);
+    } catch (const std::exception &e) {
+        logger.Error(std::string("Unhandled exception in ReceiverThread: ") + e.what());
+    } catch (...) {
+        logger.Error("Unhandled unknown exception in ReceiverThread");
+    }
+
+    CloseSocket();
+    cbor_fifo.PutBlocking(ImagePullerOutput{});
+}
+
+void TCPImagePuller::CBORThread() {
+    auto ret = cbor_fifo.GetBlocking();
+    while (ret.tcp_msg) {
+        try {
+            ret.cbor = CBORStream2Deserialize(ret.tcp_msg->payload.data(), ret.tcp_msg->payload.size());
+            outside_fifo.PutBlocking(ret);
+        } catch (const JFJochException &e) {
+            logger.ErrorException(e);
+        }
+        ret = cbor_fifo.GetBlocking();
+    }
+    outside_fifo.PutBlocking(ret);
+}
+void TCPImagePuller::Disconnect() {
+    if (disconnect.exchange(true))
+        return;
+
+    CloseSocket();
+
+    if (receiver_thread.joinable())
+        receiver_thread.join();
+    if (cbor_thread.joinable())
+        cbor_thread.join();
+}