Skip to content

Commit

Permalink
[tree,io] prevent integer overflow when dealing with cachesize
Browse files Browse the repository at this point in the history
Fixes #9292
  • Loading branch information
ferdymercury authored and pcanal committed Apr 8, 2024
1 parent 0f3645e commit 0f2bb47
Show file tree
Hide file tree
Showing 9 changed files with 37 additions and 23 deletions.
2 changes: 1 addition & 1 deletion io/io/inc/TFileCacheRead.h
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ class TFileCacheRead : public TObject {
virtual Int_t ReadBufferExtNormal(char *buf, Long64_t pos, Int_t len, Int_t &loc);
virtual Int_t ReadBufferExtPrefetch(char *buf, Long64_t pos, Int_t len, Int_t &loc);
virtual Int_t ReadBuffer(char *buf, Long64_t pos, Int_t len);
virtual Int_t SetBufferSize(Int_t buffersize);
virtual Int_t SetBufferSize(Long64_t buffersize);
virtual void SetFile(TFile *file, TFile::ECacheAction action = TFile::kDisconnect);
virtual void SetSkipZip(Bool_t /*skip*/ = kTRUE) {} // This function is only used by TTreeCacheUnzip (ignore it)
virtual void Sort();
Expand Down
5 changes: 4 additions & 1 deletion io/io/src/TFileCacheRead.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
#include "TFileCacheWrite.h"
#include "TFilePrefetch.h"
#include "TMathBase.h"
#include <limits>

ClassImp(TFileCacheRead);

Expand Down Expand Up @@ -701,15 +702,17 @@ void TFileCacheRead::WaitFinishPrefetch()
/// If the current prefetch list is too large to fit in
/// the new buffer some or all of the prefetch blocks are dropped. The
/// requested buffersize must be greater than zero.
/// It will be automatically clamped to minimum 100000 (if <= 10000) and maximum INT_MAX
/// Return values:
/// - 0 if the prefetch block lists remain unchanged
/// - 1 if some or all blocks have been removed from the prefetch list
/// - -1 on error

Int_t TFileCacheRead::SetBufferSize(Int_t buffersize)
Int_t TFileCacheRead::SetBufferSize(Long64_t buffersize)
{
if (buffersize <= 0) return -1;
if (buffersize <=10000) buffersize = 100000;
if (buffersize > std::numeric_limits<Int_t>::max()) buffersize = std::numeric_limits<Int_t>::max();

if (buffersize == fBufferSize) {
fBufferSizeMin = buffersize;
Expand Down
2 changes: 1 addition & 1 deletion tree/tree/inc/TTreeCache.h
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,7 @@ class TTreeCache : public TFileCacheRead {
virtual void ResetCache();
void ResetMissCache(); // Reset the miss cache.
void SetAutoCreated(bool val) {fAutoCreated = val;}
Int_t SetBufferSize(Int_t buffersize) override;
Int_t SetBufferSize(Long64_t buffersize) override;
virtual void SetEntryRange(Long64_t emin, Long64_t emax);
void SetFile(TFile *file, TFile::ECacheAction action=TFile::kDisconnect) override;
virtual void SetLearnPrefill(EPrefillType type = kNoPrefill);
Expand Down
2 changes: 1 addition & 1 deletion tree/tree/inc/TTreeCacheUnzip.h
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,7 @@ class TTreeCacheUnzip : public TTreeCache {
Int_t GetUnzipBuffer(char **buf, Long64_t pos, Int_t len, bool *free) override;
Int_t GetUnzipGroupSize() { return fUnzipGroupSize; }
void ResetCache() override;
Int_t SetBufferSize(Int_t buffersize) override;
Int_t SetBufferSize(Long64_t buffersize) override;
void SetUnzipBufferSize(Long64_t bufferSize);
void SetUnzipGroupSize(Int_t groupSize) { fUnzipGroupSize = groupSize; }
static void SetUnzipRelBufferSize(Float_t relbufferSize);
Expand Down
4 changes: 2 additions & 2 deletions tree/tree/inc/TTreeCloner.h
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ class TTreeCloner {
UInt_t fCloneMethod; ///< Indicates which cloning method was selected.
Long64_t fToStartEntries; ///< Number of entries in the target tree before any addition.

Int_t fCacheSize; ///< Requested size of the file cache
Long64_t fCacheSize; ///< Requested size of the file cache
TFileCacheRead *fFileCache; ///< File Cache used to reduce the number of individual reads
TFileCacheRead *fPrevCache; ///< Cache that set before the TTreeCloner ctor for the 'from' TTree if any.

Expand Down Expand Up @@ -119,7 +119,7 @@ class TTreeCloner {
bool Exec();
bool IsValid() { return fIsValid; }
bool NeedConversion() { return fNeedConversion; }
void SetCacheSize(Int_t size);
void SetCacheSize(Long64_t size);
void SortBaskets();
void WriteBaskets();

Expand Down
9 changes: 7 additions & 2 deletions tree/tree/src/TTree.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -3553,7 +3553,7 @@ Long64_t TTree::CopyEntries(TTree* tree, Long64_t nentries /* = -1 */, Option_t*
onIndexError = kBuild;
}
Ssiz_t cacheSizeLoc = opt.Index("cachesize=");
Int_t cacheSize = -1;
Long64_t cacheSize = -1;
if (cacheSizeLoc != TString::kNPOS) {
// If the parse faile, cacheSize stays at -1.
Ssiz_t cacheSizeEnd = opt.Index(" ",cacheSizeLoc+10) - (cacheSizeLoc+10);
Expand All @@ -3569,7 +3569,7 @@ Long64_t TTree::CopyEntries(TTree* tree, Long64_t nentries /* = -1 */, Option_t*
Warning("CopyEntries","The cachesize option is too large: %s (%g%s max). The default size will be used.",cacheSizeStr.String().Data(),m,munit);
}
}
if (gDebug > 0 && cacheSize != -1) Info("CopyEntries","Using Cache size: %d\n",cacheSize);
if (gDebug > 0 && cacheSize != -1) Info("CopyEntries","Using Cache size: %lld\n",cacheSize);

Long64_t nbytes = 0;
Long64_t treeEntries = tree->GetEntriesFast();
Expand Down Expand Up @@ -8670,6 +8670,8 @@ void TTree::SetBranchStyle(Int_t style)
/// - if cachesize = -1 (default) it is set to the AutoFlush value when writing
/// the Tree (default is 30 MBytes).
///
/// The cacheSize might be clamped, see TFileCacheRead::SetBufferSize
///
/// Returns:
/// - 0 size set, cache was created if possible
/// - -1 on error
Expand All @@ -8695,6 +8697,8 @@ Int_t TTree::SetCacheSize(Long64_t cacheSize)
/// this is a user requested cache. cacheSize is used to size the cache.
/// This cache should never be automatically adjusted.
///
/// The cacheSize might be clamped, see TFileCacheRead::SetBufferSize
///
/// Returns:
/// - 0 size set, or existing autosized cache almost large enough.
/// (cache was created if possible)
Expand Down Expand Up @@ -8778,6 +8782,7 @@ Int_t TTree::SetCacheSizeAux(bool autocache /* = true */, Long64_t cacheSize /*
if (res < 0) {
return -1;
}
cacheSize = pf->GetBufferSize(); // update after potential clamp
}
} else {
// no existing cache
Expand Down
23 changes: 12 additions & 11 deletions tree/tree/src/TTreeCache.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -1535,7 +1535,7 @@ bool TTreeCache::FillBuffer()
}
}

if ((ntotCurrentBuf + len) > fBufferSizeMin) {
if (((Long64_t)ntotCurrentBuf + len) > fBufferSizeMin) {
// Humm ... we are going to go over the requested size.
if (clusterIterations > 0 && cursor[i].fLoadedOnce) {
// We already have a full cluster and now we would go over the requested
Expand All @@ -1544,25 +1544,25 @@ bool TTreeCache::FillBuffer()
if (showMore || gDebug > 5) {
Info(
"FillBuffer",
"Breaking early because %d is greater than %d at cluster iteration %d will restart at %lld",
(ntotCurrentBuf + len), fBufferSizeMin, clusterIterations, minEntry);
"Breaking early because %lld is greater than %d at cluster iteration %d will restart at %lld",
((Long64_t)ntotCurrentBuf + len), fBufferSizeMin, clusterIterations, minEntry);
}
fEntryNext = minEntry;
filled = true;
break;
} else {
if (pass == kStart || !cursor[i].fLoadedOnce) {
if ((ntotCurrentBuf + len) > 4 * fBufferSizeMin) {
if (((Long64_t)ntotCurrentBuf + len) > 4LL * fBufferSizeMin) {
// Okay, so we have not even made one pass and we already have
// accumulated request for more than twice the memory size ...
// So stop for now, and will restart at the same point, hoping
// that the basket will still be in memory and not asked again ..
fEntryNext = maxReadEntry;

if (showMore || gDebug > 5) {
Info("FillBuffer", "Breaking early because %d is greater than 4*%d at cluster iteration "
Info("FillBuffer", "Breaking early because %lld is greater than 4*%d at cluster iteration "
"%d pass %d will restart at %lld",
(ntotCurrentBuf + len), fBufferSizeMin, clusterIterations, pass, fEntryNext);
((Long64_t)ntotCurrentBuf + len), fBufferSizeMin, clusterIterations, pass, fEntryNext);
}
filled = true;
break;
Expand All @@ -1571,12 +1571,12 @@ bool TTreeCache::FillBuffer()
// We have made one pass through the branches and thus already
// requested one basket per branch, let's stop prefetching
// now.
if ((ntotCurrentBuf + len) > 2 * fBufferSizeMin) {
if (((Long64_t)ntotCurrentBuf + len) > 2LL * fBufferSizeMin) {
fEntryNext = maxReadEntry;
if (showMore || gDebug > 5) {
Info("FillBuffer", "Breaking early because %d is greater than 2*%d at cluster iteration "
Info("FillBuffer", "Breaking early because %lld is greater than 2*%d at cluster iteration "
"%d pass %d will restart at %lld",
(ntotCurrentBuf + len), fBufferSizeMin, clusterIterations, pass, fEntryNext);
((Long64_t)ntotCurrentBuf + len), fBufferSizeMin, clusterIterations, pass, fEntryNext);
}
filled = true;
break;
Expand Down Expand Up @@ -1624,7 +1624,7 @@ bool TTreeCache::FillBuffer()
// Info("FillBuffer","maxCollectEntry incremented from %lld to %lld", maxReadEntry, entries[j+1]);
maxReadEntry = entries[j+1];
}
if (ntotCurrentBuf > 4 * fBufferSizeMin) {
if (ntotCurrentBuf > 4LL * fBufferSizeMin) {
// Humm something wrong happened.
Warning("FillBuffer", "There is more data in this cluster (starting at entry %lld to %lld, "
"current=%lld) than usual ... with %d %.3f%% of the branches we already have "
Expand Down Expand Up @@ -2064,12 +2064,13 @@ void TTreeCache::ResetCache()
/// Change the underlying buffer size of the cache.
/// If the change of size means some cache content is lost, or if the buffer
/// is now larger, setup for a cache refill the next time there is a read
/// Buffersize might be clamped, see TFileCacheRead::SetBufferSize
/// Returns:
/// - 0 if the buffer content is still available
/// - 1 if some or all of the buffer content has been made unavailable
/// - -1 on error

Int_t TTreeCache::SetBufferSize(Int_t buffersize)
Int_t TTreeCache::SetBufferSize(Long64_t buffersize)
{
Int_t prevsize = GetBufferSize();
Int_t res = TFileCacheRead::SetBufferSize(buffersize);
Expand Down
3 changes: 2 additions & 1 deletion tree/tree/src/TTreeCacheUnzip.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -351,12 +351,13 @@ bool TTreeCacheUnzip::FillBuffer()

////////////////////////////////////////////////////////////////////////////////
/// Change the underlying buffer size of the cache.
/// The buffersize might be clamped, see TFileCacheRead::SetBufferSize
/// Returns:
/// - 0 if the buffer content is still available
/// - 1 if some or all of the buffer content has been made unavailable
/// - -1 on error

Int_t TTreeCacheUnzip::SetBufferSize(Int_t buffersize)
Int_t TTreeCacheUnzip::SetBufferSize(Long64_t buffersize)
{
Int_t res = TTreeCache::SetBufferSize(buffersize);
if (res < 0) {
Expand Down
10 changes: 7 additions & 3 deletions tree/tree/src/TTreeCloner.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -644,12 +644,16 @@ void TTreeCloner::ImportClusterRanges()
}

////////////////////////////////////////////////////////////////////////////////
/// Set the TFile cache size to be used.
/// Set the cache size used by the matching TFile.
/// Note that the default is to use the same size as the default TTreeCache for
/// the input tree.
/// \param size Size of the cache. Zero disable the use of the cache.
/// \param size Size of the cache.
/// \note If size=0, or if it does not match the fileCache buffer size,
/// the fileCache will be deleted so that it be created later with the right size
/// (or not created if the size is 0) at the beginning of Exec.

void TTreeCloner::SetCacheSize(Int_t size)

void TTreeCloner::SetCacheSize(Long64_t size)
{
fCacheSize = size;
if (IsValid() && fFileCache) {
Expand Down

0 comments on commit 0f2bb47

Please sign in to comment.