ylib/include/util/pool.hpp

#pragma once
#include <iostream>
#include <mutex>
#include <memory>
#include <thread>
#include <atomic>
#include <condition_variable>
#include <vector>
#include <deque>
#include <algorithm>
#include <cmath>
#include <chrono>
#include <optional>

#include "base/error.h"
#include "base/exception.h"
#include "util/time.h"

namespace ylib {


    class example_core{
    public:
        virtual void recover() = 0;
        virtual void task_out() = 0;
    };
    class poolcore:public ylib::error_base
    {
    public:
        virtual ~poolcore() = default;
        virtual void recover(void* example) = 0;
    public:
        void* env = nullptr;
    protected:
        std::mutex m_mutex;
        bool m_closing = false;
        size_t m_max_size = 0;
        size_t m_pop_size = 0; // 当前借出（使用中）数量
    };
    enum EXAMPLE_START_RESULT
    {
        SR_SUCCESS,
        SR_TIMEOUT,
        SR_FAILED,
    };
    template<typename INFO>
    class example:public example_core
    {
    public:
       
    public:
        example()
        {
            m_pool = nullptr;
        }
        virtual EXAMPLE_START_RESULT start(const INFO& info) = 0;
        virtual void close() = 0;

        inline void pool(poolcore *pool)
        {
            m_pool = pool;
        }
        inline poolcore *pool()
        {
            return m_pool;
        }
    protected:
        poolcore *m_pool;
    };

    /**
     * @brief 连接池连接自动归还（可移动，不可拷贝）
     */
    template<typename EXAMPLE>
    class conn_autofree
    {
    public:
        explicit conn_autofree(EXAMPLE* conn = nullptr) : m_conn(conn) {}
        conn_autofree(const conn_autofree&) = delete;
        conn_autofree& operator=(const conn_autofree&) = delete;
        conn_autofree(conn_autofree&& other) noexcept : m_conn(other.m_conn)
        {
            other.m_conn = nullptr;
        }
        conn_autofree& operator=(conn_autofree&& other) noexcept
        {
            if (this != &other) {
                release();
                m_conn = other.m_conn;
                other.m_conn = nullptr;
            }
            return *this;
        }
        ~conn_autofree()
        {
            release();
        }

        EXAMPLE* operator->() const { return m_conn; }
        EXAMPLE* get() const { return m_conn; }
        explicit operator bool() const { return m_conn != nullptr; }
        bool empty() const { return m_conn == nullptr; }
        void reset() { release(); }

    private:
        void release()
        {
            if (m_conn == nullptr) {
                return;
            }
            if (m_conn->pool() != nullptr) {
                m_conn->pool()->recover(m_conn);
            }
            m_conn = nullptr;
        }

        EXAMPLE* m_conn = nullptr;
    };

    /// 连接池增量配置：只填需要覆盖的项，未填项使用默认值
    struct pool_options
    {
        std::optional<size_t> max_size;              // 默认 50
        std::optional<size_t> min_size;              // 默认 2
        std::optional<uint32> idle_timeout_sec;      // 默认 60；0=不按空闲超时释放
        std::optional<uint32> shrink_interval_sec;   // 默认 5
        std::optional<uint32> usage_window_sec;      // 默认 30
        std::optional<double> shrink_threshold;      // 默认 0.5
        std::optional<double> shrink_headroom;       // 默认 2.0
        std::optional<bool> enable_auto_shrink;      // 默认 true
    };

    /// 解析后实际生效的连接池配置
    struct pool_config
    {
        size_t max_size = 50;
        size_t min_size = 2;
        uint32 idle_timeout_sec = 60;
        uint32 shrink_interval_sec = 5;
        uint32 usage_window_sec = 30;
        double shrink_threshold = 0.5;
        double shrink_headroom = 2.0;
        bool enable_auto_shrink = true;
    };

    inline pool_config make_pool_config(const pool_options& opt = {})
    {
        pool_config c;
        if (opt.max_size) c.max_size = *opt.max_size;
        if (opt.min_size) c.min_size = *opt.min_size;
        if (opt.idle_timeout_sec) c.idle_timeout_sec = *opt.idle_timeout_sec;
        if (opt.shrink_interval_sec) c.shrink_interval_sec = *opt.shrink_interval_sec;
        if (opt.usage_window_sec) c.usage_window_sec = *opt.usage_window_sec;
        if (opt.shrink_threshold) c.shrink_threshold = *opt.shrink_threshold;
        if (opt.shrink_headroom) c.shrink_headroom = *opt.shrink_headroom;
        if (opt.enable_auto_shrink) c.enable_auto_shrink = *opt.enable_auto_shrink;

        // 仅做结构约束，不覆盖调用方已给出的业务参数
        if (c.max_size < 1)
            c.max_size = 1;
        if (c.min_size > c.max_size)
            c.min_size = c.max_size;
        if (c.shrink_interval_sec < 1)
            c.shrink_interval_sec = 1;
        if (c.usage_window_sec < 1)
            c.usage_window_sec = 1;
        if (c.shrink_headroom < 1.0)
            c.shrink_headroom = 1.0;
        return c;
    }

    template<typename EXAMPLE,typename INFO>
    class pool :public poolcore
    {
    private:
        struct idle_entry
        {
            void* example = nullptr;
            timestamp idle_since_msec = 0;
        };
    public:
        pool()
        {
            m_max_size = 0;
            m_pop_size = 0;
            m_closing = false;
            m_alive_size = 0;
        }
        ~pool()
        {
            close();
        }
        /// 全部使用默认配置
        bool start(const INFO& info)
        {
            return start(info, pool_options{});
        }
        /// 仅覆盖 max_size，其余默认
        bool start(const INFO& info, size_t max_size)
        {
            pool_options opt;
            opt.max_size = max_size;
            return start(info, opt);
        }
        /// 增量配置：opt 里写了的字段生效，没写的用默认值
        bool start(const INFO& info, const pool_options& opt)
        {
            close();

            {
                std::unique_lock<std::mutex> sp(m_mutex);
                m_closing = false;
                m_info = info;
                m_opt = make_pool_config(opt);

                m_max_size = m_opt.max_size;
                m_pop_size = 0;
                m_alive_size = 0;
                m_idle.clear();
                m_usage_buckets.assign(m_opt.usage_window_sec, 0);
                m_usage_bucket_idx = 0;
                m_usage_bucket_sec = 0;
            }

            // 预热到最低维持数量（创建时不长时间占锁）
            for (size_t i = 0; i < m_opt.min_size; ++i)
            {
                EXAMPLE* ex = create_example();
                if (!ex)
                    break;
                std::unique_lock<std::mutex> sp(m_mutex);
                if (m_closing)
                {
                    sp.unlock();
                    destroy_example(ex);
                    break;
                }
                m_idle.push_back(idle_entry{ ex, ylib::time::now_msec() });
                m_alive_size++;
            }

            if (m_opt.enable_auto_shrink)
            {
                m_shrink_stop = false;
                m_shrink_thread = std::thread([this]() { shrink_loop(); });
            }
            return true;
        }
        void close()
        {
            {
                std::unique_lock<std::mutex> sp(m_mutex);
                m_closing = true;
                m_shrink_stop = true;
            }
            m_shrink_cv.notify_all();
            if (m_shrink_thread.joinable())
                m_shrink_thread.join();

            std::vector<void*> doomed;
            {
                std::unique_lock<std::mutex> sp(m_mutex);
                doomed.reserve(m_idle.size());
                for (auto& e : m_idle)
                    doomed.push_back(e.example);
                m_idle.clear();
                m_alive_size = 0;
                m_pop_size = 0;
            }
            for (void* p : doomed)
                destroy_example(p);
        }
        void* get_ptr() { return (void*)get(); }
        EXAMPLE* get()
        {
            std::vector<void*> doomed;
            EXAMPLE* result = nullptr;
            {
                std::unique_lock<std::mutex> sp(m_mutex);
                if (m_closing)
                    throw ylib::exception("pool is shutting down");

                collect_idle_timeout_unlocked(doomed);

                if (!m_idle.empty())
                {
                    idle_entry entry = m_idle.front();
                    m_idle.pop_front();
                    m_pop_size++;
                    note_usage_unlocked(m_pop_size);
                    result = (EXAMPLE*)entry.example;
                    result->pool(this);
                }
                else if (m_alive_size < m_max_size)
                {
                    // 先占名额，锁外创建，失败再回滚
                    m_alive_size++;
                    m_pop_size++;
                    note_usage_unlocked(m_pop_size);
                    sp.unlock();
                    for (void* p : doomed)
                        destroy_example(p);
                    doomed.clear();

                    EXAMPLE* ex = create_example();
                    if (!ex)
                    {
                        std::unique_lock<std::mutex> sp2(m_mutex);
                        if (m_alive_size > 0) m_alive_size--;
                        if (m_pop_size > 0) m_pop_size--;
                        throw ylib::exception(m_lastErrorDesc.empty() ? "create connection failed" : m_lastErrorDesc);
                    }
                    ex->task_out();
                    return ex;
                }
                else
                {
                    sp.unlock();
                    for (void* p : doomed)
                        destroy_example(p);
                    throw ylib::exception("maximum capacity exceeded");
                }
            }
            for (void* p : doomed)
                destroy_example(p);

            result->task_out();
            return result;
        }
        void recover(void* example) override
        {
            if (example == NULL)
                return;
            ((ylib::example_core*)example)->recover();

            std::vector<void*> doomed;
            {
                std::unique_lock<std::mutex> sp(m_mutex);
                if (m_pop_size > 0)
                    m_pop_size--;

                if (m_closing)
                {
                    if (m_alive_size > 0)
                        m_alive_size--;
                    doomed.push_back(example);
                }
                else
                {
                    m_idle.push_back(idle_entry{ example, ylib::time::now_msec() });
                    collect_shrink_unlocked(doomed);
                }
            }
            for (void* p : doomed)
                destroy_example(p);
        }

        /// 空闲连接数
        size_t idle_size()
        {
            std::unique_lock<std::mutex> sp(m_mutex);
            return m_idle.size();
        }
        /// 使用中连接数
        size_t busy_size()
        {
            std::unique_lock<std::mutex> sp(m_mutex);
            return m_pop_size;
        }
        /// 已建立连接总数
        size_t alive_size()
        {
            std::unique_lock<std::mutex> sp(m_mutex);
            return m_alive_size;
        }
        /// 兼容旧接口：剩余可借出容量
        size_t size()
        {
            std::unique_lock<std::mutex> sp(m_mutex);
            if (m_max_size < m_pop_size)
                return 0;
            return m_max_size - m_pop_size;
        }
        /// 当前生效配置（已与默认值合并）
        const pool_config& options() const { return m_opt; }
    public:
        INFO m_info;
    private:
        EXAMPLE* create_example()
        {
            EXAMPLE* ex = new EXAMPLE;
            ex->pool(this);
            bool init_success = false;
            for (uint32 i = 0; i < 3; i++)
            {
                auto SR = ex->start(m_info);
                if (SR == SR_SUCCESS)
                {
                    init_success = true;
                    break;
                }
                else if (SR == SR_TIMEOUT)
                {
                    std::cout << "start failed." << ex->last_error().c_str() << std::endl;
                    std::cout << "restart " << std::to_string(i + 1).c_str() << "." << std::endl;
                }
                else if (SR == SR_FAILED)
                {
                    break;
                }
            }
            if (!init_success)
            {
                m_lastErrorDesc = ex->last_error();
                delete ex;
                return nullptr;
            }
            return ex;
        }
        void destroy_example(void* example)
        {
            if (!example)
                return;
            EXAMPLE* ex = (EXAMPLE*)example;
            try
            {
                ex->pool(nullptr);
                ex->close();
            }
            catch (...)
            {
            }
            delete ex;
        }
        void note_usage_unlocked(size_t in_use)
        {
            const size_t window = m_usage_buckets.empty() ? 1 : m_usage_buckets.size();
            timestamp sec = ylib::time::now_msec() / 1000;
            if (m_usage_bucket_sec == 0)
            {
                m_usage_bucket_sec = sec;
                m_usage_buckets.assign(window, 0);
                m_usage_bucket_idx = 0;
            }
            while (sec > m_usage_bucket_sec)
            {
                m_usage_bucket_idx = (m_usage_bucket_idx + 1) % window;
                m_usage_buckets[m_usage_bucket_idx] = 0;
                m_usage_bucket_sec++;
                if (sec > m_usage_bucket_sec && (sec - m_usage_bucket_sec) >= (timestamp)window)
                {
                    std::fill(m_usage_buckets.begin(), m_usage_buckets.end(), 0);
                    m_usage_bucket_sec = sec;
                    m_usage_bucket_idx = 0;
                    break;
                }
            }
            if (in_use > m_usage_buckets[m_usage_bucket_idx])
                m_usage_buckets[m_usage_bucket_idx] = in_use;
        }
        size_t usage_peak_unlocked() const
        {
            size_t peak = 0;
            for (size_t v : m_usage_buckets)
            {
                if (v > peak)
                    peak = v;
            }
            if (m_pop_size > peak)
                peak = m_pop_size;
            return peak;
        }
        void collect_idle_timeout_unlocked(std::vector<void*>& doomed)
        {
            if (m_opt.idle_timeout_sec == 0)
                return;
            const timestamp now = ylib::time::now_msec();
            const timestamp timeout_msec = (timestamp)m_opt.idle_timeout_sec * 1000;
            while (m_alive_size > m_opt.min_size && !m_idle.empty())
            {
                if (now - m_idle.front().idle_since_msec < timeout_msec)
                    break;
                doomed.push_back(m_idle.front().example);
                m_idle.pop_front();
                m_alive_size--;
            }
        }
        void collect_shrink_unlocked(std::vector<void*>& doomed)
        {
            if (!m_opt.enable_auto_shrink)
                return;

            collect_idle_timeout_unlocked(doomed);

            if (m_alive_size <= m_opt.min_size)
                return;

            note_usage_unlocked(m_pop_size);

            const size_t peak = usage_peak_unlocked();
            const size_t alive = m_alive_size;
            const size_t threshold_alive = (size_t)std::ceil((double)alive * m_opt.shrink_threshold);

            // 近期峰值仍占存活一半以上：不按量缩减
            if (peak >= threshold_alive)
                return;

            size_t target = (size_t)std::ceil((double)peak * m_opt.shrink_headroom);
            if (target < m_opt.min_size)
                target = m_opt.min_size;
            if (target > m_opt.max_size)
                target = m_opt.max_size;

            while (m_alive_size > target && !m_idle.empty())
            {
                doomed.push_back(m_idle.front().example);
                m_idle.pop_front();
                m_alive_size--;
            }
        }
        void shrink_loop()
        {
            while (true)
            {
                std::vector<void*> doomed;
                {
                    std::unique_lock<std::mutex> sp(m_mutex);
                    m_shrink_cv.wait_for(sp, std::chrono::seconds(m_opt.shrink_interval_sec), [this]() {
                        return m_shrink_stop || m_closing;
                    });
                    if (m_shrink_stop || m_closing)
                        return;
                    collect_shrink_unlocked(doomed);
                }
                for (void* p : doomed)
                    destroy_example(p);
            }
        }
    private:
        pool_config m_opt;
        std::deque<idle_entry> m_idle;
        size_t m_alive_size = 0;

        std::vector<size_t> m_usage_buckets;
        size_t m_usage_bucket_idx = 0;
        timestamp m_usage_bucket_sec = 0;

        std::thread m_shrink_thread;
        std::condition_variable m_shrink_cv;
        bool m_shrink_stop = true;
    };
}