Skip invalid timeframes during ROOT input

Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx

@@ -19,13 +19,16 @@
 #include "Framework/DataProcessingStats.h"
 #include "Framework/RootArrowFilesystem.h"
 #include "Framework/AlgorithmSpec.h"
+#include "Framework/ArrowContext.h"
 #include "Framework/ConfigParamRegistry.h"
 #include "Framework/ControlService.h"
 #include "Framework/CallbackService.h"
 #include "Framework/EndOfStreamContext.h"
 #include "Framework/DeviceSpec.h"
 #include "Framework/RawDeviceService.h"
 #include "Framework/DataSpecUtils.h"
+#include "Framework/MessageContext.h"
+#include "Framework/StringContext.h"
 #include "Framework/ConfigContext.h"
 #include "DataInputDirector.h"
 #include "Framework/SourceInfoHeader.h"
@@ -200,7 +203,7 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
                            numTF,
                            watchdog,
                            maxRate,
-                           didir, reportTFN, reportTFFileName, level](Monitoring& monitoring, DataAllocator& outputs, ControlService& control, DeviceSpec const& device, DataProcessingStats& dpstats) {
+                           didir, reportTFN, reportTFFileName, level](Monitoring& monitoring, DataAllocator& outputs, ControlService& control, DeviceSpec const& device, DataProcessingStats& dpstats, ArrowContext& arrowContext, MessageContext& messageContext, StringContext& stringContext) {
       // Each parallel reader device.inputTimesliceId reads the files fileCounter*device.maxInputTimeslices+device.inputTimesliceId
       // the TF to read is numTF
       assert(device.inputTimesliceId < device.maxInputTimeslices);
@@ -218,6 +221,7 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
       static size_t totalSizeUncompressed = 0;
       static size_t totalSizeCompressed = 0;
       static uint64_t totalDFSent = 0;
+      static uint64_t totalInvalidReadSkipped = 0;
 
       // check if RuntimeLimit is reached
       if (!watchdog->update()) {
@@ -232,6 +236,17 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
 
       int64_t startTime = uv_hrtime();
       int64_t startSize = totalSizeCompressed;
+      auto skipInvalidRead = [&](o2::header::DataOrigin const& origin, InvalidAODReadError const& e) {
+        auto skippedTimeframes = ++totalInvalidReadSkipped;
+        LOGP(error, "Invalid AOD read for table {}: fileCounter {}, timeFrame {}. Skipping timeframe (skipped timeframes: {}). Reason: {}",
+             origin.as<std::string>(), fcnt, ntf, skippedTimeframes, e.what());
+        arrowContext.clear();
+        messageContext.discard();
+        stringContext.clear();
+        monitoring.send(Metric{skippedTimeframes, "aod-invalid-read-skipped-timeframes"}.addTag(Key::Subsystem, monitoring::tags::Value::DPL));
+        *fileCounter = (fcnt - device.inputTimesliceId) / device.maxInputTimeslices;
+        *numTF = ntf;
+      };
       for (auto& route : requestedTables) {
         if ((device.inputTimesliceId % route.maxTimeslices) != route.timeslice) {
           continue;
@@ -242,25 +257,38 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
         auto dh = header::DataHeader(concrete.description, concrete.origin, concrete.subSpec);
         bool wasAOD = std::ranges::any_of(route.matcher.metadata, [](ConfigParamSpec const& p) { return p.name.starts_with("aod-origin-replaced"); });
 
-        if (!didir->readTree(outputs, dh, fcnt, ntf, totalSizeCompressed, totalSizeUncompressed, wasAOD)) {
-          if (first) {
-            // check if there is a next file to read
-            fcnt += device.maxInputTimeslices;
-            if (didir->atEnd(fcnt)) {
-              LOGP(info, "No input files left to read for reader {}!", device.inputTimesliceId);
-              didir->closeInputFiles();
-              monitoring.flushBuffer();
-              control.endOfStream();
-              control.readyToQuit(QuitRequest::Me);
-              return;
-            }
-            // get first folder of next file
-            ntf = 0;
-            if (!didir->readTree(outputs, dh, fcnt, ntf, totalSizeCompressed, totalSizeUncompressed, wasAOD)) {
-              LOGP(fatal, "Can not retrieve tree for table {}: fileCounter {}, timeFrame {}", concrete.origin.as<std::string>(), fcnt, ntf);
-              throw std::runtime_error("Processing is stopped!");
-            }
-          } else {
+        bool treeRead = false;
+        try {
+          treeRead = didir->readTree(outputs, dh, fcnt, ntf, totalSizeCompressed, totalSizeUncompressed, wasAOD);
+        } catch (InvalidAODReadError const& e) {
+          skipInvalidRead(concrete.origin, e);
+          return;
+        }
+
+        if (!treeRead) {
+          if (!first) {
+            LOGP(fatal, "Can not retrieve tree for table {}: fileCounter {}, timeFrame {}", concrete.origin.as<std::string>(), fcnt, ntf);
+            throw std::runtime_error("Processing is stopped!");
+          }
+          // check if there is a next file to read
+          fcnt += device.maxInputTimeslices;
+          if (didir->atEnd(fcnt)) {
+            LOGP(info, "No input files left to read for reader {}!", device.inputTimesliceId);
+            didir->closeInputFiles();
+            monitoring.flushBuffer();
+            control.endOfStream();
+            control.readyToQuit(QuitRequest::Me);
+            return;
+          }
+          // get first folder of next file
+          ntf = 0;
+          try {
+            treeRead = didir->readTree(outputs, dh, fcnt, ntf, totalSizeCompressed, totalSizeUncompressed, wasAOD);
+          } catch (InvalidAODReadError const& e) {
+            skipInvalidRead(concrete.origin, e);
+            return;
+          }
+          if (!treeRead) {
             LOGP(fatal, "Can not retrieve tree for table {}: fileCounter {}, timeFrame {}", concrete.origin.as<std::string>(), fcnt, ntf);
             throw std::runtime_error("Processing is stopped!");
           }

Framework/AnalysisSupport/src/DataInputDirector.cxx

@@ -34,6 +34,7 @@
 #include <arrow/dataset/file_base.h>
 #include <arrow/dataset/dataset.h>
 #include <uv.h>
+#include <exception>
 #include <memory>
 
 #if __has_include(<TJAlienFile.h>)
@@ -536,18 +537,25 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
   if (!format) {
     t.deactivate();
     LOGP(debug, "Could not find tree {}. Trying in parent file.", fullpath.path());
-    auto parentFile = getParentFile(counter, numTF, treename, wantedLevel, wantedOrigin);
-    if (parentFile != nullptr) {
-      int parentNumTF = parentFile->findDFNumber(0, folder.path());
-      if (parentNumTF == -1) {
-        auto parentRootFS = std::dynamic_pointer_cast<TFileFileSystem>(parentFile->mCurrentFilesystem);
-        throw std::runtime_error(fmt::format(R"(DF {} listed in parent file map but not found in the corresponding file "{}")", folder.path(), parentRootFS->GetFile()->GetName()));
-      }
-      // first argument is 0 as the parent file object contains only 1 file
-      return parentFile->readTree(outputs, dh, 0, parentNumTF, treename, totalSizeCompressed, totalSizeUncompressed);
+    std::shared_ptr<DataInputDescriptor> parentFile;
+    try {
+      parentFile = getParentFile(counter, numTF, treename, wantedLevel, wantedOrigin);
+    } catch (std::exception const& e) {
+      throw InvalidAODReadError(fmt::format("Unable to resolve parent file for tree {}: {}", treename, e.what()));
+    } catch (...) {
+      throw InvalidAODReadError(fmt::format("Unable to resolve parent file for tree {}", treename));
+    }
+    if (parentFile == nullptr) {
+      auto rootFS = std::dynamic_pointer_cast<TFileFileSystem>(mCurrentFilesystem);
+      throw std::runtime_error(fmt::format(R"(Couldn't get TTree "{}" from "{}". Please check https://aliceo2group.github.io/analysis-framework/docs/troubleshooting/#tree-not-found for more information.)", fullpath.path(), rootFS->GetFile()->GetName()));
     }
-    auto rootFS = std::dynamic_pointer_cast<TFileFileSystem>(mCurrentFilesystem);
-    throw std::runtime_error(fmt::format(R"(Couldn't get TTree "{}" from "{}". Please check https://aliceo2group.github.io/analysis-framework/docs/troubleshooting/#tree-not-found for more information.)", fullpath.path(), rootFS->GetFile()->GetName()));
+    int parentNumTF = parentFile->findDFNumber(0, folder.path());
+    if (parentNumTF == -1) {
+      auto parentRootFS = std::dynamic_pointer_cast<TFileFileSystem>(parentFile->mCurrentFilesystem);
+      throw InvalidAODReadError(fmt::format(R"(DF {} listed in parent file map but not found in the corresponding file "{}")", folder.path(), parentRootFS->GetFile()->GetName()));
+    }
+    // first argument is 0 as the parent file object contains only 1 file
+    return parentFile->readTree(outputs, dh, 0, parentNumTF, treename, totalSizeCompressed, totalSizeUncompressed);
   }
 
   auto schemaOpt = format->Inspect(fullpath);
@@ -573,7 +581,23 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
   //// add branches to read
   //// fill the table
   f2b->setLabel(treename.c_str());
-  f2b->fill(datasetSchema, format);
+  try {
+    f2b->fill(datasetSchema, format);
+  } catch (std::exception const& e) {
+    f2b.discard();
+    throw InvalidAODReadError(fmt::format("Unable to read tree {}: {}", treename, e.what()));
+  } catch (...) {
+    f2b.discard();
+    throw InvalidAODReadError(fmt::format("Unable to read tree {}", treename));
+  }
+
+  try {
+    f2b.release();
+  } catch (std::exception const& e) {
+    throw InvalidAODReadError(fmt::format("Unable to finalize tree {}: {}", treename, e.what()));
+  } catch (...) {
+    throw InvalidAODReadError(fmt::format("Unable to finalize tree {}", treename));
+  }
 
   return true;
 }

Framework/AnalysisSupport/src/DataInputDirector.h

@@ -21,6 +21,7 @@
 #include <arrow/dataset/dataset.h>
 
 #include <regex>
+#include <stdexcept>
 #include <vector>
 #include "rapidjson/fwd.h"
 
@@ -32,6 +33,12 @@ class Monitoring;
 namespace o2::framework
 {
 
+class InvalidAODReadError : public std::runtime_error
+{
+ public:
+  using std::runtime_error::runtime_error;
+};
+
 struct FileNameHolder {
   std::string fileName;
   int numberOfTimeFrames = 0;

Framework/AnalysisSupport/src/TTreePlugin.cxx

@@ -211,6 +211,9 @@ auto readBoolValues = [](uint8_t* target, ReadOps& op, TBufferFile& rootBuffer)
   while (readEntries < op.rootBranchEntries) {
     auto beginValue = readEntries;
     readLast = op.branch->GetBulkRead().GetBulkEntries(readEntries, rootBuffer);
+    if (readLast < 0) {
+      throw runtime_error_f("Error while reading branch %s starting from %d.", op.branch->GetName(), readEntries);
+    }
     int size = readLast * op.listSize;
     readEntries += readLast;
     for (int i = beginValue; i < beginValue + size; ++i) {
@@ -228,7 +231,18 @@ auto readVLAValues = [](uint8_t* target, ReadOps& op, ReadOps const& offsetOp, T
   rootBuffer.Reset();
   while (readEntries < op.rootBranchEntries) {
     auto readLast = op.branch->GetBulkRead().GetEntriesSerialized(readEntries, rootBuffer);
+    if (readLast < 0) {
+      throw runtime_error_f("Error while reading branch %s starting from %d.", op.branch->GetName(), readEntries);
+    }
+    if (readEntries + readLast > op.rootBranchEntries) {
+      throw runtime_error_f("Invalid read range for branch %s: starting from %d, read %d entries, total entries %lld.",
+                            op.branch->GetName(), readEntries, readLast, static_cast<long long>(op.rootBranchEntries));
+    }
     int size = offsets[readEntries + readLast] - offsets[readEntries];
+    if (size < 0) {
+      throw runtime_error_f("Invalid offset range for branch %s: offsets[%d]=%d, offsets[%d]=%d.",
+                            op.branch->GetName(), readEntries, offsets[readEntries], readEntries + readLast, offsets[readEntries + readLast]);
+    }
     readEntries += readLast;
     bigEndianCopy(target, rootBuffer.GetCurrent(), size, op.typeSize);
     target += (ptrdiff_t)(size * op.typeSize);

Framework/Core/include/Framework/MessageContext.h

@@ -466,6 +466,11 @@ class MessageContext
   /// discarded.
   void clear();
 
+  /// Discard pending output messages without asserting that they were sent. This
+  /// is intended for exception teardown paths where normal post-processing will
+  /// not run.
+  void discard();
+
   FairMQDeviceProxy& proxy()
   {
     return mProxy;

Framework/Core/src/MessageContext.cxx

@@ -107,6 +107,13 @@ void MessageContext::clear()
   mMessages.clear();
 }
 
+void MessageContext::discard()
+{
+  mDidDispatch = false;
+  mScheduledMessages.clear();
+  mMessages.clear();
+}
+
 int64_t MessageContext::addToCache(std::unique_ptr<fair::mq::Message>& toCache)
 {
   auto&& cached = toCache->GetTransport()->CreateMessage();