Cache parsed datasets instead of reopening per frame

src/dectris/neggia/plugin/H5ToXds.cpp

@@ -8,6 +8,7 @@
 #include <iostream>
 #include <limits>
 #include <memory>
+#include <mutex>
 #include <sstream>
 #include <stdexcept>
 #include <string>
@@ -38,6 +39,23 @@ struct H5DataCache {
 constexpr int NUM_WORKERS = 16;
 std::vector<std::unique_ptr<H5DataCache>> GLOBAL_POOL;
 
+// Parsed Datasets, indexed by dataset number - 1. Filled lazily on the
+// get_data path. EVERY access to REGISTRY -- read side and both clears
+// included -- is made with REGISTRY_MUTEX held. Entries are immutable once
+// published and are handed out as shared_ptr<const Dataset>, so a caller that
+// took a copy is unaffected by a later replacement, and a superseded entry's
+// mapping is released when the last in-flight reader drops it.
+// MAX_CACHED_DATASETS bounds both the live mapping count and the resize
+// allocation; indices past it are served uncached, exactly as before.
+constexpr size_t MAX_CACHED_DATASETS = 4096;
+std::mutex REGISTRY_MUTEX;
+std::vector<std::shared_ptr<const Dataset>> REGISTRY;
+
+void clearRegistry() {
+    std::lock_guard<std::mutex> lock(REGISTRY_MUTEX);
+    REGISTRY.clear();
+}
+
 void printVersionInfo() {
     std::cout << "This is neggia " << VERSION << " (Copyright Dectris 2020)"
               << std::endl;
@@ -385,13 +403,39 @@ void applyMaskAndTransformToInt32(const H5DataCache* dataCache,
     }
 }
 
+// Returns the Dataset holding globalFrameNumber, from REGISTRY when a usable
+// entry is cached. An entry is unusable once its file has grown past the extent
+// mapped when the entry was built: addresses parsed out of a grown file may
+// point outside that mapping. Construction runs OUTSIDE the lock, so no caller
+// blocks behind another caller's open(); two threads missing the same index may
+// both build one and the last publisher wins, which is benign because the
+// objects are value-identical.
+std::shared_ptr<const Dataset> acquireDataset(size_t globalFrameNumber,
+                                              const H5DataCache* dataCache) {
+    size_t index = globalFrameNumber / (size_t)dataCache->nframesPerDataset;
+    {
+        std::lock_guard<std::mutex> lock(REGISTRY_MUTEX);
+        if (index < REGISTRY.size() && REGISTRY[index] &&
+            !REGISTRY[index]->fileHasGrown())
+            return REGISTRY[index];
+    }
+    std::shared_ptr<const Dataset> built(new Dataset(
+            dataCache->h5File, getPathToDataset(globalFrameNumber, dataCache)));
+    std::lock_guard<std::mutex> lock(REGISTRY_MUTEX);
+    if (REGISTRY.size() <= index && index < MAX_CACHED_DATASETS)
+        REGISTRY.resize(index + 1);
+    if (index < REGISTRY.size())
+        REGISTRY[index] = built;
+    return built;
+}
+
 void readDataset(int* frame_number,
                  int data_array[],
                  const H5DataCache* dataCache) {
     size_t globalFrameNumber = correctFrameNumberOffset(*frame_number);
-    std::string pathToDataset = getPathToDataset(globalFrameNumber, dataCache);
     try {
-        Dataset dataset(dataCache->h5File, pathToDataset);
+        auto held = acquireDataset(globalFrameNumber, dataCache);
+        const Dataset& dataset = *held;
         size_t totNumberOfDatasets = dataset.dim()[0];
         size_t datasetFrameNumber =
                 getFrameNumberWithinDataset(globalFrameNumber, dataCache);
@@ -458,6 +502,9 @@ void plugin_get_header(int* nx,
     setInfoArray(info);
     try {
         H5DataCache* dataCache = getPreopenedDataCache();
+        // Drop anything a get_data-before-get_header call may have cached under
+        // an indeterminate masterFileOnly / nframesPerDataset.
+        clearRegistry();
         setXPixelSize(dataCache);
         setYPixelSize(dataCache);
         setPixelMask(dataCache);
@@ -520,7 +567,13 @@ void plugin_get_data(int* frame_number,
 }
 
 void plugin_close(int* error_flag) {
+    // Single-threaded by contract (every get_data thread joined first).
+    // Straggler shared_ptr copies keep their Datasets alive, so clearing the
+    // registry cannot dangle; the lock inside clearRegistry() keeps the
+    // container race-free even if the contract is violated.
+    clearRegistry();
     GLOBAL_POOL.clear();
+    *error_flag = 0;
 }
 
 }  // extern "C"

src/dectris/neggia/user/Dataset.cpp

@@ -71,6 +71,13 @@ std::vector<size_t> Dataset::chunkShape() const {
     return _dataLayoutMsg.chunkShape();
 }
 
+bool Dataset::fileHasGrown() const {
+    H5Superblock superblock(_h5File.fileAddress());
+    // End-of-file address: byte 40 in superblock v0, byte 28 in v2/v3.
+    size_t eofOffset = superblock.version() == 0 ? 40 : 28;
+    return superblock.read_u64(eofOffset) > _h5File.mapSize();
+}
+
 void Dataset::readRawData(ConstDataPointer rawData,
                           void* outData,
                           size_t outDataSize) const {
@@ -110,6 +117,11 @@ size_t Dataset::chunkDataSize() const {
 
 void Dataset::read(void* data, const std::vector<size_t>& chunkOffset) const {
     auto rawData = _dataLayoutMsg.getRawData(_dataSize, chunkOffset);
+    // The chunk address comes out of the file itself; refuse it when it points
+    // outside the mapping instead of dereferencing a wild pointer.
+    if (rawData.data < _h5File.fileAddress() ||
+        rawData.data >= _h5File.fileAddress() + _h5File.mapSize())
+        throw std::out_of_range("chunk address outside the mapped file");
     size_t s = chunkDataSize();
     switch (_filterId) {
         case -1:

src/dectris/neggia/user/Dataset.h

@@ -25,6 +25,12 @@ class Dataset {
     bool isChunked() const;
     std::vector<size_t> chunkShape() const;
 
+    // True when the file's own end-of-file address exceeds the extent mapped
+    // when this Dataset was built, i.e. the file has grown since. Addresses
+    // parsed out of it may then lie outside the mapping, so a Dataset that
+    // reports true must be rebuilt before it is read again.
+    bool fileHasGrown() const;
+
     // chunkOffset is ignored for contigous or raw datasets
     void read(void* data,
               const std::vector<size_t>& chunkOffset =

src/dectris/neggia/user/H5File.cpp

@@ -42,6 +42,11 @@ std::shared_ptr<char> mapFile(const std::string& fileName) {
 
 }  // namespace
 
+size_t H5File::mapSize() const {
+    const UnMap* deleter = std::get_deleter<UnMap>(_fileAddress);
+    return deleter ? deleter->size : 0;
+}
+
 H5File::H5File(const std::string& path) : _fileAddress(mapFile(path)) {
     for (ssize_t i = path.size() - 1; i > 0; i--) {
         if (path[i] == '/') {

src/dectris/neggia/user/H5File.h

@@ -12,6 +12,9 @@ class H5File {
     ~H5File();
     const char* fileAddress() const;
     std::string fileDir() const;
+    // Bytes mapped when this handle was created, i.e. the file's size at that
+    // moment. Data appended afterwards is NOT reachable through it.
+    size_t mapSize() const;
 
 private:
     std::shared_ptr<char> _fileAddress;