mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 81 Wiki Activity
Files
21f30549a740521139734c46c0e61bf761ca5d7c
Jungfraujoch/tests/LiveCollectionTest.cpp
T
leonarski_fandClaude Opus 5.5 21f30549a7 rugnux: follow a collection from the broker while it is written 
rugnux <path>/<prefix>_master.h5 http://<broker>:<port>/ - 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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
2026-10-08 15:52:22 +02:00

281 lines
11 KiB
C++
Raw Blame History

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include <catch2/catch_all.hpp>
#include <filesystem>
#include <future>
#include <thread>
#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<uint8_t> Image(const DiffractionExperiment &x, int i) {
std::vector<int16_t> image(x.GetPixelsNum());
for (size_t p = 0; p < image.size(); p++)
image[p] = static_cast<int16_t>((p * 13 + 17 * i + 3) % 20000);
JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_ZSTD);
return compressor.Compress(image);
}
DataMessage Message(const DiffractionExperiment &x, const std::vector<uint8_t> &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<uint8_t>{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<ServerSentEvent> 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<size_t>(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<std::pair<uint64_t, uint64_t>> 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<uint64_t, uint64_t>(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<std::unique_ptr<TCPImagePuller>> pullers;
std::vector<std::unique_ptr<StreamWriter>> writers;
std::vector<std::future<StreamWriterOutput>> futures;
for (int i = 0; i < 2; i++) {
pullers.push_back(std::make_unique<TCPImagePuller>(pusher.GetAddress()[0], 8 * 1024 * 1024));
writers.push_back(std::make_unique<StreamWriter>(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<uint8_t> 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<std::vector<uint8_t>> images;
for (int i = 0; i < 7; i++)
images.push_back(Image(x, i));
JFJochStartMessageReader reader;
std::vector<std::pair<uint64_t, uint64_t>> 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);
}
Reference in New Issue View Git Blame Copy Permalink