DPL: add fairmq message pointer getters to InputRecord

ehellbar · ehellbar · commit b4bf81843936 · 2026-07-06T11:35:10.000+02:00
diff --git a/Framework/Core/include/Framework/InputRecord.h b/Framework/Core/include/Framework/InputRecord.h
@@ -213,6 +213,9 @@ class InputRecord
   /// O(1) access to the part described by @a indices in slot @a pos.
   [[nodiscard]] DataRef getAtIndices(int pos, DataRefIndices indices) const;
 
+  /// Return the payload as fair::mq::Message* for the part described by @a indices in slot @a slotIdx
+  fair::mq::Message* getPayloadAtIndices(size_t slotIdx, DataRefIndices indices) const;
+
   /// O(1) advance from @a current to the next part's indices in slot @a pos.
   [[nodiscard]] DataRefIndices nextIndices(int pos, DataRefIndices current) const
   {
@@ -745,6 +748,7 @@ class InputRecord
     [[nodiscard]] DataRefIndices initialIndices() const { return {0, 1}; }
     [[nodiscard]] DataRefIndices endIndices() const { return {size_t(-1), size_t(-1)}; }
     [[nodiscard]] DataRef getAtIndices(DataRefIndices idx) const { return record->getAtIndices((int)slot, idx); }
+    [[nodiscard]] fair::mq::Message* getPayloadAtIndices(DataRefIndices idx) const { return record->getPayloadAtIndices((int)slot, idx); }
     [[nodiscard]] DataRefIndices nextIndices(DataRefIndices idx) const { return record->nextIndices((int)slot, idx); }
     [[nodiscard]] size_t size() const { return record->getNofParts((int)slot); }
 
diff --git a/Framework/Core/include/Framework/InputRecordWalker.h b/Framework/Core/include/Framework/InputRecordWalker.h
@@ -115,6 +115,11 @@ class InputRecordWalker
       return not operator==(rh);
     }
 
+    fair::mq::Message* getPayload() const
+    {
+      return mCurrentRange.getPayloadAtIndices(mCurrent.indices());
+    }
+
    private:
     bool next(bool isInitialPart = false)
     {
diff --git a/Framework/Core/include/Framework/InputSpan.h b/Framework/Core/include/Framework/InputSpan.h
@@ -13,9 +13,11 @@
 
 #include "Framework/DataRef.h"
 #include <functional>
+#include <fairmq/FwdDecls.h>
 
 extern template class std::function<o2::framework::DataRef(size_t, o2::framework::DataRefIndices)>;
 extern template class std::function<o2::framework::DataRefIndices(size_t, o2::framework::DataRefIndices)>;
+extern template class std::function<fair::mq::Message*(size_t, o2::framework::DataRefIndices)>;
 
 namespace o2::framework
 {
@@ -38,6 +40,7 @@ class InputSpan
             std::function<int(size_t)> refCountGetter,
             std::function<DataRef(size_t, DataRefIndices)> indicesGetter,
             std::function<DataRefIndices(size_t, DataRefIndices)> nextIndicesGetter,
+            std::function<fair::mq::Message*(size_t, DataRefIndices)> payloadGetter,
             size_t size);
 
   /// @a i-th element of the InputSpan (O(partidx) sequential scan via indices protocol)
@@ -56,6 +59,12 @@ class InputSpan
     return mIndicesGetter(slotIdx, indices);
   }
 
+  /// Return the payload as fair::mq::Message* for the part described by @a indices in slot @a slotIdx
+  [[nodiscard]] fair::mq::Message* getPayloadAtIndices(size_t slotIdx, DataRefIndices indices) const
+  {
+    return mPayloadGetter(slotIdx, indices);
+  }
+
   /// Advance from @a current to the indices of the next part in slot @a slotIdx in O(1).
   [[nodiscard]] DataRefIndices nextIndices(size_t slotIdx, DataRefIndices current) const
   {
@@ -179,6 +188,12 @@ class InputSpan
       return mCurrentIndices.headerIdx;
     }
 
+    // return current indices
+    [[nodiscard]] DataRefIndices indices() const
+    {
+      return mCurrentIndices;
+    }
+
     // return an iterable range over all parts in the current slot
     // only available for slot-level iterators whose parent has parts(size_t)
     [[nodiscard]] auto parts() const
@@ -201,6 +216,7 @@ class InputSpan
     [[nodiscard]] DataRefIndices initialIndices() const { return {0, 1}; }
     [[nodiscard]] DataRefIndices endIndices() const { return {size_t(-1), size_t(-1)}; }
     [[nodiscard]] DataRef getAtIndices(DataRefIndices idx) const { return span->getAtIndices(slot, idx); }
+    [[nodiscard]] fair::mq::Message* getPayloadAtIndices(DataRefIndices idx) const { return span->getPayloadAtIndices(slot, idx); }
     [[nodiscard]] DataRefIndices nextIndices(DataRefIndices idx) const { return span->nextIndices(slot, idx); }
     [[nodiscard]] size_t size() const { return span->getNofParts(slot); }
 
@@ -230,6 +246,7 @@ class InputSpan
   std::function<int(size_t)> mRefCountGetter;
   std::function<DataRef(size_t, DataRefIndices)> mIndicesGetter;
   std::function<DataRefIndices(size_t, DataRefIndices)> mNextIndicesGetter;
+  std::function<fair::mq::Message*(size_t, DataRefIndices)> mPayloadGetter;
   size_t mSize;
 };
 
diff --git a/Framework/Core/src/DataProcessingDevice.cxx b/Framework/Core/src/DataProcessingDevice.cxx
@@ -2203,7 +2203,14 @@ bool DataProcessingDevice::tryDispatchComputation(ServiceRegistryRef ref, std::v
       auto next = currentSetOfInputs[i] | get_next_pair{current};
       return next.headerIdx < currentSetOfInputs[i].size() ? next : DataRefIndices{size_t(-1), size_t(-1)};
     };
-    return InputSpan{nofPartsGetter, refCountGetter, indicesGetter, nextIndicesGetter, currentSetOfInputs.size()};
+    auto payloadGetter = [&currentSetOfInputs](size_t i, DataRefIndices current) -> fair::mq::Message* {
+      auto const& msgs = currentSetOfInputs[i];
+      if (msgs.size() <= current.payloadIdx || !msgs[current.payloadIdx]) {
+        return nullptr;
+      }
+      return msgs[current.payloadIdx].get();
+    };
+    return InputSpan{nofPartsGetter, refCountGetter, indicesGetter, nextIndicesGetter, payloadGetter, currentSetOfInputs.size()};
   };
 
   auto markInputsAsDone = [ref](TimesliceSlot slot) -> void {
diff --git a/Framework/Core/src/DataRelayer.cxx b/Framework/Core/src/DataRelayer.cxx
@@ -236,7 +236,14 @@ DataRelayer::ActivityStats DataRelayer::processDanglingInputs(std::vector<Expira
         auto next = partial[idx] | get_next_pair{current};
         return next.headerIdx < partial[idx].size() ? next : DataRefIndices{size_t(-1), size_t(-1)};
       };
-      InputSpan span{nPartsGetter, refCountGetter, indicesGetter, nextIndicesGetter, static_cast<size_t>(partial.size())};
+      auto payloadGetter = [&partial](size_t idx, DataRefIndices current) -> fair::mq::Message* {
+        auto const& msgs = partial[idx];
+        if (msgs.size() <= current.payloadIdx || !msgs[current.payloadIdx]) {
+          return nullptr;
+        }
+        return msgs[current.payloadIdx].get();
+      };
+      InputSpan span{nPartsGetter, refCountGetter, indicesGetter, nextIndicesGetter, payloadGetter, static_cast<size_t>(partial.size())};
       // Setup the input span
 
       if (expirator.checker(services, timestamp.value, span) == false) {
@@ -818,7 +825,14 @@ void DataRelayer::getReadyToProcess(std::vector<DataRelayer::RecordAction>& comp
       auto next = partial[idx] | get_next_pair{current};
       return next.headerIdx < partial[idx].size() ? next : DataRefIndices{size_t(-1), size_t(-1)};
     };
-    InputSpan span{nPartsGetter, refCountGetter, indicesGetter, nextIndicesGetter, static_cast<size_t>(partial.size())};
+    auto payloadGetter = [&partial](size_t idx, DataRefIndices current) -> fair::mq::Message* {
+      auto const& msgs = partial[idx];
+      if (msgs.size() <= current.payloadIdx || !msgs[current.payloadIdx]) {
+        return nullptr;
+      }
+      return msgs[current.payloadIdx].get();
+    };
+    InputSpan span{nPartsGetter, refCountGetter, indicesGetter, nextIndicesGetter, payloadGetter, static_cast<size_t>(partial.size())};
     CompletionPolicy::CompletionOp action = mCompletionPolicy.callbackFull(span, mInputs, mContext);
 
     auto& variables = mTimesliceIndex.getVariablesForSlot(slot);
diff --git a/Framework/Core/src/InputRecord.cxx b/Framework/Core/src/InputRecord.cxx
@@ -149,6 +149,11 @@ DataRef InputRecord::getAtIndices(int pos, DataRefIndices indices) const
   return ref;
 }
 
+fair::mq::Message* InputRecord::getPayloadAtIndices(size_t slotIdx, DataRefIndices indices) const
+{
+  return mSpan.getPayloadAtIndices(slotIdx, indices);
+}
+
 size_t InputRecord::size() const
 {
   return mSpan.size();
diff --git a/Framework/Core/src/InputSpan.cxx b/Framework/Core/src/InputSpan.cxx
@@ -13,15 +13,17 @@
 
 template class std::function<o2::framework::DataRef(size_t, o2::framework::DataRefIndices)>;
 template class std::function<o2::framework::DataRefIndices(size_t, o2::framework::DataRefIndices)>;
+template class std::function<fair::mq::Message*(size_t, o2::framework::DataRefIndices)>;
 
 namespace o2::framework
 {
 InputSpan::InputSpan(std::function<size_t(size_t)> nofPartsGetter,
                      std::function<int(size_t)> refCountGetter,
                      std::function<DataRef(size_t, DataRefIndices)> indicesGetter,
                      std::function<DataRefIndices(size_t, DataRefIndices)> nextIndicesGetter,
+                     std::function<fair::mq::Message*(size_t, DataRefIndices)> payloadGetter,
                      size_t size)
-  : mNofPartsGetter{nofPartsGetter}, mRefCountGetter(refCountGetter), mIndicesGetter{std::move(indicesGetter)}, mNextIndicesGetter{std::move(nextIndicesGetter)}, mSize{size}
+  : mNofPartsGetter{nofPartsGetter}, mRefCountGetter(refCountGetter), mIndicesGetter{std::move(indicesGetter)}, mNextIndicesGetter{std::move(nextIndicesGetter)}, mPayloadGetter{std::move(payloadGetter)}, mSize{size}
 {
 }
 
diff --git a/Framework/Core/test/benchmark_InputRecord.cxx b/Framework/Core/test/benchmark_InputRecord.cxx
@@ -52,6 +52,7 @@ static void BM_InputRecordGenericGetters(benchmark::State& state)
     nullptr,
     [](size_t, DataRefIndices) { return DataRef{nullptr, nullptr, nullptr}; },
     [](size_t, DataRefIndices) -> DataRefIndices { return {size_t(-1), size_t(-1)}; },
+    nullptr,
     0};
   ServiceRegistry registry;
   InputRecord emptyRecord(schema, span, registry);
@@ -92,6 +93,7 @@ static void BM_InputRecordGenericGetters(benchmark::State& state)
     nullptr,
     [&inputs](size_t i, DataRefIndices idx) { return DataRef{nullptr, static_cast<char const*>(inputs[2 * i + idx.headerIdx]), static_cast<char const*>(inputs[2 * i + idx.payloadIdx])}; },
     [](size_t, DataRefIndices) -> DataRefIndices { return {size_t(-1), size_t(-1)}; },
+    nullptr,
     inputs.size() / 2};
   InputRecord record{schema, span2, registry};
 
diff --git a/Framework/Core/test/test_CompletionPolicy.cxx b/Framework/Core/test/test_CompletionPolicy.cxx
@@ -60,6 +60,7 @@ TEST_CASE("TestCompletionPolicy_callback")
     nullptr,
     [&ref](size_t, DataRefIndices) -> DataRef { return ref; },
     [](size_t, DataRefIndices) -> DataRefIndices { return {size_t(-1), size_t(-1)}; },
+    nullptr,
     1};
   std::vector<InputSpec> specs;
   ServiceRegistryRef servicesRef{services};
diff --git a/Framework/Core/test/test_InputRecord.cxx b/Framework/Core/test/test_InputRecord.cxx
@@ -51,6 +51,7 @@ TEST_CASE("TestInputRecord")
     nullptr,
     [](size_t, DataRefIndices) { return DataRef{nullptr, nullptr, nullptr}; },
     [](size_t, DataRefIndices) -> DataRefIndices { return {size_t(-1), size_t(-1)}; },
+    nullptr,
     0};
   ServiceRegistry registry;
   InputRecord emptyRecord(schema, span, registry);
@@ -99,6 +100,7 @@ TEST_CASE("TestInputRecord")
     nullptr,
     [&inputs](size_t i, DataRefIndices idx) { return DataRef{nullptr, static_cast<char const*>(inputs[2 * i + idx.headerIdx]), static_cast<char const*>(inputs[2 * i + idx.payloadIdx])}; },
     [](size_t, DataRefIndices) -> DataRefIndices { return {size_t(-1), size_t(-1)}; },
+    nullptr,
     inputs.size() / 2};
   InputRecord record{schema, span2, registry};
 
diff --git a/Framework/Core/test/test_InputRecordWalker.cxx b/Framework/Core/test/test_InputRecordWalker.cxx
@@ -40,7 +40,7 @@ struct DataSet {
                                                            auto payload = static_cast<char const*>(this->messages[i].second.at(idx.payloadIdx)->data());
                                                            return DataRef{nullptr, header, payload}; }, [this](size_t i, DataRefIndices current) -> DataRefIndices {
                                                            size_t next = current.headerIdx + 2;
-                                                           return next < this->messages[i].second.size() ? DataRefIndices{next, next + 1} : DataRefIndices{size_t(-1), size_t(-1)}; }, this->messages.size()}, record{schema, span, registry}, values{std::move(v)}
+                                                           return next < this->messages[i].second.size() ? DataRefIndices{next, next + 1} : DataRefIndices{size_t(-1), size_t(-1)}; }, nullptr, this->messages.size()}, record{schema, span, registry}, values{std::move(v)}
   {
     REQUIRE(messages.size() == schema.size());
   }
diff --git a/Framework/Core/test/test_InputSpan.cxx b/Framework/Core/test/test_InputSpan.cxx
@@ -41,7 +41,7 @@ TEST_CASE("TestInputSpan")
     return next < inputs[i].size() ? DataRefIndices{next, next + 1} : DataRefIndices{size_t(-1), size_t(-1)};
   };
 
-  InputSpan span{nPartsGetter, nullptr, indicesGetter, nextIndicesGetter, inputs.size()};
+  InputSpan span{nPartsGetter, nullptr, indicesGetter, nextIndicesGetter, nullptr, inputs.size()};
   REQUIRE(span.size() == inputs.size());
   routeNo = 0;
   for (; routeNo < span.size(); ++routeNo) {
diff --git a/Framework/Utils/test/RawPageTestData.h b/Framework/Utils/test/RawPageTestData.h
@@ -53,6 +53,7 @@ struct DataSet {
              size_t next = current.headerIdx + 2;
              return next < this->messages[i].size() ? DataRefIndices{next, next + 1} : DataRefIndices{size_t(-1), size_t(-1)};
            },
+           nullptr,
            this->messages.size()},
       record{schema, span, registry},
       values{std::move(v)}
diff --git a/Framework/Utils/test/test_RootTreeWriter.cxx b/Framework/Utils/test/test_RootTreeWriter.cxx
@@ -231,6 +231,7 @@ TEST_CASE("test_RootTreeWriter")
       return DataRef{nullptr, static_cast<char const*>(store[2 * i + idx.headerIdx]->GetData()), static_cast<char const*>(store[2 * i + idx.payloadIdx]->GetData())};
     },
     [](size_t, DataRefIndices) -> DataRefIndices { return {size_t(-1), size_t(-1)}; },
+    nullptr,
     store.size() / 2};
   ServiceRegistry registry;
   InputRecord inputs{

Original file line number	Diff line number	Diff line change
`@@ -115,6 +115,11 @@ class InputRecordWalker`
`115`	`115`	`return not operator==(rh);`
`116`	`116`	`}`
`117`	`117`
	`118`	`+ fair::mq::Message* getPayload() const`
	`119`	`+ {`
	`120`	`+ return mCurrentRange.getPayloadAtIndices(mCurrent.indices());`
	`121`	`+ }`
	`122`	`+`
`118`	`123`	`private:`
`119`	`124`	`bool next(bool isInitialPart = false)`
`120`	`125`	`{`
Original file line number	Diff line number	Diff line change
`@@ -149,6 +149,11 @@ DataRef InputRecord::getAtIndices(int pos, DataRefIndices indices) const`
`149`	`149`	`return ref;`
`150`	`150`	`}`
`151`	`151`
	`152`	`+fair::mq::Message* InputRecord::getPayloadAtIndices(size_t slotIdx, DataRefIndices indices) const`
	`153`	`+{`
	`154`	`+ return mSpan.getPayloadAtIndices(slotIdx, indices);`
	`155`	`+}`
	`156`	`+`
`152`	`157`	`size_t InputRecord::size() const`
`153`	`158`	`{`
`154`	`159`	`return mSpan.size();`
Original file line number	Diff line number	Diff line change
`@@ -13,15 +13,17 @@`
`13`	`13`
`14`	`14`	`template class std::function<o2::framework::DataRef(size_t, o2::framework::DataRefIndices)>;`
`15`	`15`	`template class std::function<o2::framework::DataRefIndices(size_t, o2::framework::DataRefIndices)>;`
	`16`	`+template class std::function<fair::mq::Message*(size_t, o2::framework::DataRefIndices)>;`
`16`	`17`
`17`	`18`	`namespace o2::framework`
`18`	`19`	`{`
`19`	`20`	`InputSpan::InputSpan(std::function<size_t(size_t)> nofPartsGetter,`
`20`	`21`	`std::function<int(size_t)> refCountGetter,`
`21`	`22`	`std::function<DataRef(size_t, DataRefIndices)> indicesGetter,`
`22`	`23`	`std::function<DataRefIndices(size_t, DataRefIndices)> nextIndicesGetter,`
	`24`	`+ std::function<fair::mq::Message*(size_t, DataRefIndices)> payloadGetter,`
`23`	`25`	`size_t size)`
`24`		`- : mNofPartsGetter{nofPartsGetter}, mRefCountGetter(refCountGetter), mIndicesGetter{std::move(indicesGetter)}, mNextIndicesGetter{std::move(nextIndicesGetter)}, mSize{size}`
	`26`	`+ : mNofPartsGetter{nofPartsGetter}, mRefCountGetter(refCountGetter), mIndicesGetter{std::move(indicesGetter)}, mNextIndicesGetter{std::move(nextIndicesGetter)}, mPayloadGetter{std::move(payloadGetter)}, mSize{size}`
`25`	`27`	`{`
`26`	`28`	`}`
`27`	`29`