ECAP_Server/module/Local_Server/Data_Source/Database.h

#pragma once
#include "BaseStation.h"
#include "../Aircraft/Aircraft.h"
#include "../Aircraft/Flight_VTO.h"
#include "../External_Database/export.h"
#include <unordered_map>
#include <memory>
#include <vector>
#include <shared_mutex>
#include <functional>


template <
    typename Key_Type,
    typename Value_Type,
    size_t Shard_Count = 64
>
struct Data_Map {
    using Ptr = std::shared_ptr<Value_Type>;

    size_t size() const {
        size_t ret = 0;

        for (const auto& shard : shards) {
            std::shared_lock<std::shared_mutex> g(shard.mtx);
            ret += shard.map.size();
        }

        return ret;
    }

    Ptr get(const Key_Type& key) const {
        const auto& shard = get_shard(key);

        std::shared_lock<std::shared_mutex> g(shard.mtx);

        auto iter = shard.map.find(key);
        if (iter == shard.map.end()) {
            return nullptr;
        }

        return iter->second;
    }

    Ptr create(const Key_Type& key) {
        auto& shard = get_shard(key);

        {
            std::shared_lock<std::shared_mutex> g(shard.mtx);

            auto iter = shard.map.find(key);
            if (iter != shard.map.end()) {
                return iter->second;
            }
        }

        // 构造对象放锁外，避免构造期间占锁
        auto new_value = std::make_shared<Value_Type>(key);

        {
            std::unique_lock<std::shared_mutex> g(shard.mtx);

            auto [iter, inserted] = shard.map.emplace(key, new_value);

            // 如果别的线程已经创建了，返回已有对象
            return iter->second;
        }
    }

    void remove(const Key_Type& key) {
        auto& shard = get_shard(key);

        std::unique_lock<std::shared_mutex> g(shard.mtx);
        shard.map.erase(key);
    }

    std::vector<Ptr> values() const {
        std::vector<Ptr> ret;

        for (const auto& shard : shards) {
            std::shared_lock<std::shared_mutex> g(shard.mtx);

            for (const auto& [key, val] : shard.map) {
                ret.push_back(val);
            }
        }

        return ret;
    }

    template <typename Cond>
    std::vector<Ptr> get_cond(Cond&& cond) const {
        std::vector<Ptr> ret;

        for (const auto& shard : shards) {
            std::vector<Ptr> snapshot;

            {
                std::shared_lock<std::shared_mutex> g(shard.mtx);

                snapshot.reserve(shard.map.size());

                for (const auto& [key, val] : shard.map) {
                    snapshot.push_back(val);
                }
            }

            // cond 放锁外执行
            for (const auto& val : snapshot) {
                if (cond(val)) {
                    ret.push_back(val);
                }
            }
        }

        return ret;
    }

    template <typename Cond>
    void remove_cond(Cond&& cond) {
        for (auto& shard : shards) {
            std::vector<std::pair<Key_Type, Ptr>> snapshot;

            {
                std::shared_lock<std::shared_mutex> g(shard.mtx);

                snapshot.reserve(shard.map.size());

                for (const auto& [key, val] : shard.map) {
                    snapshot.emplace_back(key, val);
                }
            }

            std::vector<std::pair<Key_Type, Ptr>> need_remove;

            // cond 放锁外执行
            for (const auto& [key, val] : snapshot) {
                if (cond(val)) {
                    need_remove.emplace_back(key, val);
                }
            }

            if (need_remove.empty()) {
                continue;
            }

            {
                std::unique_lock<std::shared_mutex> g(shard.mtx);

                for (const auto& [key, old_val] : need_remove) {
                    auto iter = shard.map.find(key);
                    if (iter == shard.map.end()) {
                        continue;
                    }

                    // 防止同 key 已经被换成新对象，误删
                    if (iter->second == old_val) {
                        shard.map.erase(iter);
                    }
                }
            }
        }
    }

    void delete_all() {
        for (auto& shard : shards) {
            std::unique_lock<std::shared_mutex> g(shard.mtx);
            shard.map.clear();
        }
    }

protected:
    struct Shard {
        mutable std::shared_mutex mtx;
        std::unordered_map<Key_Type, Ptr> map;
    };

    Shard& get_shard(const Key_Type& key) {
        const size_t index = hasher(key) % Shard_Count;
        return shards[index];
    }

    const Shard& get_shard(const Key_Type& key) const {
        const size_t index = hasher(key) % Shard_Count;
        return shards[index];
    }

protected:
    std::hash<Key_Type> hasher;
    Shard shards[Shard_Count];
};


class DataBase : public SSR::Data_Source_Interface {
public:
    std::shared_ptr<SSR::Aircraft_Info> get_aircraft(const std::string& icao) override;
    std::shared_ptr<SSR::Aircraft_Info> create_aircraft(const std::string& icao) override;
    Psc::JSON get_aircraftlist();
    Psc::JSON get_aircraftlist(int limit_msg_num);
    BaseStation base_station;
    void delete_timeout_aircraft();
    std::string get_key() override {
        return key;
    }
    std::string key;
    std::atomic<size_t> have_pos_aircraft_num{};
    Psc::JSON get_all_aircraft_json();
    Psc::JSON get_aircraft_list_after(time_t timestamp);
#ifdef Cache_Some_Mode_S
    Psc::JSON get_limit_mode_s_msg_info();
#endif
    size_t get_aircraft_num();
    void clear_aircraft() { aircraft_map.delete_all(); }
    ~DataBase() override {
        aircraft_map.delete_all();
    }

protected:


    Data_Map<std::string, Aircraft> aircraft_map;
};


struct Mode_S_Msg;
void database_server(Global* g);