1.0.0-rc2: Fixes in preview

frame_serialize/ZMQStream2Pusher.cpp

@@ -1,66 +1,117 @@
 // Copyright (2019-2024) Paul Scherrer Institute
 
 #include "ZMQStream2Pusher.h"
+#include "CBORStream2Serializer.h"
 
-ZMQStream2Pusher::ZMQStream2Pusher(ZMQContext &context, const std::string &addr, int32_t send_buffer_high_watermark,
-                                   int32_t send_buffer_size)
-        : socket(context, ZMQSocketType::Push) {
-    Bind(addr, send_buffer_high_watermark, send_buffer_size);
-}
+ZMQStream2Pusher::ZMQStream2Pusher(const std::vector<std::string> &addr,
+                                             int32_t send_buffer_high_watermark, int32_t send_buffer_size)
+        : serialization_buffer(256*1024*1024),
+          serializer(serialization_buffer.data(), serialization_buffer.size()),
+          preview_counter(std::chrono::seconds(1)) {
+    if (addr.empty())
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No writer ZMQ address provided");
 
-ZMQStream2Pusher::ZMQStream2Pusher(const std::string &addr, int32_t send_buffer_high_watermark,
-                                   int32_t send_buffer_size)
-        : context(std::make_unique<ZMQContext>()),
-          socket(*context, ZMQSocketType::Push) {
-    Bind(addr, send_buffer_high_watermark, send_buffer_size);
-}
-
-void ZMQStream2Pusher::Bind(const std::string &addr, int32_t send_buffer_high_watermark, int32_t send_buffer_size) {
-
-    if (send_buffer_size > 0)
-        socket.SendBufferSize(send_buffer_size);
-    if (send_buffer_high_watermark > 0)
-        socket.SendWaterMark(send_buffer_high_watermark);
-    socket.SendTimeout(std::chrono::seconds(5)); // 5 seconds should be more than enough to flush buffers and to still give fast response
-    socket.Bind(addr);
-}
-
-void ZMQStream2Pusher::StartDataCollection(StartMessage &message) {
-    size_t approx_size = 1024*1024;
-    for (const auto &x : message.pixel_mask)
-        approx_size += x.size;
-
-    std::vector<uint8_t> serialization_buffer(approx_size);
-    CBORStream2Serializer serializer(serialization_buffer.data(), serialization_buffer.size());
-    serializer.SerializeSequenceStart(message);
-    if (!socket.Send(serialization_buffer.data(), serializer.GetBufferSize(), true))
-        throw JFJochException(JFJochExceptionCategory::ZeroMQ, "Timeout on pushing start message on addr " + GetAddress());
+    for (const auto &a : addr) {
+        auto s = std::make_unique<ZMQSocket>(ZMQSocketType::Push);
+        if (send_buffer_size > 0)
+            s->SendBufferSize(send_buffer_size);
+        if (send_buffer_high_watermark > 0)
+            s->SendWaterMark(send_buffer_high_watermark);
+        s->SendTimeout(std::chrono::seconds(5)); // 5 seconds should be more than enough to flush buffers and to still give fast response
+        s->Bind(a);
+        socket.emplace_back(std::move(s));
+    }
 }
 
 bool ZMQStream2Pusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number) {
-    return socket.Send(image_data, image_size, false);
+    if (preview_socket) {
+        if (preview_counter.GeneratePreview())
+            preview_socket->Send(image_data, image_size, false);
+    }
+
+    if (!socket.empty()) {
+        auto socket_number = (image_number / images_per_file) % socket.size();
+        return socket[socket_number]->Send(image_data, image_size, false);
+    } else
+        return false;
 }
 
 void ZMQStream2Pusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number,
-                                 ZeroCopyReturnValue *z) {
-    socket.SendZeroCopy(image_data,image_size, z);
+                                      ZeroCopyReturnValue *z) {
+    if (preview_socket) {
+        if (preview_counter.GeneratePreview())
+            preview_socket->Send(image_data, image_size, false);
+    }
+
+    if (!socket.empty()) {
+        auto socket_number = (image_number / images_per_file) % socket.size();
+        socket[socket_number]->SendZeroCopy(image_data, image_size, z);
+    } else
+        z->release();
 }
 
-bool ZMQStream2Pusher::EndDataCollection(const EndMessage &message) {
-    std::vector<uint8_t> serialization_buffer(80 * 1024 * 1024);
-    CBORStream2Serializer serializer(serialization_buffer.data(), serialization_buffer.size());
+void ZMQStream2Pusher::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;
 
-    serializer.SerializeSequenceEnd(message);
-    return socket.Send(serialization_buffer.data(), serializer.GetBufferSize(), true); // Blocking
+    serializer.SerializeSequenceStart(message);
+
+    for (auto &s: socket) {
+        if (!s->Send(serialization_buffer.data(), serializer.GetBufferSize(), true))
+            throw JFJochException(JFJochExceptionCategory::ZeroMQ, "Timeout on pushing start message on addr "
+                                                                   + s->GetEndpointName());
+        if (message.write_master_file) {
+            message.write_master_file = false;
+            serializer.SerializeSequenceStart(message);
+        }
+    }
+
+    if (preview_socket)
+        preview_socket->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
 }
 
 bool ZMQStream2Pusher::SendCalibration(const CompressedImage &message) {
-    std::vector<uint8_t> serialization_buffer(80 * 1024 * 1024);
-    CBORStream2Serializer serializer(serialization_buffer.data(), serialization_buffer.size());
+    if (socket.empty())
+        return false;
+
     serializer.SerializeCalibration(message);
-    return socket.Send(serialization_buffer.data(), serializer.GetBufferSize(), true); // Blocking
+
+    return socket[0]->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
 }
 
-std::string ZMQStream2Pusher::GetAddress() {
-    return socket.GetEndpointName();
+bool ZMQStream2Pusher::EndDataCollection(const EndMessage& message) {
+    serializer.SerializeSequenceEnd(message);
+
+    bool ret = true;
+    for (auto &s: socket) {
+        if (!s->Send(serialization_buffer.data(), serializer.GetBufferSize(), true))
+            ret = false;
+    }
+
+    if (preview_socket)
+        preview_socket->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
+
+    return ret;
+}
+
+std::vector<std::string> ZMQStream2Pusher::GetAddress() {
+    std::vector<std::string> ret;
+    for (auto &p: socket)
+        ret.push_back(p->GetEndpointName());
+    return ret;
+}
+
+ZMQStream2Pusher &ZMQStream2Pusher::PreviewSocket(const std::string &addr) {
+    preview_socket = std::make_unique<ZMQSocket>(ZMQSocketType::Pub);
+    preview_socket->Bind(addr);
+    return *this;
+}
+
+std::string ZMQStream2Pusher::GetPreviewAddress() {
+    if (preview_socket)
+        return preview_socket->GetEndpointName();
+    else
+        return "";
 }