Skip to content

Commit 4836d96

Browse files
Dr15JonesDr15Jones
committed
Protect constructing and setting TTreeCache
Manipulating the TTreeCache in ROOT is fairly brittle so need to have extra safe-guards.
1 parent 9632144 commit 4836d96

File tree

4 files changed

+90
-83
lines changed

4 files changed

+90
-83
lines changed

‎IOPool/Input/src/InputFile.h‎

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ Holder for an input TFile.
1010
#include "FWCore/Utilities/interface/propagate_const.h"
1111

1212
#include "TFile.h"
13+
#include "TTree.h"
14+
#include "TTreeCache.h"
1315

1416
#include <map>
1517
#include <string>
@@ -46,10 +48,30 @@ namespace edm {
4648
static void reportReadBranch(InputType inputType, std::string const& branchname);
4749

4850
TObject* Get(char const* name) { return file_->Get(name); }
49-
TFileCacheRead* GetCacheRead(TObject* iTree) const { return file_->GetCacheRead(iTree); }
50-
void SetCacheRead(TFileCacheRead* tfcr, TObject* iTree) {
51+
std::unique_ptr<TTreeCache> createCacheWithSize(TTree& iTree, unsigned int cacheSize) {
52+
iTree.SetCacheSize(static_cast<Long64_t>(cacheSize));
53+
std::unique_ptr<TTreeCache> newCache(dynamic_cast<TTreeCache*>(file_->GetCacheRead(&iTree)));
54+
file_->SetCacheRead(nullptr, &iTree, TFile::kDoNotDisconnect);
55+
return newCache;
56+
}
57+
58+
class CacheGuard {
59+
public:
60+
CacheGuard(TFile* file, TObject* tree) : file_(file), tree_(tree) {}
61+
CacheGuard() = delete;
62+
CacheGuard(CacheGuard const&) = delete;
63+
CacheGuard& operator=(CacheGuard const&) = delete;
64+
~CacheGuard() { file_->SetCacheRead(nullptr, tree_, TFile::kDoNotDisconnect); }
65+
66+
private:
67+
TFile* file_;
68+
TObject* tree_;
69+
};
70+
[[nodiscard]] CacheGuard setCacheReadTemporarily(TFileCacheRead* tfcr, TObject* iTree) {
5171
file_->SetCacheRead(tfcr, iTree, TFile::kDoNotDisconnect);
72+
return CacheGuard(file_.get(), iTree);
5273
}
74+
void clearCacheRead(TObject* iTree) { file_->SetCacheRead(nullptr, iTree, TFile::kDoNotDisconnect); }
5375
void logFileAction(char const* msg, char const* fileName) const;
5476

5577
private:

‎IOPool/Input/src/RootFile.cc‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -302,12 +302,11 @@ namespace edm {
302302
std::unique_ptr<TTreeCache> psetTreeCache =
303303
roottree::trainCache(psetTree.get(), *filePtr_, roottree::defaultNonEventCacheSize, "*");
304304
psetTreeCache->SetEnablePrefetching(false);
305-
filePtr_->SetCacheRead(psetTreeCache.get(), psetTree.get());
305+
auto guard = filePtr_->setCacheReadTemporarily(psetTreeCache.get(), psetTree.get());
306306
for (Long64_t i = 0; i != psetTree->GetEntries(); ++i) {
307307
psetTree->GetEntry(i);
308308
psetMap.insert(idToBlob);
309309
}
310-
filePtr_->SetCacheRead(nullptr, psetTree.get());
311310
}
312311

313312
// backward compatibility

‎IOPool/Input/src/RootTree.cc‎

Lines changed: 64 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -221,13 +221,15 @@ namespace edm {
221221

222222
roottree::BranchMap const& RootTree::branches() const { return branches_; }
223223

224+
std::shared_ptr<TTreeCache> RootTree::createCacheWithSize(unsigned int cacheSize) const {
225+
return filePtr_->createCacheWithSize(*tree_, cacheSize);
226+
}
227+
224228
void RootTree::setCacheSize(unsigned int cacheSize) {
225229
cacheSize_ = cacheSize;
226-
tree_->SetCacheSize(static_cast<Long64_t>(cacheSize));
227-
treeCache_.reset(dynamic_cast<TTreeCache*>(filePtr_->GetCacheRead(tree_)));
230+
treeCache_ = createCacheWithSize(cacheSize);
228231
if (treeCache_)
229232
treeCache_->SetEnablePrefetching(enablePrefetching_);
230-
filePtr_->SetCacheRead(nullptr, tree_);
231233
rawTreeCache_.reset();
232234
}
233235

@@ -245,45 +247,47 @@ namespace edm {
245247
}
246248

247249
void RootTree::setEntryNumber(EntryNumber theEntryNumber) {
248-
filePtr_->SetCacheRead(treeCache_.get(), tree_);
249-
250-
// Detect a backward skip. If the skip is sufficiently large, we roll the dice and reset the treeCache.
251-
// This will cause some amount of over-reading: we pre-fetch all the events in some prior cluster.
252-
// However, because reading one event in the cluster is supposed to be equivalent to reading all events in the cluster,
253-
// we're not incurring additional over-reading - we're just doing it more efficiently.
254-
// NOTE: Constructor guarantees treeAutoFlush_ is positive, even if TTree->GetAutoFlush() is negative.
255-
if (theEntryNumber < entryNumber_ and theEntryNumber >= 0) {
256-
//We started reading the file near the end, now we need to correct for the learning length
257-
if (switchOverEntry_ > tree_->GetEntries()) {
258-
switchOverEntry_ = switchOverEntry_ - tree_->GetEntries();
259-
if (rawTreeCache_) {
260-
rawTreeCache_->SetEntryRange(theEntryNumber, switchOverEntry_);
261-
rawTreeCache_->FillBuffer();
250+
{
251+
auto guard = filePtr_->setCacheReadTemporarily(treeCache_.get(), tree_);
252+
253+
// Detect a backward skip. If the skip is sufficiently large, we roll the dice and reset the treeCache.
254+
// This will cause some amount of over-reading: we pre-fetch all the events in some prior cluster.
255+
// However, because reading one event in the cluster is supposed to be equivalent to reading all events in the cluster,
256+
// we're not incurring additional over-reading - we're just doing it more efficiently.
257+
// NOTE: Constructor guarantees treeAutoFlush_ is positive, even if TTree->GetAutoFlush() is negative.
258+
if (theEntryNumber < entryNumber_ and theEntryNumber >= 0) {
259+
//We started reading the file near the end, now we need to correct for the learning length
260+
if (switchOverEntry_ > tree_->GetEntries()) {
261+
switchOverEntry_ = switchOverEntry_ - tree_->GetEntries();
262+
if (rawTreeCache_) {
263+
rawTreeCache_->SetEntryRange(theEntryNumber, switchOverEntry_);
264+
rawTreeCache_->FillBuffer();
265+
}
262266
}
263-
}
264-
if (performedSwitchOver_ and triggerTreeCache_) {
265-
//We are using the triggerTreeCache_ not the rawTriggerTreeCache_.
266-
//The triggerTreeCache was originally told to start from an entry further in the file.
267-
triggerTreeCache_->SetEntryRange(theEntryNumber, tree_->GetEntries());
268-
} else if (rawTriggerTreeCache_) {
269-
//move the switch point to the end of the cluster holding theEntryNumber
270-
rawTriggerSwitchOverEntry_ = -1;
271-
TTree::TClusterIterator clusterIter = tree_->GetClusterIterator(theEntryNumber);
272-
while ((rawTriggerSwitchOverEntry_ < theEntryNumber) || (rawTriggerSwitchOverEntry_ <= 0)) {
273-
rawTriggerSwitchOverEntry_ = clusterIter();
267+
if (performedSwitchOver_ and triggerTreeCache_) {
268+
//We are using the triggerTreeCache_ not the rawTriggerTreeCache_.
269+
//The triggerTreeCache was originally told to start from an entry further in the file.
270+
triggerTreeCache_->SetEntryRange(theEntryNumber, tree_->GetEntries());
271+
} else if (rawTriggerTreeCache_) {
272+
//move the switch point to the end of the cluster holding theEntryNumber
273+
rawTriggerSwitchOverEntry_ = -1;
274+
TTree::TClusterIterator clusterIter = tree_->GetClusterIterator(theEntryNumber);
275+
while ((rawTriggerSwitchOverEntry_ < theEntryNumber) || (rawTriggerSwitchOverEntry_ <= 0)) {
276+
rawTriggerSwitchOverEntry_ = clusterIter();
277+
}
278+
rawTriggerTreeCache_->SetEntryRange(theEntryNumber, rawTriggerSwitchOverEntry_);
274279
}
275-
rawTriggerTreeCache_->SetEntryRange(theEntryNumber, rawTriggerSwitchOverEntry_);
276280
}
277-
}
278-
if ((theEntryNumber < static_cast<EntryNumber>(entryNumber_ - treeAutoFlush_)) && (treeCache_) &&
279-
(!treeCache_->IsLearning()) && (entries_ > 0) && (switchOverEntry_ >= 0)) {
280-
treeCache_->SetEntryRange(theEntryNumber, entries_);
281-
treeCache_->FillBuffer();
282-
}
281+
if ((theEntryNumber < static_cast<EntryNumber>(entryNumber_ - treeAutoFlush_)) && (treeCache_) &&
282+
(!treeCache_->IsLearning()) && (entries_ > 0) && (switchOverEntry_ >= 0)) {
283+
treeCache_->SetEntryRange(theEntryNumber, entries_);
284+
treeCache_->FillBuffer();
285+
}
283286

284-
entryNumber_ = theEntryNumber;
285-
tree_->LoadTree(entryNumber_);
286-
filePtr_->SetCacheRead(nullptr, tree_);
287+
entryNumber_ = theEntryNumber;
288+
tree_->LoadTree(entryNumber_);
289+
//want guard to end here
290+
}
287291
if (treeCache_ && trainNow_ && entryNumber_ >= 0) {
288292
startTraining();
289293
trainNow_ = false;
@@ -336,8 +340,7 @@ namespace edm {
336340

337341
// ROOT will automatically expand the cache to fit one cluster; hence, we use
338342
// 5 MB as the cache size below
339-
tree_->SetCacheSize(static_cast<Long64_t>(5 * 1024 * 1024));
340-
rawTriggerTreeCache_.reset(dynamic_cast<TTreeCache*>(filePtr_->GetCacheRead(tree_)));
343+
rawTriggerTreeCache_ = createCacheWithSize(5 * 1024 * 1024);
341344
if (rawTriggerTreeCache_)
342345
rawTriggerTreeCache_->SetEnablePrefetching(false);
343346
TObjArray* branches = tree_->GetListOfBranches();
@@ -355,7 +358,6 @@ namespace edm {
355358
}
356359
performedSwitchOver_ = false;
357360
rawTriggerTreeCache_->StopLearningPhase();
358-
filePtr_->SetCacheRead(nullptr, tree_);
359361

360362
return rawTriggerTreeCache_.get();
361363
} else if (!performedSwitchOver_ and entryNumber_ < rawTriggerSwitchOverEntry_) {
@@ -369,8 +371,7 @@ namespace edm {
369371
performedSwitchOver_ = true;
370372

371373
// Train the triggerCache
372-
tree_->SetCacheSize(static_cast<Long64_t>(5 * 1024 * 1024));
373-
triggerTreeCache_.reset(dynamic_cast<TTreeCache*>(filePtr_->GetCacheRead(tree_)));
374+
triggerTreeCache_ = createCacheWithSize(5 * 1024 * 1024);
374375
triggerTreeCache_->SetEnablePrefetching(false);
375376
triggerTreeCache_->SetLearnEntries(0);
376377
triggerTreeCache_->SetEntryRange(entryNumber, tree_->GetEntries());
@@ -380,7 +381,6 @@ namespace edm {
380381
triggerTreeCache_->AddBranch(*it, kTRUE);
381382
}
382383
triggerTreeCache_->StopLearningPhase();
383-
filePtr_->SetCacheRead(nullptr, tree_);
384384
}
385385
return triggerTreeCache_.get();
386386
}
@@ -412,12 +412,10 @@ namespace edm {
412412
}
413413
TTreeCache* RootTree::getAuxCache(TBranch* auxBranch) const {
414414
if (not auxCache_ and cacheSize_ > 0) {
415-
tree_->SetCacheSize(1 * 1024 * 1024);
416-
auxCache_.reset(dynamic_cast<TTreeCache*>(filePtr_->GetCacheRead(tree_)));
415+
auxCache_ = createCacheWithSize(1 * 1024 * 1024);
417416
if (auxCache_) {
418417
auxCache_->SetEnablePrefetching(enablePrefetching_);
419418
auxCache_->SetEnablePrefetching(false);
420-
filePtr_->SetCacheRead(nullptr, tree_);
421419
auxCache_->SetLearnEntries(0);
422420
auxCache_->StartLearningPhase();
423421
auxCache_->SetEntryRange(0, tree_->GetEntries());
@@ -429,10 +427,7 @@ namespace edm {
429427
}
430428

431429
void RootTree::getEntryForAllBranches() const {
432-
filePtr_->SetCacheRead(rawTreeCache_.get(), tree_);
433-
auto ptr = filePtr_.get();
434-
auto cleanup = [ptr, tree = tree_](TTreeCache* cache) { ptr->SetCacheRead(nullptr, tree); };
435-
std::unique_ptr<TTreeCache, decltype(cleanup)> cacheGuard(rawTreeCache_.get(), cleanup);
430+
auto guard = filePtr_->setCacheReadTemporarily(rawTreeCache_.get(), tree_);
436431
tree_->GetEntry(entryNumber_);
437432
}
438433

@@ -442,24 +437,20 @@ namespace edm {
442437

443438
inline void RootTree::getEntryUsingCache(TBranch* branch, EntryNumber entryNumber, TTreeCache* cache) const {
444439
try {
445-
filePtr_->SetCacheRead(cache, tree_);
440+
auto guard = filePtr_->setCacheReadTemporarily(cache, tree_);
446441
branch->GetEntry(entryNumber);
447-
filePtr_->SetCacheRead(nullptr, tree_);
448442
} catch (cms::Exception const& e) {
449443
// We make sure the treeCache_ is detached from the file,
450444
// so that ROOT does not also delete it.
451-
filePtr_->SetCacheRead(nullptr, tree_);
452445
Exception t(errors::FileReadError, "", e);
453446
t.addContext(std::string("Reading branch ") + branch->GetName());
454447
throw t;
455448
} catch (std::exception const& e) {
456-
filePtr_->SetCacheRead(nullptr, tree_);
457449
Exception t(errors::FileReadError);
458450
t << e.what();
459451
t.addContext(std::string("Reading branch ") + branch->GetName());
460452
throw t;
461453
} catch (...) {
462-
filePtr_->SetCacheRead(nullptr, tree_);
463454
Exception t(errors::FileReadError);
464455
t << "An exception of unknown type was thrown.";
465456
t.addContext(std::string("Reading branch ") + branch->GetName());
@@ -490,8 +481,7 @@ namespace edm {
490481
assert(branchType_ == InEvent);
491482
assert(!rawTreeCache_);
492483
treeCache_->SetLearnEntries(learningEntries_);
493-
tree_->SetCacheSize(static_cast<Long64_t>(cacheSize_));
494-
rawTreeCache_.reset(dynamic_cast<TTreeCache*>(filePtr_->GetCacheRead(tree_)));
484+
rawTreeCache_ = createCacheWithSize(cacheSize_);
495485
rawTreeCache_->SetEnablePrefetching(false);
496486
rawTreeCache_->SetLearnEntries(0);
497487
if (promptRead_) {
@@ -510,7 +500,6 @@ namespace edm {
510500
rawTreeCache_->SetEntryRange(rawStart, rawEnd);
511501
rawTreeCache_->AddBranch("*", kTRUE);
512502
rawTreeCache_->StopLearningPhase();
513-
filePtr_->SetCacheRead(nullptr, tree_);
514503

515504
treeCache_->StartLearningPhase();
516505
treeCache_->SetEntryRange(treeStart, tree_->GetEntries());
@@ -522,13 +511,11 @@ namespace edm {
522511
trainedSet_.clear();
523512
triggerSet_.clear();
524513
assert(treeCache_->GetTree() == tree_);
525-
filePtr_->SetCacheRead(nullptr, tree_);
526514
}
527515

528516
void RootTree::stopTraining() {
529-
filePtr_->SetCacheRead(treeCache_.get(), tree_);
517+
auto guard = filePtr_->setCacheReadTemporarily(treeCache_.get(), tree_);
530518
treeCache_->StopLearningPhase();
531-
filePtr_->SetCacheRead(nullptr, tree_);
532519
rawTreeCache_.reset();
533520
}
534521

@@ -540,7 +527,7 @@ namespace edm {
540527
// We own the treeCache_.
541528
// We make sure the treeCache_ is detached from the file,
542529
// so that ROOT does not also delete it.
543-
filePtr_->SetCacheRead(nullptr, tree_);
530+
filePtr_->clearCacheRead(tree_);
544531
// We *must* delete the TTreeCache here because the TFilePrefetch object
545532
// references the TFile. If TFile is closed, before the TTreeCache is
546533
// deleted, the TFilePrefetch may continue to do TFile operations, causing
@@ -560,16 +547,19 @@ namespace edm {
560547
}
561548
tree_->LoadTree(0);
562549
assert(treeCache_);
563-
filePtr_->SetCacheRead(treeCache_.get(), tree_);
564-
treeCache_->StartLearningPhase();
565-
treeCache_->SetEntryRange(0, tree_->GetEntries());
566-
treeCache_->AddBranch(branchNames, kTRUE);
567-
treeCache_->StopLearningPhase();
568-
assert(treeCache_->GetTree() == tree_);
569-
// We own the treeCache_.
570-
// We make sure the treeCache_ is detached from the file,
571-
// so that ROOT does not also delete it.
572-
filePtr_->SetCacheRead(nullptr, tree_);
550+
{
551+
auto guard = filePtr_->setCacheReadTemporarily(treeCache_.get(), tree_);
552+
treeCache_->StartLearningPhase();
553+
treeCache_->SetEntryRange(0, tree_->GetEntries());
554+
treeCache_->AddBranch(branchNames, kTRUE);
555+
treeCache_->StopLearningPhase();
556+
assert(treeCache_->GetTree() == tree_);
557+
// We own the treeCache_.
558+
// We make sure the treeCache_ is detached from the file,
559+
// so that ROOT does not also delete it.
560+
561+
//want guard to end here
562+
}
573563

574564
if (branchType_ == InEvent) {
575565
// Must also manually add things to the trained set.
@@ -616,18 +606,13 @@ namespace edm {
616606
unsigned int cacheSize,
617607
char const* branchNames) {
618608
tree->LoadTree(0);
619-
tree->SetCacheSize(cacheSize);
620-
std::unique_ptr<TTreeCache> treeCache(dynamic_cast<TTreeCache*>(file.GetCacheRead(tree)));
609+
std::unique_ptr<TTreeCache> treeCache = file.createCacheWithSize(*tree, cacheSize);
621610
if (nullptr != treeCache.get()) {
622611
treeCache->StartLearningPhase();
623612
treeCache->SetEntryRange(0, tree->GetEntries());
624613
treeCache->AddBranch(branchNames, kTRUE);
625614
treeCache->StopLearningPhase();
626615
}
627-
// We own the treeCache_.
628-
// We make sure the treeCache_ is detached from the file,
629-
// so that ROOT does not also delete it.
630-
file.SetCacheRead(nullptr, tree);
631616
return treeCache;
632617
}
633618
} // namespace roottree

‎IOPool/Input/src/RootTree.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,7 @@ namespace edm {
210210
bool promptRead,
211211
InputType inputType);
212212

213+
std::shared_ptr<TTreeCache> createCacheWithSize(unsigned int cacheSize) const;
213214
void setCacheSize(unsigned int cacheSize);
214215
void setTreeMaxVirtualSize(int treeMaxVirtualSize);
215216
void startTraining();

0 commit comments

Comments
 (0)