DataAllocator.h

Go to the documentation of this file.

// Copyright 2019-2020 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 <gsl/span>
 
#include <memory>
#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>
 
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 && callback) {
      callback(*ptr);
      delete ptr;
      ptr = nullptr;
    }
  }
};
 
template <typename T>
concept VectorOfMessageableTypes = is_specialization_v<T, std::vector> &&
                                   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 TreeToTable>(std::declval<std::decay_t<T>>()); })
  decltype(auto) make(const Output& spec, Args... args)
  {
    auto t2t = std::move(LifetimeHolder<TreeToTable>(new std::decay_t<T>(args...)));
    adopt(spec, t2t);
    return t2t;
  }
 
  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 TreeToTable>&);
 
  void
    adopt(const Output& spec, LifetimeHolder<struct FragmentToBatch>&);
 
  void
    adopt(const Output& spec, std::shared_ptr<class arrow::Table>);
 
  template <typename T>
  void snapshot(const Output& spec, T const& object)
  {
    auto& proxy = mRegistry.get<MessageContext>().proxy();
    fair::mq::MessagePtr payloadMessage;
    auto serializationType = o2::header::gSerializationMethodNone;
    RouteIndex routeIndex = matchDataHeader(spec, mRegistry.get<TimingInfo>().timeslice);
    if constexpr (is_messageable<T>::value == true) {
      // Serialize a snapshot of a trivially copyable, non-polymorphic object,
      payloadMessage = proxy.createOutputMessage(routeIndex, sizeof(T));
      memcpy(payloadMessage->GetData(), &object, sizeof(T));
 
      serializationType = o2::header::gSerializationMethodNone;
    } else if constexpr (is_specialization_v<T, std::vector> == true ||
                         (gsl::details::is_span<T>::value && has_messageable_value_type<T>::value)) {
      using ElementType = typename std::remove_pointer<typename T::value_type>::type;
      if constexpr (is_messageable<ElementType>::value) {
        // 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();
        payloadMessage = proxy.createOutputMessage(routeIndex, sizeInBytes);
 
        if constexpr (std::is_pointer<typename T::value_type>::value == false) {
          // vector of elements
          if (object.data() && sizeInBytes) {
            memcpy(payloadMessage->GetData(), object.data(), sizeInBytes);
          }
        } else {
          // 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;
          }
        }
 
        serializationType = o2::header::gSerializationMethodNone;
      } else if constexpr (has_root_dictionary<ElementType>::value) {
        return snapshot(spec, ROOTSerialized<T const>(object));
      } else {
        static_assert(always_static_assert_v<T>,
                      "value type of std::vector not supported by API, supported types:"
                      "\n - messageable tyeps (trivially copyable, non-polymorphic structures)"
                      "\n - pointers to those"
                      "\n - types with ROOT dictionary and implementing ROOT ClassDef interface");
      }
    } else if constexpr (is_container<T>::value == true && has_messageable_value_type<T>::value == true) {
      // 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();
      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;
      }
      serializationType = o2::header::gSerializationMethodNone;
    } else if constexpr (has_root_dictionary<T>::value == true || is_specialization_v<T, ROOTSerialized> == true) {
      // Serialize a snapshot of an object with root dictionary
      payloadMessage = proxy.createOutputMessage(routeIndex);
      payloadMessage->Rebuild(4096, {64});
      if constexpr (is_specialization_v<T, ROOTSerialized> == true) {
        // 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;
        static_assert(std::is_same<typename T::hint_type, const char>::value ||
                        std::is_same<typename T::hint_type, TClass>::value ||
                        std::is_void<typename T::hint_type>::value,
                      "class hint must be of type TClass or const char");
 
        const TClass* cl = nullptr;
        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);
      } else {
        typename root_serializer<T>::serializer().Serialize(*payloadMessage, &object, TClass::GetClass(typeid(T)));
      }
      serializationType = o2::header::gSerializationMethodROOT;
    } else {
      static_assert(always_static_assert_v<T>,
                    "data type T not supported by API, \n specializations available for"
                    "\n - trivially copyable, non-polymorphic structures"
                    "\n - std::vector of messageable structures or pointers to those"
                    "\n - types with ROOT dictionary and implementing ROOT ClassDef interface");
    }
    addPartToContext(routeIndex, std::move(payloadMessage), spec, serializationType);
  }
 
  void snapshot(const Output& spec, const char* payload, size_t payloadSize,
                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;
  };
 
  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};
  }
 
  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);
 
  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}; //
  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);
  }
 
  context.add<MessageContext::TrivialObject>(std::move(headerMessage), std::move(payloadMessage), routeIndex);
  return cacheId;
}
 
} // namespace o2::framework
 
#endif // O2_FRAMEWORK_DATAALLOCATOR_H_