Files
Jungfraujoch/image_pusher/ZMQStream2Pusher.cpp
2025-07-16 20:19:48 +02:00

128 lines
4.9 KiB
C++

// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "ZMQStream2Pusher.h"
#include "../frame_serialize/CBORStream2Serializer.h"
ZMQStream2Pusher::ZMQStream2Pusher(const std::vector<std::string> &addr,
std::optional<int32_t> send_buffer_high_watermark,
std::optional<int32_t> send_buffer_size)
: serialization_buffer(256*1024*1024),
serializer(serialization_buffer.data(), serialization_buffer.size())
{
if (addr.empty())
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No writer ZMQ address provided");
for (const auto &a : addr) {
auto s = std::make_unique<ZMQSocket>(ZMQSocketType::Push);
if (send_buffer_size)
s->SendBufferSize(send_buffer_size.value());
if (send_buffer_high_watermark)
s->SendWaterMark(send_buffer_high_watermark.value());
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) {
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 zmq_socket_free(void *data, void *hint) {
auto z = (ZeroCopyReturnValue *) hint;
z->release();
}
void ZMQStream2Pusher::SendImage(ZeroCopyReturnValue &z) {
if (!socket.empty()) {
auto socket_number = (z.GetImageNumber() / images_per_file) % socket.size();
socket[socket_number]->SendZeroCopy(z.GetImage(), z.GetImageSize(), zmq_socket_free, &z);
} else
z.release();
}
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;
run_number = message.run_number;
run_name = message.run_name;
for (int i = 0; i < socket.size(); i++) {
message.socket_number = i;
if (i > 0)
message.write_master_file = false; // Only writer on first socket is asked to write master file
serializer.SerializeSequenceStart(message);
if (!socket[i]->Send(serialization_buffer.data(), serializer.GetBufferSize(), true))
throw JFJochException(JFJochExceptionCategory::ZeroMQ, "Timeout on pushing start message on addr "
+ socket[i]->GetEndpointName());
}
}
bool ZMQStream2Pusher::SendCalibration(const CompressedImage &message) {
if (socket.empty())
return false;
serializer.SerializeCalibration(message);
return socket[0]->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
}
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;
}
return ret;
}
std::vector<std::string> ZMQStream2Pusher::GetAddress() {
std::vector<std::string> ret;
for (auto &p: socket)
ret.push_back(p->GetEndpointName());
return ret;
}
std::string ZMQStream2Pusher::Finalize() {
std::string ret;
if (writer_notification_socket) {
for (int 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 (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 ZMQStream2Pusher::GetWriterNotificationSocketAddress() const {
if (writer_notification_socket)
return writer_notification_socket->GetEndpointName();
else
return "";
}
ZMQStream2Pusher &ZMQStream2Pusher::WriterNotificationSocket(const std::string &addr) {
writer_notification_socket = std::make_unique<ZMQWriterNotificationPuller>(addr, std::chrono::minutes(1));
return *this;
}
std::string ZMQStream2Pusher::PrintSetup() const {
std::string output = "ZMQStream2Pusher: Sending images to the following ZeroMQ sockets: ";
for (const auto &s: socket)
output += s->GetEndpointName() + " ";
return output;
}