Common: TreeStream optional upper limit on file size

Signed-off-by: Felix Schlepper <[email protected]>

Author: f3sch

Common/Utils/include/CommonUtils/TreeStream.h

@@ -63,6 +63,8 @@ class TreeStream
   const char* getName() const { return mTree.GetName(); }
   void setID(int id) { mID = id; }
   int getID() const { return mID; }
+  void setMaxFileSize(Long64_t max) { mMaxFileSize = max; }
+  Long64_t getMaxFileSize() const { return mMaxFileSize; }
 
   TreeStream& operator<<(const Bool_t& b)
   {
@@ -177,6 +179,7 @@ class TreeStream
   int mNextNameCounter = 0;        ///< next name counter
   int mStatus = 0;                 ///< status of the layout
   TString mNextName;               ///< name for next entry
+  Long64_t mMaxFileSize{0};        ///< optional upper limit on file size
 
   ClassDefNV(TreeStream, 0);
 };

Common/Utils/include/CommonUtils/TreeStreamRedirector.h

@@ -44,21 +44,25 @@ class TreeStreamRedirector
   TreeStreamRedirector(const char* fname = "", const char* option = "recreate");
   virtual ~TreeStreamRedirector();
   void Close();
-  TFile* GetFile() { return mDirectory->GetFile(); }
   TDirectory* GetDirectory() { return mDirectory; }
-  virtual TreeStream& operator<<(Int_t id);
-  virtual TreeStream& operator<<(const char* name);
+  TFile* GetFile() { return mDirectory->GetFile(); }
+  Long64_t GetMaxFileSize() const noexcept { return mMaxFileSize; }
   void SetDirectory(TDirectory* sfile);
   void SetFile(TFile* sfile);
+  void SetMaxFileSize(Long64_t max) { mMaxFileSize = max; }
   static void FixLeafNameBug(TTree* tree);
 
+  virtual TreeStream& operator<<(Int_t id);
+  virtual TreeStream& operator<<(const char* name);
+
  private:
   TreeStreamRedirector(const TreeStreamRedirector& tsr);
   TreeStreamRedirector& operator=(const TreeStreamRedirector& tsr);
 
   std::unique_ptr<TDirectory> mOwnDirectory;             // own directory of the redirector
   TDirectory* mDirectory = nullptr;                      // output directory
   std::vector<std::unique_ptr<TreeStream>> mDataLayouts; // array of data layouts
+  Long64_t mMaxFileSize{0};                              // optional upper limit on file size
 
   ClassDefNV(TreeStreamRedirector, 0);
 };

Common/Utils/src/TreeStream.cxx

@@ -13,6 +13,7 @@
 
 #include "CommonUtils/TreeStream.h"
 #include <TBranch.h>
+#include <TFile.h>
 
 using namespace o2::utils;
 
@@ -108,6 +109,11 @@ void TreeStream::Fill()
 {
   // Fill the tree
 
+  // Stop filling if maximum file size is exceeded
+  if (mMaxFileSize != 0 && mTree.GetDirectory()->GetFile()->GetEND() > mMaxFileSize) {
+    return;
+  }
+
   int entries = mElements.size();
   if (entries > mTree.GetNbranches()) {
     BuildTree();

Common/Utils/src/TreeStreamRedirector.cxx

@@ -77,6 +77,7 @@ TreeStream& TreeStreamRedirector::operator<<(Int_t id)
   mDataLayouts.emplace_back(std::unique_ptr<TreeStream>(new TreeStream(Form("Tree%d", id))));
   auto layout = mDataLayouts.back().get();
   layout->setID(id);
+  layout->setMaxFileSize(mMaxFileSize);
   if (backup) {
     backup->cd();
   }
@@ -101,6 +102,7 @@ TreeStream& TreeStreamRedirector::operator<<(const char* name)
   mDataLayouts.emplace_back(std::unique_ptr<TreeStream>(new TreeStream(name)));
   auto layout = mDataLayouts.back().get();
   layout->setID(-1);
+  layout->setMaxFileSize(mMaxFileSize);
   if (backup) {
     backup->cd();
   }