From 21f30549a740521139734c46c0e61bf761ca5d7c Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 15:52:22 +0200 Subject: [PATCH] rugnux: follow a collection from the broker while it is written rugnux /_master.h5 http://:/ - when the master file is not on disk yet - processes the broker's current collection as it is written: - the start message comes from GET /live/start.cbor; the given path has to end in the collection's master file name, and the data files are read next to it; - JFJochStartMessageReader builds the dataset from the start message (detector, geometry, goniometer, mask) and reads the images straight from the data files; a read of an image whose file the broker has not reported waits for it; - BrokerFeed reads GET /live/events in a thread of its own and hands the reader each closed file. On end, the run finishes with the images the collection delivered (a read past them returns no image). A broken stream, another collection, or a refused follower makes every waiting read throw LiveCollectionLost, which ends the run like a fatal resource error; - the token is taken from JUNGFRAUJOCH_HTTP_TOKEN (the variable the viewer uses) and sent only in the Authorization header; without it, and with no master file on disk, rugnux stops at once; - rotation data and --mode mx only; plain http only (TLS is not linked into rugnux). With the master file on disk, the URL is ignored and the run is offline as before. The streamed run does not read the master file and is not compared with it: its result can differ from an offline run of the same files in the last bits of the metadata. Tests ([Live]): LiveCollection, the server-sent-event parser, the FileWriter callback, the TCP pusher's closed-file reports with two real writers, and JFJochStartMessageReader against the master the writer produced. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- reader/CMakeLists.txt | 2 + reader/JFJochStartMessageReader.cpp | 196 +++++++++++++++++++ reader/JFJochStartMessageReader.h | 56 ++++++ rugnux/BrokerFeed.cpp | 142 ++++++++++++++ rugnux/BrokerFeed.h | 54 ++++++ rugnux/CMakeLists.txt | 5 +- rugnux/Rugnux.cpp | 4 + rugnux/rugnux_cli.cpp | 48 ++++- tests/CMakeLists.txt | 1 + tests/LiveCollectionTest.cpp | 280 ++++++++++++++++++++++++++++ 10 files changed, 785 insertions(+), 3 deletions(-) create mode 100644 reader/JFJochStartMessageReader.cpp create mode 100644 reader/JFJochStartMessageReader.h create mode 100644 rugnux/BrokerFeed.cpp create mode 100644 rugnux/BrokerFeed.h create mode 100644 tests/LiveCollectionTest.cpp diff --git a/reader/CMakeLists.txt b/reader/CMakeLists.txt index b98c248d8..de234d23f 100644 --- a/reader/CMakeLists.txt +++ b/reader/CMakeLists.txt @@ -17,6 +17,8 @@ ADD_LIBRARY(JFJochReader STATIC JFJochSMVReader.cpp JFJochSMVReader.h HDF5ImageLocator.cpp + JFJochStartMessageReader.cpp + JFJochStartMessageReader.h HDF5ImageLocator.h HDF5ImageSource.cpp HDF5ImageSource.h diff --git a/reader/JFJochStartMessageReader.cpp b/reader/JFJochStartMessageReader.cpp new file mode 100644 index 000000000..eb8f768f4 --- /dev/null +++ b/reader/JFJochStartMessageReader.cpp @@ -0,0 +1,196 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "JFJochStartMessageReader.h" + +#include + +#include "spdlog/fmt/fmt.h" +#include "../common/JFJochException.h" + +namespace { + // The data file name the writer gives file `file_number` (HDF5Metadata::DataFileName), without + // the directory: the data files sit next to the master file. + std::string DataFileName(const StartMessage &start, uint64_t file_number) { + const std::string prefix = std::filesystem::path(start.file_prefix).filename().string(); + return fmt::format("{:s}_data_{:06d}.h5", prefix, file_number + 1); + } + + // The dataset as the start message describes it - the same quantities an HDF5 master states, taken + // from the message the writer wrote it from. + std::shared_ptr Dataset(const StartMessage &start, const DiffractionExperiment &defaults) { + auto dataset = std::make_shared(); + dataset->arm_date = start.arm_date; + dataset->jfjoch_release = start.jfjoch_release; + dataset->error_value = start.error_value; + dataset->file_detect_ice_rings = start.detect_ice_rings; + + auto &x = dataset->experiment; + x = defaults; + // The reader hands every image out as signed 32-bit, as JFJochHDF5Reader does. + x.BitDepthImage(32); + x.PixelSigned(true); + x.FilePrefix(std::filesystem::path(start.file_prefix).filename().string()); + x.BeamX_pxl(start.beam_center_x); + x.BeamY_pxl(start.beam_center_y); + x.DetectorDistance_mm(start.detector_distance * 1000.0f); + x.PoniRot1_rad(start.poni_rot1.value_or(0.0f)); + x.PoniRot2_rad(start.poni_rot2.value_or(0.0f)); + x.PoniRot3_rad(start.poni_rot3.value_or(0.0f)); + x.IncidentEnergy_keV(WVL_1A_IN_KEV / start.incident_wavelength); + if (start.incident_wavelength_spread) + x.BandwidthFWHM(start.incident_wavelength_spread.value() / start.incident_wavelength); + x.DetectIceRings(start.detect_ice_rings.value_or(false)); + + DetectorSetup detector = DetDECTRIS(start.image_size_x, start.image_size_y, start.detector_description, {}); + detector.PixelSize_um(start.pixel_size_x * 1e6f); + detector.MirrorY(start.mirror_y); + detector.ImageOrientation(DetectorOrientation(start.detector_orientation_mirror_y, + start.detector_orientation_quarter_turns)); + detector.SensorThickness_um(start.sensor_thickness * 1e6f); + detector.SensorMaterial(start.sensor_material); + detector.SaturationLimit(SaturationLimitFromValue(start.saturation_value)); + detector.BitDepthImage(32); + detector.MinFrameTime(std::chrono::microseconds(0)); + detector.MinCountTime(std::chrono::microseconds(0)); + detector.ReadOutTime(std::chrono::nanoseconds(0)); + x.Detector(detector); + + x.FrameTime(std::chrono::duration_cast(std::chrono::duration(start.frame_time)), + std::chrono::duration_cast(std::chrono::duration(start.count_time))); + + x.Goniometer(start.goniometer); + x.GridScan(start.grid_scan); + x.Smargon(start.smargon_position); + x.NumTriggers(1); + x.ImagesPerTrigger(static_cast(start.number_of_images)); + + InstrumentMetadata instrument; + instrument.SourceName(start.source_name); + instrument.SourceType(start.source_type); + instrument.InstrumentName(start.instrument_name); + x.ImportInstrumentMetadata(instrument); + x.SampleName(start.sample_name); + x.SampleTemperature_K(start.sample_temperature_K); + x.RingCurrent_mA(start.ring_current_mA); + x.TotalFlux(start.total_flux); + x.AttenuatorTransmission(start.attenuator_transmission); + if (start.beam_size_x && start.beam_size_y) { + x.BeamSizeX_um(start.beam_size_x.value() * 1e6f); + x.BeamSizeY_um(start.beam_size_y.value() * 1e6f); + } + + if (!start.pixel_mask.empty()) + dataset->pixel_mask = std::make_shared(start.pixel_mask.begin()->second); + else + dataset->pixel_mask = std::make_shared( + std::vector(start.image_size_x * start.image_size_y)); + return dataset; + } +} + +void JFJochStartMessageReader::Open(const StartMessage &start, const std::string &master_path) { + if (start.images_per_file <= 0) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "The start message gives no images per file"); + + number_of_images_ = start.number_of_images; + images_per_file_ = static_cast(start.images_per_file); + const uint64_t files = (number_of_images_ + images_per_file_ - 1) / images_per_file_; + + const std::filesystem::path directory = std::filesystem::path(master_path).parent_path(); + HDF5ImageLocator::Layout layout; + layout.format = FileWriterFormat::NXmxLegacy; + layout.images_per_file = images_per_file_; + for (uint64_t f = 0; f < files; f++) + layout.legacy_files.push_back({(directory / DataFileName(start, f)).string(), "/entry/data/data"}); + { + std::unique_lock ul(hdf5_mutex); + image_source_.Configure(std::move(layout)); + } + { + std::unique_lock ul(m_); + file_images_.assign(files, -1); + ended_ = false; + lost_.clear(); + } + SetStartMessage(Dataset(start, default_experiment)); +} + +void JFJochStartMessageReader::FileClosed(uint64_t file_number, uint64_t image_count) { + { + std::unique_lock ul(m_); + if (file_number < file_images_.size()) + file_images_[file_number] = static_cast(image_count); + } + cv_.notify_all(); +} + +void JFJochStartMessageReader::Ended() { + { + std::unique_lock ul(m_); + ended_ = true; + } + cv_.notify_all(); +} + +void JFJochStartMessageReader::Lost(const std::string &why) { + { + std::unique_lock ul(m_); + if (!ended_) + lost_ = why; + } + cv_.notify_all(); +} + +uint64_t JFJochStartMessageReader::GetNumberOfImages() const { + return number_of_images_; +} + +void JFJochStartMessageReader::Close() { + std::unique_lock ul(hdf5_mutex); + image_source_.Clear(); + SetStartMessage({}); +} + +bool JFJochStartMessageReader::WaitForImage(int64_t image_number) { + if (image_number < 0 || static_cast(image_number) >= number_of_images_) + throw JFJochException(JFJochExceptionCategory::HDF5, "Image out of bounds"); + const uint64_t f = static_cast(image_number) / images_per_file_; + std::unique_lock ul(m_); + cv_.wait(ul, [&] { return file_images_[f] >= 0 || ended_ || !lost_.empty(); }); + if (file_images_[f] < 0) { + if (!lost_.empty()) + throw LiveCollectionLost(lost_); + return false; // the collection ended without this file + } + return static_cast(static_cast(image_number) % images_per_file_) < file_images_[f]; +} + +bool JFJochStartMessageReader::ReadRawImage(int64_t image_number, JFJochReaderRawImage &ret) { + if (!WaitForImage(image_number)) + return false; + // As JFJochHDF5Reader: only the chunk lookup under the HDF5 lock, the read after it. + std::optional chunk; + { + std::unique_lock ul(hdf5_mutex); + const auto loc = image_source_.Resolve(image_number); + chunk = image_source_.PrepareDirectRead(loc); + if (!chunk) { + ret.image = image_source_.ReadImageAt(ret.image_buffer, loc); + return true; + } + } + ret.image = HDF5ImageSource::ReadDirect(ret.image_buffer, *chunk); + return true; +} + +bool JFJochStartMessageReader::LoadImage_i(std::shared_ptr &dataset, DataMessage &message, + std::vector &buffer, int64_t image_number, bool) { + if (!dataset || !WaitForImage(image_number)) + return false; + std::unique_lock ul(hdf5_mutex); + message.image = image_source_.ReadImageAt(buffer, image_source_.Resolve(image_number)); + message.number = image_number; + return true; +} diff --git a/reader/JFJochStartMessageReader.h b/reader/JFJochStartMessageReader.h new file mode 100644 index 000000000..75d943f69 --- /dev/null +++ b/reader/JFJochStartMessageReader.h @@ -0,0 +1,56 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include +#include +#include + +#include "JFJochReader.h" +#include "HDF5ImageSource.h" + +// The broker that reported a collection's data files can no longer be heard from, and an image this +// reader was asked for has not been reported: it will never be read. +class LiveCollectionLost : public std::runtime_error { +public: + explicit LiveCollectionLost(const std::string &what) : std::runtime_error(what) {} +}; + +// Reads a collection while it is being written: the dataset comes from its start message, as the +// broker sends it, and the images straight from the data files - /entry/data/data in each - once the +// broker reports a file closed and renamed into place (FileClosed). A read of an image whose file has +// not been reported waits for it; after Ended it returns false for an image no file holds, and after +// Lost it throws. The master file is never read. +class JFJochStartMessageReader : public JFJochReader { +public: + // master_path: where the collection's master file will be; the data files are next to it. + void Open(const StartMessage &start, const std::string &master_path); + + void FileClosed(uint64_t file_number, uint64_t image_count); + void Ended(); + void Lost(const std::string &why); + + [[nodiscard]] uint64_t GetNumberOfImages() const override; + void Close() override; + bool ReadRawImage(int64_t image_number, JFJochReaderRawImage &image) override; + std::vector ReadSpots(int64_t image) const override { return {}; } + +private: + bool LoadImage_i(std::shared_ptr &dataset, DataMessage &message, + std::vector &buffer, int64_t image_number, bool update_dataset) override; + // Waits until the file of this image is reported; false when the collection ended without it. + bool WaitForImage(int64_t image_number); + + HDF5ImageSource image_source_; + uint64_t number_of_images_ = 0; + uint64_t images_per_file_ = 1; + + std::mutex m_; + std::condition_variable cv_; + std::vector file_images_; // per data file: images it holds, -1 while not reported + bool ended_ = false; + std::string lost_; +}; diff --git a/rugnux/BrokerFeed.cpp b/rugnux/BrokerFeed.cpp new file mode 100644 index 000000000..c5b02dec3 --- /dev/null +++ b/rugnux/BrokerFeed.cpp @@ -0,0 +1,142 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "BrokerFeed.h" + +#include +#include + +#include "../common/JFJochException.h" +#include "../common/Logger.h" +#include "../frame_serialize/CBORStream2Deserializer.h" + +std::vector ServerSentEventParser::Feed(const char *bytes, size_t size) { + pending.append(bytes, size); + std::vector events; + size_t end; + while ((end = pending.find("\n\n")) != std::string::npos) { + const std::string block = pending.substr(0, end); + pending.erase(0, end + 2); + ServerSentEvent e; + size_t pos = 0; + while (pos <= block.size()) { + size_t eol = block.find('\n', pos); + if (eol == std::string::npos) + eol = block.size(); + const std::string line = block.substr(pos, eol - pos); + if (line.starts_with("event: ")) + e.event = line.substr(7); + else if (line.starts_with("data: ")) + e.data = line.substr(6); + pos = eol + 1; + } + if (!e.event.empty()) + events.push_back(std::move(e)); + } + return events; +} + +namespace { + httplib::Headers Authorization(const std::string &token) { + return {{"Authorization", "Bearer " + token}}; + } + + std::string Refused(const httplib::Result &res, const std::string &what) { + if (!res) + return what + ": no answer from the broker (" + httplib::to_string(res.error()) + ")"; + if (res->status == 401) + return what + ": the broker refused the token in JUNGFRAUJOCH_HTTP_TOKEN"; + return what + ": the broker answered " + std::to_string(res->status) + " " + res->body; + } +} + +BrokerFeed::BrokerFeed(const std::string &url, std::string in_token) : token(std::move(in_token)) { + if (url.starts_with("https://")) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "Following a collection over https is not supported by this build - use http://"); + if (!url.starts_with("http://")) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "The broker address must start with http://"); + const size_t path = url.find('/', 7); + host = url.substr(0, path); + if (host.find('@') != std::string::npos) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "The broker address must not carry credentials"); + if (path != std::string::npos) + prefix = url.substr(path); + while (prefix.ends_with('/')) + prefix.pop_back(); +} + +BrokerFeed::~BrokerFeed() { + stop = true; + if (thread.joinable()) + thread.join(); +} + +StartMessage BrokerFeed::FetchStart() { + httplib::Client client(host); + auto res = client.Get(prefix + "/live/start.cbor", Authorization(token)); + if (!res || res->status != 200) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + Refused(res, "No collection to follow at " + host)); + auto msg = CBORStream2Deserialize(res->body); + if (!msg || !msg->start_message) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "The broker at " + host + " sent no start message"); + return msg->start_message.value(); +} + +void BrokerFeed::Follow(JFJochStartMessageReader &reader, uint64_t run_number, const std::string &file_prefix) { + thread = std::thread([this, &reader, run_number, file_prefix] { + Logger logger("BrokerFeed"); + httplib::Client client(host); + // The broker writes at least a keepalive every 10 s; three missed ones mean it is gone. + client.set_read_timeout(std::chrono::seconds(30)); + ServerSentEventParser parser; + bool ended = false; + std::string lost; + size_t files = 0; + int status = 0; + auto res = client.Get(prefix + "/live/events", Authorization(token), + [&](const httplib::Response &response) { + status = response.status; + return status == 200; + }, + [&](const char *data, size_t size) { + for (const auto &e: parser.Feed(data, size)) { + const auto j = nlohmann::json::parse(e.data); + if (e.event == "collection") { + if (j.at("run_number").get() != run_number + || j.at("file_prefix").get() != file_prefix) + lost = "the broker is now on another collection (" + j.at("file_prefix").get() + ")"; + } else if (e.event == "file") { + reader.FileClosed(j.at("file_number").get(), j.at("image_count").get()); + files++; + } else if (e.event == "end") { + if (j.contains("error")) + logger.Warning("The collection ended with an error: {}", j.at("error").get()); + logger.Info("The collection ended with {} images in {} data files", + j.at("images").get(), files); + ended = true; + } else if (e.event == "superseded") { + lost = "a new collection began before this one ended"; + } + if (!lost.empty()) + return false; + } + return !stop; + }); + if (ended) + reader.Ended(); + else if (!lost.empty()) + reader.Lost(lost); + else if (status == 401) + reader.Lost("Following the collection at " + host + ": the broker refused the token in JUNGFRAUJOCH_HTTP_TOKEN"); + else if (status != 200 && status != 0) + reader.Lost("Following the collection at " + host + ": the broker answered " + std::to_string(status) + + (status == 409 ? " (another reader is following it)" : "")); + else if (!stop) + reader.Lost("Lost the broker at " + host + " before the collection ended"); + }); +} diff --git a/rugnux/BrokerFeed.h b/rugnux/BrokerFeed.h new file mode 100644 index 000000000..0e21b859e --- /dev/null +++ b/rugnux/BrokerFeed.h @@ -0,0 +1,54 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include +#include + +#include "../common/JFJochMessages.h" +#include "../reader/JFJochStartMessageReader.h" + +// One server-sent event: its type and its data line. +struct ServerSentEvent { + std::string event; + std::string data; +}; + +// Splits a server-sent event stream into events as it arrives. Feed it the bytes in any pieces; +// comments (keepalives) are dropped. +class ServerSentEventParser { + std::string pending; +public: + std::vector Feed(const char *bytes, size_t size); +}; + +// rugnux's side of following a collection while it is written: the start message from the broker +// (GET /live/start.cbor), then its event stream (GET /live/events) handed to a JFJochStartMessageReader. +// The token goes in the Authorization header and nowhere else. +class BrokerFeed { +public: + // url: http://host:port[/prefix][/] + BrokerFeed(const std::string &url, std::string token); + ~BrokerFeed(); + + // Throws with the broker's answer when there is no collection to follow. + StartMessage FetchStart(); + + // Reads the event stream in a thread of its own until the collection ends, reporting every closed + // data file to `reader`. The reader is told the broker is lost when the stream breaks before the + // end, or when it describes another collection than run_number / file_prefix. + void Follow(JFJochStartMessageReader &reader, uint64_t run_number, const std::string &file_prefix); + + // scheme://host:port, for messages - never the token. + [[nodiscard]] const std::string &Host() const { return host; } + +private: + std::string host; // http://host:port + std::string prefix; // path before /live, without a trailing slash + std::string token; + std::atomic stop{false}; + std::thread thread; +}; diff --git a/rugnux/CMakeLists.txt b/rugnux/CMakeLists.txt index 3efa2a0cd..6b437d83b 100644 --- a/rugnux/CMakeLists.txt +++ b/rugnux/CMakeLists.txt @@ -21,9 +21,12 @@ ADD_LIBRARY(JFJochRugnux STATIC ReportDocument.h DiagnosticOutput.cpp DiagnosticOutput.h + BrokerFeed.cpp + BrokerFeed.h ) -TARGET_LINK_LIBRARIES(JFJochRugnux JFJochReader JFJochImageAnalysis JFJochWriter JFJochModelWriter JFJochPreview gemmi fftw3f) +TARGET_LINK_LIBRARIES(JFJochRugnux JFJochReader JFJochImageAnalysis JFJochWriter JFJochModelWriter JFJochPreview gemmi fftw3f + httplib::httplib) # rugnux is the single offline analysis CLI; --mode picks what it does: mx (the full pipeline, the # default), azint, scale (re-scale/merge stored reflections) or calibration (powder-ring geometry). diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 919861e05..e04cb0658 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -31,6 +31,7 @@ #include #include "../reader/JFJochHDF5Reader.h" +#include "../reader/JFJochStartMessageReader.h" #include "../common/JFJochMath.h" #include "../common/ParallelFor.h" #include "../common/Logger.h" @@ -149,6 +150,9 @@ namespace rugnux_internal { // Not a resource, but just as fatal: a memo verification (verify_first_pass_memo) that found a // difference must end the run, not pass for a probe that scored nothing. if (dynamic_cast(&e)) return true; + // Nor a collection followed from a broker that is gone: the frames it had not reported will never + // be read, and a run that skipped them would look complete. + if (dynamic_cast(&e)) return true; const auto *jf = dynamic_cast(&e); return jf != nullptr && (jf->Category() == JFJochExceptionCategory::GPUCUDAError || jf->Category() == JFJochExceptionCategory::MemAllocFailed); diff --git a/rugnux/rugnux_cli.cpp b/rugnux/rugnux_cli.cpp index 90dab4842..e9dc2ee2c 100644 --- a/rugnux/rugnux_cli.cpp +++ b/rugnux/rugnux_cli.cpp @@ -27,6 +27,9 @@ #include "../reader/JFJochCBFReader.h" #include "../reader/JFJochMarCCDReader.h" #include "../reader/JFJochSMVReader.h" +#include "../reader/JFJochStartMessageReader.h" +#include "BrokerFeed.h" +#include "../writer/HDF5NXmx.h" #include #include #include @@ -93,7 +96,13 @@ enum class RugnuxMode { MX, Azint, Scale, Calibration }; constexpr int CALIBRATION_AZIM_BINS_DEFAULT = 32; void print_usage() { - std::cout << "Usage rugnux {} " << std::endl; + std::cout << "Usage rugnux {} []" << std::endl; + std::cout << " With a broker URL (http://host:port/) and no on disk yet, follows the broker's current " + "collection while it is written: must be the path its master file will have, the data " + "files are read next to it as the broker reports them closed, and the run ends with the images the " + "collection delivered. Rotation data, --mode mx only. The collection's token is read from the " + "environment variable JUNGFRAUJOCH_HTTP_TOKEN and sent only to the broker. With on disk, " + "the URL is ignored" << std::endl; std::cout << "Options:" << std::endl; std::cout << " -o, --output-prefix Output file prefix (default: output). mx and scale runs always write " "_report.txt, the results report; every run writes _log.txt, the full timestamped log " @@ -1504,13 +1513,14 @@ static int RunRugnux(int argc, char **argv) { } } - if (optind != argc - 1) { + if (optind != argc - 1 && optind != argc - 2) { logger.Error("Input file not specified"); print_usage(); exit(EXIT_FAILURE); } input_file = argv[optind]; + const std::string broker_url = (optind == argc - 2) ? argv[optind + 1] : ""; // -N defaults to 0 = "use all hardware threads"; resolve it to a concrete count here so every mode // behaves the same. The scale/merge engines expand 0 on their own, but the per-image processing @@ -1566,6 +1576,8 @@ static int RunRugnux(int argc, char **argv) { JFJochCBFReader cbf_reader; JFJochMarCCDReader marccd_reader; JFJochSMVReader smv_reader; + JFJochStartMessageReader live_reader; + std::unique_ptr broker_feed; JFJochReader *reader_ptr = nullptr; // Which format this is. Asked INSIDE the try: these open and read the file - a directory whose // first plausible member cannot be read throws here - and an input the program cannot read has to @@ -1597,6 +1609,38 @@ static int RunRugnux(int argc, char **argv) { logger.Error("{} holds no CBF, marCCD or SMV sweep this program can read (a directory is " "read as the sweep of frames it holds)", input_file); exit(EXIT_FAILURE); + } else if (!broker_url.empty() && !std::filesystem::exists(input_file)) { + // The collection is still being written: follow it from the broker. + if (mode != RugnuxMode::MX) { + logger.Error("Following a collection from the broker is for --mode mx only"); + exit(EXIT_FAILURE); + } + const char *token = std::getenv("JUNGFRAUJOCH_HTTP_TOKEN"); + if (token == nullptr) { + logger.Error("{} does not exist yet, and following it from the broker needs the collection's token " + "in the environment variable JUNGFRAUJOCH_HTTP_TOKEN", input_file); + exit(EXIT_FAILURE); + } + broker_feed = std::make_unique(broker_url, token); + const StartMessage start = broker_feed->FetchStart(); + // The start message names the files relative to the collection's root; the path given here + // has to be where its master file will be. + const std::string master = std::filesystem::path(HDF5Metadata::MasterFileName(start)).generic_string(); + const std::string given = std::filesystem::path(input_file).lexically_normal().generic_string(); + if (given != master && !given.ends_with("/" + master)) { + logger.Error("The broker at {} is writing {}, not {}", broker_feed->Host(), master, input_file); + exit(EXIT_FAILURE); + } + if (!start.goniometer.has_value() || !start.goniometer->IsScanning()) { + logger.Error("Following a collection from the broker is for rotation data, and {} is not", master); + exit(EXIT_FAILURE); + } + live_reader.Open(start, input_file); + broker_feed->Follow(live_reader, start.run_number, start.file_prefix); + reader_ptr = &live_reader; + logger.Info("Following collection {} (run {}) from the broker at {}: {} images in files of {}", + master, start.run_number, broker_feed->Host(), start.number_of_images, start.images_per_file); + console->OnNote("still being written - following it from the broker"); } else { hdf5_reader.ReadFile(input_file); reader_ptr = &hdf5_reader; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 78249a265..2283f4395 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -92,6 +92,7 @@ ADD_EXECUTABLE(jfjoch_test ModuleSummationTest.cpp ZMQMetadataSocketTest.cpp JFJochReaderTest.cpp + LiveCollectionTest.cpp SweepLayoutTest.cpp ShadowFinderTest.cpp HotPixelFinderTest.cpp diff --git a/tests/LiveCollectionTest.cpp b/tests/LiveCollectionTest.cpp new file mode 100644 index 000000000..47f4210c4 --- /dev/null +++ b/tests/LiveCollectionTest.cpp @@ -0,0 +1,280 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#include +#include +#include + +#include "../common/DiffractionExperiment.h" +#include "../compression/JFJochCompressor.h" +#include "../image_puller/TCPImagePuller.h" +#include "../image_pusher/LiveCollection.h" +#include "../image_pusher/TCPStreamPusher.h" +#include "../reader/JFJochHDF5Reader.h" +#include "../reader/JFJochStartMessageReader.h" +#include "../rugnux/BrokerFeed.h" +#include "../writer/FileWriter.h" +#include "../writer/StreamWriter.h" + +// Following a collection while it is written: the writer reports each closed data file, the broker +// tracks them (LiveCollection), and rugnux reads the images from the start message and the files. + +namespace { + const std::string live_dir = "live_collection_test"; + + DiffractionExperiment LiveExperiment(const std::string &prefix) { + DiffractionExperiment x(DetJF(1)); + x.FilePrefix(prefix).ImagesPerTrigger(7).OverwriteExistingFiles(true); + x.BitDepthImage(16).ImagesPerFile(2).SetFileWriterFormat(FileWriterFormat::NXmxVDS).PixelSigned(true); + x.Compression(CompressionAlgorithm::BSHUF_ZSTD); + x.BeamX_pxl(512.5).BeamY_pxl(256.25).DetectorDistance_mm(120.0).IncidentEnergy_keV(12.4); + x.Goniometer(GoniometerAxis("omega", 10.0f, 0.25f, Coord(-1, 0, 0), {})); + return x; + } + + std::vector Image(const DiffractionExperiment &x, int i) { + std::vector image(x.GetPixelsNum()); + for (size_t p = 0; p < image.size(); p++) + image[p] = static_cast((p * 13 + 17 * i + 3) % 20000); + JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_ZSTD); + return compressor.Compress(image); + } + + DataMessage Message(const DiffractionExperiment &x, const std::vector &compressed, int i) { + DataMessage message{}; + message.image = CompressedImage(compressed, x.GetXPixelsNum(), x.GetYPixelsNum(), + CompressedImageMode::Int16, CompressionAlgorithm::BSHUF_ZSTD); + message.number = i; + return message; + } +} + +TEST_CASE("LiveCollection_Lifecycle", "[Live]") { + LiveCollection live; + CHECK(live.Get().generation == 0); + CHECK(!live.StartCBOR()); + + StartMessage start{}; + start.file_prefix = "sub/run"; + start.images_per_file = 100; + start.number_of_images = 250; + start.run_number = 7; + live.Begin(start, {1, 2, 3}); + CHECK(live.StartCBOR() == std::vector{1, 2, 3}); + + auto s = live.Wait(1, 0, std::chrono::milliseconds(0)); + CHECK(s.generation == 1); + CHECK(s.run_number == 7); + CHECK(s.files.empty()); + + // A waiter wakes when a file is reported. + auto waiter = std::async(std::launch::async, [&] { return live.Wait(1, 0, std::chrono::seconds(10)); }); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + live.FileClosed(1, 100); + s = waiter.get(); + REQUIRE(s.files.size() == 1); + CHECK(s.files[0].file == "sub/run_data_000002.h5"); + CHECK(s.files[0].first_image == 100); + CHECK(s.files[0].image_count == 100); + + live.FileClosed(2, 50); + live.End(250); + s = live.Get(); + CHECK(s.files.size() == 2); + CHECK(s.ended); + CHECK(s.end_images == 250); + live.FileClosed(0, 100); // nothing after the end + CHECK(live.Get().files.size() == 2); + + // One follower at a time. + CHECK(live.AcquireFollower()); + CHECK_FALSE(live.AcquireFollower()); + live.ReleaseFollower(); + CHECK(live.AcquireFollower()); + + // A new collection starts from nothing. + live.Begin(start, {}); + s = live.Get(); + CHECK(s.generation == 2); + CHECK(s.files.empty()); + CHECK_FALSE(s.ended); +} + +TEST_CASE("ServerSentEventParser", "[Live]") { + ServerSentEventParser parser; + const std::string stream = ": keepalive\n\nevent: file\ndata: {\"file_number\":3}\n\nevent: end\ndata: {}\n\n"; + std::vector events; + // Fed in pieces, as the network delivers it. + for (size_t i = 0; i < stream.size(); i += 5) { + auto e = parser.Feed(stream.data() + i, std::min(5, stream.size() - i)); + events.insert(events.end(), e.begin(), e.end()); + } + REQUIRE(events.size() == 2); + CHECK(events[0].event == "file"); + CHECK(events[0].data == "{\"file_number\":3}"); + CHECK(events[1].event == "end"); +} + +TEST_CASE("FileWriter_FileClosedCallback", "[Live]") { + RegisterHDF5Filter(); + std::filesystem::remove_all(live_dir); + const auto x = LiveExperiment(live_dir + "/cb"); + StartMessage start; + x.FillMessage(start); + std::vector> closed; // file number (from 1), images + { + FileWriter writer(start); + writer.FileClosedCallback([&](const HDF5DataFileStatistics &s) { + // Reported once the file is in place under its own name. + CHECK(std::filesystem::exists(s.filename)); + closed.emplace_back(s.file_number, s.max_image_number + 1); + }); + for (int i = 0; i < 7; i++) + writer.Write(Message(x, Image(x, i), i)); + CHECK(closed.size() == 3); // the full files close as they fill + EndMessage end; + end.max_image_number = 7; + writer.WriteHDF5(end); + writer.Finalize(); + } + REQUIRE(closed.size() == 4); + CHECK(closed[3] == std::make_pair(4, 1)); // the partial last file, at END + std::filesystem::remove_all(live_dir); +} + +TEST_CASE("TCPStreamPusher_ReportsClosedFiles", "[Live][TCP]") { + RegisterHDF5Filter(); + Logger logger("TCPStreamPusher_ReportsClosedFiles"); + std::filesystem::remove_all(live_dir); + std::filesystem::create_directories(live_dir); + const auto x = LiveExperiment(live_dir + "/tcp"); + + TCPStreamPusher pusher("tcp://127.0.0.1:*", 2); + std::vector> pullers; + std::vector> writers; + std::vector> futures; + for (int i = 0; i < 2; i++) { + pullers.push_back(std::make_unique(pusher.GetAddress()[0], 8 * 1024 * 1024)); + writers.push_back(std::make_unique(logger, *pullers.back())); + futures.push_back(std::async(std::launch::async, [w = writers.back().get()] { return w->Run(); })); + } + for (int i = 0; i < 100 && pusher.GetConnectedWriters() < 2; i++) + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + REQUIRE(pusher.GetConnectedWriters() == 2); + + StartMessage start; + x.FillMessage(start); + start.write_master_file = true; + pusher.StartDataCollection(start); + LiveCollection *live = pusher.GetLiveCollection(); + REQUIRE(live != nullptr); + REQUIRE(live->StartCBOR()); + + std::vector buffer(16 * 1024 * 1024); + CBORStream2Serializer serializer(buffer.data(), buffer.size()); + for (int i = 0; i < 7; i++) { + const auto compressed = Image(x, i); + serializer.SerializeImage(Message(x, compressed, i)); + REQUIRE(pusher.SendImage(buffer.data(), serializer.GetBufferSize(), i)); + } + EndMessage end; + end.max_image_number = 7; + REQUIRE(pusher.EndDataCollection(end)); + + const auto s = live->Get(); + CHECK(s.ended); + CHECK(s.end_images == 7); + REQUIRE(s.files.size() == 4); + uint64_t images = 0; + for (const auto &f: s.files) { + CHECK(std::filesystem::exists(std::filesystem::path(live_dir) / std::filesystem::path(f.file).filename())); + CHECK(f.first_image == f.file_number * 2); + images += f.image_count; + } + CHECK(images == 7); + CHECK(std::filesystem::exists(live_dir + "/tcp_master.h5")); + + for (auto &w: writers) + w->Cancel(); + for (auto &f: futures) + REQUIRE_NOTHROW(f.get()); + std::filesystem::remove_all(live_dir); +} + +TEST_CASE("JFJochStartMessageReader", "[Live]") { + RegisterHDF5Filter(); + std::filesystem::remove_all(live_dir); + const auto x = LiveExperiment(live_dir + "/sm"); + StartMessage start; + x.FillMessage(start); + std::vector> images; + for (int i = 0; i < 7; i++) + images.push_back(Image(x, i)); + + JFJochStartMessageReader reader; + std::vector> closed; + { + FileWriter writer(start); + writer.FileClosedCallback([&](const HDF5DataFileStatistics &s) { + closed.emplace_back(s.file_number - 1, s.max_image_number + 1); + }); + reader.Open(start, live_dir + "/sm_master.h5"); + CHECK(reader.GetNumberOfImages() == 7); + + // A read waits for its file. + auto read = std::async(std::launch::async, [&] { return reader.GetRawImage(2); }); + for (int i = 0; i < 4; i++) + writer.Write(Message(x, images[i], i)); + CHECK(read.wait_for(std::chrono::milliseconds(200)) == std::future_status::timeout); + for (const auto &[f, n]: closed) + reader.FileClosed(f, n); + auto raw = read.get(); + REQUIRE(raw->image_buffer.size() == images[2].size()); + CHECK(memcmp(raw->image_buffer.data(), images[2].data(), images[2].size()) == 0); + + // Stopped after 5 images: the rest is never read. + writer.Write(Message(x, images[4], 4)); + EndMessage end; + end.max_image_number = 5; + writer.WriteHDF5(end); + writer.Finalize(); + } + for (const auto &[f, n]: closed) + reader.FileClosed(f, n); + reader.Ended(); + JFJochReaderRawImage raw; + CHECK(reader.ReadRawImage(4, raw)); + CHECK_FALSE(reader.ReadRawImage(5, raw)); + CHECK_FALSE(reader.ReadRawImage(6, raw)); + + // The dataset is the one the master file describes. + JFJochHDF5Reader offline; + offline.ReadFile(live_dir + "/sm_master.h5"); + const auto &a = reader.GetDataset()->experiment; + const auto &b = offline.GetDataset()->experiment; + CHECK(a.GetBeamX_pxl() == b.GetBeamX_pxl()); + CHECK(a.GetBeamY_pxl() == b.GetBeamY_pxl()); + CHECK(a.GetDetectorDistance_mm() == Catch::Approx(b.GetDetectorDistance_mm())); + CHECK(a.GetWavelength_A() == Catch::Approx(b.GetWavelength_A())); + CHECK(a.GetPixelSize_mm() == Catch::Approx(b.GetPixelSize_mm())); + CHECK(a.GetXPixelsNum() == b.GetXPixelsNum()); + CHECK(a.GetYPixelsNum() == b.GetYPixelsNum()); + REQUIRE(a.GetGoniometer()); + REQUIRE(b.GetGoniometer()); + CHECK(a.GetGoniometer()->GetStart_deg() == Catch::Approx(b.GetGoniometer()->GetStart_deg())); + CHECK(a.GetGoniometer()->GetIncrement_deg() == Catch::Approx(b.GetGoniometer()->GetIncrement_deg())); + CHECK(reader.GetDataset()->pixel_mask->GetMask() == offline.GetDataset()->pixel_mask->GetMask()); + offline.Close(); + + // A broker that is gone fails a read that would wait. + JFJochStartMessageReader lost; + lost.Open(start, live_dir + "/sm_master.h5"); + lost.Lost("gone"); + CHECK_THROWS_AS(lost.ReadRawImage(0, raw), LiveCollectionLost); + + reader.Close(); + std::filesystem::remove_all(live_dir); + REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0); +}