DPL: properly reserve memory when writing (#15510)

Author: ktf

Detectors/AOD/src/AODProducerWorkflowSpec.cxx

@@ -602,20 +602,20 @@ void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
     int end = start + trackRef.getEntriesOfSource(src);
     int nToReserve = end - start; // + last index for a given table
     if (src == GIndex::Source::MFT) {
-      mftTracksCursor.reserve(nToReserve + mftTracksCursor.lastIndex());
+      mftTracksCursor.reserve(nToReserve + mftTracksCursor.lastIndex() + 1);
       if (mStoreAllMFTCov) {
-        mftTracksCovCursor.reserve(nToReserve + mftTracksCovCursor.lastIndex());
+        mftTracksCovCursor.reserve(nToReserve + mftTracksCovCursor.lastIndex() + 1);
       }
     } else if (src == GIndex::Source::MCH || src == GIndex::Source::MFTMCH || src == GIndex::Source::MCHMID) {
-      fwdTracksCursor.reserve(nToReserve + fwdTracksCursor.lastIndex());
-      fwdTracksCovCursor.reserve(nToReserve + fwdTracksCovCursor.lastIndex());
+      fwdTracksCursor.reserve(nToReserve + fwdTracksCursor.lastIndex() + 1);
+      fwdTracksCovCursor.reserve(nToReserve + fwdTracksCovCursor.lastIndex() + 1);
       if (!mStoreAllMFTCov && src == GIndex::Source::MFTMCH) {
-        mftTracksCovCursor.reserve(nToReserve + mftTracksCovCursor.lastIndex());
+        mftTracksCovCursor.reserve(nToReserve + mftTracksCovCursor.lastIndex() + 1);
       }
     } else {
-      tracksCursor.reserve(nToReserve + tracksCursor.lastIndex());
-      tracksCovCursor.reserve(nToReserve + tracksCovCursor.lastIndex());
-      tracksExtraCursor.reserve(nToReserve + tracksExtraCursor.lastIndex());
+      tracksCursor.reserve(nToReserve + tracksCursor.lastIndex() + 1);
+      tracksCovCursor.reserve(nToReserve + tracksCovCursor.lastIndex() + 1);
+      tracksExtraCursor.reserve(nToReserve + tracksExtraCursor.lastIndex() + 1);
     }
     for (int ti = start; ti < end; ti++) {
       const auto& trackIndex = GIndices[ti];
@@ -715,9 +715,9 @@ void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
   }
   /// Add strangeness tracks to the table
   auto sTracks = data.getStrangeTracks();
-  tracksCursor.reserve(mVertexStrLUT[collisionID + 1] + tracksCursor.lastIndex());
-  tracksCovCursor.reserve(mVertexStrLUT[collisionID + 1] + tracksCovCursor.lastIndex());
-  tracksExtraCursor.reserve(mVertexStrLUT[collisionID + 1] + tracksExtraCursor.lastIndex());
+  tracksCursor.reserve(mVertexStrLUT[collisionID + 1] + tracksCursor.lastIndex() + 1);
+  tracksCovCursor.reserve(mVertexStrLUT[collisionID + 1] + tracksCovCursor.lastIndex() + 1);
+  tracksExtraCursor.reserve(mVertexStrLUT[collisionID + 1] + tracksExtraCursor.lastIndex() + 1);
   for (int iS{mVertexStrLUT[collisionID]}; iS < mVertexStrLUT[collisionID + 1]; ++iS) {
     auto& collStrTrk = mCollisionStrTrk[iS];
     auto& sTrk = sTracks[collStrTrk.second];
@@ -1236,9 +1236,9 @@ void AODProducerWorkflowDPL::fillMCTrackLabelsTable(MCTrackLabelCursorType& mcTr
   for (int src = GIndex::NSources; src--;) {
     int start = trackRef.getFirstEntryOfSource(src);
     int end = start + trackRef.getEntriesOfSource(src);
-    mcMFTTrackLabelCursor.reserve(end - start + mcMFTTrackLabelCursor.lastIndex());
-    mcFwdTrackLabelCursor.reserve(end - start + mcFwdTrackLabelCursor.lastIndex());
-    mcTrackLabelCursor.reserve(end - start + mcTrackLabelCursor.lastIndex());
+    mcMFTTrackLabelCursor.reserve(end - start + mcMFTTrackLabelCursor.lastIndex() + 1);
+    mcFwdTrackLabelCursor.reserve(end - start + mcFwdTrackLabelCursor.lastIndex() + 1);
+    mcTrackLabelCursor.reserve(end - start + mcTrackLabelCursor.lastIndex() + 1);
     for (int ti = start; ti < end; ti++) {
       const auto trackIndex = primVerGIs[ti];
 
@@ -1320,7 +1320,7 @@ void AODProducerWorkflowDPL::fillMCTrackLabelsTable(MCTrackLabelCursorType& mcTr
   auto sTrackLabels = data.getStrangeTracksMCLabels();
   // check if vertexId and vertexId + 1 maps into mVertexStrLUT
   if (!(vertexId < 0 || vertexId >= mVertexStrLUT.size() - 1)) {
-    mcTrackLabelCursor.reserve(mVertexStrLUT[vertexId + 1] + mcTrackLabelCursor.lastIndex());
+    mcTrackLabelCursor.reserve(mVertexStrLUT[vertexId + 1] + mcTrackLabelCursor.lastIndex() + 1);
     for (int iS{mVertexStrLUT[vertexId]}; iS < mVertexStrLUT[vertexId + 1]; ++iS) {
       auto& collStrTrk = mCollisionStrTrk[iS];
       auto& label = sTrackLabels[collStrTrk.second];
@@ -1448,9 +1448,9 @@ void AODProducerWorkflowDPL::addClustersToFwdTrkClsTable(const o2::globaltrackin
 
   if (mchTrackID > -1 && mchTrackID < mchTracks.size()) {
     const auto& mchTrack = mchTracks[mchTrackID];
-    fwdTrkClsCursor.reserve(mchTrack.getNClusters() + fwdTrkClsCursor.lastIndex());
     int first = mchTrack.getFirstClusterIdx();
     int last = mchTrack.getLastClusterIdx();
+    fwdTrkClsCursor.reserve(last - first + 1 + fwdTrkClsCursor.lastIndex() + 1);
     for (int i = first; i <= last; i++) {
       const auto& cluster = mchClusters[i];
       fwdTrkClsCursor(fwdTrackId,
@@ -1678,10 +1678,10 @@ void AODProducerWorkflowDPL::addToCaloTable(TCaloHandler& caloHandler, TCaloCurs
   auto inputEvent = caloHandler.buildEvent(eventID);
   auto cellsInEvent = inputEvent.mCells;        // get cells belonging to current event
   auto cellMClabels = inputEvent.mMCCellLabels; // get MC labels belonging to current event (only implemented for EMCal currently!)
-  caloCellCursor.reserve(cellsInEvent.size() + caloCellCursor.lastIndex());
-  caloTRGCursor.reserve(cellsInEvent.size() + caloTRGCursor.lastIndex());
+  caloCellCursor.reserve(cellsInEvent.size() + caloCellCursor.lastIndex() + 1);
+  caloTRGCursor.reserve(cellsInEvent.size() + caloTRGCursor.lastIndex() + 1);
   if (mUseMC) {
-    mcCaloCellLabelCursor.reserve(cellsInEvent.size() + mcCaloCellLabelCursor.lastIndex());
+    mcCaloCellLabelCursor.reserve(cellsInEvent.size() + mcCaloCellLabelCursor.lastIndex() + 1);
   }
   for (auto iCell = 0U; iCell < cellsInEvent.size(); iCell++) {
     caloCellCursor(bcID,

Framework/Core/test/test_ASoA.cxx

@@ -1376,3 +1376,49 @@ TEST_CASE("TestCombinedGetter")
     ++count;
   }
 }
+
+TEST_CASE("TestWritingCursorLastIndexAndReserve")
+{
+  // Nails down the WritingCursor semantics the AOD-producer reserves depend on:
+  // lastIndex() returns the *last index* (rows - 1), not the row count, and
+  // reserve(newRows + lastIndex() + 1) reserves exactly the post-batch total so a
+  // fully-filled, no-skip batch neither overruns (the fwdTrkCls crash) nor trips
+  // the release() / per-row UnsafeAppend guard.
+  Produces<o2::aod::Points> cursor; // Points has two persistent columns: X, Y
+  auto* builder = new TableBuilder();
+  cursor.resetCursor(LifetimeHolder<TableBuilder>(builder));
+
+  // Empty cursor: no row written, so the last index is -1 and rows == lastIndex()+1 == 0.
+  REQUIRE(cursor.lastIndex() == -1);
+
+  // operator() increments before the append, but only to the index of the row it
+  // writes: after N writes lastIndex() == N - 1, NOT N.
+  cursor(10, 20);
+  REQUIRE(cursor.lastIndex() == 0);
+  cursor(11, 21);
+  REQUIRE(cursor.lastIndex() == 1);
+  cursor(12, 22);
+  REQUIRE(cursor.lastIndex() == 2);
+  REQUIRE(cursor.lastIndex() + 1 == 3); // rows-so-far == last index + 1
+
+  // Reserve a second batch the correct way: total = newRows + rowsSoFar
+  //                                               = newRows + (lastIndex() + 1).
+  // The (buggy) newRows + lastIndex() would reserve 4 here and under-reserve the
+  // 5th row; the + 1 makes it exactly 5.
+  int64_t const newRows = 2;
+  int64_t const reserved = newRows + cursor.lastIndex() + 1; // correct total -> reserve(5)
+  cursor.reserve(reserved);
+  cursor(13, 23); // row index 3
+  cursor(14, 24); // row index 4 — fills the batch exactly (5 rows total)
+  REQUIRE(cursor.lastIndex() == 4);
+
+  // The contract release() enforces: rows filled (lastIndex()+1) must not exceed
+  // what was reserved. Correct (+1) gives reserved == 5 -> 5 <= 5 (green); the buggy
+  // newRows + lastIndex() reserves only 4 -> 5 <= 4 fails (red).
+  REQUIRE(cursor.lastIndex() + 1 <= reserved);
+
+  auto table = builder->finalize();
+  REQUIRE(table->num_rows() == 5);
+  REQUIRE(table->num_columns() == 2);
+  delete builder;
+}