PubSubService.h

Go to the documentation of this file.

 
#pragma once
 
#include <trantor/utils/NonCopyable.h>
#include <functional>
#include <mutex>
#include <shared_mutex>
#include <string>
#include <memory>
#include <unordered_map>
 
namespace drogon
{
using SubscriberID = uint64_t;

template <typename MessageType>
class Topic : public trantor::NonCopyable
{
  public:
    using MessageHandler = std::function<void(const MessageType &)>;
#if __cplusplus >= 201703L | defined _WIN32
    using SharedMutex = std::shared_mutex;
#else
    using SharedMutex = std::shared_timed_mutex;
#endif
    void publish(const MessageType &message) const
    {
        std::shared_lock<SharedMutex> lock(mutex_);
        for (auto &pair : handlersMap_)
        {
            pair.second(message);
        }
    }

    SubscriberID subscribe(const MessageHandler &handler)
    {
        std::unique_lock<SharedMutex> lock(mutex_);
        handlersMap_[++id_] = handler;
        return id_;
    }

    SubscriberID subscribe(MessageHandler &&handler)
    {
        std::unique_lock<SharedMutex> lock(mutex_);
        handlersMap_[++id_] = std::move(handler);
        return id_;
    }

    void unsubscribe(SubscriberID id)
    {
        std::unique_lock<SharedMutex> lock(mutex_);
        handlersMap_.erase(id);
    }

    bool empty() const
    {
        std::shared_lock<SharedMutex> lock(mutex_);
        return handlersMap_.empty();
    }

    void clear()
    {
        std::unique_lock<SharedMutex> lock(mutex_);
        handlersMap_.clear();
    }
 
  private:
    std::unordered_map<SubscriberID, MessageHandler> handlersMap_;
    mutable SharedMutex mutex_;
    SubscriberID id_{0};
};

template <typename MessageType>
class PubSubService : public trantor::NonCopyable
{
  public:
    using MessageHandler =
        std::function<void(const std::string &, const MessageType &)>;
#if __cplusplus >= 201703L | defined _WIN32
    using SharedMutex = std::shared_mutex;
#else
    using SharedMutex = std::shared_timed_mutex;
#endif

    void publish(const std::string &topicName, const MessageType &message) const
    {
        std::shared_ptr<Topic<MessageType>> topicPtr;
        {
            std::shared_lock<SharedMutex> lock(mutex_);
            auto iter = topicMap_.find(topicName);
            if (iter != topicMap_.end())
            {
                topicPtr = iter->second;
            }
            else
            {
                return;
            }
        }
        topicPtr->publish(message);
    }

    SubscriberID subscribe(const std::string &topicName,
                           const MessageHandler &handler)
    {
        auto topicHandler = [topicName, handler](const MessageType &message) {
            handler(topicName, message);
        };
        return subscribeToTopic(topicName, std::move(topicHandler));
    }

    SubscriberID subscribe(const std::string &topicName,
                           MessageHandler &&handler)
    {
        auto topicHandler = [topicName, handler = std::move(handler)](
                                const MessageType &message) {
            handler(topicName, message);
        };
        return subscribeToTopic(topicName, std::move(topicHandler));
    }

    void unsubscribe(const std::string &topicName, SubscriberID id)
    {
        {
            std::shared_lock<SharedMutex> lock(mutex_);
            auto iter = topicMap_.find(topicName);
            if (iter == topicMap_.end())
            {
                return;
            }
            iter->second->unsubscribe(id);
            if (!iter->second->empty())
                return;
        }
        std::unique_lock<SharedMutex> lock(mutex_);
        auto iter = topicMap_.find(topicName);
        if (iter == topicMap_.end())
        {
            return;
        }
        if (iter->second->empty())
            topicMap_.erase(iter);
    }

    size_t size() const
    {
        std::shared_lock<SharedMutex> lock(mutex_);
        return topicMap_.size();
    }

    void clear()
    {
        std::unique_lock<SharedMutex> lock(mutex_);
        topicMap_.clear();
    }

    void removeTopic(const std::string &topicName)
    {
        std::unique_lock<SharedMutex> lock(mutex_);
        topicMap_.erase(topicName);
    }

    bool isTopicEmpty(const std::string &topicName) const
    {
        std::shared_ptr<Topic<MessageType>> topicPtr;
        {
            std::shared_lock<SharedMutex> lock(mutex_);
            auto iter = topicMap_.find(topicName);
            if (iter != topicMap_.end())
            {
                topicPtr = iter->second;
            }
            else
            {
                return true;
            }
        }
        return topicPtr->empty();
    }
 
  private:
    std::unordered_map<std::string, std::shared_ptr<Topic<MessageType>>>
        topicMap_;
    mutable SharedMutex mutex_;
    SubscriberID subID_ = 0;
 
    SubscriberID subscribeToTopic(
        const std::string &topicName,
        typename Topic<MessageType>::MessageHandler &&handler)
    {
        {
            std::shared_lock<SharedMutex> lock(mutex_);
            auto iter = topicMap_.find(topicName);
            if (iter != topicMap_.end())
            {
                return iter->second->subscribe(std::move(handler));
            }
        }
        std::unique_lock<SharedMutex> lock(mutex_);
        auto iter = topicMap_.find(topicName);
        if (iter != topicMap_.end())
        {
            return iter->second->subscribe(std::move(handler));
        }
        auto topicPtr = std::make_shared<Topic<MessageType>>();
        auto id = topicPtr->subscribe(std::move(handler));
        topicMap_[topicName] = std::move(topicPtr);
        return id;
    }
};
}  // namespace drogon