Optimise basket sizes created by RootCnvSvc

See merge request gaudi/Gaudi!1381

GaudiExamples/tests/qmtest/gaudiexamples.qms/return_codes.qms/corrupted_input.qmt

@@ -22,10 +22,17 @@ from Gaudi.Configuration import *
 import shutil
 shutil.copy('ROOTIO.dst', 'ROOTIO_corrupted.dst')
 
-f = open('ROOTIO_corrupted.dst', 'rb+')
-f.seek(1024)
-f.write(b'corruption!')
-f.close()
+# find a specific byte sequence to determine where to corrupt the file
+magic_bytes = b"\xa6\xbc\x8c\xe9"
+with open('ROOTIO_corrupted.dst', 'rb') as f:
+    data = f.read()
+    if data.count(magic_bytes) != 1:
+        raise NotImplementedError("The test file is not as expected")
+    seek_index = data.index(magic_bytes) - 4
+
+with open('ROOTIO_corrupted.dst', 'rb+') as f:
+    f.seek(seek_index)
+    f.write(b'corruption!')
 
 esel = EventSelector()
 esel.Input = ["DATAFILE='PFN:ROOTIO_corrupted.dst' SVC='Gaudi::RootEvtSelector' OPT='READ'"]

RootCnv/include/RootCnv/RootCnvSvc.h

@@ -70,12 +70,12 @@ namespace Gaudi {
                                                "Records name to fire incident for file records" };
 
     // ROOT Write parameters
-    Gaudi::Property<int> m_autoFlush{ this, "AutoFlush", 100,
-                                      "AutoFlush parameter for ROOT TTree (Number of events between auto flushes)" };
-    Gaudi::Property<int> m_basketSize{ this, "BasketSize", 2 * 1024 * 1024 /*MBYTE*/,
-                                       "Basket optimization parameter for ROOT TTree (total basket size)" };
-    Gaudi::Property<int> m_bufferSize{ this, "BufferSize", 2 * 1024 /*kBYTE*/,
-                                       "Buffer size optimization parameter for ROOT TTree" };
+    Gaudi::Property<int> m_minBufferSize{ this, "MinBufferSize", 2 * 1024 /*kBYTE*/,
+                                          "Minimum buffer size to use for writing TTree baskets" };
+    Gaudi::Property<int> m_maxBufferSize{ this, "MaxBufferSize", 100 * 1024 * 1024 /*MBYTE*/,
+                                          "Maximum buffer size to use for writing TTree baskets" };
+    Gaudi::Property<int> m_approxEventsPerBasket{ this, "ApproxEventsPerBasket", 1000,
+                                                  "Resize TBasket buffers to fit approximately this many events" };
     Gaudi::Property<int> m_splitLevel{ this, "SplitLevel", 0, "Split level optimization parameter for ROOT TTree" };
     Gaudi::Property<std::string> m_compression{ this, "GlobalCompression", "",
                                                 "Compression-algorithm:compression-level,  empty: do nothing" };

RootCnv/include/RootCnv/RootDataConnection.h

@@ -288,10 +288,12 @@ namespace Gaudi {
 
     /// Save object of a given class to section and container
     std::pair<int, unsigned long> saveObj( std::string_view section, std::string_view cnt, TClass* cl, DataObject* pObj,
-                                           int buff_siz, int split_lvl, bool fill_missing = false );
+                                           int minBufferSize, int maxBufferSize, int approxEventsPerBasket,
+                                           int split_lvl, bool fill_missing = false );
     /// Save object of a given class to section and container
     std::pair<int, unsigned long> save( std::string_view section, std::string_view cnt, TClass* cl, void* pObj,
-                                        int buff_siz, int split_lvl, bool fill_missing = false );
+                                        int minBufferSize, int maxBufferSize, int approxEventsPerBasket, int split_lvl,
+                                        bool fill_missing = false );
 
     /// Open data stream in read mode
     StatusCode connectRead() override;

RootCnv/src/RootCnvSvc.cpp

@@ -329,15 +329,6 @@ StatusCode RootCnvSvc::commitOutput( CSTR dsn, bool /* doCommit */ ) {
       }
 
       b->GetTree()->SetEntries( evt );
-      if ( evt == 1 ) { b->GetTree()->OptimizeBaskets( m_basketSize, 1.1, "" ); }
-      if ( evt > 0 && ( evt % m_autoFlush ) == 0 ) {
-        if ( evt == m_autoFlush ) {
-          b->GetTree()->SetAutoFlush( m_autoFlush );
-          b->GetTree()->OptimizeBaskets( m_basketSize, 1., "" );
-        } else {
-          b->GetTree()->FlushBaskets();
-        }
-      }
       if ( log().level() <= MSG::DEBUG )
         log() << MSG::DEBUG << "Set section entries of " << m_currSection << " to " << long( evt ) << " entries."
               << endmsg;
@@ -365,7 +356,8 @@ StatusCode RootCnvSvc::createAddress( long typ, const CLID& clid, const string*
 StatusCode RootCnvSvc::createNullRep( const std::string& path ) {
   size_t len     = path.find( '/', 1 );
   string section = path.substr( 1, len == string::npos ? string::npos : len - 1 );
-  m_current->saveObj( section, path, nullptr, nullptr, m_bufferSize, m_splitLevel );
+  m_current->saveObj( section, path, nullptr, nullptr, m_minBufferSize, m_maxBufferSize, m_approxEventsPerBasket,
+                      m_splitLevel );
   return S_OK;
 }
 
@@ -374,7 +366,8 @@ StatusCode RootCnvSvc::createNullRef( const std::string& path ) {
   RootObjectRefs*          refs    = nullptr;
   size_t                   len     = path.find( '/', 1 );
   string                   section = path.substr( 1, len == string::npos ? string::npos : len - 1 );
-  pair<int, unsigned long> ret     = m_current->save( section, path + "#R", nullptr, refs, m_bufferSize, m_splitLevel );
+  pair<int, unsigned long> ret = m_current->save( section, path + "#R", nullptr, refs, m_minBufferSize, m_maxBufferSize,
+                                                  m_approxEventsPerBasket, m_splitLevel );
   if ( log().level() <= MSG::VERBOSE )
     log() << MSG::VERBOSE << "Writing object:" << path << " " << ret.first << " " << hex << ret.second << dec
           << " [NULL]" << endmsg;
@@ -391,7 +384,8 @@ StatusCode RootCnvSvc::i__createRep( DataObject* pObj, IOpaqueAddress*& refpAddr
   TClass*                  cl   = ( clid == CLID_DataObject ) ? m_classDO : getClass( pObj );
   size_t                   len  = p[1].find( '/', 1 );
   string                   sect = p[1].substr( 1, len == string::npos ? string::npos : len - 1 );
-  pair<int, unsigned long> ret  = m_current->saveObj( sect, p[1], cl, pObj, m_bufferSize, m_splitLevel, true );
+  pair<int, unsigned long> ret  = m_current->saveObj( sect, p[1], cl, pObj, m_minBufferSize, m_maxBufferSize,
+                                                      m_approxEventsPerBasket, m_splitLevel, true );
   if ( ret.first > 1 || ( clid == CLID_DataObject && ret.first == 1 ) ) {
     unsigned long ip[2] = { 0, ret.second };
     if ( m_currSection.empty() ) m_currSection = p[1];
@@ -428,8 +422,8 @@ StatusCode RootCnvSvc::i__fillRepRefs( IOpaqueAddress* /* pA */, DataObject* pOb
           int                      link_id = m_current->makeLink( lnk->path() );
           refs.links.push_back( link_id );
         }
-        pair<int, unsigned long> ret =
-            m_current->save( sect, id + "#R", m_classRefs, &refs, m_bufferSize, m_splitLevel, true );
+        pair<int, unsigned long> ret = m_current->save( sect, id + "#R", m_classRefs, &refs, m_minBufferSize,
+                                                        m_maxBufferSize, m_approxEventsPerBasket, m_splitLevel, true );
         if ( ret.first > 1 ) {
           if ( log().level() <= MSG::DEBUG )
             log() << MSG::DEBUG << "Writing object:" << id << " " << ret.first << " " << hex << ret.second << dec

RootCnv/src/RootDataConnection.cpp

@@ -31,6 +31,7 @@
 #include "TClass.h"
 #include "TFile.h"
 #include "TLeaf.h"
+#include "TMemFile.h"
 #include "TROOT.h"
 #include "TTree.h"
 #if ROOT_VERSION_CODE >= ROOT_VERSION( 5, 33, 0 )
@@ -41,6 +42,7 @@ static int s_compressionLevel = 1;
 #endif
 
 // C/C++ include files
+#include <limits>
 #include <numeric>
 #include <stdexcept>
 
@@ -474,23 +476,27 @@ CSTR RootDataConnection::empty() const { return s_empty; }
 
 /// Save object of a given class to section and container
 pair<int, unsigned long> RootDataConnection::saveObj( std::string_view section, std::string_view cnt, TClass* cl,
-                                                      DataObject* pObj, int buff_siz, int split_lvl, bool fill ) {
+                                                      DataObject* pObj, int minBufferSize, int maxBufferSize,
+                                                      int approxEventsPerBasket, int split_lvl, bool fill ) {
   DataObjectPush push( pObj );
-  return save( section, cnt, cl, pObj, buff_siz, split_lvl, fill );
+  return save( section, cnt, cl, pObj, minBufferSize, maxBufferSize, approxEventsPerBasket, split_lvl, fill );
 }
 
 /// Save object of a given class to section and container
 pair<int, unsigned long> RootDataConnection::save( std::string_view section, std::string_view cnt, TClass* cl,
-                                                   void* pObj, int buff_siz, int split_lvl, bool fill_missing ) {
+                                                   void* pObj, int minBufferSize, int maxBufferSize,
+                                                   int approxEventsPerBasket, int split_lvl, bool fill_missing ) {
   split_lvl  = 0;
-  TBranch* b = getBranch( section, cnt, cl, pObj ? &pObj : nullptr, buff_siz, split_lvl );
+  TBranch* b = getBranch( section, cnt, cl, pObj ? &pObj : nullptr, minBufferSize, split_lvl );
   if ( b ) {
     Long64_t evt = b->GetEntries();
     // msgSvc() << MSG::DEBUG << cnt.c_str() << " Obj:" << (void*)pObj
-    //         << " Split:" << split_lvl << " Buffer size:" << buff_siz << endl;
+    //         << " Split:" << split_lvl << " Buffer size:" << minBufferSize << endl;
+    bool set_buffer_size = ( evt == 0 );
     if ( fill_missing ) {
       Long64_t num, nevt = b->GetTree()->GetEntries();
       if ( nevt > evt ) {
+        set_buffer_size = true;
         b->SetAddress( nullptr );
         num = nevt - evt;
         while ( num > 0 ) {
@@ -502,6 +508,20 @@ pair<int, unsigned long> RootDataConnection::save( std::string_view section, std
         evt = b->GetEntries();
       }
     }
+    if ( set_buffer_size ) {
+      auto     dummy_file = make_unique<TMemFile>( "dummy.root", "CREATE" );
+      auto     dummy_tree = make_unique<TTree>( "DummyTree", "DummyTree", split_lvl, dummy_file->GetDirectory( "/" ) );
+      TBranch* dummy_branch = dummy_tree->Branch( "DummyBranch", cl->GetName(), &pObj, minBufferSize, split_lvl );
+      Int_t    nWritten     = dummy_branch->Fill();
+      if ( nWritten < 0 ) return { nWritten, evt };
+      Int_t newBasketSize = nWritten * approxEventsPerBasket;
+      // Ensure that newBasketSize doesn't wrap around
+      if ( std::numeric_limits<Int_t>::max() / approxEventsPerBasket < nWritten ) {
+        newBasketSize = std::numeric_limits<Int_t>::max();
+      }
+      b->SetBasketSize( std::min( maxBufferSize, std::max( minBufferSize, newBasketSize ) ) );
+      msgSvc() << MSG::DEBUG << "Setting basket size to " << newBasketSize << " for " << cnt << endmsg;
+    }
     b->SetAddress( &pObj );
     return { b->Fill(), evt };
   }