diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index 171cbfb4..7e4d503a 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -122,6 +122,7 @@ ADD_LIBRARY(JFJochCommon STATIC TopPixels.cpp TopPixels.h hkl_key.h + JfjochTCP.h ) TARGET_LINK_LIBRARIES(JFJochCommon JFJochLogger Compression JFCalibration gemmi Threads::Threads -lrt ) diff --git a/common/JfjochTCP.h b/common/JfjochTCP.h new file mode 100644 index 00000000..b5fe36a7 --- /dev/null +++ b/common/JfjochTCP.h @@ -0,0 +1,28 @@ +// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include + +constexpr uint32_t JFJOCH_TCP_MAGIC = 0x4A464A54; // JFJT +constexpr uint32_t JFJOCH_TCP_VERSION = 1; + +enum class TCPFrameType : uint16_t { + START = 1, + DATA = 2, + CALIBRATION = 3, + END = 4 +}; + +struct alignas(64) TcpFrameHeader { + uint32_t magic = JFJOCH_TCP_MAGIC; + uint16_t version = JFJOCH_TCP_VERSION ; + uint16_t type = 0; + uint64_t image_number = 0; + uint64_t payload_size = 0; + uint32_t socket_number = 0; + uint32_t flags = 0; + uint64_t run_number = 0; + uint64_t reserved[4] = {0, 0, 0, 0}; +}; \ No newline at end of file diff --git a/image_puller/CMakeLists.txt b/image_puller/CMakeLists.txt index d1bb95d2..2bcabd63 100644 --- a/image_puller/CMakeLists.txt +++ b/image_puller/CMakeLists.txt @@ -2,5 +2,7 @@ ADD_LIBRARY(JFJochImagePuller ZMQImagePuller.cpp ZMQImagePuller.h ImagePuller.cpp ImagePuller.h TestImagePuller.cpp - TestImagePuller.h) + TestImagePuller.h + TcpImagePuller.cpp + TcpImagePuller.h) TARGET_LINK_LIBRARIES(JFJochImagePuller JFJochZMQ) diff --git a/image_puller/ImagePuller.h b/image_puller/ImagePuller.h index d0ea6e62..c6ef4971 100644 --- a/image_puller/ImagePuller.h +++ b/image_puller/ImagePuller.h @@ -10,9 +10,17 @@ #include "../common/ZMQWrappers.h" #include "../frame_serialize/CBORStream2Deserializer.h" #include "../common/ThreadSafeFIFO.h" +#include "../common/JfjochTCP.h" + +struct RawFrame { + TcpFrameHeader header{}; + std::vector payload; + bool end = false; +}; struct ImagePullerOutput { - std::shared_ptr msg; + std::shared_ptr zmq_msg; + std::shared_ptr tcp_msg; std::shared_ptr cbor; }; diff --git a/image_puller/TcpImagePuller.cpp b/image_puller/TcpImagePuller.cpp new file mode 100644 index 00000000..cfa216cb --- /dev/null +++ b/image_puller/TcpImagePuller.cpp @@ -0,0 +1,232 @@ +// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "TcpImagePuller.h" + +#include +#include +#include +#include +#include +#include + +static std::pair 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(std::numeric_limits::max())) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "TCP port out of range in address: " + addr); + + return {host, static_cast(port_i)}; +} + +TCPImagePuller::TCPImagePuller(const std::string &tcp_addr, + std::optional 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(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(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(); + 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(); +} \ No newline at end of file diff --git a/image_puller/TcpImagePuller.h b/image_puller/TcpImagePuller.h new file mode 100644 index 00000000..5f4b5b33 --- /dev/null +++ b/image_puller/TcpImagePuller.h @@ -0,0 +1,41 @@ +// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include +#include +#include "ImagePuller.h" +#include "../common/Logger.h" +#include "../common/ThreadSafeFIFO.h" + +class TCPImagePuller : public ImagePuller { + int fd = -1; + std::mutex fd_mutex; + + std::string addr; + std::string host; + uint16_t port = 0; + std::optional receive_buffer_size; + std::atomic disconnect{false}; + + ThreadSafeFIFO cbor_fifo{200}; + + std::thread receiver_thread; + std::thread cbor_thread; + + Logger logger{"TCPImagePuller"}; + + bool ReadExact(void *buf, size_t size); + bool EnsureConnected(); + void CloseSocket(); + void ReceiverThread(); + void CBORThread(); +public: + explicit TCPImagePuller(const std::string &tcp_addr, std::optional rcv_buffer_size = {}); + + ~TCPImagePuller() override; + void Disconnect() override; +}; \ No newline at end of file diff --git a/image_puller/ZMQImagePuller.cpp b/image_puller/ZMQImagePuller.cpp index db21ef05..33c8e943 100644 --- a/image_puller/ZMQImagePuller.cpp +++ b/image_puller/ZMQImagePuller.cpp @@ -48,7 +48,7 @@ void ZMQImagePuller::Disconnect() { void ZMQImagePuller::PullerThread() { while (true) { ImagePullerOutput ret; - ret.msg = std::make_shared(); + ret.zmq_msg = std::make_shared(); bool received = false; while (!received) { if (disconnect) { @@ -56,7 +56,7 @@ void ZMQImagePuller::PullerThread() { return; } try { - received = socket.Receive(*ret.msg, false); + received = socket.Receive(*ret.zmq_msg, false); if (!received) std::this_thread::sleep_for(std::chrono::milliseconds(1)); } catch (const JFJochException &e) { @@ -70,9 +70,9 @@ void ZMQImagePuller::PullerThread() { void ZMQImagePuller::CBORThread() { auto ret = cbor_fifo.GetBlocking(); - while (ret.msg) { + while (ret.zmq_msg) { try { - ret.cbor = CBORStream2Deserialize(ret.msg->data(), ret.msg->size()); + ret.cbor = CBORStream2Deserialize(ret.zmq_msg->data(), ret.zmq_msg->size()); if (ret.cbor->msg_type == CBORImageType::END) logger.Info("Received END"); @@ -98,7 +98,7 @@ void ZMQImagePuller::RepubThread() { auto ret = repub_fifo.GetBlocking(); bool repub_active = false; - while (ret.msg) { + while (ret.zmq_msg) { try { if (ret.cbor->msg_type == CBORImageType::START) { // Start message needs to be cleaned when running republish @@ -112,7 +112,7 @@ void ZMQImagePuller::RepubThread() { logger.Info("Republish active"); } else { if (repub_active) - repub_socket->Send(ret.msg->data(), ret.msg->size(), true); + repub_socket->Send(ret.zmq_msg->data(), ret.zmq_msg->size(), true); } } catch (const JFJochException &e) { logger.ErrorException(e); diff --git a/image_pusher/CMakeLists.txt b/image_pusher/CMakeLists.txt index f5cb4ccd..0e644aac 100644 --- a/image_pusher/CMakeLists.txt +++ b/image_pusher/CMakeLists.txt @@ -9,6 +9,10 @@ ADD_LIBRARY(ImagePusher STATIC NonePusher.h ZMQStream2PusherSocket.cpp ZMQStream2PusherSocket.h + TcpStreamPusherSocket.cpp + TcpStreamPusherSocket.h + TcpStreamPusher.cpp + TcpStreamPusher.h ) TARGET_LINK_LIBRARIES(ImagePusher JFJochZMQ CBORStream2FrameSerialize JFJochCommon Compression JFJochWriter) \ No newline at end of file diff --git a/image_pusher/TcpStreamPusher.cpp b/image_pusher/TcpStreamPusher.cpp new file mode 100644 index 00000000..5472638b --- /dev/null +++ b/image_pusher/TcpStreamPusher.cpp @@ -0,0 +1,113 @@ +// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "TcpStreamPusher.h" + + +TCPStreamPusher::TCPStreamPusher(const std::vector &addr, + std::optional send_buffer_size, + std::optional zerocopy_threshold, + size_t send_queue_size) + : serialization_buffer(256 * 1024 * 1024), + serializer(serialization_buffer.data(), serialization_buffer.size()) { + if (addr.empty()) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No TCP writer address provided"); + + for (size_t i = 0; i < addr.size(); i++) { + socket.emplace_back(std::make_unique( + addr[i], static_cast(i), send_buffer_size, zerocopy_threshold, send_queue_size)); + } +} + + +void TCPStreamPusher::StartDataCollection(StartMessage &message) { + if (message.images_per_file < 1) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "Images per file cannot be zero or negative"); + images_per_file = message.images_per_file; + run_number = message.run_number; + run_name = message.run_name; + + for (size_t i = 0; i < socket.size(); i++) { + if (!socket[i]->AcceptConnection(std::chrono::seconds(5))) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "TCP accept timeout/failure on socket " + socket[i]->GetEndpointName()); + } + + for (size_t i = 0; i < socket.size(); i++) { + message.socket_number = static_cast(i); + if (i > 0) + message.write_master_file = false; + + serializer.SerializeSequenceStart(message); + socket[i]->SetRunNumber(run_number); + + if (!socket[i]->Send(serialization_buffer.data(), serializer.GetBufferSize(), TCPFrameType::START)) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Timeout/failure sending START"); + } + + for (auto &s : socket) + s->StartWriterThread(); +} + +bool TCPStreamPusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number) { + if (socket.empty()) + return false; + + auto socket_number = (image_number / images_per_file) % socket.size(); + if (socket[socket_number]->IsBroken()) + return false; + + return socket[socket_number]->Send(image_data, image_size, TCPFrameType::DATA, image_number); +} + +void TCPStreamPusher::SendImage(ZeroCopyReturnValue &z) { + if (socket.empty()) { + z.release(); + return; + } + + auto socket_number = (z.GetImageNumber() / images_per_file) % socket.size(); + if (socket[socket_number]->IsBroken()) { + z.release(); + return; + } + + socket[socket_number]->SendImage(z); +} + +bool TCPStreamPusher::EndDataCollection(const EndMessage &message) { + serializer.SerializeSequenceEnd(message); + + bool ret = true; + for (auto &s : socket) { + s->StopWriterThread(); + if (s->IsBroken()) + ret = false; + else if (!s->Send(serialization_buffer.data(), serializer.GetBufferSize(), TCPFrameType::END)) + ret = false; + } + return ret; +} + +std::string TCPStreamPusher::PrintSetup() const { + std::string output = "TCPStream2Pusher: Sending images to sockets: "; + for (const auto &s : socket) + output += s->GetEndpointName() + " "; + return output; +} + +std::string TCPStreamPusherSocket::GetEndpointName() const { + return endpoint; +} + +void TCPStreamPusherSocket::SetRunNumber(uint64_t in_run_number) { + run_number = in_run_number; +} + +bool TCPStreamPusher::SendCalibration(const CompressedImage &message) { + if (socket.empty()) + return false; + serializer.SerializeCalibration(message); + return socket[0]->Send(serialization_buffer.data(), serializer.GetBufferSize(), TCPFrameType::CALIBRATION); +} diff --git a/image_pusher/TcpStreamPusher.h b/image_pusher/TcpStreamPusher.h new file mode 100644 index 00000000..3d01abae --- /dev/null +++ b/image_pusher/TcpStreamPusher.h @@ -0,0 +1,30 @@ +// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include "TcpStreamPusherSocket.h" + +class TCPStreamPusher : public ImagePusher { + std::vector serialization_buffer; + CBORStream2Serializer serializer; + std::vector> socket; + + int64_t images_per_file = 1; + uint64_t run_number = 0; + std::string run_name; + +public: + explicit TCPStreamPusher(const std::vector& addr, + std::optional send_buffer_size = {}, + std::optional zerocopy_threshold = {}, + size_t send_queue_size = 4096); + + void StartDataCollection(StartMessage& message) override; + bool EndDataCollection(const EndMessage& message) override; + bool SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number) override; + void SendImage(ZeroCopyReturnValue &z) override; + bool SendCalibration(const CompressedImage& message) override; + + std::string PrintSetup() const override; +}; \ No newline at end of file diff --git a/image_pusher/TcpStreamPusherSocket.cpp b/image_pusher/TcpStreamPusherSocket.cpp new file mode 100644 index 00000000..54e80e6a --- /dev/null +++ b/image_pusher/TcpStreamPusherSocket.cpp @@ -0,0 +1,415 @@ +// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "TcpStreamPusherSocket.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static std::pair ParseTcpAddress(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(std::numeric_limits::max())) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "TCP port out of range in address: " + addr); + + return {host, static_cast(port_i)}; +} + +TCPStreamPusherSocket::TCPStreamPusherSocket(const std::string &addr, + uint32_t in_socket_number, + std::optional in_send_buffer_size, + std::optional in_zerocopy_threshold, + size_t send_queue_size) + : queue(send_queue_size), + endpoint(addr), + socket_number(in_socket_number), + zerocopy_threshold(in_zerocopy_threshold), + send_buffer_size(in_send_buffer_size) { + auto [host, port] = ParseTcpAddress(addr); + + listen_fd = ::socket(AF_INET, SOCK_STREAM, 0); + if (listen_fd < 0) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "socket(listen) failed"); + + int one = 1; + setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)); + + sockaddr_in sin{}; + sin.sin_family = AF_INET; + sin.sin_port = htons(port); + if (host == "*" || host == "0.0.0.0") + sin.sin_addr.s_addr = htonl(INADDR_ANY); + else if (inet_pton(AF_INET, host.c_str(), &sin.sin_addr) != 1) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "inet_pton failed for " + host); + + if (bind(listen_fd, reinterpret_cast(&sin), sizeof(sin)) != 0) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "bind() failed to " + addr); + + if (listen(listen_fd, 16) != 0) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "listen() failed on " + addr); +} + +TCPStreamPusherSocket::~TCPStreamPusherSocket() { + try { + StopWriterThread(); + } catch (...) {} + CloseDataSocket(); + if (listen_fd >= 0) + close(listen_fd); +} + +void TCPStreamPusherSocket::CloseDataSocket() { + int old_fd = fd.exchange(-1); + if (old_fd >= 0) { + shutdown(old_fd, SHUT_RDWR); + close(old_fd); + } +} + +bool TCPStreamPusherSocket::AcceptConnection(std::chrono::milliseconds timeout) { + std::unique_lock ul(send_mutex); + + if (broken) + return false; + + if (fd.load() >= 0) + return true; + + if (ever_connected) { + broken = true; // session policy: no reconnect + return false; + } + + pollfd pfd{}; + pfd.fd = listen_fd; + pfd.events = POLLIN; + + const int prc = ::poll(&pfd, 1, static_cast(timeout.count())); + if (prc == 0) { + logger.Error("TCP accept timeout (" + std::to_string(timeout.count()) + " ms) on " + endpoint); + return false; + } + if (prc < 0) { + if (errno == EINTR) + return false; + logger.Error("TCP poll() failed on " + endpoint); + return false; + } + if ((pfd.revents & (POLLERR | POLLHUP | POLLNVAL)) != 0) { + logger.Error("TCP listen socket error on " + endpoint); + return false; + } + + int new_fd = ::accept(listen_fd, nullptr, nullptr); + if (new_fd < 0) + return false; + + int one = 1; + setsockopt(new_fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one)); + if (send_buffer_size) + setsockopt(new_fd, SOL_SOCKET, SO_SNDBUF, &send_buffer_size.value(), sizeof(int32_t)); +#ifdef SO_ZEROCOPY + setsockopt(new_fd, SOL_SOCKET, SO_ZEROCOPY, &one, sizeof(one)); +#endif + + fd.store(new_fd); + ever_connected = true; + logger.Info("TCP peer connected on " + endpoint); + return true; +} + +bool TCPStreamPusherSocket::IsConnectionAlive() const { + if (broken) + return false; + + int local_fd = fd.load(); + if (local_fd < 0) + return false; + + pollfd pfd{}; + pfd.fd = local_fd; + pfd.events = POLLOUT; + if (::poll(&pfd, 1, 0) < 0) + return false; + if ((pfd.revents & (POLLERR | POLLHUP | POLLNVAL)) != 0) + return false; + + int so_error = 0; + socklen_t len = sizeof(so_error); + if (::getsockopt(local_fd, SOL_SOCKET, SO_ERROR, &so_error, &len) != 0) + return false; + + return so_error == 0; +} + +bool TCPStreamPusherSocket::EnsureAccepted() { + if (fd.load() >= 0) + return true; + return AcceptConnection(std::chrono::duration_cast(AcceptTimeout)); +} + +bool TCPStreamPusherSocket::SendAll(const void *buf, size_t len) { + const uint8_t *p = static_cast(buf); + size_t sent = 0; + while (sent < len) { + int local_fd = fd.load(); + if (local_fd < 0 || broken) + return false; + + ssize_t rc = ::send(local_fd, p + sent, len - sent, MSG_NOSIGNAL); + if (rc < 0) { + if (errno == EINTR) + continue; + + if (errno == EPIPE || errno == ECONNRESET || errno == ENOTCONN) { + CloseDataSocket(); + broken = true; + logger.Error("TCP peer disconnected on " + endpoint + ", stopping this stream"); + return false; + } + return false; + } + sent += static_cast(rc); + } + return true; +} + +bool TCPStreamPusherSocket::SendPayloadZC(const uint8_t *data, size_t size, ZeroCopyReturnValue *z) { +#if defined(MSG_ZEROCOPY) && defined(SO_ZEROCOPY) + int local_fd = fd.load(); + if (local_fd < 0) + return false; + + msghdr msg{}; + iovec iov{}; + iov.iov_base = const_cast(data); + iov.iov_len = size; + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + + while (true) { + ssize_t rc = ::sendmsg(local_fd, &msg, MSG_ZEROCOPY | MSG_NOSIGNAL); + if (rc < 0) { + if (errno == EINTR) + continue; + if (errno == EAGAIN) + return SendAll(data, size); + return false; + } + if (static_cast(rc) != size) + return false; + break; + } + + std::unique_lock ul(inflight_mutex); + inflight.push_back(InflightZC{.z = z, .tx_id = next_tx_id.fetch_add(1)}); + return true; +#else + (void) z; + return SendAll(data, size); +#endif +} + +bool TCPStreamPusherSocket::SendFrame(const uint8_t *data, size_t size, TCPFrameType type, int64_t image_number, ZeroCopyReturnValue *z) { + TcpFrameHeader h{}; + h.type = static_cast(type); + h.payload_size = size; + h.image_number = image_number >= 0 ? static_cast(image_number) : 0; + h.socket_number = socket_number; + h.run_number = run_number; + + if (!SendAll(&h, sizeof(h))) { + if (z) + z->release(); + return false; + } + + if (size == 0) { + if (z) + z->release(); + return true; + } + + if (z && zerocopy_threshold && size >= zerocopy_threshold.value()) { + bool ok = SendPayloadZC(data, size, z); + if (!ok) + z->release(); + return ok; + } + + bool ok = SendAll(data, size); + if (z) + z->release(); + return ok; +} + +void TCPStreamPusherSocket::WriterThread() { + while (active) { + const auto e = queue.GetBlocking(); + if (e.end) + break; + + if (!e.z) + continue; + + if (!EnsureAccepted()) + broken = true; + + if (broken) { + e.z->release(); + continue; + } + + bool ok = false; + { + std::unique_lock ul(send_mutex); + ok = SendFrame(static_cast(e.z->GetImage()), + e.z->GetImageSize(), + TCPFrameType::DATA, + e.z->GetImageNumber(), + e.z); + } + + if (!ok) { + broken = true; + logger.Error("TCP send failed on " + endpoint + ", stopping this stream"); + } + } +} + +bool TCPStreamPusherSocket::IsBroken() const { + return broken; +} + +void TCPStreamPusherSocket::CompletionThread() { +#if defined(MSG_ZEROCOPY) && defined(SO_ZEROCOPY) + while (active) { + int local_fd = fd.load(); + if (local_fd < 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + continue; + } + + char cmsgbuf[512]; + char dummy[1]; + iovec iov{.iov_base = dummy, .iov_len = sizeof(dummy)}; + msghdr msg{}; + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + msg.msg_control = cmsgbuf; + msg.msg_controllen = sizeof(cmsgbuf); + + ssize_t rc = ::recvmsg(local_fd, &msg, MSG_ERRQUEUE | MSG_DONTWAIT); + if (rc < 0) { + if (errno == EAGAIN || errno == EINTR) { + std::this_thread::sleep_for(std::chrono::microseconds(100)); + continue; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + continue; + } + + for (cmsghdr *cmsg = CMSG_FIRSTHDR(&msg); cmsg; cmsg = CMSG_NXTHDR(&msg, cmsg)) { + if (cmsg->cmsg_level != SOL_IP || cmsg->cmsg_type != IP_RECVERR) + continue; + auto *serr = reinterpret_cast(CMSG_DATA(cmsg)); + if (!serr || serr->ee_origin != SO_EE_ORIGIN_ZEROCOPY) + continue; + + uint64_t first = serr->ee_info; + uint64_t last = serr->ee_data; + + std::unique_lock ul(inflight_mutex); + while (!inflight.empty() && inflight.front().tx_id <= last) { + auto item = inflight.front(); + inflight.pop_front(); + if (item.tx_id >= first && item.z) + item.z->release(); + } + } + } + + std::unique_lock ul(inflight_mutex); + while (!inflight.empty()) { + if (inflight.front().z) + inflight.front().z->release(); + inflight.pop_front(); + } +#endif +} + +void TCPStreamPusherSocket::StartWriterThread() { + active = true; + send_future = std::async(std::launch::async, &TCPStreamPusherSocket::WriterThread, this); + completion_future = std::async(std::launch::async, &TCPStreamPusherSocket::CompletionThread, this); +} + +void TCPStreamPusherSocket::StopWriterThread() { + if (!active) + return; + active = false; + queue.PutBlocking({.end = true}); + + if (send_future.valid()) + send_future.get(); + if (completion_future.valid()) + completion_future.get(); + + // Keep fd open: END frame may still be sent after writer thread stops. + // Socket is closed in destructor / explicit close path. +} + +void TCPStreamPusherSocket::SendImage(ZeroCopyReturnValue &z) { + queue.PutBlocking(ImagePusherQueueElement{ + .image_data = static_cast(z.GetImage()), + .z = &z, + .end = false + }); +} + +bool TCPStreamPusherSocket::Send(const uint8_t *data, size_t size, TCPFrameType type, int64_t image_number) { + if (broken) + return false; + + std::unique_lock ul(send_mutex); + + if (fd.load() < 0) + return false; + + if (!IsConnectionAlive()) { + broken = true; + return false; + } + + return SendFrame(data, size, type, image_number, nullptr); +} diff --git a/image_pusher/TcpStreamPusherSocket.h b/image_pusher/TcpStreamPusherSocket.h new file mode 100644 index 00000000..e2338c19 --- /dev/null +++ b/image_pusher/TcpStreamPusherSocket.h @@ -0,0 +1,84 @@ +// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "ImagePusher.h" +#include "../common/ThreadSafeFIFO.h" +#include "../common/Logger.h" +#include "../common/JfjochTCP.h" + +class TCPStreamPusherSocket { + struct InflightZC { + ZeroCopyReturnValue *z = nullptr; + uint64_t tx_id = 0; + }; + + std::mutex send_mutex; + std::atomic active = false; + std::future send_future; + std::future completion_future; + + ThreadSafeFIFO queue; + + std::atomic fd{-1}; + int listen_fd = -1; + std::string endpoint; + uint32_t socket_number = 0; + uint64_t run_number = 0; + std::optional zerocopy_threshold; + std::optional send_buffer_size; + + constexpr static auto AcceptTimeout = std::chrono::seconds(5); + + std::atomic ever_connected{false}; + std::atomic broken{false}; + + std::atomic next_tx_id{1}; + std::mutex inflight_mutex; + std::deque inflight; + + Logger logger{"TCPStream2PusherSocket"}; + + void WriterThread(); + void CompletionThread(); + + bool EnsureAccepted(); + void CloseDataSocket(); + + bool SendAll(const void *buf, size_t len); + bool SendFrame(const uint8_t *data, size_t size, TCPFrameType type, int64_t image_number, ZeroCopyReturnValue *z); + bool SendPayloadZC(const uint8_t *data, size_t size, ZeroCopyReturnValue *z); +public: + explicit TCPStreamPusherSocket(const std::string& addr, + uint32_t in_socket_number, + std::optional send_buffer_size, + std::optional in_zerocopy_threshold, + size_t send_queue_size = 4096); + + ~TCPStreamPusherSocket(); + + std::string GetEndpointName() const; + + bool Send(const uint8_t *data, size_t size, TCPFrameType type, int64_t image_number = -1); + + bool AcceptConnection(std::chrono::milliseconds timeout = std::chrono::duration_cast(AcceptTimeout)); + bool IsConnectionAlive() const; + + void StartWriterThread(); + void StopWriterThread(); + + void SetRunNumber(uint64_t in_run_number); + + void SendImage(ZeroCopyReturnValue &z); + + bool IsBroken() const; +}; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 248d7dac..7eb1cbfe 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -64,6 +64,7 @@ ADD_EXECUTABLE(jfjoch_test RotationIndexerTest.cpp TopPixelsTest.cpp HKLKeyTest.cpp + TcpTest.cpp ) target_link_libraries(jfjoch_test Catch2WithMain JFJochBroker JFJochReceiver JFJochReader JFJochWriter JFJochImageAnalysis JFJochCommon JFJochHLSSimulation JFJochPreview) diff --git a/tests/TcpTest.cpp b/tests/TcpTest.cpp new file mode 100644 index 00000000..90fedefe --- /dev/null +++ b/tests/TcpTest.cpp @@ -0,0 +1,98 @@ +// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include +#include + +#include "../image_pusher/TcpStreamPusher.h" +#include "../image_puller/TcpImagePuller.h" + +TEST_CASE("TCPImageCommTest_2Writers", "[TCP]") { + const size_t nframes = 128; + const int64_t npullers = 2; + const int64_t images_per_file = 16; + + DiffractionExperiment x(DetJF(1)); + x.Raw(); + x.PedestalG0Frames(0).NumTriggers(1).UseInternalPacketGenerator(false).IncidentEnergy_keV(12.4) + .ImagesPerTrigger(nframes).Compression(CompressionAlgorithm::NO_COMPRESSION); + + std::mt19937 g1(1387); + std::uniform_int_distribution dist; + std::vector image1(x.GetPixelsNum() * nframes); + for (auto &i : image1) i = dist(g1); + + std::vector addr{ + "tcp://127.0.0.1:19001", + "tcp://127.0.0.1:19002" +}; + + std::vector> puller; + for (int i = 0; i < npullers; i++) { + puller.push_back(std::make_unique( + addr[i], 64 * 1024 * 1024)); // decoded cbor ring + } + + TCPStreamPusher pusher( + addr, + 64 * 1024 * 1024, + 128 * 1024, // zerocopy threshold + 8192 // sender queue + ); + + std::vector received(npullers, 0); + + std::thread sender([&] { + std::vector serialization_buffer(16 * 1024 * 1024); + CBORStream2Serializer serializer(serialization_buffer.data(), serialization_buffer.size()); + + StartMessage start{ + .images_per_file = images_per_file, + .write_master_file = true + }; + EndMessage end{}; + + pusher.StartDataCollection(start); + + for (int64_t i = 0; i < static_cast(nframes); i++) { + DataMessage data_message; + data_message.number = i; + data_message.image = CompressedImage(image1.data() + i * x.GetPixelsNum(), + x.GetPixelsNum() * sizeof(uint16_t), + x.GetXPixelsNum(), + x.GetYPixelsNum(), + x.GetImageMode(), + x.GetCompressionAlgorithm()); + serializer.SerializeImage(data_message); + REQUIRE(pusher.SendImage(serialization_buffer.data(), serializer.GetBufferSize(), i)); + } + + REQUIRE(pusher.EndDataCollection(end)); + }); + + for (int w = 0; w < npullers; w++) { + bool seen_end = false; + while (!seen_end) { + auto out = puller[w]->PollImage(std::chrono::seconds(10)); + REQUIRE(out.has_value()); + REQUIRE(out->cbor != nullptr); + if (out->cbor->end_message) { + seen_end = true; + continue; + } + if (out->cbor->data_message) { + auto n = out->cbor->data_message->number; + REQUIRE(((n / images_per_file) % npullers) == w); + received[w]++; + } + } + } + + sender.join(); + + REQUIRE(received[0] == nframes / 2); + REQUIRE(received[1] == nframes / 2); + + for (auto &p : puller) + p->Disconnect(); +} \ No newline at end of file