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); +}