Jungfraujoch/broker/JFJochServices.cpp

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

#include "JFJochServices.h"

#include <filesystem>
#include <fstream>
#include <cstdlib>
#include "../common/JFJochException.h"
#include "../detector_control/SLSDetectorWrapper.h"
#include "../detector_control/DectrisDetectorWrapper.h"

JFJochServices::JFJochServices(Logger &in_logger) : logger(in_logger) {}

void JFJochServices::Preflight(const DiffractionExperiment& experiment) {
    if (receiver == nullptr)
        return;
    receiver->Preflight(experiment);
}

void JFJochServices::Start(const DiffractionExperiment& experiment,
                           const PixelMask &pixel_mask,
                           const JFCalibration *calibration) {
    logger.Info("Measurement start for: {}", experiment.GetFilePrefix());

    cannot_stop_detector = false;

    if (receiver == nullptr) {
        logger.Info("   Done!");
        return;
    }

    logger.Info("   ... receiver start");
    if (image_puller)
        image_puller->ResumeAndClear();
    receiver->Start(experiment, pixel_mask, calibration, image_puller);

    // From here the receiver is running asynchronously. If starting the detector fails, stop the
    // receiver again so the service returns to idle and the run can be retried without
    // re-initialising. The detector failure is then propagated (a critical detector error stays
    // critical, so the broker still goes to the Error state).
    try {
        std::shared_lock ul(detector_mutex);
        if (detector && !experiment.IsUsingInternalPacketGen()) {
            logger.Info("   ... detector start");
            detector->Start(experiment);
        }
    } catch (const std::exception &e) {
        logger.Error("   ... detector failed to start ({}) - stopping detector and receiver", e.what());
        // Best-effort cleanup - it must not replace the original detector exception (which may be
        // critical), so swallow anything it throws and let the throw below re-raise e. The detector
        // may have partially armed before failing, so stop it too; otherwise it lingers in a BUSY
        // state and the next start fails with "detector busy" until a manual re-initialisation.
        try {
            {
                std::shared_lock ul(detector_mutex);
                if (detector)
                    detector->Stop();
            }
            receiver->Cancel(false);
            receiver->Stop();
        } catch (const std::exception &stop_error) {
            logger.Warning("Stop after failed start reported: {}", stop_error.what());
        }
        throw;
    }

    logger.Info("   Done!");
}

void JFJochServices::Off() {
    // Disconnect explicitly rather than relying on the destructor: the receiver service keeps its
    // own shared_ptr to the puller of the last run, so dropping this one does not necessarily close
    // the socket, and a puller left connected keeps taking its share of the detector stream.
    if (image_puller)
        image_puller->Disconnect();
    image_puller.reset();

    std::unique_ptr<DetectorWrapper> old_detector;
    {
        std::unique_lock ul(detector_mutex);
        old_detector = std::move(detector);
    }
    if (old_detector) {
        old_detector->Deactivate();
        // destroyed here, outside lock
    }
}

void JFJochServices::On(DiffractionExperiment &x) {
    if (x.IsUsingInternalPacketGen() || (receiver == nullptr))  {
        std::unique_lock ul(detector_mutex);
        detector.reset();
    } else {
        logger.Info("Detector on");

        // The previous puller has to be disconnected before the new one connects. Dropping the
        // shared_ptr is not enough: the receiver service keeps its own reference to the puller of
        // the last run, so the socket would stay open - and two PULL sockets on the same DECTRIS
        // stream share the messages round-robin, so the old one silently swallows half of the
        // stream, the start message included. The run then never starts and the receiver waits for
        // a start message that has already been thrown away.
        if (image_puller)
            image_puller->Disconnect();

        std::unique_ptr<DetectorWrapper> new_detector;
        switch (x.GetDetectorType()) {
            case DetectorType::EIGER:
            case DetectorType::JUNGFRAU:
                new_detector = std::make_unique<SLSDetectorWrapper>();
                image_puller.reset();
                break;
            case DetectorType::DECTRIS:
                new_detector = std::make_unique<DectrisDetectorWrapper>();
                image_puller = std::make_shared<ZMQImagePuller>(x.GetDetectorSetup().GetDECTRISStream2Addr());
                image_puller->Suspend();
                break;
        }
        new_detector->Initialize(x, receiver->GetNetworkConfig());
        {
            std::unique_lock ul(detector_mutex);
            detector = std::move(new_detector);
        }
        logger.Info("   ... done");
    }
}

JFJochServicesOutput JFJochServices::Stop() {
    JFJochServicesOutput ret;

    // Captured as an exception_ptr (not a copied JFJochException) so the dynamic type survives: a
    // critical hardware fault (e.g. PCIeDeviceException) must stay critical when re-thrown, otherwise
    // slicing it to a plain JFJochException would send the state machine to Idle instead of Error.
    std::exception_ptr exception;

    bool detector_error = false;

    if (receiver != nullptr) {
        try {
            {
                std::shared_lock ul(detector_mutex);
                if (detector) {
                    logger.Info("Wait for detector idle");
                    DetectorState state = detector->GetState();
                    while ((!cannot_stop_detector)
                           && ((state == DetectorState::WAITING) || (state == DetectorState::BUSY))) {
                        // check detector state every 100 ms - every check is an HTTP request for DECTRIS
                        std::this_thread::sleep_for(std::chrono::milliseconds(100));
                        state = detector->GetState();
                    }
                    if (state == DetectorState::IDLE) {
                        logger.Info("   ... detector idle");
                        receiver->Cancel(true); // cancel silently
                    } else {
                        logger.Error("   ... detector not idle: state {}{}; {}", to_string(state),
                                     cannot_stop_detector ? " (stop command failed)" : "",
                                     detector->GetErrorDescription());
                        receiver->Cancel(false);
                        detector_error = true;
                    }
                }
            }

            logger.Info("Wait for receiver done");
            ret.receiver_output = receiver->Stop();
            if (image_puller)
                image_puller->Suspend();

            logger.Info("    ... Receiver efficiency: {} %     Max delay: {}     Compression ratio {}x",
                    static_cast<int>(ret.receiver_output.efficiency * 100.0),
                    ret.receiver_output.status.max_receive_delay.value_or(0),
                    static_cast<int>(std::round(ret.receiver_output.status.compressed_ratio.value_or(1))));

            if (ret.receiver_output.efficiency < 1.0) {
                for (int i = 0; i < ret.receiver_output.received_packets.size(); i++) {
                    if (ret.receiver_output.received_packets[i] != ret.receiver_output.expected_packets[i])
                        logger.Info("    ...   Module: {}     Packets received: {} out of {}", i,
                                    ret.receiver_output.received_packets[i], ret.receiver_output.expected_packets[i]);
                }
            }

            // A writer that broke mid-run (e.g. a lost connection) leaves a truncated file. This is
            // reported to the caller via receiver_output.writer_err (the state machine surfaces it as
            // an error message) - it is not a detector fault, so it does not need re-initialisation.

            logger.Info("    ... finished with success");
        } catch (const JFJochException &e) {
            logger.Error("   ... finished with error {}", e.what());
            exception = std::current_exception();
        }
    } else {
        logger.Info("No receiver - sleeping for 30 seconds");
        std::this_thread::sleep_for(std::chrono::seconds(30));
        logger.Info("Sleep done");
    }


    // A detector fault must win over an ordinary receiver/writer error: it needs re-initialisation,
    // so raise the critical exception first. If only the receiver failed, the run can be retried from
    // Idle, so that ordinary error is raised second.
    if (detector_error)
        throw JFJochCriticalException("Error in detector operation");

    if (exception)
        std::rethrow_exception(exception);

    return ret;
}

void JFJochServices::Cancel() {
    std::shared_lock ul(detector_mutex);
    if (detector) {
        // Best effort - if detector cannot be stopped, this is OK, important to still stop receiver
        try {
            detector->Stop();
        } catch (const std::exception &e) {
            logger.Error("Cannot stop detector: {}", e.what());
            cannot_stop_detector = true;
        } catch (...) {
            cannot_stop_detector = true;
        }
    }
    if (receiver != nullptr)
        receiver->Cancel(false);
}

JFJochServices &JFJochServices::Receiver(JFJochReceiverService *input) {
    receiver = input;
    return *this;
}

std::optional<JFJochReceiverStatus> JFJochServices::GetReceiverStatus() const {
    if (receiver == nullptr)
        return {};
    return receiver->GetStatus();
}

std::optional<float> JFJochServices::GetReceiverProgress() const {
    if (receiver == nullptr)
        return {};
    return receiver->GetProgress();
}

MultiLinePlot JFJochServices::GetPlots(const PlotRequest &request) {
    if (receiver == nullptr)
        return {};
    return receiver->GetDataProcessingPlot(request);
}

void JFJochServices::GetPlotRaw(std::vector<float> &v, PlotType type, const std::string &roi) {
    if (receiver != nullptr)
        receiver->GetPlotRaw(v, type, roi);
}

void JFJochServices::SetSpotFindingSettings(const SpotFindingSettings &settings) {
    if (receiver)
        receiver->SetSpotFindingSettings(settings);
}

void JFJochServices::Trigger() {
    std::shared_lock ul(detector_mutex);
    if (detector && (receiver != nullptr))
        detector->Trigger();
}

std::optional<DetectorStatus> JFJochServices::GetDetectorStatus() const {
    std::shared_lock ul(detector_mutex, std::defer_lock);
    if (ul.try_lock_for(std::chrono::milliseconds(500)) && detector)
        return detector->GetStatus();
    return {};
}

std::string JFJochServices::GetPreviewJPEG(const PreviewImageSettings &settings, int64_t image_number) const {
    if (receiver != nullptr)
        return receiver->GetJPEGFromBuffer(settings, image_number);
    else
        return {};
}

std::string JFJochServices::GetPreviewTIFF(int64_t image_number) const {
    if (receiver != nullptr)
        return receiver->GetTIFFFromBuffer(image_number);
    else
        return "";
}

void JFJochServices::ConfigureDetector(const DiffractionExperiment &experiment) {
    std::unique_lock ul(detector_mutex); // While configuring detector ensure exclusive access (even though pointer is not modified here)
    if (detector)
        detector->Configure(experiment);
}

void JFJochServices::LoadInternalGeneratorImage(const DiffractionExperiment &experiment,
                                                const std::vector<uint16_t> &image,
                                                uint64_t image_number) {
    if (receiver)
        receiver->LoadInternalGeneratorImage(experiment, image, image_number);
}

void JFJochServices::GetXFELPulseID(std::vector<uint64_t> &v) const {
    if (receiver)
        receiver->GetXFELPulseID(v);
}

void JFJochServices::GetXFELEventCode(std::vector<uint64_t> &v) const {
    if (receiver)
        receiver->GetXFELEventCode(v);
}

std::vector<DeviceStatus> JFJochServices::GetDeviceStatus() const {
    std::vector<DeviceStatus> ret;
    if (receiver)
        ret = receiver->GetDeviceStatus();
    return ret;
}

ZMQPreviewSettings JFJochServices::GetPreviewSocketSettings() {
    if (receiver)
        return receiver->GetPreviewSocketSettings();
    return {};
}

ZMQMetadataSettings JFJochServices::GetMetadataSocketSettings() {
    if (receiver)
        return receiver->GetMetadataSocketSettings();
    return {};
}

void JFJochServices::SetPreviewSocketSettings(const ZMQPreviewSettings &input) {
    if (receiver)
        receiver->PreviewSocketSettings(input);
}

void JFJochServices::SetMetadataSocketSettings(const ZMQMetadataSettings &input) {
    if (receiver)
        receiver->MetadataSocketSettings(input);
}

void JFJochServices::GetStartMessageFromBuffer(std::vector<uint8_t> &v) {
    if (receiver)
        return receiver->GetStartMessageFromBuffer(v);
}

bool JFJochServices::GetImageFromBuffer(std::vector<uint8_t> &v, int64_t image_number) {
    if (receiver)
        return receiver->GetImageFromBuffer(v, image_number);
    return false;
}

ImageBufferStatus JFJochServices::GetImageBufferStatus() const {
    if (receiver)
        return receiver->GetImageBufferStatus();
    else return ImageBufferStatus{.total_slots = 0, .available_slots = 0};
}

void JFJochServices::ClearImageBuffer() const {
    if (receiver)
        receiver->ClearImageBuffer();
}

void JFJochServices::ClearResults() const {
    if (receiver)
        receiver->ClearResults();
}

void JFJochServices::LoadDetectorPixelMask(PixelMask &mask) {
    std::shared_lock ul(detector_mutex);
    if (detector)
        detector->LoadPixelMask(mask);
}

void JFJochServices::SetupIndexing(const IndexingSettings &input) {
    if (receiver)
        receiver->Indexing(input);
}

void JFJochServices::SetRugnuxLive(const RugnuxLiveSettings &input) {
    if (receiver)
        receiver->RugnuxLive(input.enable, input.model_path, input.reserved_cpus, input.deadline_s);
    rugnux = input;
}

RugnuxLiveSettings JFJochServices::GetRugnuxLive() const {
    return rugnux;
}

RugnuxLiveSnapshot JFJochServices::GetRugnuxStatus() const {
    if (receiver)
        return receiver->GetRugnuxStatus();
    return {};
}

void JFJochServices::SetRugnuxModel(const std::string &model) {
    RugnuxLiveSettings s = rugnux;
    if (!s.model_path.empty())
        std::filesystem::remove(s.model_path);
    s.model_path.clear();

    if (!model.empty()) {
        if (rugnux_model_dir.empty()) {
            char dir[] = "/tmp/jfjoch_rugnux_XXXXXX";   // mkdtemp creates it 0700
            if (mkdtemp(dir) == nullptr)
                throw JFJochException(JFJochExceptionCategory::FileWriteError,
                                      "Cannot create directory for the rugnux model");
            rugnux_model_dir = dir;
        }
        // gemmi detects PDB/mmCIF from the content, but gzip only from the extension
        const bool gzipped = model.size() >= 2 && model[0] == '\x1f' && model[1] == '\x8b';
        s.model_path = rugnux_model_dir + (gzipped ? "/model.gz" : "/model");
        {
            std::ofstream f(s.model_path, std::ios::binary);
            f.write(model.data(), static_cast<std::streamsize>(model.size()));
            if (!f)
                throw JFJochException(JFJochExceptionCategory::FileWriteError, "Cannot write the rugnux model");
        }
        std::filesystem::permissions(s.model_path,
                                     std::filesystem::perms::owner_read | std::filesystem::perms::owner_write);
    }
    SetRugnuxLive(s);
}

LiveCollection *JFJochServices::GetLiveCollection() {
    return receiver ? receiver->GetLiveCollection() : nullptr;
}

ImagePusherStatus JFJochServices::GetImagePusherStatus() const {
    if (receiver)
        return receiver->GetImagePusherStatus();
    return {
        .pusher_type = ImagePusherType::None,
        .address = {},
        .connected_writers = 0
    };
}