session management optimization

kv_store.hpp

@@ -1,12 +1,13 @@
 #pragma once
 
 #include <unordered_map>
-#include <mutex>
+#include <shared_mutex>
 #include <fstream>
 #include <optional>
 #include <string>
 #include <variant>
 #include <sstream>
+#include <mutex>
 
 using std::string;
 
@@ -15,13 +16,13 @@ public:
 	using Value = std::variant<int, double, string, bool>;
 
 	void set(const string& key, const Value& value) {
-		std::lock_guard<std::mutex> lock(mutex_);
+		std::unique_lock<std::shared_mutex> lock(mutex_);
 		data_[key] = value;
 	}
 
 	template<typename T>
 	std::optional<T> get(const string& key) {
-		std::lock_guard<std::mutex> lock(mutex_);
+		std::shared_lock<std::shared_mutex> lock(mutex_);
 		auto it = data_.find(key);
 		if (it == data_.end()) return std::nullopt;
 		if (auto* ptr = std::get_if<T>(&it->second)) {
@@ -30,17 +31,23 @@ public:
 		return std::nullopt;
 	}
 
+	const std::unordered_map<string, Value>& get_all() {
+		std::shared_lock<std::shared_mutex> lock(mutex_);
+		return data_;
+	}
+
 	bool del(const string& key) {
-		std::lock_guard<std::mutex> lock(mutex_);
+		std::unique_lock<std::shared_mutex> lock(mutex_);
 		return data_.erase(key) > 0;
 	}
 
 	size_t size() {
-		std::lock_guard<std::mutex> lock(mutex_);
+		std::shared_lock<std::shared_mutex> lock(mutex_);
 		return data_.size();
 	}
 
 	void set_file(const string& filename) {
+		std::unique_lock<std::shared_mutex> lock(mutex_);
 		filename_ = filename;
 	}
 
@@ -48,7 +55,7 @@ public:
 		std::ifstream file(filename_);
 		if (!file.is_open()) return;
 
-		std::lock_guard<std::mutex> lock(mutex_);
+		std::unique_lock<std::shared_mutex> lock(mutex_);
 		string line;
 		while (std::getline(file, line)) {
 			size_t eq = line.find('=');
@@ -68,7 +75,7 @@ public:
 	}
 
 	void save() {
-		std::lock_guard<std::mutex> lock(mutex_);
+		std::shared_lock<std::shared_mutex> lock(mutex_);
 		std::ofstream file(filename_);
 		if (!file.is_open()) return;
 
@@ -94,6 +101,6 @@ private:
 	}
 
 	std::unordered_map<string, Value> data_;
-	mutable std::mutex mutex_;
+	mutable std::shared_mutex mutex_;
 	string filename_ = "store.txt";
 };

server.hpp

@@ -151,16 +151,19 @@ private:
 			}
 
 			// Handle session
-			std::string session_id = get_session_id(req);
-			if (session_id.empty()) {
-				session_id = sessions.create();
+			std::string session_id;
+			bool needs_session = sessions.needs_session(req);
+
+			if (needs_session) {
+				std::string_view existing_id = sessions.extract_session_id(req);
+				session_id = existing_id.empty() ? sessions.create() : std::string(existing_id);
 			}
 
 			Response response;
 
 			// Try router first
 			if (router.handle(req, response)) {
-				set_session_cookie(response, session_id);
+				if (needs_session) set_session_cookie(response, session_id);
 				send_response(client_fd, response, req.version);
 				return;
 			}
@@ -172,26 +175,11 @@ private:
 			} else {
 				response.status = 404;
 				response.set_text("Not Found");
-				set_session_cookie(response, session_id);
+				if (needs_session) set_session_cookie(response, session_id);
 				send_response(client_fd, response, req.version);
 			}
 		}
 
-		std::string get_session_id(const Request& req) {
-			auto it = req.headers.find("Cookie");
-			if (it == req.headers.end()) return "";
-
-			std::string_view cookies = it->second;
-			size_t pos = cookies.find("session_id=");
-			if (pos == std::string_view::npos) return "";
-
-			pos += 11; // length of "session_id="
-			size_t end = cookies.find(';', pos);
-			if (end == std::string_view::npos) end = cookies.length();
-
-			return std::string(cookies.substr(pos, end - pos));
-		}
-
 		void set_session_cookie(Response& response, const std::string& session_id) {
 			response.cookies.push_back("session_id=" + session_id + "; HttpOnly; Path=/; SameSite=Strict");
 		}

session_store.hpp

@@ -1,45 +1,168 @@
 #pragma once
 
 #include "kv_store.hpp"
+#include "request.hpp"
 #include <random>
 #include <unordered_map>
 #include <mutex>
+#include <array>
+#include <chrono>
+#include <optional>
 
 class SessionStore {
+private:
+	static constexpr size_t SHARD_COUNT = 16;
+	static constexpr auto SESSION_TIMEOUT = std::chrono::hours(24);
+
+	struct SessionData {
+		std::unordered_map<std::string, std::string> data;
+		std::chrono::steady_clock::time_point last_access;
+		bool dirty = false;
+	};
+
+	std::array<std::mutex, SHARD_COUNT> mutexes_;
+	std::array<std::unordered_map<std::string, SessionData>, SHARD_COUNT> shards_;
+	KeyValueStore store_;
+	std::mt19937 rng_;
+
+	size_t get_shard(const std::string& id) const {
+		return std::hash<std::string>{}(id) % SHARD_COUNT;
+	}
+
+	void lazy_load_session(const std::string& id, size_t shard) {
+		auto& shard_map = shards_[shard];
+		if (shard_map.find(id) != shard_map.end()) return;
+
+		// Load from persistent store
+		SessionData session;
+		session.last_access = std::chrono::steady_clock::now();
+
+		// Try to load existing session data
+		for (const auto& [key, value] : store_.get_all()) {
+			if (key.starts_with(id + ":")) {
+				std::string session_key = key.substr(id.length() + 1);
+				if (auto* str_val = std::get_if<std::string>(&value)) {
+					session.data[session_key] = *str_val;
+				}
+			}
+		}
+
+		shard_map[id] = std::move(session);
+	}
+
+	void cleanup_expired_sessions() {
+		auto now = std::chrono::steady_clock::now();
+
+		for (size_t i = 0; i < SHARD_COUNT; ++i) {
+			std::lock_guard<std::mutex> lock(mutexes_[i]);
+			auto& shard = shards_[i];
+
+			for (auto it = shard.begin(); it != shard.end();) {
+				if (now - it->second.last_access > SESSION_TIMEOUT) {
+					it = shard.erase(it);
+				} else {
+					++it;
+				}
+			}
+		}
+	}
+
 public:
 	SessionStore() : rng_(std::random_device{}()) {
 		store_.set_file("sessions.txt");
 	}
 
+	bool needs_session(const Request& req) const {
+		// Skip sessions for static content
+		if (req.path.starts_with("/static") ||
+		    req.path.starts_with("/assets") ||
+		    req.path.ends_with(".css") ||
+		    req.path.ends_with(".js") ||
+		    req.path.ends_with(".png") ||
+		    req.path.ends_with(".jpg") ||
+		    req.path.ends_with(".jpeg") ||
+		    req.path.ends_with(".gif") ||
+		    req.path.ends_with(".ico") ||
+		    req.path.ends_with(".svg") ||
+		    req.path.ends_with(".woff") ||
+		    req.path.ends_with(".woff2") ||
+		    req.path.ends_with(".ttf")) {
+			return false;
+		}
+
+		// Need sessions for API and admin routes
+		return req.path.starts_with("/api") ||
+		       req.path.starts_with("/admin") ||
+		       req.method == HttpMethod::POST ||
+		       req.method == HttpMethod::PUT ||
+		       req.method == HttpMethod::DELETE;
+	}
+
 	std::string create() {
 		std::string id = generate_id();
-		std::lock_guard<std::mutex> lock(mutex_);
-		sessions_[id] = {};
+		size_t shard = get_shard(id);
+
+		std::lock_guard<std::mutex> lock(mutexes_[shard]);
+		SessionData session;
+		session.last_access = std::chrono::steady_clock::now();
+		session.dirty = true;
+		shards_[shard][id] = std::move(session);
+
 		return id;
 	}
 
 	std::optional<std::string> get(const std::string& id, const std::string& key) {
-		std::lock_guard<std::mutex> lock(mutex_);
-		auto it = sessions_.find(id);
-		if (it == sessions_.end()) return std::nullopt;
+		size_t shard = get_shard(id);
+		std::lock_guard<std::mutex> lock(mutexes_[shard]);
 
-		auto data_it = it->second.find(key);
-		return data_it != it->second.end() ? std::optional<std::string>{data_it->second} : std::nullopt;
+		lazy_load_session(id, shard);
+
+		auto it = shards_[shard].find(id);
+		if (it == shards_[shard].end()) return std::nullopt;
+
+		it->second.last_access = std::chrono::steady_clock::now();
+		auto data_it = it->second.data.find(key);
+		return data_it != it->second.data.end() ?
+			std::optional<std::string>{data_it->second} : std::nullopt;
 	}
 
 	void set(const std::string& id, const std::string& key, const std::string& value) {
-		std::lock_guard<std::mutex> lock(mutex_);
-		sessions_[id][key] = value;
+		size_t shard = get_shard(id);
+		std::lock_guard<std::mutex> lock(mutexes_[shard]);
+
+		lazy_load_session(id, shard);
+
+		auto& session = shards_[shard][id];
+		session.data[key] = value;
+		session.last_access = std::chrono::steady_clock::now();
+		session.dirty = true;
 	}
 
 	bool del(const std::string& id) {
-		std::lock_guard<std::mutex> lock(mutex_);
-		return sessions_.erase(id) > 0;
+		size_t shard = get_shard(id);
+		std::lock_guard<std::mutex> lock(mutexes_[shard]);
+
+		bool found = shards_[shard].erase(id) > 0;
+
+		// Remove from persistent store
+		if (found) {
+			for (const auto& [key, value] : store_.get_all()) {
+				if (key.starts_with(id + ":")) {
+					store_.del(key);
+				}
+			}
+		}
+
+		return found;
 	}
 
 	size_t size() {
-		std::lock_guard<std::mutex> lock(mutex_);
-		return sessions_.size();
+		size_t total = 0;
+		for (size_t i = 0; i < SHARD_COUNT; ++i) {
+			std::lock_guard<std::mutex> lock(mutexes_[i]);
+			total += shards_[i].size();
+		}
+		return total;
 	}
 
 	void set_file(const std::string& filename) {
@@ -48,24 +171,34 @@ public:
 
 	void load() {
 		store_.load();
+		// Sessions loaded lazily on first access
 	}
 
 	void save() {
-		std::lock_guard<std::mutex> lock(mutex_);
-		for (const auto& [id, data] : sessions_) {
-			for (const auto& [key, value] : data) {
-				store_.set(id + ":" + key, value);
+		// Save only dirty sessions
+		for (size_t i = 0; i < SHARD_COUNT; ++i) {
+			std::lock_guard<std::mutex> lock(mutexes_[i]);
+
+			for (auto& [id, session] : shards_[i]) {
+				if (!session.dirty) continue;
+
+				for (const auto& [key, value] : session.data) {
+					store_.set(id + ":" + key, value);
+				}
+				session.dirty = false;
 			}
 		}
+
+		cleanup_expired_sessions();
 		store_.save();
 	}
 
-private:
-	std::unordered_map<std::string, std::unordered_map<std::string, std::string>> sessions_;
-	mutable std::mutex mutex_;
-	KeyValueStore store_;
-	std::mt19937 rng_;
+	// Use existing cookie helpers
+	static std::string_view extract_session_id(const Request& req) {
+		return req.get_cookie("session_id");
+	}
 
+private:
 	std::string generate_id() {
 		static const char chars[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
 		std::string id;