Skip to content

Commit d7ded9d

Browse files
authored
DPL Analysis: cap the read buffers of the native AOD reader (#16032)
1 parent 451f837 commit d7ded9d

9 files changed

Lines changed: 148 additions & 15 deletions

File tree

‎Framework/AODNative/include/AODNative/ByteSource.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,11 @@ struct Options {
3636
int threads = 1;
3737
int prefetchDepth = 1;
3838
size_t coalesceGap = 64 * 1024;
39+
/// Most bytes the read plans of a file hold at once, 0 for no limit. A DF
40+
/// larger than this is read in groups of trees. kAutoBufferBudget: 16 MB for
41+
/// a local file, 64 MB for a remote one, where each read is a round trip.
42+
static constexpr size_t kAutoBufferBudget = SIZE_MAX;
43+
size_t bufferBudget = kAutoBufferBudget;
3944
double spanFraction = 0.8;
4045
int streamWindows = 4;
4146
size_t streamWindowBytes = 2 * 1024 * 1024;

‎Framework/AODNative/include/AODNative/ReadPlan.h‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,16 +28,26 @@ namespace o2::aod::native
2828
class ReadBufferPool
2929
{
3030
public:
31+
/// Mapped directly rather than taken from malloc, so that a buffer which is
32+
/// dropped always goes back to the system, whatever its size.
33+
struct Unmap {
34+
size_t size = 0;
35+
void operator()(uint8_t* p) const;
36+
};
3137
struct Buffer {
32-
std::unique_ptr<uint8_t[]> data;
38+
std::unique_ptr<uint8_t[], Unmap> data;
3339
size_t capacity = 0;
3440
};
41+
/// Keeps at most @a maxFreeBytes in free buffers, dropping the largest
42+
/// first; with 0, at most kMaxFree buffers, dropping the smallest first.
43+
explicit ReadBufferPool(size_t maxFreeBytes = 0) : mMaxFreeBytes(maxFreeBytes) {}
3544
/// The smallest free buffer of at least @a size bytes, or a new one.
3645
Buffer get(size_t size);
3746
void put(Buffer buffer);
3847

3948
private:
4049
static constexpr size_t kMaxFree = 2;
50+
size_t mMaxFreeBytes;
4151
std::mutex mMutex;
4252
std::vector<Buffer> mFree;
4353
};
@@ -58,6 +68,8 @@ class ReadPlan
5868
void add(uint64_t offset, size_t length);
5969
/// Registers every basket of @a tree which is not embedded in its TBranch.
6070
void add(TreeInfo const& tree);
71+
/// The bytes add(@a tree) registers.
72+
static uint64_t bytesOf(TreeInfo const& tree);
6173
/// The bytes of [offset, offset + length), if a registered range covers
6274
/// them, fetching the pending ranges first if needed. Empty otherwise. The
6375
/// span stays valid as long as the plan.

‎Framework/AODNative/src/ReadPlan.cxx‎

Lines changed: 39 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
#include "AODNative/ReadPlan.h"
1313
#include <algorithm>
1414
#include <cstring>
15+
#include <new>
16+
#include <sys/mman.h>
1517

1618
namespace o2::aod::native
1719
{
@@ -34,17 +36,38 @@ ReadBufferPool::Buffer ReadBufferPool::get(size_t size)
3436
// resident next to the new one. Some headroom, so that the next DF, a bit
3537
// larger, still fits.
3638
mFree.clear();
37-
auto capacity = size + size / 8;
38-
return {std::make_unique_for_overwrite<uint8_t[]>(capacity), capacity};
39+
auto capacity = std::max<size_t>(size + size / 8, 1);
40+
void* p = mmap(nullptr, capacity, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
41+
if (p == MAP_FAILED) {
42+
throw std::bad_alloc();
43+
}
44+
return {std::unique_ptr<uint8_t[], Unmap>((uint8_t*)p, Unmap{capacity}), capacity};
45+
}
46+
47+
void ReadBufferPool::Unmap::operator()(uint8_t* p) const
48+
{
49+
munmap(p, size);
3950
}
4051

4152
void ReadBufferPool::put(Buffer buffer)
4253
{
4354
std::lock_guard lock{mMutex};
4455
mFree.push_back(std::move(buffer));
45-
if (mFree.size() > kMaxFree) {
56+
if (mMaxFreeBytes == 0 && mFree.size() > kMaxFree) {
4657
mFree.erase(std::ranges::min_element(mFree, {}, &Buffer::capacity));
4758
}
59+
auto freeBytes = [this] {
60+
size_t total = 0;
61+
for (auto const& b : mFree) {
62+
total += b.capacity;
63+
}
64+
return total;
65+
};
66+
// The same headroom as get() allocates, so that a buffer for a group just
67+
// under the budget can still be kept.
68+
while (mMaxFreeBytes > 0 && !mFree.empty() && freeBytes() > mMaxFreeBytes + mMaxFreeBytes / 8) {
69+
mFree.erase(std::ranges::max_element(mFree, {}, &Buffer::capacity));
70+
}
4871
}
4972

5073
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)
80103
}
81104
}
82105

106+
uint64_t ReadPlan::bytesOf(TreeInfo const& tree)
107+
{
108+
uint64_t total = 0;
109+
for (auto const& branch : tree.branches) {
110+
for (auto const& basket : branch.baskets) {
111+
if (!basket.embedded) {
112+
total += basket.nbytes;
113+
}
114+
}
115+
}
116+
return total;
117+
}
118+
83119
std::span<const uint8_t> ReadPlan::lookup(uint64_t offset, size_t length) const
84120
{
85121
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‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,3 +185,27 @@ TEST_CASE("buffers of finished plans are reused by the next one")
185185
auto big = pool->get(1 << 16);
186186
REQUIRE(big.capacity >= 1 << 16);
187187
}
188+
189+
TEST_CASE("a pool with a budget drops the largest free buffers beyond it")
190+
{
191+
ReadBufferPool pool{1000};
192+
auto small = pool.get(400);
193+
auto medium = pool.get(500);
194+
auto large = pool.get(2000);
195+
auto* smallData = small.data.get();
196+
auto* mediumData = medium.data.get();
197+
pool.put(std::move(small));
198+
pool.put(std::move(medium));
199+
pool.put(std::move(large)); // over the budget: the largest goes
200+
REQUIRE(pool.get(300).data.get() == smallData);
201+
REQUIRE(pool.get(480).data.get() == mediumData);
202+
}
203+
204+
TEST_CASE("bytesOf counts the baskets a plan reads, not the embedded ones")
205+
{
206+
TreeInfo tree;
207+
tree.branches.resize(2);
208+
tree.branches[0].baskets = {BasketRef{100, 300, 0, 10, false, {}}, BasketRef{400, 200, 10, 5, false, {}}};
209+
tree.branches[1].baskets = {BasketRef{600, 50, 0, 10, false, {}}, BasketRef{0, 0, 10, 5, true, {1, 2, 3}}};
210+
REQUIRE(ReadPlan::bytesOf(tree) == 550);
211+
}

‎Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,9 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
206206
}
207207
context.nativeOptions.threads = std::max(1, options.get<int>("aod-reader-threads"));
208208
context.nativeOptions.prefetchDepth = std::max(0, options.get<int>("aod-reader-prefetch"));
209+
if (auto mb = options.get<int>("aod-reader-buffer-mb"); mb >= 0) {
210+
context.nativeOptions.bufferBudget = mb * size_t{1024 * 1024};
211+
}
209212
auto localIO = options.get<std::string>("aod-reader-local-io");
210213
if (localIO != "mmap" && localIO != "pread") {
211214
LOGP(fatal, "Invalid --aod-reader-local-io \"{}\": use mmap or pread.", localIO);

‎Framework/AnalysisSupport/src/NativeAOD.cxx‎

Lines changed: 46 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,19 @@ namespace o2::framework
2020

2121
using namespace o2::aod::native;
2222

23+
namespace
24+
{
25+
Options withBudget(Options options, RootFile& file)
26+
{
27+
if (options.bufferBudget == Options::kAutoBufferBudget) {
28+
options.bufferBudget = (file.source().isRemote() ? 64 : 16) * size_t{1024 * 1024};
29+
}
30+
return options;
31+
}
32+
} // namespace
33+
2334
NativeAODFileSystem::NativeAODFileSystem(std::shared_ptr<RootFile> file, Options options, RootObjectReadingFactory& factory)
24-
: mFile(std::move(file)), mOptions(options), mObjectFactory(factory)
35+
: mFile(std::move(file)), mOptions(withBudget(options, *mFile)), mObjectFactory(factory), mPool(std::make_shared<ReadBufferPool>(mOptions.bufferBudget))
2536
{
2637
}
2738

@@ -63,12 +74,20 @@ std::shared_ptr<ReadPlan> NativeAODFileSystem::planFor(std::string const& path)
6374
if (df != mPlanDF) {
6475
auto previous = std::move(mPlanTrees);
6576
mPlanTrees.clear();
77+
mPlanOf.clear();
78+
mPlanBytes = 0;
6679
mPlanDF = df;
6780
// Also when the directory is already known, e.g. the first DF, read at open.
6881
adoptLookahead(df);
6982
directory(df);
7083
if (mAdoptedDF == df) {
7184
mPlan = std::move(mAdoptedPlan);
85+
for (auto& name : mAdoptedTrees) {
86+
mPlanOf[name] = mPlan;
87+
if (auto* handle = findTree(df, name)) {
88+
mPlanBytes += ReadPlan::bytesOf(handle->tree);
89+
}
90+
}
7291
mPlanTrees = std::move(mAdoptedTrees);
7392
} else {
7493
mPlan = std::make_shared<ReadPlan>(mFile->sourcePtr(), mOptions, mPool);
@@ -90,7 +109,15 @@ std::shared_ptr<ReadPlan> NativeAODFileSystem::planFor(std::string const& path)
90109
}
91110
}
92111
addToPlan(treeName);
93-
return mPlan;
112+
// Let go of the plan once its tree is handed out: the plan, and its buffers,
113+
// then go away when the reads of its trees are done.
114+
auto it = mPlanOf.find(treeName);
115+
if (it == mPlanOf.end()) {
116+
return nullptr;
117+
}
118+
auto plan = std::move(it->second);
119+
mPlanOf.erase(it);
120+
return plan;
94121
}
95122

96123
void NativeAODFileSystem::releaseMappedDF(std::string const& df)
@@ -182,10 +209,19 @@ void NativeAODFileSystem::startLookahead(std::string const& afterDF, std::vector
182209
Lookahead lookahead{df, {}, {}, std::make_shared<ReadPlan>(file->sourcePtr(), options, pool)};
183210
file->setTreesToPrefetch(trees);
184211
lookahead.keys = file->keysOf(df);
212+
// Only the trees which fit in the budget; the others are read when asked for.
213+
uint64_t bytes = 0;
185214
for (auto const& name : trees) {
186215
for (auto const& key : lookahead.keys) {
187216
if (key.name == name && key.className == "TTree") {
188-
lookahead.trees.emplace_back(name, file->readTree(key));
217+
auto tree = file->readTree(key);
218+
auto treeBytes = ReadPlan::bytesOf(tree);
219+
if (options.bufferBudget > 0 && bytes > 0 && bytes + treeBytes > options.bufferBudget) {
220+
lookahead.plan->fetch();
221+
return lookahead;
222+
}
223+
bytes += treeBytes;
224+
lookahead.trees.emplace_back(name, std::move(tree));
189225
lookahead.plan->add(lookahead.trees.back().second);
190226
}
191227
}
@@ -210,7 +246,14 @@ void NativeAODFileSystem::addToPlan(std::string const& treeName)
210246
// Reported when the tree itself is asked for.
211247
}
212248
if (handle) {
249+
auto bytes = ReadPlan::bytesOf(handle->tree);
250+
if (mOptions.bufferBudget > 0 && mPlanBytes > 0 && mPlanBytes + bytes > mOptions.bufferBudget) {
251+
mPlan = std::make_shared<ReadPlan>(mFile->sourcePtr(), mOptions, mPool);
252+
mPlanBytes = 0;
253+
}
213254
mPlan->add(handle->tree);
255+
mPlanBytes += bytes;
256+
mPlanOf[treeName] = mPlan;
214257
mPlanTrees.push_back(treeName);
215258
}
216259
}

‎Framework/AnalysisSupport/src/NativeAOD.h‎

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

‎Framework/Core/src/WorkflowHelpers.cxx‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,7 @@ void WorkflowHelpers::injectServiceDevices(WorkflowSpec& workflow, ConfigContext
227227
ConfigParamSpec{"aod-reader-threads", VariantType::Int, 1, {"Decompression threads of the native AOD reader"}},
228228
ConfigParamSpec{"aod-reader-prefetch", VariantType::Int, 1, {"DFs the native AOD reader prepares ahead for remote files: 0 (off) or 1"}},
229229
ConfigParamSpec{"aod-reader-local-io", VariantType::String, "pread", {"How the native AOD reader reads local files: pread or mmap"}},
230+
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"}},
230231
ConfigParamSpec{"aod-reader-json", VariantType::String, {"json configuration file"}},
231232
ConfigParamSpec{"time-limit", VariantType::Int64, 0ll, {"Maximum run time limit in seconds"}},
232233
ConfigParamSpec{"orbit-offset-enumeration", VariantType::Int64, 0ll, {"initial value for the orbit"}},

‎Framework/Core/src/runDataProcessing.cxx‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2089,6 +2089,7 @@ int runStateMachine(DataProcessorSpecs const& workflow,
20892089
"--aod-reader-threads",
20902090
"--aod-reader-prefetch",
20912091
"--aod-reader-local-io",
2092+
"--aod-reader-buffer-mb",
20922093
"--aod-parent-access-level",
20932094
"--aod-parent-base-path-replacement",
20942095
"--aod-origin-level-mapping",

0 commit comments

Comments
 (0)