Jungfraujoch/writer/StreamWriter.cpp

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

#include "../common/JFJochException.h"
#include "StreamWriter.h"

#include <utility>
#include "FileWriter.h"

StreamWriter::StreamWriter(Logger &in_logger,
                           ImagePuller &in_image_puller,
                           std::string in_file_done_address,
                           bool in_verbose)
        : verbose(in_verbose),
          image_puller(in_image_puller),
          logger(in_logger),
          file_done_address(std::move(in_file_done_address)),
          socket_number(0),
          run_number(0),
          max_image_number(0) {
}

void StreamWriter::NotifyTcpAck(TCPFrameType ack_for, bool ok, bool fatal, TCPAckCode code, const std::string &error_text) {
    if (!image_puller.SupportsAck())
        return;

    PullerAckMessage ack;
    ack.ack_for = ack_for;
    ack.ok = ok;
    ack.fatal = fatal;
    ack.error_code = code;
    ack.error_text = error_text;
    ack.run_number = run_number;
    ack.socket_number = static_cast<uint32_t>(socket_number);
    ack.processed_images = processed_images.load();

    if (image_puller_output.cbor && image_puller_output.cbor->data_message)
        ack.image_number = image_puller_output.cbor->data_message->number;

    if (!image_puller.SendAck(ack))
        logger.Warning("Failed to send TCP ACK");
}

void StreamWriter::ProcessStartMessage() {
    if (state == StreamWriterState::Finalized)
        return; // Should not happen (?)

    if (state != StreamWriterState::Idle)
        FinalizeDataCollection();

    err = "";
    tcp_data_fatal_sent = false;

    max_image_number = 0;

    processed_images = 0;
    processed_image_size = 0;

    file_prefix = image_puller_output.cbor->start_message->file_prefix;
    run_number = image_puller_output.cbor->start_message->run_number;
    run_name = image_puller_output.cbor->start_message->run_name;

    socket_number = 0;
    if (image_puller_output.cbor->start_message->socket_number)
        socket_number = image_puller_output.cbor->start_message->socket_number.value();
    writer_notification_zmq_addr = image_puller_output.cbor->start_message->writer_notification_zmq_addr;

    try {
        // Fail fast on an overwrite conflict only when the transport can report it
        // back to the broker (TCP ACK). The ZeroMQ path has no back-channel, so it
        // keeps writing .tmp files and fails at the final rename instead.
        file_writer = std::make_unique<FileWriter>(*image_puller_output.cbor->start_message,
                                                   image_puller.SupportsAck());
        if (!file_done_address.empty())
            file_writer->SetupFinalizedFileSocket(file_done_address);
        logger.Info("Starting writing for dataset {} of {} images",
                    image_puller_output.cbor->start_message->file_prefix,
                    image_puller_output.cbor->start_message->number_of_images);
        state = StreamWriterState::Started;
        NotifyTcpAck(TCPFrameType::START, true, false, TCPAckCode::None);
    } catch (const JFJochException &e) {
        logger.ErrorException(e);
        logger.Error("Error writing start message - switching to error state");
        state = StreamWriterState::Error;
        err = e.what();
        NotifyTcpAck(TCPFrameType::START, false, true, TCPAckCode::StartFailed, err);
    }
}

// Dry run for a run that has not started: check that the output could be written and answer,
// without opening anything and without touching state - the writer stays Idle and ready for the
// START that follows.
void StreamWriter::ProcessPreflight() {
    const StartMessage &msg = image_puller_output.cbor->start_message.value();

    // The pusher's ACK reader discards frames whose run_number does not match the run it is about
    // to start, so the ACK has to carry the run number of the message being checked - it is not yet
    // the one ProcessStartMessage will store.
    run_number = msg.run_number;
    socket_number = msg.socket_number.value_or(0);

    try {
        FileWriter::Preflight(msg);
        NotifyTcpAck(TCPFrameType::PREFLIGHT, true, false, TCPAckCode::None);
    } catch (const JFJochException &e) {
        logger.Warning("Pre-flight check failed: {}", e.what());
        // Not fatal: nothing was started, so the connection stays usable for the next attempt.
        NotifyTcpAck(TCPFrameType::PREFLIGHT, false, false, TCPAckCode::StartFailed, e.what());
    }
}

void StreamWriter::ProcessCalibrationImage() {
    switch (state) {
        case StreamWriterState::Started:
            try {
                file_writer->WriteHDF5(*image_puller_output.cbor->calibration);
            } catch (const std::exception &e) {
                logger.Warning(e.what());
                logger.Warning("Error during writing calibration data - skipping");
                NotifyTcpAck(TCPFrameType::CALIBRATION, false, false, TCPAckCode::DataWriteFailed, e.what());
            }
            break;
        case StreamWriterState::Receiving:
            logger.Warning("Unexpected calibration message");
            break;
        case StreamWriterState::Error:
        case StreamWriterState::Idle:
        case StreamWriterState::Finalized:
            break;
    }
}

void StreamWriter::ProcessDataImage() {
    switch (state) {
        case StreamWriterState::Idle:
            logger.Warning("Missing meaningful image while waiting for START");
            mute_data_msg_in_idle = true;
            break;
        case StreamWriterState::Started:
            start_time = std::chrono::system_clock::now();
            state = StreamWriterState::Receiving;
            // Follow through to receiving - no brake!
        case StreamWriterState::Receiving:
            try {
                if (verbose)
                    logger.Info("Received data message {}",
                                image_puller_output.cbor->data_message->number);

                file_writer->Write(*image_puller_output.cbor->data_message);
                if (max_image_number < image_puller_output.cbor->data_message->number + 1)
                    max_image_number = image_puller_output.cbor->data_message->number + 1;

                processed_images++;
                processed_image_size += image_puller_output.cbor->data_message->image.GetCompressedSize();
                if (verbose)
                    logger.Info("Written");
                NotifyTcpAck(TCPFrameType::DATA, true, false, TCPAckCode::None);
            } catch (const JFJochException &e) {
                logger.ErrorException(e);
                logger.Warning("Error writing image - switching to error state");
                state = StreamWriterState::Error;
                err = e.what();
                NotifyTcpAck(TCPFrameType::DATA, false, true, TCPAckCode::DataWriteFailed, err);
            }
            break;
        case StreamWriterState::Error:
            // Error state => Wait till end only
        case StreamWriterState::Finalized:
            break;
    }
}

void StreamWriter::ProcessEndMessage() {
    // Ignore end message when idle state!
    if (state == StreamWriterState::Idle || state == StreamWriterState::Finalized)
        return;

    if (verbose)
        logger.Info("Received end message");

    if (state != StreamWriterState::Error) {
        try {
            if ((image_puller_output.cbor->end_message->max_image_number == 0) && (max_image_number > 0))
                image_puller_output.cbor->end_message->max_image_number = max_image_number;
            file_writer->WriteHDF5(*image_puller_output.cbor->end_message);
        } catch (const JFJochException &e) {
            logger.ErrorException(e);
            logger.Error("Error writing end message - switching to error state");
            state = StreamWriterState::Error;
            err = e.what();
        }
    }

    FinalizeDataCollection();

    const bool error_state = (state == StreamWriterState::Error);

    NotifyReceiverOnFinalizedWrite(writer_notification_zmq_addr);
    NotifyTcpAck(TCPFrameType::END, !error_state, error_state,
                 error_state ? TCPAckCode::EndFailed : TCPAckCode::None,
                 error_state ? err : "");

    // To exit main image loop in CollectImages(), state must be finalized
    state = StreamWriterState::Finalized;
}

void StreamWriter::FinalizeDataCollection() {
    end_time = std::chrono::system_clock::now();

    bool finalize_error = false;

    if (file_writer && (state != StreamWriterState::Error)) {
        try {
            hdf5_data_file_statistics = file_writer->Finalize();
        } catch (const JFJochException &e) {
            finalize_error = true;
            state = StreamWriterState::Error;
            err = e.what();
            logger.ErrorException(e);
            logger.Error("Error finalizing writing - switching to error state");
        } catch (const std::exception &e) {
            finalize_error = true;
            state = StreamWriterState::Error;
            err = e.what();
            logger.Error("Error finalizing writing - switching to error state: {}", e.what());
        }
    } else {
        hdf5_data_file_statistics.clear();
    }

    file_writer.reset();
    logger.Info("Data writing finished");

    if (!finalize_error && state != StreamWriterState::Error)
        state = StreamWriterState::Finalized;
}

void StreamWriter::CollectImages() {
    state = StreamWriterState::Idle;
    mute_data_msg_in_idle = false;

    bool run = true;
    while (run && state != StreamWriterState::Finalized) {
        run = WaitForImage();
        // No new frame: WaitForImage leaves the previous one in image_puller_output, and the body
        // below would go on to process it a second time - writing the last image twice, or
        // repeating the last check - before the loop condition noticed.
        if (!run)
            break;

        if (image_puller_output.tcp_msg &&
            static_cast<TCPFrameType>(image_puller_output.tcp_msg->header.type) == TCPFrameType::CANCEL) {
            logger.Warning("Received TCP CANCEL, finalizing data collection");
            if (state != StreamWriterState::Idle)
                FinalizeDataCollection();
            const bool error_state = (state == StreamWriterState::Error);
            NotifyTcpAck(TCPFrameType::CANCEL, !error_state, error_state,
                         error_state ? TCPAckCode::EndFailed : TCPAckCode::None,
                         error_state ? err : "");

            state = StreamWriterState::Finalized;
            continue;
        }

        if (!image_puller_output.cbor) {
            logger.Warning("Missing CBOR payload for non-CANCEL TCP frame");
            continue;
        }

        // A PREFLIGHT carries a StartMessage payload, so it has to be told apart by frame type
        // before the payload-driven dispatch below would take it for a real START.
        if (image_puller_output.tcp_msg &&
            static_cast<TCPFrameType>(image_puller_output.tcp_msg->header.type) == TCPFrameType::PREFLIGHT) {
            if (image_puller_output.cbor->start_message)
                ProcessPreflight();
            else
                logger.Warning("PREFLIGHT frame without start message");
            continue;
        }

        if (image_puller_output.cbor->start_message)
            ProcessStartMessage();
        else if (image_puller_output.cbor->calibration)
            ProcessCalibrationImage();
        else if (image_puller_output.cbor->data_message)
            ProcessDataImage();
        else if (image_puller_output.cbor->end_message)
            ProcessEndMessage();
        else
            logger.Warning("Unknown message type");
    }
}

void StreamWriter::Cancel() {
    logger.Info("Cancel requested");
    abort = true;
}

StreamWriterOutput StreamWriter::Run() {
    hdf5_data_file_statistics.clear();
    try {
        CollectImages();
    } catch (std::exception &e) {
        // Error during collecting images will skip to end data collection
        // End data collection will consume all images till the end
        logger.ErrorException(e);
        logger.Error("Exception not properly handled by CollectImages()");
    }

    StreamWriterOutput ret;
    ret.data_file_stats = hdf5_data_file_statistics;
    ret.stats = GetStatistics();
    logger.Info("Write task done. Images = {} Throughput = {:.0f} MB/s Frame rate = {:.0f} Hz max occupation of FIFO {}",
                ret.stats.processed_images, ret.stats.performance_MBs, ret.stats.performance_Hz,
                ret.stats.max_puller_fifo_utilization);
    return ret;
}

bool StreamWriter::WaitForImage() {
    try {
        std::optional<ImagePullerOutput> ret;
        while (!ret && !abort)
            ret = image_puller.PollImage();

        if (ret.has_value())
            image_puller_output = ret.value();
        return ret.has_value();
    } catch (const JFJochException &e) {
        logger.ErrorException(e);
        return false;
    }
}

StreamWriterStatistics StreamWriter::GetStatistics() const {
    float perf_MBs = 0.0f, perf_Hz = 0.0f;

    if ((state != StreamWriterState::Started) && (processed_images > 0)) {
        int64_t time_us;
        if (state == StreamWriterState::Idle || state == StreamWriterState::Finalized)
            time_us = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time).count();
        else
            time_us = std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now() - start_time).count();

        // MByte/s ==> Byte/us
        perf_MBs = static_cast<float>(processed_image_size) / static_cast<float>(time_us);
        perf_Hz = static_cast<float>(processed_images) * 1e6f / static_cast<float>(time_us);
    }

    return {
            .processed_images = processed_images,
            .performance_MBs = perf_MBs,
            .performance_Hz = perf_Hz,
            .file_prefix = file_prefix,
            .run_name = run_name,
            .run_number = run_number,
            .socket_number = socket_number,
            .state = state,
            .max_puller_fifo_utilization = image_puller.GetMaxFifoUtilization()
    };
}

void StreamWriter::NotifyReceiverOnFinalizedWrite(const std::string &detector_update_zmq_addr) {
    if (debug_skip_write_notification) {
        logger.Info("StreamWriter: Skipping notification");
        return;
    }

    if (detector_update_zmq_addr.empty())
        return;

    nlohmann::json j;

    auto stats = GetStatistics();
    j["socket_number"] = socket_number;
    j["processed_images"] = processed_images.load();
    j["socket_number"] = stats.socket_number;
    j["run_number"] = stats.run_number;
    j["run_name"] = stats.run_name;
    j["performance_MBs"] = stats.performance_MBs;
    if (state == StreamWriterState::Error) {
        j["ok"] = false;
        j["error"] = err;
    } else
        j["ok"] = true;
    try {
        logger.Info("Sending notification to {}", detector_update_zmq_addr);
        ZMQSocket s(ZMQSocketType::Push);
        s.SendTimeout(std::chrono::seconds(1));
        s.Connect(detector_update_zmq_addr);
        s.Send(j.dump());
    } catch (const JFJochException &e) {
        logger.ErrorException(e);
        logger.Error("Error sending notification to detector update socket");
    }
}

void StreamWriter::DebugSkipWriteNotification(bool input) {
    debug_skip_write_notification = input;
}