MultiServicePool.h

Go to the documentation of this file.

 
#pragma once
#include "IoThreadPool.h"
#include "catapult/utils/Logging.h"
#include "catapult/functions.h"
#include <memory>
#include <thread>
#include <vector>
 
namespace catapult { namespace thread {
 
    class MultiServicePool {
    public:
        enum class IsolatedPoolMode {
            Enabled,
 
            Disabled
        };
 
        static constexpr size_t DefaultPoolConcurrency() {
            return 0;
        }
 
    public:
        // region ServiceGroup
 
        class ServiceGroup {
        public:
            ServiceGroup(const std::shared_ptr<thread::IoThreadPool>& pPool, const std::string& name)
                    : m_pPool(pPool)
                    , m_name(name)
            {}
 
        public:
            size_t numServices() const {
                return m_shutdownFunctions.size();
            }
 
        public:
            template<typename TService>
            auto registerService(const std::shared_ptr<TService>& pService) {
                m_services.push_back(pService);
 
                auto serviceName = m_name + " service " + std::to_string(m_shutdownFunctions.size() + 1);
                m_shutdownFunctions.push_back([pService, serviceName]() {
                    // shutdown the service
                    CATAPULT_LOG(info) << "shutting down " << serviceName;
                    pService->shutdown();
                });
 
                CATAPULT_LOG(debug) << "registered " << serviceName;
                return pService;
            }
 
            template<typename TFactory, typename... TArgs>
            auto pushService(TFactory factory, TArgs&&... args) {
                // create a service around the pool
                auto pService = factory(m_pPool, std::forward<TArgs>(args)...);
                return registerService(pService);
            }
 
            void shutdown() {
                // 1. shutdown the services
                for (auto iter = m_shutdownFunctions.rbegin(); m_shutdownFunctions.rend() != iter; ++iter)
                    (*iter)();
 
                m_shutdownFunctions.clear();
 
                // 2. wait for all service references to be released
                for (const auto& pService : m_services)
                    WaitForLastReference(pService);
            }
 
        private:
            std::shared_ptr<thread::IoThreadPool> m_pPool;
            std::string m_name;
            std::vector<std::shared_ptr<void>> m_services;
            std::vector<action> m_shutdownFunctions;
        };
 
        // endregion
 
    public:
        MultiServicePool(const std::string& name, size_t numWorkerThreads, IsolatedPoolMode isolatedPoolMode = IsolatedPoolMode::Enabled)
                : m_name(name)
                , m_isolatedPoolMode(isolatedPoolMode)
                , m_numTotalIsolatedPoolThreads(0)
                , m_numServiceGroups(0)
                , m_pPool(CreateThreadPool(numWorkerThreads, name))
        {}
 
        ~MultiServicePool() {
            shutdown();
        }
 
    public:
        size_t numWorkerThreads() const {
            return (m_pPool ? m_pPool->numWorkerThreads() : 0) + m_numTotalIsolatedPoolThreads;
        }
 
        size_t numServiceGroups() const {
            return m_numServiceGroups;
        }
 
        size_t numServices() const {
            // don't include the service groups in the count
            auto numServices = m_shutdownFunctions.size() - numServiceGroups();
            for (const auto& pGroup : m_serviceGroups)
                numServices += pGroup->numServices();
 
            return numServices;
        }
 
    private:
        template<typename TService>
        auto registerService(const std::shared_ptr<TService>& pService, const std::string& serviceName) {
            m_shutdownFunctions.push_back([pService, serviceName]() {
                CATAPULT_LOG(info) << "shutting down " << serviceName;
 
                // shutdown the service and wait for all callbacks to complete
                pService->shutdown();
                WaitForLastReference(pService);
            });
 
            return pService;
        }
 
    public:
        std::shared_ptr<ServiceGroup> pushServiceGroup(const std::string& name) {
            auto pGroup = std::make_shared<ServiceGroup>(m_pPool, m_name + "::" + name);
            m_serviceGroups.push_back(pGroup);
            registerService(pGroup, name + " (service group)");
            ++m_numServiceGroups;
            return pGroup;
        }
 
        std::shared_ptr<thread::IoThreadPool> pushIsolatedPool(const std::string& name) {
            return pushIsolatedPool(name, DefaultPoolConcurrency());
        }
 
        std::shared_ptr<thread::IoThreadPool> pushIsolatedPool(const std::string& name, size_t numWorkerThreads) {
            class PoolServiceAdapter {
            public:
                explicit PoolServiceAdapter(const std::shared_ptr<thread::IoThreadPool>& pPool) : m_pPool(pPool)
                {}
 
            public:
                void shutdown() {
                    WaitForLastReference(m_pPool);
                    m_pPool->join();
                }
 
            private:
                std::shared_ptr<thread::IoThreadPool> m_pPool;
            };
 
            // when isolated pool mode is disabled, use the main pool for everything
            if (IsolatedPoolMode::Disabled == m_isolatedPoolMode)
                return m_pPool;
 
            auto pPool = CreateThreadPool(numWorkerThreads, name);
            registerService(std::make_shared<PoolServiceAdapter>(pPool), name + " (isolated pool)");
            m_numTotalIsolatedPoolThreads += pPool->numWorkerThreads();
            return pPool;
        }
 
    public:
        void shutdown() {
            // if the pool has already been destroyed, don't do anything
            if (!m_pPool)
                return;
 
            // 1. clear the dependent entities
            m_serviceGroups.clear();
 
            // 2. shutdown the services
            for (auto iter = m_shutdownFunctions.rbegin(); m_shutdownFunctions.rend() != iter; ++iter)
                (*iter)();
 
            m_shutdownFunctions.clear();
            m_numTotalIsolatedPoolThreads = 0;
            m_numServiceGroups = 0;
 
            // 3. after the services are destroyed, the thread pool can be safely destroyed
            WaitForLastReference(m_pPool);
            m_pPool.reset();
        }
 
    private:
        static std::shared_ptr<thread::IoThreadPool> CreateThreadPool(size_t numWorkerThreads, const std::string& name) {
            numWorkerThreads = DefaultPoolConcurrency() == numWorkerThreads ? std::thread::hardware_concurrency() : numWorkerThreads;
            auto pPool = thread::CreateIoThreadPool(numWorkerThreads, name.c_str());
            pPool->start();
            return std::move(pPool);
        }
 
        template<typename T>
        static void WaitForLastReference(const std::shared_ptr<T>& pVoid) {
            volatile long useCount;
            while (1 < (useCount = pVoid.use_count()))
                std::this_thread::yield();
        }
 
    private:
        std::string m_name;
        IsolatedPoolMode m_isolatedPoolMode;
        size_t m_numTotalIsolatedPoolThreads;
        size_t m_numServiceGroups;
        std::shared_ptr<thread::IoThreadPool> m_pPool;
        std::vector<std::shared_ptr<ServiceGroup>> m_serviceGroups;
        std::vector<action> m_shutdownFunctions;
    };
}}