subscriber.hpp

Go to the documentation of this file.

#ifndef KAGOME_SUBSCRIPTION_SUBSCRIBER_HPP
#define KAGOME_SUBSCRIPTION_SUBSCRIBER_HPP

#include <atomic>
#include <functional>
#include <memory>
#include <mutex>

#include "subscription/subscription_engine.hpp"

namespace kagome::subscription {

  inline SubscriptionSetId generateNextId() {
    static std::atomic<SubscriptionSetId> id{0ll};
    SubscriptionSetId result = ++id;
    return result;
  }

  template <typename EventKey, typename Receiver, typename... Arguments>
  class Subscriber final : public std::enable_shared_from_this<
                               Subscriber<EventKey, Receiver, Arguments...>> {
   public:
    using EventType = EventKey;
    using ReceiverType = Receiver;
    using Hash = size_t;

    using SubscriptionEngineType =
        SubscriptionEngine<EventType, ReceiverType, Arguments...>;
    using SubscriptionEnginePtr = std::shared_ptr<SubscriptionEngineType>;

    using CallbackFnType = std::function<void(SubscriptionSetId,
                                              ReceiverType &,
                                              const EventType &,
                                              const Arguments &...)>;

   private:
    using SubscriptionsContainer =
        std::unordered_map<EventType,
                           typename SubscriptionEngineType::IteratorType>;
    using SubscriptionsSets =
        std::unordered_map<SubscriptionSetId, SubscriptionsContainer>;
    SubscriptionEnginePtr engine_;
    ReceiverType object_;

    std::mutex subscriptions_cs_;
    SubscriptionsSets subscriptions_sets_;

    CallbackFnType on_notify_callback_;

   public:
    template <typename... SubscriberConstructorArgs>
    explicit Subscriber(SubscriptionEnginePtr &ptr,
                        SubscriberConstructorArgs &&...args)
        : engine_(ptr),
          object_(std::forward<SubscriberConstructorArgs>(args)...) {}

    ~Subscriber() {
      // Unsubscribe all
      for (auto &[_, subscriptions] : subscriptions_sets_)
        for (auto &[key, it] : subscriptions) engine_->unsubscribe(key, it);
    }

    Subscriber(const Subscriber &) = delete;
    Subscriber &operator=(const Subscriber &) = delete;

    Subscriber(Subscriber &&) = default;             // NOLINT
    Subscriber &operator=(Subscriber &&) = default;  // NOLINT

    void setCallback(CallbackFnType &&f) {
      on_notify_callback_ = std::move(f);
    }

    SubscriptionSetId generateSubscriptionSetId() {
      return generateNextId();
    }

    void subscribe(SubscriptionSetId id, const EventType &key) {
      std::lock_guard lock(subscriptions_cs_);
      auto &&[it, inserted] = subscriptions_sets_[id].emplace(
          key, typename SubscriptionEngineType::IteratorType{});

      if (inserted)
        it->second = engine_->subscribe(id, key, this->weak_from_this());
    }

    bool unsubscribe(SubscriptionSetId id, const EventType &key) {
      std::lock_guard<std::mutex> lock(subscriptions_cs_);
      if (auto set_it = subscriptions_sets_.find(id);
          set_it != subscriptions_sets_.end()) {
        auto &subscriptions = set_it->second;
        auto it = subscriptions.find(key);
        if (subscriptions.end() != it) {
          engine_->unsubscribe(key, it->second);
          subscriptions.erase(it);
          return true;
        }
      }
      return false;
    }

    bool unsubscribe(SubscriptionSetId id) {
      std::lock_guard<std::mutex> lock(subscriptions_cs_);
      if (auto set_it = subscriptions_sets_.find(id);
          set_it != subscriptions_sets_.end()) {
        auto &subscriptions = set_it->second;
        for (auto &[key, it] : subscriptions) engine_->unsubscribe(key, it);

        subscriptions_sets_.erase(set_it);
        return true;
      }
      return false;
    }

    void unsubscribe() {
      std::lock_guard<std::mutex> lock(subscriptions_cs_);
      for (auto &[_, subscriptions] : subscriptions_sets_)
        for (auto &[key, it] : subscriptions) engine_->unsubscribe(key, it);

      subscriptions_sets_.clear();
    }

    void on_notify(SubscriptionSetId set_id,
                   const EventType &key,
                   const Arguments &...args) {
      if (nullptr != on_notify_callback_)
        on_notify_callback_(set_id, object_, key, args...);
    }

    ReceiverType &get() {
      return object_;
    }
  };

}  // namespace kagome::subscription

#endif  // KAGOME_SUBSCRIPTION_SUBSCRIBER_HPP