DPL Analysis: cap the read buffers of the native AOD reader

The native reader fetches the compressed baskets of all the trees a DF is
read from with one coalesced read into one buffer, kept until the DF is
done, and the buffer pool keeps up to two of them for reuse. On Run 2 AO2Ds
(about 62 MB of compressed data per DF) this holds about 75 MB of private
memory per reader for the whole job: IgProf shows the largest allocation of
the reader, 72 MB, in ReadBufferPool::get, and smaps shows it as anonymous
memory, unaffected by fewer threads, no prefetch or MALLOC_ARENA_MAX=1.

--aod-reader-buffer-mb N caps that. The trees of a DF, in the order they
are read, are packed into read plans of at most N MB of compressed data;
a plan is freed when the reads of its trees are done, and the buffer pool
keeps at most N MB of free buffers. A DF which fits is still read in one
go, so small DFs are not affected, and a single tree larger than N gets a
plan of its own. The lookahead for remote files prepares only the trees
which fit. Only the trees the reader asks for are read, as before.

The default is 0, no limit, which keeps the current behaviour until the
default is chosen from measurements.

Author: ktf

Framework/AODNative/include/AODNative/ByteSource.h

@@ -36,6 +36,9 @@ struct Options {
   int threads = 1;
   int prefetchDepth = 1;
   size_t coalesceGap = 64 * 1024;
+  /// Most bytes the read plans of a file hold at once, 0 for no limit. A DF
+  /// larger than this is read in groups of trees.
+  size_t bufferBudget = 0;
   double spanFraction = 0.8;
   int streamWindows = 4;
   size_t streamWindowBytes = 2 * 1024 * 1024;

Framework/AODNative/include/AODNative/ReadPlan.h

@@ -32,12 +32,16 @@ class ReadBufferPool
     std::unique_ptr<uint8_t[]> data;
     size_t capacity = 0;
   };
+  /// Keeps at most @a maxFreeBytes in free buffers, dropping the largest
+  /// first; with 0, at most kMaxFree buffers, dropping the smallest first.
+  explicit ReadBufferPool(size_t maxFreeBytes = 0) : mMaxFreeBytes(maxFreeBytes) {}
   /// The smallest free buffer of at least @a size bytes, or a new one.
   Buffer get(size_t size);
   void put(Buffer buffer);
 
  private:
   static constexpr size_t kMaxFree = 2;
+  size_t mMaxFreeBytes;
   std::mutex mMutex;
   std::vector<Buffer> mFree;
 };
@@ -58,6 +62,8 @@ class ReadPlan
   void add(uint64_t offset, size_t length);
   /// Registers every basket of @a tree which is not embedded in its TBranch.
   void add(TreeInfo const& tree);
+  /// The bytes add(@a tree) registers.
+  static uint64_t bytesOf(TreeInfo const& tree);
   /// The bytes of [offset, offset + length), if a registered range covers
   /// them, fetching the pending ranges first if needed. Empty otherwise. The
   /// span stays valid as long as the plan.

Framework/AODNative/src/ReadPlan.cxx

@@ -42,9 +42,21 @@ void ReadBufferPool::put(Buffer buffer)
 {
   std::lock_guard lock{mMutex};
   mFree.push_back(std::move(buffer));
-  if (mFree.size() > kMaxFree) {
+  if (mMaxFreeBytes == 0 && mFree.size() > kMaxFree) {
     mFree.erase(std::ranges::min_element(mFree, {}, &Buffer::capacity));
   }
+  auto freeBytes = [this] {
+    size_t total = 0;
+    for (auto const& b : mFree) {
+      total += b.capacity;
+    }
+    return total;
+  };
+  // The same headroom as get() allocates, so that a buffer for a group just
+  // under the budget can still be kept.
+  while (mMaxFreeBytes > 0 && !mFree.empty() && freeBytes() > mMaxFreeBytes + mMaxFreeBytes / 8) {
+    mFree.erase(std::ranges::max_element(mFree, {}, &Buffer::capacity));
+  }
 }
 
 ReadPlan::ReadPlan(std::shared_ptr<ByteSource> source, Options const& options, std::shared_ptr<ReadBufferPool> pool)
@@ -80,6 +92,19 @@ void ReadPlan::add(TreeInfo const& tree)
   }
 }
 
+uint64_t ReadPlan::bytesOf(TreeInfo const& tree)
+{
+  uint64_t total = 0;
+  for (auto const& branch : tree.branches) {
+    for (auto const& basket : branch.baskets) {
+      if (!basket.embedded) {
+        total += basket.nbytes;
+      }
+    }
+  }
+  return total;
+}
+
 std::span<const uint8_t> ReadPlan::lookup(uint64_t offset, size_t length) const
 {
   auto it = std::upper_bound(mChunks.begin(), mChunks.end(), offset, [](uint64_t off, Chunk const& c) { return off < c.offset; });

Framework/AODNative/test/test_ReadPlan.cxx

@@ -185,3 +185,27 @@ TEST_CASE("buffers of finished plans are reused by the next one")
   auto big = pool->get(1 << 16);
   REQUIRE(big.capacity >= 1 << 16);
 }
+
+TEST_CASE("a pool with a budget drops the largest free buffers beyond it")
+{
+  ReadBufferPool pool{1000};
+  auto small = pool.get(400);
+  auto medium = pool.get(500);
+  auto large = pool.get(2000);
+  auto* smallData = small.data.get();
+  auto* mediumData = medium.data.get();
+  pool.put(std::move(small));
+  pool.put(std::move(medium));
+  pool.put(std::move(large)); // over the budget: the largest goes
+  REQUIRE(pool.get(300).data.get() == smallData);
+  REQUIRE(pool.get(480).data.get() == mediumData);
+}
+
+TEST_CASE("bytesOf counts the baskets a plan reads, not the embedded ones")
+{
+  TreeInfo tree;
+  tree.branches.resize(2);
+  tree.branches[0].baskets = {BasketRef{100, 300, 0, 10, false, {}}, BasketRef{400, 200, 10, 5, false, {}}};
+  tree.branches[1].baskets = {BasketRef{600, 50, 0, 10, false, {}}, BasketRef{0, 0, 10, 5, true, {1, 2, 3}}};
+  REQUIRE(ReadPlan::bytesOf(tree) == 550);
+}

Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx

@@ -206,6 +206,7 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
     }
     context.nativeOptions.threads = std::max(1, options.get<int>("aod-reader-threads"));
     context.nativeOptions.prefetchDepth = std::max(0, options.get<int>("aod-reader-prefetch"));
+    context.nativeOptions.bufferBudget = std::max(0, options.get<int>("aod-reader-buffer-mb")) * size_t{1024 * 1024};
     auto localIO = options.get<std::string>("aod-reader-local-io");
     if (localIO != "mmap" && localIO != "pread") {
       LOGP(fatal, "Invalid --aod-reader-local-io \"{}\": use mmap or pread.", localIO);

Framework/AnalysisSupport/src/NativeAOD.cxx

@@ -21,7 +21,7 @@ namespace o2::framework
 using namespace o2::aod::native;
 
 NativeAODFileSystem::NativeAODFileSystem(std::shared_ptr<RootFile> file, Options options, RootObjectReadingFactory& factory)
-  : mFile(std::move(file)), mOptions(options), mObjectFactory(factory)
+  : mFile(std::move(file)), mOptions(options), mObjectFactory(factory), mPool(std::make_shared<ReadBufferPool>(options.bufferBudget))
 {
 }
 
@@ -63,12 +63,20 @@ std::shared_ptr<ReadPlan> NativeAODFileSystem::planFor(std::string const& path)
   if (df != mPlanDF) {
     auto previous = std::move(mPlanTrees);
     mPlanTrees.clear();
+    mPlanOf.clear();
+    mPlanBytes = 0;
     mPlanDF = df;
     // Also when the directory is already known, e.g. the first DF, read at open.
     adoptLookahead(df);
     directory(df);
     if (mAdoptedDF == df) {
       mPlan = std::move(mAdoptedPlan);
+      for (auto& name : mAdoptedTrees) {
+        mPlanOf[name] = mPlan;
+        if (auto* handle = findTree(df, name)) {
+          mPlanBytes += ReadPlan::bytesOf(handle->tree);
+        }
+      }
       mPlanTrees = std::move(mAdoptedTrees);
     } else {
       mPlan = std::make_shared<ReadPlan>(mFile->sourcePtr(), mOptions, mPool);
@@ -90,7 +98,15 @@ std::shared_ptr<ReadPlan> NativeAODFileSystem::planFor(std::string const& path)
     }
   }
   addToPlan(treeName);
-  return mPlan;
+  // Let go of the plan once its tree is handed out: the plan, and its buffers,
+  // then go away when the reads of its trees are done.
+  auto it = mPlanOf.find(treeName);
+  if (it == mPlanOf.end()) {
+    return nullptr;
+  }
+  auto plan = std::move(it->second);
+  mPlanOf.erase(it);
+  return plan;
 }
 
 void NativeAODFileSystem::releaseMappedDF(std::string const& df)
@@ -182,10 +198,19 @@ void NativeAODFileSystem::startLookahead(std::string const& afterDF, std::vector
       Lookahead lookahead{df, {}, {}, std::make_shared<ReadPlan>(file->sourcePtr(), options, pool)};
       file->setTreesToPrefetch(trees);
       lookahead.keys = file->keysOf(df);
+      // Only the trees which fit in the budget; the others are read when asked for.
+      uint64_t bytes = 0;
       for (auto const& name : trees) {
         for (auto const& key : lookahead.keys) {
           if (key.name == name && key.className == "TTree") {
-            lookahead.trees.emplace_back(name, file->readTree(key));
+            auto tree = file->readTree(key);
+            auto treeBytes = ReadPlan::bytesOf(tree);
+            if (options.bufferBudget > 0 && bytes > 0 && bytes + treeBytes > options.bufferBudget) {
+              lookahead.plan->fetch();
+              return lookahead;
+            }
+            bytes += treeBytes;
+            lookahead.trees.emplace_back(name, std::move(tree));
             lookahead.plan->add(lookahead.trees.back().second);
           }
         }
@@ -210,7 +235,14 @@ void NativeAODFileSystem::addToPlan(std::string const& treeName)
     // Reported when the tree itself is asked for.
   }
   if (handle) {
+    auto bytes = ReadPlan::bytesOf(handle->tree);
+    if (mOptions.bufferBudget > 0 && mPlanBytes > 0 && mPlanBytes + bytes > mOptions.bufferBudget) {
+      mPlan = std::make_shared<ReadPlan>(mFile->sourcePtr(), mOptions, mPool);
+      mPlanBytes = 0;
+    }
     mPlan->add(handle->tree);
+    mPlanBytes += bytes;
+    mPlanOf[treeName] = mPlan;
     mPlanTrees.push_back(treeName);
   }
 }

Framework/AnalysisSupport/src/NativeAOD.h

@@ -52,13 +52,17 @@ class NativeAODFileSystem : public VirtualRootFileSystemBase
   /// The tree at @a path, parsed on first use. nullptr if there is no such
   /// tree. Throws o2::aod::native::FormatError if it cannot be parsed.
   NativeTreeHandle* getTreeHandle(std::string const& path);
-  /// The read plan of the DF of the tree at @a path, with that tree in it;
-  /// nullptr when the file is mapped and nothing needs reading. A new DF gets a
-  /// new plan, with the trees read in the previous DF already in it: a reader
-  /// asks for the same tables in every DF, so the first read of a DF fetches
-  /// all of them together. For a remote file, the next DF is prepared in the
-  /// background meanwhile (keys, TTrees and data of the same trees), so that
-  /// its round trips overlap with the decoding of this one.
+  /// The read plan with the tree at @a path in it; nullptr when the file is
+  /// mapped and nothing needs reading. A new DF gets new plans, with the trees
+  /// read in the previous DF already in them: a reader asks for the same tables
+  /// in every DF, so the first read of a plan fetches all of its trees
+  /// together. With Options::bufferBudget, the trees of a DF are split into
+  /// plans of at most that many bytes, in the order they are read, and a plan
+  /// is freed once the reads of its trees are done; a tree larger than the
+  /// budget gets a plan of its own. For a remote file, the next DF is prepared
+  /// in the background meanwhile (keys, TTrees and data of the same trees, up
+  /// to the budget), so that its round trips overlap with the decoding of this
+  /// one.
   std::shared_ptr<o2::aod::native::ReadPlan> planFor(std::string const& path);
   o2::aod::native::RootFile& file() { return *mFile; }
   /// The name the file was asked for by (an LFN for a grid file), which is
@@ -97,8 +101,12 @@ class NativeAODFileSystem : public VirtualRootFileSystemBase
   std::mutex mMutex;
   std::map<std::string, std::vector<o2::aod::native::KeyInfo>> mDirectories;
   std::map<std::string, std::unique_ptr<NativeTreeHandle>> mTrees;
-  std::shared_ptr<o2::aod::native::ReadBufferPool> mPool = std::make_shared<o2::aod::native::ReadBufferPool>();
+  std::shared_ptr<o2::aod::native::ReadBufferPool> mPool;
+  // The plan trees are being added to, and how many bytes it holds.
   std::shared_ptr<o2::aod::native::ReadPlan> mPlan;
+  uint64_t mPlanBytes = 0;
+  // The plan of each tree of the DF, until planFor() hands it out.
+  std::map<std::string, std::shared_ptr<o2::aod::native::ReadPlan>> mPlanOf;
   std::string mPlanDF;
   // The DF whose pages are mapped, released when the next one starts.
   std::string mMappedDF;

Framework/Core/src/WorkflowHelpers.cxx

@@ -227,6 +227,7 @@ void WorkflowHelpers::injectServiceDevices(WorkflowSpec& workflow, ConfigContext
                 ConfigParamSpec{"aod-reader-threads", VariantType::Int, 1, {"Decompression threads of the native AOD reader"}},
                 ConfigParamSpec{"aod-reader-prefetch", VariantType::Int, 1, {"DFs the native AOD reader prepares ahead for remote files: 0 (off) or 1"}},
                 ConfigParamSpec{"aod-reader-local-io", VariantType::String, "pread", {"How the native AOD reader reads local files: pread or mmap"}},
+                ConfigParamSpec{"aod-reader-buffer-mb", VariantType::Int, 0, {"Most MB the native AOD reader holds in read buffers, 0 for no limit"}},
                 ConfigParamSpec{"aod-reader-json", VariantType::String, {"json configuration file"}},
                 ConfigParamSpec{"time-limit", VariantType::Int64, 0ll, {"Maximum run time limit in seconds"}},
                 ConfigParamSpec{"orbit-offset-enumeration", VariantType::Int64, 0ll, {"initial value for the orbit"}},

Framework/Core/src/runDataProcessing.cxx

@@ -2089,6 +2089,7 @@ int runStateMachine(DataProcessorSpecs const& workflow,
             "--aod-reader-threads",
             "--aod-reader-prefetch",
             "--aod-reader-local-io",
+            "--aod-reader-buffer-mb",
             "--aod-parent-access-level",
             "--aod-parent-base-path-replacement",
             "--aod-origin-level-mapping",