ktf · Nov 3, 2026 · Nov 2, 2026
@@ -36,6 +36,11 @@ 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. kAutoBufferBudget: 16 MB for
+  /// a local file, 64 MB for a remote one, where each read is a round trip.
+  static constexpr size_t kAutoBufferBudget = SIZE_MAX;
+  size_t bufferBudget = kAutoBufferBudget;
   double spanFraction = 0.8;
   int streamWindows = 4;
   size_t streamWindowBytes = 2 * 1024 * 1024;

@@ -28,16 +28,26 @@ namespace o2::aod::native
 class ReadBufferPool
 {
  public:
+  /// Mapped directly rather than taken from malloc, so that a buffer which is
+  /// dropped always goes back to the system, whatever its size.
+  struct Unmap {
+    size_t size = 0;
+    void operator()(uint8_t* p) const;
+  };
   struct Buffer {
-    std::unique_ptr<uint8_t[]> data;
+    std::unique_ptr<uint8_t[], Unmap> 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 +68,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.

@@ -12,6 +12,8 @@
 #include "AODNative/ReadPlan.h"
 #include <algorithm>
 #include <cstring>
+#include <new>
+#include <sys/mman.h>
 
 namespace o2::aod::native
 {
@@ -34,17 +36,38 @@ ReadBufferPool::Buffer ReadBufferPool::get(size_t size)
   // resident next to the new one. Some headroom, so that the next DF, a bit
   // larger, still fits.
   mFree.clear();
-  auto capacity = size + size / 8;
-  return {std::make_unique_for_overwrite<uint8_t[]>(capacity), capacity};
+  auto capacity = std::max<size_t>(size + size / 8, 1);
+  void* p = mmap(nullptr, capacity, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  if (p == MAP_FAILED) {
+    throw std::bad_alloc();
+  }
+  return {std::unique_ptr<uint8_t[], Unmap>((uint8_t*)p, Unmap{capacity}), capacity};
+}
+
+void ReadBufferPool::Unmap::operator()(uint8_t* p) const
+{
+  munmap(p, size);
 }
 
 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 +103,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; });

@@ -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);
+}
@@ -206,6 +206,9 @@ 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"));
+    if (auto mb = options.get<int>("aod-reader-buffer-mb"); mb >= 0) {
+      context.nativeOptions.bufferBudget = 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);

@@ -20,8 +20,19 @@ namespace o2::framework
 
 using namespace o2::aod::native;
 
+namespace
+{
+Options withBudget(Options options, RootFile& file)
+{
+  if (options.bufferBudget == Options::kAutoBufferBudget) {
+    options.bufferBudget = (file.source().isRemote() ? 64 : 16) * size_t{1024 * 1024};
+  }
+  return options;
+}
+} // namespace
+
 NativeAODFileSystem::NativeAODFileSystem(std::shared_ptr<RootFile> file, Options options, RootObjectReadingFactory& factory)
-  : mFile(std::move(file)), mOptions(options), mObjectFactory(factory)
+  : mFile(std::move(file)), mOptions(withBudget(options, *mFile)), mObjectFactory(factory), mPool(std::make_shared<ReadBufferPool>(mOptions.bufferBudget))
 {
 }
 
@@ -63,12 +74,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 +109,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 +209,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 +246,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);
   }
 }

@@ -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;

@@ -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, -1, {"Most MB the native AOD reader holds in read buffers: 0 for no limit, -1 for 16 MB for local files and 64 MB for remote ones"}},
                 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"}},

@@ -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",