MessageContext.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_MESSAGECONTEXT_H_
#define O2_FRAMEWORK_MESSAGECONTEXT_H_
 
#include "Framework/DispatchControl.h"
#include "Framework/FairMQDeviceProxy.h"
#include "Framework/OutputRoute.h"
#include "Framework/RouteState.h"
#include "Framework/RoutingIndices.h"
#include "Framework/RuntimeError.h"
#include "Framework/DataProcessingHeader.h"
#include "Framework/TypeTraits.h"
 
#include "Headers/DataHeader.h"
#include "Headers/Stack.h"
#include "MemoryResources/MemoryResources.h"
 
#include <fairmq/Message.h>
#include <fairmq/Parts.h>
 
#include <cassert>
#include <functional>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <vector>
 
#include <fairmq/FwdDecls.h>
 
namespace o2::framework
{
 
template <typename T, typename = void>
struct enable_root_serialization : std::false_type {
  using debug_type = T;
};
 
template <typename T, typename = void>
struct root_serializer : std::false_type {
};
 
struct Output;
 
class MessageContext
{
 public:
  enum class DispatchState {
    NotDispatched,
    Dispatched,
    Discarded,
  };
 
  constexpr static ServiceKind service_kind = ServiceKind::Stream;
 
  // so far we are only using one instance per named channel
  static constexpr int DefaultChannelIndex = 0;
 
  MessageContext(FairMQDeviceProxy& proxy)
    : mProxy{proxy}
  {
  }
 
  MessageContext(FairMQDeviceProxy& proxy, DispatchControl&& dispatcher)
    : mProxy{proxy}, mDispatchControl{dispatcher}
  {
  }
 
  void init(DispatchControl&& dispatcher)
  {
    mDispatchControl = dispatcher;
  }
 
  // this is the virtual interface for context objects
  class ContextObject
  {
   public:
    ContextObject() = delete;
    ContextObject(fair::mq::MessagePtr&& headerMsg, fair::mq::MessagePtr&& payloadMsg, RouteIndex routeIndex)
      : mParts{}, mRouteIndex{routeIndex}
    {
      mParts.AddPart(std::move(headerMsg));
      mParts.AddPart(std::move(payloadMsg));
    }
 
    ContextObject(fair::mq::MessagePtr&& headerMsg, RouteIndex routeIndex)
      : mParts{}, mRouteIndex{routeIndex}
    {
      mParts.AddPart(std::move(headerMsg));
    }
 
    virtual ~ContextObject() = default;
 
    virtual fair::mq::Parts finalize()
    {
      fair::mq::Parts parts = std::move(mParts);
      assert(parts.Size() == 2);
      auto* header = o2::header::get<o2::header::DataHeader*>(parts.At(0)->GetData());
      if (header == nullptr) {
        throw std::logic_error("No valid header message found");
      } else {
        // o2::header::get returns const pointer, but here we can change the message
        const_cast<o2::header::DataHeader*>(header)->payloadSize = parts.At(1)->GetSize();
      }
      // return value optimization returns by move
      return parts;
    }
 
    [[nodiscard]] RouteIndex route() const
    {
      return mRouteIndex;
    }
 
    [[nodiscard]] bool empty() const
    {
      return mParts.Size() == 0;
    }
 
    o2::header::DataHeader const* header()
    {
      // we would expect this function to be const but the fair::mq::Parts API does not allow this
      if (empty() || mParts.At(0) == nullptr) {
        return nullptr;
      }
      return o2::header::get<o2::header::DataHeader*>(mParts.At(0)->GetData());
    }
 
    o2::framework::DataProcessingHeader const* dataProcessingHeader()
    {
      if (empty() || mParts.At(0) == nullptr) {
        return nullptr;
      }
      return o2::header::get<o2::framework::DataProcessingHeader*>(mParts.At(0)->GetData());
    }
 
    o2::header::Stack const* headerStack()
    {
      // we would expect this function to be const but the fair::mq::Parts API does not allow this
      if (empty() || mParts.At(0) == nullptr) {
        return nullptr;
      }
      return o2::header::get<o2::header::DataHeader*>(mParts.At(0)->GetData()) ? reinterpret_cast<o2::header::Stack*>(mParts.At(0)->GetData()) : nullptr;
    }
 
   protected:
    fair::mq::Parts mParts;
    RouteIndex mRouteIndex{-1};
  };
 
  class TrivialObject : public ContextObject
  {
   public:
    TrivialObject() = delete;
    template <typename ContextType>
    TrivialObject(ContextType* context, fair::mq::MessagePtr&& headerMsg, fair::mq::MessagePtr&& payloadMsg, RouteIndex routeIndex)
      : ContextObject(std::forward<fair::mq::MessagePtr>(headerMsg), std::forward<fair::mq::MessagePtr>(payloadMsg), routeIndex)
    {
    }
    template <typename ContextType, typename... Args>
    TrivialObject(ContextType* context, fair::mq::MessagePtr&& headerMsg, RouteIndex routeIndex, int index, Args... args)
      : ContextObject(std::forward<fair::mq::MessagePtr>(headerMsg), context->createMessage(routeIndex, index, std::forward<Args>(args)...), routeIndex)
    {
    }
    ~TrivialObject() override = default;
 
    auto* data()
    {
      assert(mParts.Size() == 2);
      return mParts[1].GetData();
    }
  };
 
  // A memory resource which can force a minimum alignment, so that
  // the whole polymorphic allocator business is happy...
  class AlignedMemoryResource : public pmr::FairMQMemoryResource
  {
   public:
    AlignedMemoryResource(fair::mq::MemoryResource* other)
      : mUpstream(other)
    {
    }
 
    AlignedMemoryResource(AlignedMemoryResource const& other)
      : mUpstream(other.mUpstream)
    {
    }
 
    bool isValid()
    {
      return mUpstream != nullptr;
    }
    fair::mq::MessagePtr getMessage(void* p) override
    {
      return mUpstream->getMessage(p);
    }
 
    void* setMessage(fair::mq::MessagePtr fmm) override
    {
      return mUpstream->setMessage(std::move(fmm));
    }
 
    fair::mq::TransportFactory* getTransportFactory() noexcept override
    {
      return mUpstream->getTransportFactory();
    }
 
    [[nodiscard]] size_t getNumberOfMessages() const noexcept override
    {
      return mUpstream->getNumberOfMessages();
    }
 
   protected:
    void* do_allocate(size_t bytes, size_t alignment) override
    {
      return mUpstream->allocate(bytes, alignment < 64 ? 64 : alignment);
    }
 
    void do_deallocate(void* p, size_t bytes, size_t alignment) override
    {
      return mUpstream->deallocate(p, bytes, alignment < 64 ? 64 : alignment);
    }
 
    [[nodiscard]] bool do_is_equal(const pmr::memory_resource& other) const noexcept override
    {
      return this == &other;
    }
 
   private:
    fair::mq::MemoryResource* mUpstream = nullptr;
  };
 
  template <typename T>
  class ContainerRefObject : public ContextObject
  {
   public:
    using value_type = typename T::value_type;
    using return_type = T;
    using buffer_type = return_type;
    static_assert(std::is_base_of<o2::pmr::polymorphic_allocator<value_type>, typename T::allocator_type>::value, "container must have polymorphic allocator");
    ContainerRefObject() = delete;
    template <typename ContextType, typename... Args>
    ContainerRefObject(ContextType* context, fair::mq::MessagePtr&& headerMsg, RouteIndex routeIndex, int index, Args&&... args)
      : ContextObject(std::forward<fair::mq::MessagePtr>(headerMsg), routeIndex),
        // the transport factory
        mFactory{context->proxy().getOutputTransport(routeIndex)},
        // the memory resource takes ownership of the message
        mResource{mFactory ? AlignedMemoryResource(mFactory->GetMemoryResource()) : AlignedMemoryResource(nullptr)},
        // create the vector with apropriate underlying memory resource for the message
        mData{std::forward<Args>(args)..., pmr::polymorphic_allocator<value_type>(&mResource)}
    {
      // FIXME: drop this repeated check and make sure at initial setup of devices that everything is fine
      // introduce error policy
      if (mFactory == nullptr) {
        throw runtime_error_f("failed to get transport factory for route %d", routeIndex);
      }
      if (mResource.isValid() == false) {
        throw runtime_error_f("no memory resource for channel %d", routeIndex);
      }
    }
    ~ContainerRefObject() override = default;
 
    fair::mq::Parts finalize() final
    {
      assert(mParts.Size() == 1);
      auto payloadMsg = o2::pmr::getMessage(std::move(mData));
      mParts.AddPart(std::move(payloadMsg));
      return ContextObject::finalize();
    }
 
    operator return_type&()
    {
      return mData;
    }
 
    return_type& get()
    {
      return mData;
    }
 
    value_type* data()
    {
      return mData.data();
    }
 
   private:
    fair::mq::TransportFactory* mFactory = nullptr; 
    AlignedMemoryResource mResource;                
    buffer_type mData;                              
  };
 
  template <typename T, typename _BASE = ContainerRefObject<std::vector<T, o2::pmr::polymorphic_allocator<T>>>>
  class VectorObject : public _BASE
  {
   public:
    template <typename... Args>
    VectorObject(Args&&... args) : _BASE(std::forward<Args>(args)...)
    {
    }
  };
 
  // SpanObject creates a trivial binary object for an array of elements of
  // type T and holds a span over the elements
  // FIXME: probably obsolete after introducing of vector with polymorphic_allocator
  template <typename T>
  class SpanObject : public ContextObject
  {
   public:
    static_assert(is_messageable<T>::value == true, "unconsistent type");
    using value_type = gsl::span<T>;
    SpanObject() = delete;
    template <typename ContextType>
    SpanObject(ContextType* context, fair::mq::MessagePtr&& headerMsg, RouteIndex routeIndex, int index, size_t nElements)
      : ContextObject(std::forward<fair::mq::MessagePtr>(headerMsg), routeIndex)
    {
      // create the span object for the memory of the payload message
      // TODO: we probably also want to check consistency of the header message, i.e. payloadSize member
      auto payloadMsg = context->createMessage(routeIndex, index, nElements * sizeof(T));
      mValue = value_type(reinterpret_cast<T*>(payloadMsg->GetData()), nElements);
      assert(mParts.Size() == 1);
      mParts.AddPart(std::move(payloadMsg));
    }
    ~SpanObject() override = default;
 
    operator value_type&()
    {
      return mValue;
    }
 
    value_type& get()
    {
      return mValue;
    }
 
   private:
    value_type mValue;
  };
 
  using Messages = std::vector<std::unique_ptr<ContextObject>>;
 
  template <typename T, typename BASE = std::default_delete<T>>
  class ScopeHook : public BASE
  {
   public:
    using base = std::default_delete<T>;
    using self_type = ScopeHook<T>;
    ScopeHook() = default;
    ScopeHook(MessageContext* context)
      : mContext(context)
    {
    }
    ~ScopeHook() = default;
 
    // forbid assignment operator to prohibid changing the Deleter
    // resource control property once used in the unique_ptr
    self_type& operator=(const self_type&) = delete;
 
    void operator()(T* ptr) const
    {
      if (!mContext) {
        // TODO: decide whether this is an error or not
        // can also check if the standard constructor can be dropped to make sure that
        // the ScopeHook is always set up with a context
        throw runtime_error("No context available to schedule the context object");
        return base::operator()(ptr);
      }
      // keep the object alive and add to message list of the context
      mContext->schedule(Messages::value_type(ptr));
    }
 
   private:
    MessageContext* mContext = nullptr;
  };
 
  template <typename T>
  using ContextObjectScope = std::unique_ptr<T, ScopeHook<T>>;
 
  template <typename T, typename... Args>
  auto& add(Args&&... args)
  {
    mMessages.push_back(std::move(make<T>(std::forward<Args>(args)...)));
    // return a reference to the element in the vector of unique pointers
    return *dynamic_cast<T*>(mMessages.back().get());
  }
 
  template <typename T, typename... Args>
  Messages::value_type make(Args&&... args)
  {
    static_assert(std::is_base_of<ContextObject, T>::value == true, "type must inherit ContextObject interface");
    return std::make_unique<T>(this, std::forward<Args>(args)...);
  }
 
  template <typename T, typename... Args>
  ContextObjectScope<T> make_scoped(Args&&... args)
  {
    ContextObjectScope<T> scope(dynamic_cast<T*>(make<T>(std::forward<Args>(args)...).release()), ScopeHook<T>(this));
    return scope;
  }
 
  void schedule(Messages::value_type&& message);
 
  Messages getMessagesForSending()
  {
    // before starting iteration, message lists are merged
    for (auto& message : mScheduledMessages) {
      mMessages.emplace_back(std::move(message));
    }
    mScheduledMessages.clear();
    return std::move(mMessages);
  }
 
  size_t size()
  {
    return mMessages.size();
  }
 
  void clear();
 
  void discard();
 
  FairMQDeviceProxy& proxy()
  {
    return mProxy;
  }
 
  // Add a message to cache and returns a unique identifier for
  // such cached message.
  int64_t addToCache(std::unique_ptr<fair::mq::Message>& message);
  // Clone a message from cache so that it can be added to the context
  [[nodiscard]] std::unique_ptr<fair::mq::Message> cloneFromCache(int64_t id) const;
  // Prune a message from cache
  void pruneFromCache(int64_t id);
 
  // FIXME: can that be const?
  fair::mq::MessagePtr createMessage(RouteIndex routeIndex, int index, size_t size);
  fair::mq::MessagePtr createMessage(RouteIndex routeIndex, int index, void* data, size_t size, fair::mq::FreeFn* ffn, void* hint);
 
  o2::header::DataHeader* findMessageHeader(const Output& spec);
  o2::header::Stack* findMessageHeaderStack(const Output& spec);
  [[nodiscard]] int countDeviceOutputs(bool excludeDPLOrigin = false) const;
  void fakeDispatch() { mDispatchState = DispatchState::Dispatched; }
  [[nodiscard]] DispatchState dispatchState() const { return mDispatchState; }
  o2::framework::DataProcessingHeader* findMessageDataProcessingHeader(const Output& spec);
  std::pair<o2::header::DataHeader*, o2::framework::DataProcessingHeader*> findMessageHeaders(const Output& spec);
 
 private:
  FairMQDeviceProxy& mProxy;
  Messages mMessages;
  Messages mScheduledMessages;
  DispatchState mDispatchState = DispatchState::NotDispatched;
  DispatchControl mDispatchControl;
  std::unordered_map<int64_t, std::unique_ptr<fair::mq::Message>> mMessageCache;
};
} // namespace o2::framework
#endif // O2_FRAMEWORK_MESSAGECONTEXT_H_