From 652053128dceb1d15274635aa135659e796ba12f Mon Sep 17 00:00:00 2001
From: Sky Johnson <skylearj@gmail.com>
Date: Fri, 13 Jun 2025 18:55:13 -0500
Subject: [PATCH] shard kv_store

---
 kv_store.hpp | 80 +++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 54 insertions(+), 26 deletions(-)

diff --git a/kv_store.hpp b/kv_store.hpp
index 276d4a0..4fc95d1 100644
--- a/kv_store.hpp
+++ b/kv_store.hpp
@@ -6,7 +6,7 @@
 #include <optional>
 #include <string>
 #include <variant>
-#include <sstream>
+#include <array>
 #include <mutex>
 
 using std::string;
@@ -15,39 +15,64 @@ 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) {
-		std::unique_lock<std::shared_mutex> lock(mutex_);
-		data_[key] = 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) {
-		std::shared_lock<std::shared_mutex> lock(mutex_);
-		auto it = data_.find(key);
-		if (it == data_.end()) return std::nullopt;
+		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;
 	}
 
-	const std::unordered_map<string, Value>& get_all() {
-		std::shared_lock<std::shared_mutex> lock(mutex_);
-		return data_;
+	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) {
-		std::unique_lock<std::shared_mutex> lock(mutex_);
-		return data_.erase(key) > 0;
+		size_t shard = get_shard(key);
+		std::unique_lock<std::shared_mutex> lock(mutexes_[shard]);
+		return shards_[shard].erase(key) > 0;
 	}
 
 	size_t size() {
-		std::shared_lock<std::shared_mutex> lock(mutex_);
-		return data_.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) {
-		std::unique_lock<std::shared_mutex> lock(mutex_);
 		filename_ = filename;
 	}
 
@@ -55,7 +80,6 @@ public:
 		std::ifstream file(filename_);
 		if (!file.is_open()) return;
 
-		std::unique_lock<std::shared_mutex> lock(mutex_);
 		string line;
 		while (std::getline(file, line)) {
 			size_t eq = line.find('=');
@@ -65,22 +89,30 @@ public:
 			string val_str = line.substr(eq + 1);
 			char type = line[0];
 
+			Value value;
 			switch (type) {
-				case 'i': data_[key] = std::stoi(val_str); break;
-				case 'd': data_[key] = std::stod(val_str); break;
-				case 's': data_[key] = val_str; break;
-				case 'b': data_[key] = (val_str == "1"); break;
+				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::shared_lock<std::shared_mutex> lock(mutex_);
 		std::ofstream file(filename_);
 		if (!file.is_open()) return;
 
-		for (const auto& [key, value] : data_) {
-			file << serialize_entry(key, value) << "\n";
+		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";
+			}
 		}
 	}
 
@@ -99,8 +131,4 @@ private:
 			}
 		}, value);
 	}
-
-	std::unordered_map<string, Value> data_;
-	mutable std::shared_mutex mutex_;
-	string filename_ = "store.txt";
 };