diff --git a/Framework/AODNative/include/AODNative/ByteSource.h b/Framework/AODNative/include/AODNative/ByteSource.h index 113bc3adaef93..26376923328c2 100644 --- a/Framework/AODNative/include/AODNative/ByteSource.h +++ b/Framework/AODNative/include/AODNative/ByteSource.h @@ -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; diff --git a/Framework/AODNative/include/AODNative/ReadPlan.h b/Framework/AODNative/include/AODNative/ReadPlan.h index b018adbacaaaa..0fbc235244d28 100644 --- a/Framework/AODNative/include/AODNative/ReadPlan.h +++ b/Framework/AODNative/include/AODNative/ReadPlan.h @@ -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 data; + std::unique_ptr 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 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. diff --git a/Framework/AODNative/src/ReadPlan.cxx b/Framework/AODNative/src/ReadPlan.cxx index 27b8c58695b2b..a093a4504b3bc 100644 --- a/Framework/AODNative/src/ReadPlan.cxx +++ b/Framework/AODNative/src/ReadPlan.cxx @@ -12,6 +12,8 @@ #include "AODNative/ReadPlan.h" #include #include +#include +#include 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(capacity), capacity}; + auto capacity = std::max(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*)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 source, Options const& options, std::shared_ptr 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 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; }); diff --git a/Framework/AODNative/test/test_ReadPlan.cxx b/Framework/AODNative/test/test_ReadPlan.cxx index 2efc1b0d101b8..fc08e28580d8d 100644 --- a/Framework/AODNative/test/test_ReadPlan.cxx +++ b/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); +} diff --git a/Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx b/Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx index 86b7779855bb0..4f0e438651cc2 100644 --- a/Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx +++ b/Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx @@ -206,6 +206,9 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const } context.nativeOptions.threads = std::max(1, options.get("aod-reader-threads")); context.nativeOptions.prefetchDepth = std::max(0, options.get("aod-reader-prefetch")); + if (auto mb = options.get("aod-reader-buffer-mb"); mb >= 0) { + context.nativeOptions.bufferBudget = mb * size_t{1024 * 1024}; + } auto localIO = options.get("aod-reader-local-io"); if (localIO != "mmap" && localIO != "pread") { LOGP(fatal, "Invalid --aod-reader-local-io \"{}\": use mmap or pread.", localIO); diff --git a/Framework/AnalysisSupport/src/NativeAOD.cxx b/Framework/AnalysisSupport/src/NativeAOD.cxx index b5d92015b718f..3d88439c158df 100644 --- a/Framework/AnalysisSupport/src/NativeAOD.cxx +++ b/Framework/AnalysisSupport/src/NativeAOD.cxx @@ -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 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(mOptions.bufferBudget)) { } @@ -63,12 +74,20 @@ std::shared_ptr 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(mFile->sourcePtr(), mOptions, mPool); @@ -90,7 +109,15 @@ std::shared_ptr 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(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(mFile->sourcePtr(), mOptions, mPool); + mPlanBytes = 0; + } mPlan->add(handle->tree); + mPlanBytes += bytes; + mPlanOf[treeName] = mPlan; mPlanTrees.push_back(treeName); } } diff --git a/Framework/AnalysisSupport/src/NativeAOD.h b/Framework/AnalysisSupport/src/NativeAOD.h index a477873875ccb..ec62e6347cff3 100644 --- a/Framework/AnalysisSupport/src/NativeAOD.h +++ b/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 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> mDirectories; std::map> mTrees; - std::shared_ptr mPool = std::make_shared(); + std::shared_ptr mPool; + // The plan trees are being added to, and how many bytes it holds. std::shared_ptr mPlan; + uint64_t mPlanBytes = 0; + // The plan of each tree of the DF, until planFor() hands it out. + std::map> mPlanOf; std::string mPlanDF; // The DF whose pages are mapped, released when the next one starts. std::string mMappedDF; diff --git a/Framework/Core/src/WorkflowHelpers.cxx b/Framework/Core/src/WorkflowHelpers.cxx index c87330f9dbbca..59052efb8bcc1 100644 --- a/Framework/Core/src/WorkflowHelpers.cxx +++ b/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, -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"}}, diff --git a/Framework/Core/src/runDataProcessing.cxx b/Framework/Core/src/runDataProcessing.cxx index c4b2cb892f79d..2877cc58ccf88 100644 --- a/Framework/Core/src/runDataProcessing.cxx +++ b/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",