DataAllocator.h

Go to the documentation of this file.

// Copyright 2019-2026 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
#ifndef O2_FRAMEWORK_DATAALLOCATOR_H_
#define O2_FRAMEWORK_DATAALLOCATOR_H_
 
#include "Framework/MessageContext.h"
#include "Framework/RootMessageContext.h"
#include "Framework/StringContext.h"
#include "Framework/Output.h"
#include "Framework/OutputRef.h"
#include "Framework/OutputRoute.h"
#include "Framework/DataChunk.h"
#include "Framework/FairMQDeviceProxy.h"
#include "Framework/TimingInfo.h"
#include "Framework/TypeTraits.h"
#include "Framework/Traits.h"
#include "Framework/SerializationMethods.h"
#include "Framework/ServiceRegistry.h"
#include "Framework/RuntimeError.h"
#include "Framework/RouteState.h"
 
#include "Headers/DataHeader.h"
#include <TClass.h>
 
#include <memory>
#include <ranges>
#include <vector>
#include <map>
#include <string>
#include <utility>
#include <type_traits>
#include <utility>
#include <cstddef>
 
// Do not change this for a full inclusion of fair::mq::Device.
#include <fairmq/FwdDecls.h>
#include <fairmq/Version.h>
#if (FAIRMQ_VERSION_DEC >= 111000)
#include <fairmq/shmem/Message.h>
#endif
 
namespace arrow
{
class Schema;
class Table;
 
namespace ipc
{
class RecordBatchWriter;
} // namespace ipc
} // namespace arrow
 
namespace o2::framework
{
struct ServiceRegistry;
 
template <typename T>
struct LifetimeHolder {
  using type = T;
  T* ptr = nullptr;
  std::function<void(T&)> callback = nullptr;
  LifetimeHolder(T* ptr_) : ptr(ptr_),
                            callback(nullptr)
  {
  }
  LifetimeHolder() = delete;
  // Never copy it, because there is only one LifetimeHolder pointer
  // created object.
  LifetimeHolder(const LifetimeHolder&) = delete;
  LifetimeHolder& operator=(const LifetimeHolder&) = delete;
  LifetimeHolder(LifetimeHolder&& other)
  {
    this->ptr = other.ptr;
    other.ptr = nullptr;
    if (other.callback) {
      this->callback = std::move(other.callback);
    } else {
      this->callback = nullptr;
    }
    other.callback = nullptr;
  }
  LifetimeHolder& operator=(LifetimeHolder&& other)
  {
    this->ptr = other.ptr;
    other.ptr = nullptr;
    if (other.callback) {
      this->callback = std::move(other.callback);
    } else {
      this->callback = nullptr;
    }
    other.callback = nullptr;
    return *this;
  }
 
  // On deletion we invoke the callback and then delete the object,
  // when prensent.
  ~LifetimeHolder()
  {
    release();
  }
 
  T* operator->() { return ptr; }
  T& operator*() { return *ptr; }
 
  // release the owned object, if any. This allows to
  // invoke the callback early (e.g. for the Product<> case)
  void release()
  {
    if (!ptr) {
      return;
    }
 
    std::unique_ptr<T> released{ptr};
    ptr = nullptr;
    auto releaseCallback = std::move(callback);
    if (!releaseCallback) {
      return;
    }
    releaseCallback(*released);
  }
 
  // Delete the owned object without invoking the release callback. This is used
  // when a partially filled object must be abandoned.
  void discard()
  {
    if (!ptr) {
      return;
    }
 
    callback = nullptr;
    delete ptr;
    ptr = nullptr;
  }
};
 
template <typename T>
concept VectorOfMessageableTypes = is_specialization_v<T, std::vector> &&
                                   is_messageable<typename T::value_type>::value;
 
template <typename T>
concept ContiguousMessageablesRange = std::ranges::contiguous_range<T> &&
                                      is_messageable<typename T::value_type>::value;
 
class DataAllocator
{
 public:
  constexpr static ServiceKind service_kind = ServiceKind::Stream;
  using AllowedOutputRoutes = std::vector<OutputRoute>;
  using DataHeader = o2::header::DataHeader;
  using DataOrigin = o2::header::DataOrigin;
  using DataDescription = o2::header::DataDescription;
  using SubSpecificationType = o2::header::DataHeader::SubSpecificationType;
 
  template <typename T>
    requires std::is_fundamental_v<T>
  struct UninitializedVector {
    using value_type = T;
  };
 
  DataAllocator(ServiceRegistryRef ref);
 
  DataChunk& newChunk(const Output&, size_t);
 
  inline DataChunk& newChunk(OutputRef&& ref, size_t size) { return newChunk(getOutputByBind(std::move(ref)), size); }
 
  void adoptChunk(const Output&, char*, size_t, fair::mq::FreeFn*, void*);
 
  // This method can be used to send a 0xdeadbeef message associated to a given
  // output. The @a spec will be used to determine the channel to which the
  // output will need to be sent, however the actual message will be empty
  // and with subspecification 0xdeadbeef.
  void cookDeadBeef(const Output& spec);
 
  template <typename T, typename... Args>
    requires is_specialization_v<T, o2::framework::DataAllocator::UninitializedVector>
  decltype(auto) make(const Output& spec, Args... args)
  {
    auto& timingInfo = mRegistry.get<TimingInfo>();
    auto& context = mRegistry.get<MessageContext>();
 
    auto routeIndex = matchDataHeader(spec, timingInfo.timeslice);
    // plain buffer as polymorphic spectator std::vector, which does not run constructors / destructors
    using ValueType = typename T::value_type;
 
    // Note: initial payload size is 0 and will be set by the context before sending
    fair::mq::MessagePtr headerMessage = headerMessageFromOutput(spec, routeIndex, o2::header::gSerializationMethodNone, 0);
    return context.add<MessageContext::VectorObject<ValueType, MessageContext::ContainerRefObject<std::vector<ValueType, o2::pmr::NoConstructAllocator<ValueType>>>>>(
                    std::move(headerMessage), routeIndex, 0, std::forward<Args>(args)...)
      .get();
  }
 
  template <typename T, typename... Args>
    requires VectorOfMessageableTypes<T>
  decltype(auto) make(const Output& spec, Args... args)
  {
    auto& timingInfo = mRegistry.get<TimingInfo>();
    auto& context = mRegistry.get<MessageContext>();
 
    auto routeIndex = matchDataHeader(spec, timingInfo.timeslice);
    // this catches all std::vector objects with messageable value type before checking if is also
    // has a root dictionary, so non-serialized transmission is preferred
    using ValueType = typename T::value_type;
 
    // Note: initial payload size is 0 and will be set by the context before sending
    fair::mq::MessagePtr headerMessage = headerMessageFromOutput(spec, routeIndex, o2::header::gSerializationMethodNone, 0);
    return context.add<MessageContext::VectorObject<ValueType>>(std::move(headerMessage), routeIndex, 0, std::forward<Args>(args)...).get();
  }
 
  template <typename T, typename... Args>
    requires(!VectorOfMessageableTypes<T> && has_root_dictionary<T>::value == true && is_messageable<T>::value == false)
  decltype(auto) make(const Output& spec, Args... args)
  {
    auto& timingInfo = mRegistry.get<TimingInfo>();
    auto& context = mRegistry.get<MessageContext>();
 
    auto routeIndex = matchDataHeader(spec, timingInfo.timeslice);
    // Extended support for types implementing the Root ClassDef interface, both TObject
    // derived types and others
    if constexpr (enable_root_serialization<T>::value) {
      fair::mq::MessagePtr headerMessage = headerMessageFromOutput(spec, routeIndex, o2::header::gSerializationMethodROOT, 0);
 
      return context.add<typename enable_root_serialization<T>::object_type>(std::move(headerMessage), routeIndex, std::forward<Args>(args)...).get();
    } else {
      static_assert(enable_root_serialization<T>::value, "Please make sure you include RootMessageContext.h");
    }
  }
 
  template <typename T, typename... Args>
    requires std::is_base_of_v<std::string, T>
  decltype(auto) make(const Output& spec, Args... args)
  {
    auto* s = new std::string(args...);
    adopt(spec, s);
    return *s;
  }
 
  template <typename T, typename... Args>
    requires(requires { static_cast<struct TableBuilder>(std::declval<std::decay_t<T>>()); })
  decltype(auto) make(const Output& spec, Args... args)
  {
    auto tb = std::move(LifetimeHolder<TableBuilder>(new std::decay_t<T>(args...)));
    adopt(spec, tb);
    return tb;
  }
 
  template <typename T, typename... Args>
    requires(requires { static_cast<struct FragmentToBatch>(std::declval<std::decay_t<T>>()); })
  decltype(auto) make(const Output& spec, Args... args)
  {
    auto f2b = std::move(LifetimeHolder<FragmentToBatch>(new std::decay_t<T>(args...)));
    adopt(spec, f2b);
    return f2b;
  }
 
  template <typename T>
    requires is_messageable<T>::value && (!is_specialization_v<T, UninitializedVector>)
  decltype(auto) make(const Output& spec)
  {
    return *reinterpret_cast<T*>(newChunk(spec, sizeof(T)).data());
  }
 
  template <typename T>
    requires is_messageable<T>::value && (!is_specialization_v<T, UninitializedVector>)
  decltype(auto) make(const Output& spec, std::integral auto nElements)
  {
    auto& timingInfo = mRegistry.get<TimingInfo>();
    auto& context = mRegistry.get<MessageContext>();
    auto routeIndex = matchDataHeader(spec, timingInfo.timeslice);
 
    fair::mq::MessagePtr headerMessage = headerMessageFromOutput(spec, routeIndex, o2::header::gSerializationMethodNone, nElements * sizeof(T));
    return context.add<MessageContext::SpanObject<T>>(std::move(headerMessage), routeIndex, 0, nElements).get();
  }
 
  template <typename T, typename Arg>
  decltype(auto) make(const Output& spec, std::same_as<std::shared_ptr<arrow::Schema>> auto schema)
  {
    std::shared_ptr<arrow::ipc::RecordBatchWriter> writer;
    create(spec, &writer, schema);
    return writer;
  }
 
  void
    adopt(const Output& spec, std::string*);
 
  void
    adopt(const Output& spec, LifetimeHolder<struct TableBuilder>&);
 
  void
    adopt(const Output& spec, LifetimeHolder<struct FragmentToBatch>&);
 
  void
    adopt(const Output& spec, std::shared_ptr<class arrow::Table>);
 
  template <typename T>
    requires(!std::ranges::contiguous_range<T> && is_messageable<T>::value)
  void snapshot(const Output& spec, T const& object)
  {
    return snapshot(spec, std::span<T const>(&object, &object + 1));
  }
 
  void snapshot(const Output& spec, std::string_view const& object)
  {
    return snapshot(spec, std::span<char const>(object.data(), object.size()));
  }
 
  // This is for snapshotting a range of contiguous messageable types
  template <typename T>
    requires(ContiguousMessageablesRange<T> && !std::is_pointer_v<typename T::value_type>)
  void snapshot(const Output& spec, T const& object)
  {
    auto& proxy = mRegistry.get<MessageContext>().proxy();
    RouteIndex routeIndex = matchDataHeader(spec, mRegistry.get<TimingInfo>().timeslice);
    using ElementType = typename std::remove_pointer<typename T::value_type>::type;
    // Serialize a snapshot of a std::vector of trivially copyable, non-polymorphic elements
    // Note: in most cases it is better to use the `make` function und work with the provided
    // reference object
    constexpr auto elementSizeInBytes = sizeof(ElementType);
    auto sizeInBytes = elementSizeInBytes * object.size();
    fair::mq::MessagePtr payloadMessage = proxy.createOutputMessage(routeIndex, sizeInBytes);
 
    // vector of elements
    if (object.data() && sizeInBytes) {
      memcpy(payloadMessage->GetData(), object.data(), sizeInBytes);
    }
 
    addPartToContext(routeIndex, std::move(payloadMessage), spec, o2::header::gSerializationMethodNone);
  }
 
  // A random access range of pointers we can serialise by storing the contens one after the other.
  // On the receiving side you will have to retrieve it via a span
  template <typename T>
    requires(std::ranges::random_access_range<T> && is_messageable<typename std::remove_pointer_t<typename T::value_type>>::value && std::is_pointer_v<typename T::value_type>)
  void snapshot(const Output& spec, T const& object)
  {
    auto& proxy = mRegistry.get<MessageContext>().proxy();
    RouteIndex routeIndex = matchDataHeader(spec, mRegistry.get<TimingInfo>().timeslice);
    using ElementType = typename std::remove_pointer_t<typename T::value_type>;
    // Serialize a snapshot of a std::vector of trivially copyable, non-polymorphic elements
    // Note: in most cases it is better to use the `make` function und work with the provided
    // reference object
    constexpr auto elementSizeInBytes = sizeof(ElementType);
    auto sizeInBytes = elementSizeInBytes * object.size();
    fair::mq::MessagePtr payloadMessage = proxy.createOutputMessage(routeIndex, sizeInBytes);
 
    // serialize vector of pointers to elements
    auto target = reinterpret_cast<unsigned char*>(payloadMessage->GetData());
    for (auto const& pointer : object) {
      memcpy(target, pointer, elementSizeInBytes);
      target += elementSizeInBytes;
    }
 
    addPartToContext(routeIndex, std::move(payloadMessage), spec, o2::header::gSerializationMethodNone);
  }
 
  // This is for a range where we can know upfront how many elements there are,
  // so that we can preallocate the final size by simply multipling sizeof(T) x N elements
  template <typename T>
    requires(!std::ranges::contiguous_range<T> && std::ranges::sized_range<T> && has_messageable_value_type<T>::value)
  void snapshot(const Output& spec, T const& object)
  {
    auto& proxy = mRegistry.get<MessageContext>().proxy();
    RouteIndex routeIndex = matchDataHeader(spec, mRegistry.get<TimingInfo>().timeslice);
    // Serialize a snapshot of a std::container of trivially copyable, non-polymorphic elements
    // Note: in most cases it is better to use the `make` function und work with the provided
    // reference object
    constexpr auto elementSizeInBytes = sizeof(typename T::value_type);
    auto sizeInBytes = elementSizeInBytes * object.size();
    fair::mq::MessagePtr payloadMessage = proxy.createOutputMessage(routeIndex, sizeInBytes);
 
    // serialize vector of pointers to elements
    auto target = reinterpret_cast<unsigned char*>(payloadMessage->GetData());
    for (auto const& entry : object) {
      memcpy(target, (void*)&entry, elementSizeInBytes);
      target += elementSizeInBytes;
    }
    addPartToContext(routeIndex, std::move(payloadMessage), spec, o2::header::gSerializationMethodNone);
  }
 
  template <typename T>
    requires(is_specialization_v<T, ROOTSerialized>)
  void snapshot(const Output& spec, T const& object)
  {
    auto& proxy = mRegistry.get<MessageContext>().proxy();
    RouteIndex routeIndex = matchDataHeader(spec, mRegistry.get<TimingInfo>().timeslice);
    // Serialize a snapshot of an object with root dictionary
    fair::mq::MessagePtr payloadMessage = proxy.createOutputMessage(routeIndex);
    payloadMessage->Rebuild(4096, {64});
    const TClass* cl = nullptr;
    // Explicitely ROOT serialize a snapshot of object.
    // An object wrapped into type `ROOTSerialized` is explicitely marked to be ROOT serialized
    // and is expected to have a ROOT dictionary. Availability can not be checked at compile time
    // for all cases.
    using WrappedType = typename T::wrapped_type;
 
    if (object.getHint() == nullptr) {
      // get TClass info by wrapped type
      cl = TClass::GetClass(typeid(WrappedType));
    } else if (std::is_same<typename T::hint_type, TClass>::value) {
      // the class info has been passed directly
      cl = reinterpret_cast<const TClass*>(object.getHint());
    } else if (std::is_same<typename T::hint_type, const char>::value) {
      // get TClass info by optional name
      cl = TClass::GetClass(reinterpret_cast<const char*>(object.getHint()));
    }
    if (has_root_dictionary<WrappedType>::value == false && cl == nullptr) {
      if (std::is_same<typename T::hint_type, const char>::value) {
        throw runtime_error_f("ROOT serialization not supported, dictionary not found for type %s",
                              reinterpret_cast<const char*>(object.getHint()));
      } else {
        throw runtime_error_f("ROOT serialization not supported, dictionary not found for type %s",
                              typeid(WrappedType).name());
      }
    }
    typename root_serializer<T>::serializer().Serialize(*payloadMessage, &object(), cl);
    addPartToContext(routeIndex, std::move(payloadMessage), spec, o2::header::gSerializationMethodROOT);
  }
 
  template <typename T>
    requires(!is_messageable<T>::value && !ContiguousMessageablesRange<T> && has_root_dictionary<T>::value && !is_specialization_v<T, ROOTSerialized>)
  void snapshot(const Output& spec, T const& object)
  {
    auto& proxy = mRegistry.get<MessageContext>().proxy();
    RouteIndex routeIndex = matchDataHeader(spec, mRegistry.get<TimingInfo>().timeslice);
    // Serialize a snapshot of an object with root dictionary
    fair::mq::MessagePtr payloadMessage = proxy.createOutputMessage(routeIndex);
    payloadMessage->Rebuild(4096, {64});
    typename root_serializer<T>::serializer().Serialize(*payloadMessage, &object, TClass::GetClass(typeid(T)));
    addPartToContext(routeIndex, std::move(payloadMessage), spec, o2::header::gSerializationMethodROOT);
  }
 
  void snapshot(const Output& spec, const char* payload, size_t payloadSize,
                o2::header::SerializationMethod serializationMethod = o2::header::gSerializationMethodNone);
 
  void forwardPayload(const Output& spec, fair::mq::Message& inputPayload,
                      o2::header::SerializationMethod serializationMethod = o2::header::gSerializationMethodNone);
 
  template <typename T, typename... Args>
  decltype(auto) make(OutputRef&& ref, Args&&... args)
  {
    return make<T>(getOutputByBind(std::move(ref)), std::forward<Args>(args)...);
  }
 
  template <typename T>
  void adopt(OutputRef&& ref, T* obj)
  {
    return adopt(getOutputByBind(std::move(ref)), obj);
  }
 
  // get the memory resource associated with an output
  o2::pmr::FairMQMemoryResource* getMemoryResource(const Output& spec)
  {
    auto& timingInfo = mRegistry.get<TimingInfo>();
    auto& proxy = mRegistry.get<FairMQDeviceProxy>();
    RouteIndex routeIndex = matchDataHeader(spec, timingInfo.timeslice);
    return *proxy.getOutputTransport(routeIndex);
  }
 
  // make a stl (pmr) vector
  template <typename T, typename... Args>
  o2::pmr::vector<T> makeVector(const Output& spec, Args&&... args)
  {
    o2::pmr::FairMQMemoryResource* targetResource = getMemoryResource(spec);
    return o2::pmr::vector<T>{targetResource, std::forward<Args>(args)...};
  }
 
  struct CacheId {
    int64_t value;
    int64_t handle;
    int64_t segment;
  };
 
  enum struct CacheStrategy : int {
    Never = 0,
    Always = 1
  };
 
  template <typename ContainerT>
  CacheId adoptContainer(const Output& /*spec*/, ContainerT& /*container*/, CacheStrategy /* cache  = false */, o2::header::SerializationMethod /* method = header::gSerializationMethodNone*/)
  {
    static_assert(always_static_assert_v<ContainerT>, "Container cannot be moved. Please make sure it is backed by a o2::pmr::FairMQMemoryResource");
    return {0, 0, 0};
  }
 
  template <typename ContainerT>
  CacheId adoptContainer(const Output& spec, ContainerT&& container, CacheStrategy cache = CacheStrategy::Never, o2::header::SerializationMethod method = header::gSerializationMethodNone);
 
  void adoptFromCache(Output const& spec, CacheId id, header::SerializationMethod method = header::gSerializationMethodNone);
 
  void pruneFromCache(CacheId id);
 
  template <typename... Args>
  auto snapshot(OutputRef&& ref, Args&&... args)
  {
    return snapshot(getOutputByBind(std::move(ref)), std::forward<Args>(args)...);
  }
 
  bool isAllowed(Output const& query);
 
  o2::header::DataHeader* findMessageHeader(const Output& spec)
  {
    return mRegistry.get<MessageContext>().findMessageHeader(spec);
  }
 
  o2::header::DataHeader* findMessageHeader(OutputRef&& ref)
  {
    return mRegistry.get<MessageContext>().findMessageHeader(getOutputByBind(std::move(ref)));
  }
 
  o2::header::Stack* findMessageHeaderStack(const Output& spec)
  {
    return mRegistry.get<MessageContext>().findMessageHeaderStack(spec);
  }
 
  o2::header::Stack* findMessageHeaderStack(OutputRef&& ref)
  {
    return mRegistry.get<MessageContext>().findMessageHeaderStack(getOutputByBind(std::move(ref)));
  }
 
  int countDeviceOutputs(bool excludeDPLOrigin = false)
  {
    return mRegistry.get<MessageContext>().countDeviceOutputs(excludeDPLOrigin);
  }
 
 private:
  ServiceRegistryRef mRegistry;
 
  RouteIndex matchDataHeader(const Output& spec, size_t timeframeId);
  fair::mq::MessagePtr headerMessageFromOutput(Output const& spec,                                  //
                                               RouteIndex index,                                    //
                                               o2::header::SerializationMethod serializationMethod, //
                                               size_t payloadSize);                                 //
 
  Output getOutputByBind(OutputRef&& ref);
  void addPartToContext(RouteIndex routeIndex, fair::mq::MessagePtr&& payload,
                        const Output& spec,
                        o2::header::SerializationMethod serializationMethod);
};
 
template <typename ContainerT>
DataAllocator::CacheId DataAllocator::adoptContainer(const Output& spec, ContainerT&& container, CacheStrategy cache, header::SerializationMethod method)
{
  // Find a matching channel, extract the message for it form the container
  // and put it in the queue to be sent at the end of the processing
  auto& timingInfo = mRegistry.get<TimingInfo>();
  auto routeIndex = matchDataHeader(spec, timingInfo.timeslice);
 
  auto& context = mRegistry.get<MessageContext>();
  auto* transport = mRegistry.get<FairMQDeviceProxy>().getOutputTransport(routeIndex);
  fair::mq::MessagePtr payloadMessage = o2::pmr::getMessage(std::forward<ContainerT>(container), *transport);
  fair::mq::MessagePtr headerMessage = headerMessageFromOutput(spec, routeIndex,         //
                                                               method,                   //
                                                               payloadMessage->GetSize() //
  );
 
  CacheId cacheId{0, 0, 0}; //
  if (cache == CacheStrategy::Always) {
    // The message will be shallow cloned in the cache. Since the
    // clone is indistinguishable from the original, we can keep sending
    // the original.
    cacheId.value = context.addToCache(payloadMessage);
#if (FAIRMQ_VERSION_DEC >= 111000)
    auto meta = dynamic_cast<fair::mq::shmem::Message*>(payloadMessage.get())->GetMeta();
    cacheId.handle = meta.fHandle;
    cacheId.segment = meta.fSegmentId;
#endif
  }
 
  context.add<MessageContext::TrivialObject>(std::move(headerMessage), std::move(payloadMessage), routeIndex);
  return cacheId;
}
 
} // namespace o2::framework
 
#endif // O2_FRAMEWORK_DATAALLOCATOR_H_