135 lines
3.4 KiB
C++
135 lines
3.4 KiB
C++
#pragma once
|
|
|
|
#include <unordered_map>
|
|
#include <shared_mutex>
|
|
#include <fstream>
|
|
#include <optional>
|
|
#include <string>
|
|
#include <variant>
|
|
#include <array>
|
|
#include <mutex>
|
|
|
|
using std::string;
|
|
|
|
class KeyValueStore {
|
|
public:
|
|
using Value = std::variant<int, double, string, bool>;
|
|
|
|
private:
|
|
static constexpr size_t SHARD_COUNT = 16;
|
|
|
|
std::array<std::shared_mutex, SHARD_COUNT> mutexes_;
|
|
std::array<std::unordered_map<string, Value>, SHARD_COUNT> shards_;
|
|
string filename_ = "store.txt";
|
|
|
|
size_t get_shard(const string& key) const {
|
|
return std::hash<string>{}(key) % SHARD_COUNT;
|
|
}
|
|
|
|
public:
|
|
void set(const string& key, const Value& value) {
|
|
size_t shard = get_shard(key);
|
|
std::unique_lock<std::shared_mutex> lock(mutexes_[shard]);
|
|
shards_[shard][key] = value;
|
|
}
|
|
|
|
template<typename T>
|
|
std::optional<T> get(const string& key) {
|
|
size_t shard = get_shard(key);
|
|
std::shared_lock<std::shared_mutex> lock(mutexes_[shard]);
|
|
auto it = shards_[shard].find(key);
|
|
if (it == shards_[shard].end()) return std::nullopt;
|
|
if (auto* ptr = std::get_if<T>(&it->second)) {
|
|
return *ptr;
|
|
}
|
|
return std::nullopt;
|
|
}
|
|
|
|
std::unordered_map<string, Value> get_all() {
|
|
std::unordered_map<string, Value> result;
|
|
|
|
for (size_t i = 0; i < SHARD_COUNT; ++i) {
|
|
std::shared_lock<std::shared_mutex> lock(mutexes_[i]);
|
|
for (const auto& [key, value] : shards_[i]) {
|
|
result[key] = value;
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
bool del(const string& key) {
|
|
size_t shard = get_shard(key);
|
|
std::unique_lock<std::shared_mutex> lock(mutexes_[shard]);
|
|
return shards_[shard].erase(key) > 0;
|
|
}
|
|
|
|
size_t size() {
|
|
size_t total = 0;
|
|
for (size_t i = 0; i < SHARD_COUNT; ++i) {
|
|
std::shared_lock<std::shared_mutex> lock(mutexes_[i]);
|
|
total += shards_[i].size();
|
|
}
|
|
return total;
|
|
}
|
|
|
|
void set_file(const string& filename) {
|
|
filename_ = filename;
|
|
}
|
|
|
|
void load() {
|
|
std::ifstream file(filename_);
|
|
if (!file.is_open()) return;
|
|
|
|
string line;
|
|
while (std::getline(file, line)) {
|
|
size_t eq = line.find('=');
|
|
if (eq == string::npos || eq < 2) continue;
|
|
|
|
string key = line.substr(2, eq - 2);
|
|
string val_str = line.substr(eq + 1);
|
|
char type = line[0];
|
|
|
|
Value value;
|
|
switch (type) {
|
|
case 'i': value = std::stoi(val_str); break;
|
|
case 'd': value = std::stod(val_str); break;
|
|
case 's': value = val_str; break;
|
|
case 'b': value = (val_str == "1"); break;
|
|
default: continue;
|
|
}
|
|
|
|
size_t shard = get_shard(key);
|
|
std::unique_lock<std::shared_mutex> lock(mutexes_[shard]);
|
|
shards_[shard][key] = std::move(value);
|
|
}
|
|
}
|
|
|
|
void save() {
|
|
std::ofstream file(filename_);
|
|
if (!file.is_open()) return;
|
|
|
|
for (size_t i = 0; i < SHARD_COUNT; ++i) {
|
|
std::shared_lock<std::shared_mutex> lock(mutexes_[i]);
|
|
for (const auto& [key, value] : shards_[i]) {
|
|
file << serialize_entry(key, value) << "\n";
|
|
}
|
|
}
|
|
}
|
|
|
|
private:
|
|
string serialize_entry(const string& key, const Value& value) {
|
|
return std::visit([&key](const auto& val) -> string {
|
|
using T = std::decay_t<decltype(val)>;
|
|
if constexpr (std::is_same_v<T, int>) {
|
|
return "i:" + key + "=" + std::to_string(val);
|
|
} else if constexpr (std::is_same_v<T, double>) {
|
|
return "d:" + key + "=" + std::to_string(val);
|
|
} else if constexpr (std::is_same_v<T, string>) {
|
|
return "s:" + key + "=" + val;
|
|
} else if constexpr (std::is_same_v<T, bool>) {
|
|
return "b:" + key + "=" + (val ? "1" : "0");
|
|
}
|
|
}, value);
|
|
}
|
|
};
|