cpp_server/kv_store.hpp
2025-06-13 18:55:13 -05:00

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);
}
};