DPL: allow plugins to account for bytes

Author: ktf

Framework/AnalysisSupport/src/DataInputDirector.cxx

@@ -537,6 +537,14 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
     if (handle) {
       format = capability.factory().format();
       creator = capability.factory().deferredOutputStreamer;
+      // Account for the bytes we are about to read. This used to sit further down, where
+      // the TTree was opened by hand; moving the reading to the arrow::Dataset API left
+      // the accounting behind, which is why aod-bytes-read-* and the --aod-max-read-rate
+      // pacing that derives from them both read zero. Each format reports its own size,
+      // so we just ask; here is where the object is resolved and its size is known.
+      if (capability.accountBytes) {
+        capability.accountBytes(handle, totalSizeCompressed, totalSizeUncompressed);
+      }
       break;
     }
   }

Framework/AnalysisSupport/src/TTreePlugin.cxx

@@ -392,8 +392,10 @@ class TTreeFileFormat : public arrow::dataset::FileFormat
 class SingleTreeFileSystem : public TTreeFileSystem
 {
  public:
-  SingleTreeFileSystem(TTree* tree)
+  SingleTreeFileSystem(TTree* tree, size_t& totalCompressedSize, size_t& totalUncompressedSize)
     : TTreeFileSystem(),
+      mTotUncompressedSize(totalUncompressedSize),
+      mTotCompressedSize(totalCompressedSize),
       mTree(tree)
   {
   }
@@ -417,8 +419,11 @@ class SingleTreeFileSystem : public TTreeFileSystem
   }
 
  private:
-  size_t mTotUncompressedSize;
-  size_t mTotCompressedSize;
+  // References, not values: a TTreeFileFormat built in GetObjectHandler binds to these,
+  // so by-value members would have it accumulate into copies that are thrown away (and,
+  // being uninitialised here, read as indeterminate).
+  size_t& mTotUncompressedSize;
+  size_t& mTotCompressedSize;
   std::unique_ptr<TTree> mTree;
 };
 

Framework/Core/include/Framework/RootArrowFilesystem.h

@@ -123,6 +123,11 @@ struct RootObjectReadingCapability {
   // Wether or not this actually supports reading an object of the following class
   std::function<bool(char const*)> checkSupport;
 
+  // Accounts the bytes of the object behind `handle` against the two counters, so that
+  // the generic reading code need not know how a given format reports its size. Left
+  // null by formats which cannot report it.
+  std::function<void(void* handle, size_t& compressed, size_t& uncompressed)> accountBytes = nullptr;
+
   // This must be implemented to load the actual RootArrowFactory plugin which
   // implements this capability. This way the detection of the file format
   // (via get handle) does not need to know about the actual code which performs

Framework/Core/src/Plugin.cxx

@@ -19,6 +19,7 @@
 #include "Framework/PluginManager.h"
 #include <TBufferFile.h>
 #include <TDirectory.h>
+#include <TTree.h>
 #include <TDirectoryFile.h>
 #include <TClass.h>
 #include <arrow/dataset/file_base.h>
@@ -272,6 +273,10 @@ struct TTreeObjectReadingCapability : o2::framework::RootObjectReadingCapability
       .lfn2objectPath = [](std::string s) { return s; },
       .getHandle = getHandleByClass("TTree"),
       .checkSupport = matchClassByName("TTree"),
+      .accountBytes = [](void* handle, size_t& compressed, size_t& uncompressed) {
+        auto* tree = (TTree*)handle;
+        compressed += tree->GetZipBytes();
+        uncompressed += tree->GetTotBytes(); },
       .factory = [context]() -> RootArrowFactory& {
         lazyLoadFactory(context->implementations, "O2FrameworkAnalysisTTreeSupport:TTreeObjectReadingImplementation");
         return context->implementations.back();