Jungfraujoch/image_pusher/HDF5FilePusher.cpp

// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "HDF5FilePusher.h"
#include "../frame_serialize/CBORStream2Deserializer.h"
#include "../frame_serialize/CBORStream2Serializer.h"

HDF5FilePusher::HDF5FilePusher(const std::string &repub_address,
                               const std::optional<int32_t> &repub_watermark) {
    if (!repub_address.empty()) {
        repub_socket = std::make_unique<ZMQSocket>(ZMQSocketType::Push);
        repub_socket->SendWaterMark(repub_watermark.value_or(default_repub_watermark));
        repub_socket->SendTimeout(RepubTimeout);
        repub_socket->Bind(repub_address);
    }
}

void HDF5FilePusher::Preflight(StartMessage &message) {
    // In-process writer, so the check is the same call the FileWriter constructor will make.
    FileWriter::Preflight(message);
}

void HDF5FilePusher::StartDataCollection(StartMessage &message) {
    if (writer)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState, "Image pusher is already writing images");
    writer = std::make_unique<FileWriter>(message);
    writer->FileClosedCallback([this](const HDF5DataFileStatistics &s) {
        live.FileClosed(s.file_number - 1, s.max_image_number + 1);   // the statistics count files from 1
    });
    writer_future = std::async(std::launch::async, &HDF5FilePusher::WriterThread, this);
    images_written = 0;
    {
        serializer.SerializeSequenceStart(message);
        live.Begin(message, std::vector<uint8_t>(serialization_buffer.begin(),
                                                 serialization_buffer.begin() + serializer.GetBufferSize()));
    }

    if (repub_socket) {
        try {
            StartMessage repub_message = message;
            repub_message.writer_notification_zmq_addr = "";

            serializer.SerializeSequenceStart(repub_message);

            repub_active = repub_socket->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
            if (repub_active)
                logger.Info("Republish active");
        } catch (const JFJochException &e) {
            logger.ErrorException(e);
        }
    }
}

bool HDF5FilePusher::EndDataCollection(const EndMessage &message) {
    if (!writer_future.valid())
        throw JFJochException(JFJochExceptionCategory::WrongDAQState, "Image pusher wasn't writing files");

    writer_queue.PutBlocking({.end = true});

    // Release the writer on every exit path. If finalization throws (e.g. a target
    // file already exists and overwrite is off) the writer must still be torn down,
    // otherwise the next collection fails with "already writing images".
    try {
        writer_future.get();

        if (!writer)
            throw JFJochException(JFJochExceptionCategory::WrongDAQState, "Image pusher not ready for writing images");
        writer->WriteHDF5(message);
        writer->WriteResultFiles(message, logger);
        writer->Finalize();
    } catch (...) {
        writer.reset();
        live.End(message.max_image_number, "the writer failed to finish the collection");
        throw;
    }
    writer.reset();
    live.End(message.max_image_number);

    if (repub_socket) {
        try {
            serializer.SerializeSequenceEnd(message);

            if (repub_active)
                repub_socket->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
        } catch (const JFJochException &e) {
            logger.ErrorException(e);
        }

        if (repub_active)
            logger.Info("Republish finished");
        repub_active = false;
    }

    return true;
}

bool HDF5FilePusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number) {
    if (!writer)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState, "Image pusher not ready for sending");

    auto deserialized = CBORStream2Deserialize(image_data, image_size);
    if (deserialized->data_message) {
        writer->Write(*deserialized->data_message);
        images_written++;

        if (repub_socket && repub_active) {
            try {
                repub_socket->Send(image_data, image_size, false);
            } catch (const JFJochException &e) {
                logger.ErrorException(e);
            }
        }
    } else
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "HDF5FilePusher::SendImage accepts only data image");
    return true;
}

bool HDF5FilePusher::SendImage(ZeroCopyReturnValue &z) {
    writer_queue.PutBlocking(ImagePusherQueueElement{
        .image_data = (uint8_t *) z.GetImage(),
        .z = &z,
        .end = false
    });
    return true;
}

bool HDF5FilePusher::SendCalibration(const CompressedImage &message) {
    if (!writer)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState, "Image pusher not ready for sending");

    writer->WriteHDF5(message);
    return true;
}

void HDF5FilePusher::WriterThread() {
    ImagePusherQueueElement e = writer_queue.GetBlocking();
    while (!e.end) {
        SendImage(e.image_data, e.z->GetImageSize(), e.z->GetImageNumber());
        e.z->release();
        e = writer_queue.GetBlocking();
    }
}

std::string HDF5FilePusher::PrintSetup() const {
    std::filesystem::path currentPath = std::filesystem::current_path();
    return "HDF5FilePusher: Images are written directly to file in base directory " + currentPath.string();
}

std::optional<uint64_t> HDF5FilePusher::GetImagesWritten() const {
    return images_written;
}

size_t HDF5FilePusher::GetConnectedWriters() const {
    return 1;
}