Jungfraujoch/writer/FileWriter.h

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

#pragma once

#include <numeric>

#include "HDF5DataFile.h"
#include "../common/JFJochMessages.h"
#include "../common/ZMQWrappers.h"
#include "HDF5NXmx.h"
#include <functional>
#include <set>
#include <unordered_set>

class FileWriter {
    FileWriterFormat format = FileWriterFormat::NXmxLegacy;

    StartMessage start_message;
    std::unique_ptr<NXmx> master_file;
    std::vector<std::unique_ptr<HDF5DataFile> > files;
    std::vector<HDF5DataFileStatistics> stats;
    std::unique_ptr<ZMQSocket> finalized_file_socket;
    std::function<void(const HDF5DataFileStatistics &)> file_closed_callback;

    std::unordered_set<uint64_t> closed_files;
    // File numbers with a live HDF5DataFile, in order. CloseOldFiles walks this and not `files`, whose
    // slots stay (empty) for every file of the run - a thousand of them for a long run, visited on every
    // image.
    std::set<uint64_t> open_files;
    constexpr static uint64_t close_file_lag_images = 1000;
    constexpr static uint64_t default_images_per_file = 1000;
    void CreateHDF5MasterFile(const StartMessage& msg);
    static void CheckOutputFilesAvailable(const StartMessage &msg, FileWriterFormat format);
    void AddStats(const std::optional<HDF5DataFileStatistics>& s);
    void CloseFile(uint64_t file_number);
    void CloseOldFiles(uint64_t current_image_number);

public:
    // check_overwrite_at_start: fail the constructor if an output file already
    // exists and overwrite is off. Only safe for transports with a back-channel
    // (direct HDF5 pusher, TCP) that can report the failure to the broker before
    // acquisition arms. The ZeroMQ pusher is fire-and-forget with no back-channel,
    // so it must pass false: the writer then falls back to writing .tmp files and
    // failing at the final rename (see docs/JFJOCH_WRITER.md).
    // trusted_path: the caller supplied the output prefix itself (offline rugnux / the viewer's
    // processing), so the multi-user CheckPath guard - which forbids absolute paths and '..' traversal
    // for remotely-supplied prefixes - is skipped, allowing an absolute output path. The broker/writer
    // never set it, so their behaviour is unchanged.
    explicit FileWriter(const StartMessage &request, bool check_overwrite_at_start = true,
                        bool trusted_path = false);
    // Everything the constructor does before it opens anything: the path guard, creating the
    // output directory and checking that no output file is in the way. Sent to the writer as a
    // PREFLIGHT frame before the detector is armed, so a run that cannot be written is refused
    // while refusing it is still free. Shares its code with the constructor, so a pre-flight that
    // passes cannot be contradicted by the start that follows it.
    static void Preflight(const StartMessage &request, bool trusted_path = false);
    void Write(const DataMessage& msg);
    void WriteHDF5(const DataMessage& msg);
    void WriteHDF5(const CompressedImage& msg);
    void WriteHDF5(const EndMessage& msg);
    std::vector<HDF5DataFileStatistics> Finalize();
    void SetupFinalizedFileSocket(const std::string &addr);
    // Called for every data file once it is closed and renamed into place.
    void FileClosedCallback(std::function<void(const HDF5DataFileStatistics &)> callback);
    std::optional<std::string> GetZMQAddr();
};