MCKinematicsReader: load and release the track references per event

The track references of a source were loaded for all events at the first access and never released. They are
now loaded per event on demand (with the baskets dropped at the end of a TTree cluster, as for the tracks) and
releaseTracksForSourceAndEvent frees them together with the tracks. For a PbPb TF (10 signal events, 1.75M
references): +304 MB held from the first access until the end before, now +135 MB at peak and +55 MB after
releasing all events.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: matthias-kleiner

Steer/include/Steer/MCKinematicsReader.h

@@ -87,7 +87,7 @@ class MCKinematicsReader
   /// variant returning all tracks for source and event at once
   std::vector<MCTrack> const& getTracks(int source, int event) const;
 
-  /// API to ask releasing tracks (freeing memory) for source + event
+  /// API to ask releasing tracks and track references (freeing memory) for source + event
   void releaseTracksForSourceAndEvent(int source, int event);
 
   /// variant returning all tracks for an event id (source = 0) at once
@@ -129,7 +129,8 @@ class MCKinematicsReader
   void initTracksForSource(int source) const;
   void loadTracksForSourceAndEvent(int source, int eventID) const;
   void loadHeadersForSource(int source) const;
-  void loadTrackRefsForSource(int source) const;
+  void initTrackRefsForSource(int source) const;
+  void loadTrackRefsForSourceAndEvent(int source, int event) const;
   void initIndexedTrackRefs(std::vector<o2::TrackReference>& refs, o2::dataformats::MCTruthContainer<o2::TrackReference>& indexedrefs) const;
 
   DigitizationContext const* mDigitizationContext = nullptr;
@@ -142,6 +143,7 @@ class MCKinematicsReader
   mutable std::vector<std::vector<std::vector<o2::MCTrack>*>> mTracks;                                       // the in-memory track container
   mutable std::vector<std::vector<o2::dataformats::MCEventHeader>> mHeaders;                                 // the in-memory header container
   mutable std::vector<std::vector<o2::dataformats::MCTruthContainer<o2::TrackReference>>> mIndexedTrackRefs; // the in-memory track ref container
+  mutable std::vector<std::vector<bool>> mTrackRefsLoaded;                                                   // whether the track refs of a source/event are in memory
 
   bool mInitialized = false; // whether initialized
 };
@@ -206,23 +208,29 @@ inline gsl::span<o2::TrackReference> MCKinematicsReader::getTrackRefs(int source
   }
   auto& perEvent = mIndexedTrackRefs[source];
   if (perEvent.size() == 0) {
-    loadTrackRefsForSource(source);
+    initTrackRefsForSource(source);
   }
   if (static_cast<size_t>(event) >= perEvent.size()) {
     return {};
   }
+  if (!mTrackRefsLoaded[source][event]) {
+    loadTrackRefsForSourceAndEvent(source, event);
+  }
   return perEvent[event].getLabels(track);
 }
 
 inline const std::vector<o2::TrackReference>& MCKinematicsReader::getTrackRefsByEvent(int source, int event) const
 {
   auto const& perEvent = mIndexedTrackRefs.at(source);
   if (perEvent.size() == 0) {
-    loadTrackRefsForSource(source);
+    initTrackRefsForSource(source);
   }
   if (static_cast<size_t>(event) >= perEvent.size()) {
     reportMissingEvent("events of track references", source, event, perEvent.size());
   }
+  if (!mTrackRefsLoaded[source][event]) {
+    loadTrackRefsForSourceAndEvent(source, event);
+  }
   return perEvent[event].getTruthArray();
 }
 

Steer/src/MCKinematicsReader.cxx

@@ -111,6 +111,11 @@ void MCKinematicsReader::releaseTracksForSourceAndEvent(int source, int eventID)
     delete mTracks[source][eventID];
     mTracks[source][eventID] = nullptr;
   }
+  // the track references of this event as well (reloaded on demand)
+  if (static_cast<size_t>(eventID) < mTrackRefsLoaded.at(source).size() && mTrackRefsLoaded[source][eventID]) {
+    mIndexedTrackRefs[source][eventID] = o2::dataformats::MCTruthContainer<o2::TrackReference>();
+    mTrackRefsLoaded[source][eventID] = false;
+  }
 }
 
 void MCKinematicsReader::loadHeadersForSource(int source) const
@@ -135,31 +140,43 @@ void MCKinematicsReader::loadHeadersForSource(int source) const
   }
 }
 
-void MCKinematicsReader::loadTrackRefsForSource(int source) const
+void MCKinematicsReader::initTrackRefsForSource(int source) const
 {
   auto chain = mInputChains[source];
   if (chain) {
     // todo: get name from NameConfig
     auto br = chain->GetBranch("TrackRefs");
     if (br) {
-      std::vector<o2::TrackReference>* refs = nullptr;
-      br->SetAddress(&refs);
       mIndexedTrackRefs[source].resize(br->GetEntries());
-      for (int event = 0; event < br->GetEntries(); ++event) {
-        br->GetEntry(event);
-        if (refs) {
-          // we convert the original flat vector into an indexed structure
-          initIndexedTrackRefs(*refs, mIndexedTrackRefs[source][event]);
-          delete refs;
-          refs = nullptr;
-        }
-      }
+      mTrackRefsLoaded[source].assign(br->GetEntries(), false);
     } else {
       LOG(warn) << "TrackRefs branch not found";
     }
   }
 }
 
+void MCKinematicsReader::loadTrackRefsForSourceAndEvent(int source, int event) const
+{
+  // todo: get name from NameConfig
+  auto br = mInputChains[source]->GetBranch("TrackRefs");
+  std::vector<o2::TrackReference>* refs = nullptr; // allocated by ROOT, owned by us
+  br->SetAddress(&refs);
+  br->GetEntry(event);
+  if (refs) {
+    // we convert the original flat vector into an indexed structure
+    initIndexedTrackRefs(*refs, mIndexedTrackRefs[source][event]);
+    delete refs;
+  }
+  br->ResetAddress();
+  // free the decompressed baskets if no later entry reads them, i.e. at the end of the cluster of this event
+  auto clusterIt = br->GetTree()->GetClusterIterator(event);
+  clusterIt.Next();
+  if (event + 1 >= clusterIt.GetNextEntry()) {
+    br->DropBaskets("all");
+  }
+  mTrackRefsLoaded[source][event] = true;
+}
+
 bool MCKinematicsReader::initFromDigitContext(o2::steer::DigitizationContext const* context)
 {
   if (mInitialized) {
@@ -177,6 +194,7 @@ bool MCKinematicsReader::initFromDigitContext(o2::steer::DigitizationContext con
   mTracks.resize(mInputChains.size());
   mHeaders.resize(mInputChains.size());
   mIndexedTrackRefs.resize(mInputChains.size());
+  mTrackRefsLoaded.resize(mInputChains.size());
 
   // actual loading will be done only if someone asks
   // the first time for a particular source ...
@@ -210,6 +228,7 @@ bool MCKinematicsReader::initFromKinematics(std::string_view name)
   mTracks.resize(1);
   mHeaders.resize(1);
   mIndexedTrackRefs.resize(1);
+  mTrackRefsLoaded.resize(1);
   mInitialized = true;
 
   return true;