Implement raw TCP/IP in jfjoch_broker and jfjoch_writer

image_pusher/TCPStreamPusher.cpp

@@ -0,0 +1,141 @@
+// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#include "TCPStreamPusher.h"
+
+
+TCPStreamPusher::TCPStreamPusher(const std::vector<std::string> &addr,
+                                 std::optional<int32_t> send_buffer_size,
+                                 std::optional<size_t> 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<TCPStreamPusherSocket>(
+                addr[i], static_cast<uint32_t>(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<int64_t>(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::Finalize() {
+    std::string ret;
+    if (writer_notification_socket) {
+        for (size_t i = 0; i < socket.size(); i++) {
+            auto n = writer_notification_socket->Receive(run_number, run_name);
+            if (!n)
+                ret += "Writer " + std::to_string(i) + ": no end notification received within 1 minute from collection end";
+            else if (static_cast<size_t>(n->socket_number) >= socket.size())
+                ret += "Writer " + std::to_string(i) + ": mismatch in socket number";
+            else if (!n->ok)
+                ret += "Writer " + std::to_string(i) + ": " + n->error;
+        }
+    }
+    return ret;
+}
+
+std::string TCPStreamPusher::GetWriterNotificationSocketAddress() const {
+    if (writer_notification_socket)
+        return writer_notification_socket->GetEndpointName();
+    else
+        return "";
+}
+
+TCPStreamPusher &TCPStreamPusher::WriterNotificationSocket(const std::string &addr) {
+    writer_notification_socket = std::make_unique<ZMQWriterNotificationPuller>(addr, std::chrono::minutes(1));
+    return *this;
+}
+
+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);
+}