DPL: add allocator to forward a fair::mq::Message payload by shallow copy (#15577)

Author: ehellbar

Framework/Core/include/Framework/DataAllocator.h

@@ -455,6 +455,11 @@ class DataAllocator
   void snapshot(const Output& spec, const char* payload, size_t payloadSize,
                 o2::header::SerializationMethod serializationMethod = o2::header::gSerializationMethodNone);
 
+  /// create a shallow copy of the @a inputPayload and forward it to the output route of @a spec
+  /// if the transport types of the input and output routes are different, a real copy is created as fallback
+  void forwardPayload(const Output& spec, fair::mq::Message& inputPayload,
+                      o2::header::SerializationMethod serializationMethod = o2::header::gSerializationMethodNone);
+
   /// make an object of type T and route to output specified by OutputRef
   /// The object is owned by the framework, returned reference can be used to fill the object.
   ///

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); }
 

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)
     {

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;
 };
 

Framework/Core/src/DataAllocator.cxx

@@ -342,6 +342,26 @@ void DataAllocator::snapshot(const Output& spec, const char* payload, size_t pay
   addPartToContext(routeIndex, std::move(payloadMessage), spec, serializationMethod);
 }
 
+void DataAllocator::forwardPayload(const Output& spec, fair::mq::Message& inputPayload,
+                                   o2::header::SerializationMethod serializationMethod)
+{
+  auto& proxy = mRegistry.get<FairMQDeviceProxy>();
+  auto& timingInfo = mRegistry.get<TimingInfo>();
+
+  RouteIndex routeIndex = matchDataHeader(spec, timingInfo.timeslice);
+  auto* transport = proxy.getOutputTransport(routeIndex);
+
+  if (inputPayload.GetTransport() == transport) {
+    auto payloadMessage = transport->CreateMessage();
+    payloadMessage->Copy(inputPayload);
+    addPartToContext(routeIndex, std::move(payloadMessage), spec, serializationMethod);
+  } else {
+    auto payloadMessage = transport->CreateMessage(inputPayload.GetSize(), fair::mq::Alignment{64});
+    memcpy(payloadMessage->GetData(), inputPayload.GetData(), inputPayload.GetSize());
+    addPartToContext(routeIndex, std::move(payloadMessage), spec, serializationMethod);
+  }
+}
+
 Output DataAllocator::getOutputByBind(OutputRef&& ref)
 {
   if (ref.label.empty()) {

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 {

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);

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();

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}
 {
 }
 

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};