Jungfraujoch/image_pusher/ZMQStream2Pusher.cpp

// 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)
{
    if (addr.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No writer ZMQ address provided");

    for (const auto &a : addr) {
        auto s = std::make_unique<ZMQStream2PusherSocket>(a, send_buffer_high_watermark, send_buffer_size);
        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);
    }
    return false;
}


bool ZMQStream2Pusher::SendImage(ZeroCopyReturnValue &z) {
    if (!socket.empty()) {
        auto socket_number = (z.GetImageNumber() / images_per_file) % socket.size();
        socket[socket_number]->SendImage(z);
        return true;
    } else {
        z.release();
        return false;
    }
}

void ZMQStream2Pusher::StartDataCollection(StartMessage& message) {
    {
        std::unique_lock ul(images_written_mutex);
        images_written = std::nullopt;
    }
    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;
    early_notifications.clear();
    run_number = message.run_number;
    run_name = message.run_name;
    transmission_error = false;

    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()))
            throw JFJochException(JFJochExceptionCategory::ZeroMQ, "Timeout on pushing start message on addr "
                                                                   + socket[i]->GetEndpointName());
    }

    for (const auto & i : socket)
        i->StartWriterThread();
}

bool ZMQStream2Pusher::SendCalibration(const CompressedImage &message) {
    if (socket.empty())
        return false;

    serializer.SerializeCalibration(message);

    return socket[0]->Send(serialization_buffer.data(), serializer.GetBufferSize());
}

bool ZMQStream2Pusher::EndDataCollection(const EndMessage& message) {
    serializer.SerializeSequenceEnd(message);

    bool ret = true;

    // The writer on the first socket writes the master file, so it gets END last, and - where the
    // writers report back - only once every other writer has reported that it closed (renamed into
    // place) its data files: the master then appears only after every data file it links to.
    early_notifications.clear();
    for (size_t i = 1; i < socket.size(); i++) {
        socket[i]->StopWriterThread();
        if (!socket[i]->Send(serialization_buffer.data(), serializer.GetBufferSize()))
            ret = false;
    }
    if (writer_notification_socket)
        for (size_t i = 1; i < socket.size(); i++)
            early_notifications.push_back(writer_notification_socket->Receive(run_number, run_name));
    if (!socket.empty()) {
        socket[0]->StopWriterThread();
        if (!socket[0]->Send(serialization_buffer.data(), serializer.GetBufferSize()))
            ret = false;
    }
    transmission_error = !ret;
    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 (transmission_error)
        ret += "Timeout sending images (e.g., writer disabled during data collection);";
    uint64_t images = 0;
    bool images_saved = false;
    if (writer_notification_socket) {
        for (int i = 0; i < socket.size(); i++) {
            // The notifications EndDataCollection already waited for, then the rest.
            auto n = (i < early_notifications.size()) ? early_notifications[i]
                                                     : 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;
                images += n->processed_images;
                images_saved = true;
            }
        }
    }

    {
        std::unique_lock ul(images_written_mutex);
        if (images_saved)
            images_written = images;
        else
            images_written = std::nullopt;
    }
    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;
}

std::optional<uint64_t> ZMQStream2Pusher::GetImagesWritten() const {
    std::unique_lock ul(images_written_mutex);
    return images_written;
}

size_t ZMQStream2Pusher::GetConnectedWriters() const {
    return socket.size();
}