pool.hpp

/*
    SPDX-License-Identifier: MIT
    Copyright © 2023-2024 Amebis
*/
 
#pragma once
 
#include "compat.hpp"
#include "spinlock.hpp"
#ifdef _WIN32
#include "windows.h"
#endif
#include <list>
#include <map>
#include <mutex>
 
namespace stdex
{
    template <class T>
    class pool
    {
    public:
#ifdef _WIN32
        using numaid_t = USHORT;
#else
        using numaid_t = int;
#endif
 
    private:
        struct numaentry_t {
            mutable spinlock lock;
            std::list<T> list;
        };
 
        mutable std::mutex m_mutex;
        std::map<numaid_t, numaentry_t> m_available;
 
    private:
        static numaid_t numa_node()
        {
#if defined(_WIN32)
            PROCESSOR_NUMBER Processor;
            GetCurrentProcessorNumberEx(&Processor);
            USHORT NodeNumber = 0;
            return GetNumaProcessorNodeEx(&Processor, &NodeNumber) ? NodeNumber : 0;
#elif defined(__APPLE__)
            return 0;
#else
            return numa_node_of_cpu(sched_getcpu());
#endif
        }
 
        numaentry_t& numa_entry(numaid_t numa = numa_node())
        {
            const std::lock_guard<std::mutex> guard(m_mutex);
            return m_available[numa];
        }
 
    public:
        T pop(_In_ numaid_t numa = numa_node())
        {
            auto& ne = numa_entry(numa);
            const std::lock_guard<spinlock> guard(ne.lock);
            if (!ne.list.empty()) {
                auto r = std::move(ne.list.front());
                ne.list.pop_front();
                return r;
            }
            return T();
        }
 
        void push(_Inout_ T&& r, _In_ numaid_t numa = numa_node())
        {
            auto& ne = numa_entry(numa);
            const std::lock_guard<spinlock> guard(ne.lock);
            ne.list.push_front(std::move(r));
        }
 
        void clear()
        {
            const std::lock_guard<std::mutex> guard_m(m_mutex);
            for (auto& ne : m_available) {
                const std::lock_guard<spinlock> guard_l(ne.second.lock);
                while (!ne.second.list.empty())
                    ne.second.list.pop_front();
            }
        }
    };
}