Input/src/RootTree.h

0001 #ifndef IOPool_Input_RootTree_h
0002 #define IOPool_Input_RootTree_h
0003
0004 /*----------------------------------------------------------------------
0005
0006 RootTree.h // used by ROOT input sources
0007
0008 ----------------------------------------------------------------------*/
0009
0010 #include "DataFormats/Provenance/interface/ProductDescription.h"
0011 #include "DataFormats/Provenance/interface/BranchID.h"
0012 #include "DataFormats/Provenance/interface/IndexIntoFile.h"
0013 #include "FWCore/Framework/interface/Frameworkfwd.h"
0014 #include "FWCore/ServiceRegistry/interface/ServiceRegistryfwd.h"
0015 #include "FWCore/Utilities/interface/BranchType.h"
0016 #include "FWCore/Utilities/interface/InputType.h"
0017 #include "FWCore/Utilities/interface/Signal.h"
0018 #include "FWCore/Utilities/interface/thread_safety_macros.h"
0019
0020 #include "Rtypes.h"
0021 #include "TBranch.h"
0022
0023 #include <memory>
0024 #include <string>
0025 #include <vector>
0026 #include <unordered_set>
0027 #include <unordered_map>
0028
0029 class TClass;
0030 class TTree;
0031 class TTreeCache;
0032
0033 namespace edm {
0034   class RootDelayedReaderBase;
0035   class InputFile;
0036
0037   namespace roottree {
0038     unsigned int const defaultCacheSize = 20U * 1024 * 1024;
0039     unsigned int const defaultNonEventCacheSize = 1U * 1024 * 1024;
0040     unsigned int const defaultLearningEntries = 20U;
0041     unsigned int const defaultNonEventLearningEntries = 1U;
0042     using EntryNumber = IndexIntoFile::EntryNumber_t;
0043     struct BranchInfo {
0044       BranchInfo(ProductDescription const& prod)
0045           : productDescription_(prod), productBranch_(nullptr), classCache_(nullptr), offsetToWrapperBase_(0) {}
0046       ProductDescription const productDescription_;
0047       void setBranch(TBranch* branch, TClass const* wrapperBaseTClass);
0048       std::unique_ptr<WrapperBase> newWrapper() const;
0049       TBranch* productBranch_;
0050
0051     private:
0052       //All access to a ROOT file is serialized
0053       TClass* classCache_;
0054       Int_t offsetToWrapperBase_;
0055     };
0056
0057     class BranchMap {
0058     public:
0059       using Map = std::unordered_map<unsigned int, BranchInfo>;
0060
0061       void reserve(Map::size_type iSize) { map_.reserve(iSize); }
0062       void insert(edm::BranchID const& iKey, BranchInfo const& iInfo) { map_.emplace(iKey.id(), iInfo); }
0063       BranchInfo const* find(BranchID const& iKey) const { return find(iKey.id()); }
0064       BranchInfo const* find(unsigned int iKey) const {
0065         auto itFound = map_.find(iKey);
0066         if (itFound == map_.end()) {
0067           return nullptr;
0068         }
0069         return &itFound->second;
0070       }
0071
0072       using const_iterator = Map::const_iterator;
0073       const_iterator begin() const { return map_.cbegin(); }
0074       const_iterator end() const { return map_.cend(); }
0075       Map::size_type size() const { return map_.size(); }
0076
0077     private:
0078       Map map_;
0079     };
0080
0081     Int_t getEntry(TBranch* branch, EntryNumber entryNumber);
0082     Int_t getEntry(TTree* tree, EntryNumber entryNumber);
0083     std::unique_ptr<TTreeCache> trainCache(TTree* tree,
0084                                            InputFile& file,
0085                                            unsigned int cacheSize,
0086                                            char const* branchNames);
0087   }  // namespace roottree
0088
0089   class RootTree {
0090   public:
0091     using BranchMap = roottree::BranchMap;
0092     using EntryNumber = roottree::EntryNumber;
0093     struct Options {
0094       unsigned int treeCacheSize = 0U;
0095       int treeMaxVirtualSize;
0096       bool enablePrefetching;
0097       bool promptReading = false;
0098
0099       Options usingDefaultNonEventOptions() const {
0100         return {roottree::defaultNonEventCacheSize, treeMaxVirtualSize, enablePrefetching, false};
0101       }
0102     };
0103
0104     RootTree(std::shared_ptr<InputFile> filePtr,
0105              BranchType const& branchType,
0106              unsigned int nIndexes,
0107              Options const& options,
0108              unsigned int learningEntries,
0109              InputType inputType);
0110
0111     RootTree(std::shared_ptr<InputFile> filePtr,
0112              BranchType const& branchType,
0113              std::string const& processName,
0114              unsigned int nIndexes,
0115              Options const& options,
0116              unsigned int learningEntries,
0117              InputType inputType);
0118
0119     void init(std::string const& productTreeName, unsigned int maxVirtualSize, unsigned int cacheSize);
0120
0121     ~RootTree();
0122
0123     RootTree(RootTree const&) = delete;             // Disallow copying and moving
0124     RootTree& operator=(RootTree const&) = delete;  // Disallow copying and moving
0125
0126     bool isValid() const;
0127     void numberOfBranchesToAdd(BranchMap::Map::size_type iSize) { branches_.reserve(iSize); }
0128     void addBranch(ProductDescription const& prod, std::string const& oldBranchName);
0129     void dropBranch(std::string const& oldBranchName);
0130     void getEntry(TBranch* branch, EntryNumber entry) const;
0131     void getEntryForAllBranches() const;
0132     void setPresence(ProductDescription& prod, std::string const& oldBranchName);
0133
0134     bool next() { return ++entryNumber_ < entries_; }
0135     bool nextWithCache();
0136     bool current() const { return entryNumber_ < entries_ && entryNumber_ >= 0; }
0137     bool current(EntryNumber entry) const { return entry < entries_ && entry >= 0; }
0138     void rewind() { entryNumber_ = 0; }
0139     void rewindToInvalid() { entryNumber_ = IndexIntoFile::invalidEntry; }
0140     void close();
0141     bool skipEntries(unsigned int& offset);
0142     EntryNumber const& entryNumber() const { return entryNumber_; }
0143     EntryNumber const& entryNumberForIndex(unsigned int index) const;
0144     EntryNumber const& entries() const { return entries_; }
0145     void setEntryNumber(EntryNumber theEntryNumber);
0146     void insertEntryForIndex(unsigned int index);
0147     std::vector<std::string> const& branchNames() const { return branchNames_; }
0148     RootDelayedReaderBase* rootDelayedReader() const;
0149     DelayedReader* resetAndGetRootDelayedReader() const;
0150     template <typename T>
0151     void fillAux(T*& pAux) {
0152       auxBranch_->SetAddress(&pAux);
0153       auto cache = getAuxCache(auxBranch_);
0154       getEntryUsingCache(auxBranch_, entryNumber_, cache);
0155       auxBranch_->SetAddress(nullptr);
0156     }
0157
0158     template <typename T>
0159     void fillBranchEntry(TBranch* branch, T*& pbuf) {
0160       branch->SetAddress(&pbuf);
0161       getEntry(branch, entryNumber_);
0162       branch->SetAddress(nullptr);
0163     }
0164
0165     template <typename T>
0166     void fillBranchEntryMeta(TBranch* branch, EntryNumber entryNumber, T*& pbuf) {
0167       if (metaTree_ != nullptr) {
0168         // Metadata was in separate tree.  Not cached.
0169         branch->SetAddress(&pbuf);
0170         roottree::getEntry(branch, entryNumber);
0171       } else {
0172         fillBranchEntry<T>(branch, entryNumber, pbuf);
0173       }
0174     }
0175
0176     template <typename T>
0177     void fillBranchEntry(TBranch* branch, EntryNumber entryNumber, T*& pbuf) {
0178       branch->SetAddress(&pbuf);
0179       getEntry(branch, entryNumber);
0180     }
0181
0182     TTree const* tree() const { return tree_; }
0183     TTree* tree() { return tree_; }
0184     TTree const* metaTree() const { return metaTree_; }
0185     TTree* metaTree() { return metaTree_; }
0186     BranchMap const& branches() const;
0187
0188     //For backwards compatibility
0189     TBranch* branchEntryInfoBranch() const { return branchEntryInfoBranch_; }
0190
0191     inline TTreeCache* checkTriggerCache(TBranch* branch, EntryNumber entryNumber) const;
0192     TTreeCache* checkTriggerCacheImpl(TBranch* branch, EntryNumber entryNumber) const;
0193     inline TTreeCache* selectCache(TBranch* branch, EntryNumber entryNumber) const;
0194     void trainCache(char const* branchNames);
0195     void resetTraining() { trainNow_ = true; }
0196
0197     BranchType branchType() const { return branchType_; }
0198     std::string const& processName() const { return processName_; }
0199
0200     void setSignals(
0201         signalslot::Signal<void(StreamContext const&, ModuleCallingContext const&)> const* preEventReadSource,
0202         signalslot::Signal<void(StreamContext const&, ModuleCallingContext const&)> const* postEventReadSource);
0203
0204   private:
0205     RootTree(std::shared_ptr<InputFile> filePtr,
0206              BranchType const& branchType,
0207              unsigned int nIndexes,
0208              unsigned int learningEntries,
0209              bool enablePrefetching,
0210              bool promptRead,
0211              InputType inputType);
0212
0213     std::shared_ptr<TTreeCache> createCacheWithSize(unsigned int cacheSize) const;
0214     void setCacheSize(unsigned int cacheSize);
0215     void setTreeMaxVirtualSize(int treeMaxVirtualSize);
0216     void startTraining();
0217     void stopTraining();
0218     void getEntryUsingCache(TBranch* branch, EntryNumber entry, TTreeCache*) const;
0219     TTreeCache* getAuxCache(TBranch* auxBranch) const;
0220
0221     std::shared_ptr<InputFile> filePtr_;
0222     // We use bare pointers for pointers to some ROOT entities.
0223     // Root owns them and uses bare pointers internally.
0224     // Therefore,using smart pointers here will do no good.
0225     TTree* tree_ = nullptr;
0226     TTree* metaTree_ = nullptr;
0227     BranchType branchType_;
0228     std::string processName_;
0229     TBranch* auxBranch_ = nullptr;
0230     // We use a smart pointer to own the TTreeCache.
0231     // Unfortunately, ROOT owns it when attached to a TFile, but not after it is detached.
0232     // So, we make sure to it is detached before closing the TFile so there is no double delete.
0233     std::shared_ptr<TTreeCache> treeCache_;
0234     std::shared_ptr<TTreeCache> rawTreeCache_;
0235     CMS_SA_ALLOW mutable std::shared_ptr<TTreeCache> auxCache_;
0236     //All access to a ROOT file is serialized
0237     CMS_SA_ALLOW mutable std::shared_ptr<TTreeCache> triggerTreeCache_;
0238     CMS_SA_ALLOW mutable std::shared_ptr<TTreeCache> rawTriggerTreeCache_;
0239     CMS_SA_ALLOW mutable std::unordered_set<TBranch*> trainedSet_;
0240     CMS_SA_ALLOW mutable std::unordered_set<TBranch*> triggerSet_;
0241     EntryNumber entries_ = 0;
0242     EntryNumber entryNumber_ = IndexIntoFile::invalidEntry;
0243     std::unique_ptr<std::vector<EntryNumber> > entryNumberForIndex_;
0244     std::vector<std::string> branchNames_;
0245     BranchMap branches_;
0246     bool trainNow_ = false;
0247     EntryNumber switchOverEntry_ = -1;
0248     CMS_SA_ALLOW mutable EntryNumber rawTriggerSwitchOverEntry_ = -1;
0249     CMS_SA_ALLOW mutable bool performedSwitchOver_ = false;
0250     unsigned int learningEntries_;
0251     unsigned int cacheSize_ = 0;
0252     unsigned long treeAutoFlush_ = 0;
0253     // Enable asynchronous I/O in ROOT (done in a separate thread).  Only takes
0254     // effect on the primary treeCache_; all other caches have this explicitly disabled.
0255     bool enablePrefetching_;
0256     bool enableTriggerCache_;
0257     bool promptRead_;
0258     std::unique_ptr<RootDelayedReaderBase> rootDelayedReader_;
0259
0260     TBranch* branchEntryInfoBranch_ = nullptr;  //backwards compatibility
0261     // below for backward compatibility
0262     TTree* infoTree_ = nullptr;  // backward compatibility
0263   };
0264 }  // namespace edm
0265 #endif