From fe9a8a60cb0898eb8f845fcdc2acd1dd423199c1 Mon Sep 17 00:00:00 2001 From: Philippe Canal Date: Tue, 29 Sep 2020 18:41:24 -0500 Subject: [PATCH] TBranch::GetBulkEntries proper handling of already loaded basket This fixes Issue #6416 and Issue #6417 --- tree/tree/src/TBranch.cxx | 89 +++++++++++++++++++++++++++++---------- 1 file changed, 66 insertions(+), 23 deletions(-) diff --git a/tree/tree/src/TBranch.cxx b/tree/tree/src/TBranch.cxx index 02654e1c046f5..c69a3a95e3253 100644 --- a/tree/tree/src/TBranch.cxx +++ b/tree/tree/src/TBranch.cxx @@ -1322,6 +1322,10 @@ const char* TBranch::GetIconName() const /// the user_buffer will back the memory of the newly-constructed basket. /// /// Assumes that this branch is enabled. +/// +/// Returns -1 if the entry does not exist +/// Returns -2 in case of error +/// Returns the index of the basket in case of success. Int_t TBranch::GetBasketAndFirst(TBasket *&basket, Long64_t &first, TBuffer *user_buffer) { @@ -1332,9 +1336,10 @@ Int_t TBranch::GetBasketAndFirst(TBasket *&basket, Long64_t &first, // make sure basket buffers are in memory. basket = fCurrentBasket; first = fFirstBasketEntry; + return fReadBasket; } else { if ((entry < fFirstEntry) || (entry >= fEntryNumber)) { - return 0; + return -1; } first = fFirstBasketEntry; Long64_t last = fNextBasketEntry - 1; @@ -1344,7 +1349,7 @@ Int_t TBranch::GetBasketAndFirst(TBasket *&basket, Long64_t &first, if (fReadBasket < 0) { fNextBasketEntry = -1; Error("GetBasketAndFirst", "In the branch %s, no basket contains the entry %lld\n", GetName(), entry); - return -1; + return -2; } if (fReadBasket == fWriteBasket) { fNextBasketEntry = fEntryNumber; @@ -1363,7 +1368,7 @@ Int_t TBranch::GetBasketAndFirst(TBasket *&basket, Long64_t &first, fCurrentBasket = 0; fFirstBasketEntry = -1; fNextBasketEntry = -1; - return -1; + return -2; } if (fTree->GetClusterPrefetch()) { TTree::TClusterIterator clusterIterator = fTree->GetClusterIterator(entry); @@ -1377,18 +1382,18 @@ Int_t TBranch::GetBasketAndFirst(TBasket *&basket, Long64_t &first, // cause a reset of the first / next basket entries back to -1. fFirstBasketEntry = first; fNextBasketEntry = updatedNext; - } - if (user_buffer) { - // Disassociate basket from memory buffer for bulk IO - // When the user provides a memory buffer (i.e., for bulk IO), we should - // make sure to drop all references to that buffer in the TTree afterward. - fCurrentBasket = nullptr; - fBaskets[fReadBasket] = nullptr; + if (user_buffer) { + // Disassociate basket from memory buffer for bulk IO + // When the user provides a memory buffer (i.e., for bulk IO), we should + // make sure to drop all references to that buffer in the TTree afterward. + fCurrentBasket = nullptr; + fBaskets[fReadBasket] = nullptr; + } } else { fCurrentBasket = basket; } + return fReadBasket; } - return 1; } //////////////////////////////////////////////////////////////////////////////// @@ -1437,7 +1442,7 @@ Int_t TBranch::GetBulkEntries(Long64_t entry, TBuffer &user_buf) TBasket *basket = nullptr; Long64_t first; Int_t result = GetBasketAndFirst(basket, first, &user_buf); - if (R__unlikely(result <= 0)) return -1; + if (R__unlikely(result < 0)) return -1; // Only support reading from full clusters. if (R__unlikely(entry != first)) { //printf("Failed to read from full cluster; first entry is %ld; requested entry is %ld.\n", first, entry); @@ -1458,22 +1463,41 @@ Int_t TBranch::GetBulkEntries(Long64_t entry, TBuffer &user_buf) return -1; } + if (&user_buf != buf) { + // The basket was already in memory and might (and might not) be backed by persistent + // storage. + R__ASSERT(result == fReadBasket); + if (fBasketSeek[fReadBasket]) { + // It is backed, so we can be destructive + user_buf.SetBuffer(buf->Buffer(), buf->BufferSize()); + buf->ResetBit(TBufferIO::kIsOwner); + fCurrentBasket = nullptr; + fBaskets[fReadBasket] = nullptr; + } else { + // This is the only copy, we can't return it as is to the user, just make a copy. + if (user_buf.BufferSize() < buf->BufferSize()) { + user_buf.AutoExpand(buf->BufferSize()); + } + memcpy(user_buf.Buffer(), buf->Buffer(), buf->BufferSize()); + } + } + Int_t bufbegin = basket->GetKeylen(); - buf->SetBufferOffset(bufbegin); + user_buf.SetBufferOffset(bufbegin); Int_t N = ((fNextBasketEntry < 0) ? fEntryNumber : fNextBasketEntry) - first; //printf("Requesting %d events; fNextBasketEntry=%lld; first=%lld.\n", N, fNextBasketEntry, first); - if (R__unlikely(!leaf->ReadBasketFast(*buf, N))) { + if (R__unlikely(!leaf->ReadBasketFast(user_buf, N))) { Error("GetBulkEntries", "Leaf failed to read.\n"); return -1; } user_buf.SetBufferOffset(bufbegin); - fCurrentBasket = nullptr; - fBaskets[fReadBasket] = nullptr; - R__ASSERT(fExtraBasket == nullptr && "fExtraBasket should have been set to nullptr by GetFreshBasket"); - fExtraBasket = basket; - basket->DisownBuffer(); + if (fCurrentBasket == nullptr) { + R__ASSERT(fExtraBasket == nullptr && "fExtraBasket should have been set to nullptr by GetFreshBasket"); + fExtraBasket = basket; + basket->DisownBuffer(); + } return N; } @@ -1498,7 +1522,7 @@ Int_t TBranch::GetEntriesSerialized(Long64_t entry, TBuffer &user_buf, TBuffer * TBasket *basket = nullptr; Long64_t first; Int_t result = GetBasketAndFirst(basket, first, &user_buf); - if (R__unlikely(result <= 0)) { return -1; } + if (R__unlikely(result < 0)) { return -1; } // Only support reading from full clusters. if (R__unlikely(entry != first)) { Error("GetEntriesSerialized", "Failed to read from full cluster; first entry is %lld; requested entry is %lld.\n", first, entry); @@ -1519,13 +1543,32 @@ Int_t TBranch::GetEntriesSerialized(Long64_t entry, TBuffer &user_buf, TBuffer * return -1; } + if (&user_buf != buf) { + // The basket was already in memory and might (and might not) be backed by persistent + // storage. + R__ASSERT(result == fReadBasket); + if (fBasketSeek[fReadBasket]) { + // It is backed, so we can be destructive + user_buf.SetBuffer(buf->Buffer(), buf->BufferSize()); + buf->ResetBit(TBufferIO::kIsOwner); + fCurrentBasket = nullptr; + fBaskets[fReadBasket] = nullptr; + } else { + // This is the only copy, we can't return it as is to the user, just make a copy. + if (user_buf.BufferSize() < buf->BufferSize()) { + user_buf.AutoExpand(buf->BufferSize()); + } + memcpy(user_buf.Buffer(), buf->Buffer(), buf->BufferSize()); + } + } + Int_t bufbegin = basket->GetKeylen(); - buf->SetBufferOffset(bufbegin); + user_buf.SetBufferOffset(bufbegin); Int_t N = ((fNextBasketEntry < 0) ? fEntryNumber : fNextBasketEntry) - first; //Info("GetEntriesSerialized", "Requesting %d events; fNextBasketEntry=%lld; first=%lld.\n", N, fNextBasketEntry, first); - if (R__unlikely(!leaf->ReadBasketSerialized(*buf, N))) { + if (R__unlikely(!leaf->ReadBasketSerialized(user_buf, N))) { Error("GetEntriesSerialized", "Leaf failed to read.\n"); return -1; } @@ -1587,7 +1630,7 @@ Int_t TBranch::GetEntry(Long64_t entry, Int_t getall) Long64_t first; Int_t result = GetBasketAndFirst(basket, first, nullptr); - if (R__unlikely(result <= 0)) { return result; } + if (R__unlikely(result < 0)) { return result + 1; } basket->PrepareBasket(entry); TBuffer* buf = basket->GetBufferRef();