305 lines
7.6 KiB
C++
305 lines
7.6 KiB
C++
#pragma once
|
|
|
|
#include "kv_store.hpp"
|
|
#include "request.hpp"
|
|
#include <random>
|
|
#include <unordered_map>
|
|
#include <mutex>
|
|
#include <array>
|
|
#include <chrono>
|
|
#include <optional>
|
|
#include <atomic>
|
|
#include <list>
|
|
|
|
class SessionStore {
|
|
private:
|
|
static constexpr size_t SHARD_COUNT = 16;
|
|
static constexpr size_t MAX_SESSIONS_PER_SHARD = 64;
|
|
static constexpr size_t MAX_SESSION_DATA_SIZE = 16;
|
|
static constexpr size_t MAX_VALUE_SIZE = 1024;
|
|
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;
|
|
std::list<std::string>::iterator lru_iter;
|
|
bool dirty = false;
|
|
};
|
|
|
|
std::array<std::mutex, SHARD_COUNT> mutexes_;
|
|
std::array<std::unordered_map<std::string, SessionData>, SHARD_COUNT> shards_;
|
|
std::array<std::list<std::string>, SHARD_COUNT> lru_lists_;
|
|
KeyValueStore store_;
|
|
std::mt19937 rng_;
|
|
|
|
size_t get_shard(const std::string& id) const {
|
|
return std::hash<std::string>{}(id) % SHARD_COUNT;
|
|
}
|
|
|
|
void evict_lru_session(size_t shard) {
|
|
auto& shard_map = shards_[shard];
|
|
auto& lru_list = lru_lists_[shard];
|
|
|
|
if (lru_list.empty()) return;
|
|
|
|
std::string oldest_id = lru_list.back();
|
|
lru_list.pop_back();
|
|
|
|
auto it = shard_map.find(oldest_id);
|
|
if (it != shard_map.end()) {
|
|
// Save dirty session before eviction
|
|
if (it->second.dirty) {
|
|
for (const auto& [key, value] : it->second.data) {
|
|
store_.set(oldest_id + ":" + key, value);
|
|
}
|
|
}
|
|
shard_map.erase(it);
|
|
}
|
|
}
|
|
|
|
void touch_session(const std::string& id, size_t shard) {
|
|
auto& session = shards_[shard][id];
|
|
auto& lru_list = lru_lists_[shard];
|
|
|
|
// Move to front of LRU
|
|
lru_list.erase(session.lru_iter);
|
|
lru_list.push_front(id);
|
|
session.lru_iter = lru_list.begin();
|
|
session.last_access = std::chrono::steady_clock::now();
|
|
}
|
|
|
|
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;
|
|
|
|
// Check session limit
|
|
if (shard_map.size() >= MAX_SESSIONS_PER_SHARD) {
|
|
evict_lru_session(shard);
|
|
}
|
|
|
|
SessionData session;
|
|
session.last_access = std::chrono::steady_clock::now();
|
|
|
|
// Load from persistent store
|
|
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)) {
|
|
if (session.data.size() < MAX_SESSION_DATA_SIZE) {
|
|
session.data[session_key] = *str_val;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Add to LRU
|
|
lru_lists_[shard].push_front(id);
|
|
session.lru_iter = lru_lists_[shard].begin();
|
|
|
|
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];
|
|
auto& lru_list = lru_lists_[i];
|
|
|
|
for (auto it = shard.begin(); it != shard.end();) {
|
|
if (now - it->second.last_access > SESSION_TIMEOUT) {
|
|
lru_list.erase(it->second.lru_iter);
|
|
it = shard.erase(it);
|
|
} else {
|
|
++it;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
public:
|
|
SessionStore() : rng_(std::random_device{}()) {
|
|
store_.set_file("sessions.txt");
|
|
}
|
|
|
|
std::string create() {
|
|
std::string id = generate_id();
|
|
size_t shard = get_shard(id);
|
|
|
|
std::lock_guard<std::mutex> lock(mutexes_[shard]);
|
|
|
|
// Check session limit
|
|
if (shards_[shard].size() >= MAX_SESSIONS_PER_SHARD) {
|
|
evict_lru_session(shard);
|
|
}
|
|
|
|
SessionData session;
|
|
session.last_access = std::chrono::steady_clock::now();
|
|
session.dirty = true;
|
|
|
|
// Add to LRU
|
|
lru_lists_[shard].push_front(id);
|
|
session.lru_iter = lru_lists_[shard].begin();
|
|
|
|
shards_[shard][id] = std::move(session);
|
|
return id;
|
|
}
|
|
|
|
std::optional<std::string> get(const std::string& id, const std::string& key) {
|
|
size_t shard = get_shard(id);
|
|
std::lock_guard<std::mutex> lock(mutexes_[shard]);
|
|
|
|
lazy_load_session(id, shard);
|
|
|
|
auto it = shards_[shard].find(id);
|
|
if (it == shards_[shard].end()) return std::nullopt;
|
|
|
|
touch_session(id, shard);
|
|
|
|
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) {
|
|
if (value.size() > MAX_VALUE_SIZE) return;
|
|
|
|
size_t shard = get_shard(id);
|
|
std::lock_guard<std::mutex> lock(mutexes_[shard]);
|
|
|
|
lazy_load_session(id, shard);
|
|
|
|
auto& session = shards_[shard][id];
|
|
|
|
// Check data limit
|
|
if (session.data.size() >= MAX_SESSION_DATA_SIZE && session.data.find(key) == session.data.end()) {
|
|
return;
|
|
}
|
|
|
|
session.data[key] = value;
|
|
touch_session(id, shard);
|
|
session.dirty = true;
|
|
}
|
|
|
|
bool del(const std::string& id) {
|
|
size_t shard = get_shard(id);
|
|
std::lock_guard<std::mutex> lock(mutexes_[shard]);
|
|
|
|
auto it = shards_[shard].find(id);
|
|
if (it == shards_[shard].end()) return false;
|
|
|
|
lru_lists_[shard].erase(it->second.lru_iter);
|
|
shards_[shard].erase(it);
|
|
|
|
// Remove from persistent store
|
|
for (const auto& [key, value] : store_.get_all()) {
|
|
if (key.starts_with(id + ":")) {
|
|
store_.del(key);
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
std::string refresh(const std::string& old_id) {
|
|
size_t old_shard = get_shard(old_id);
|
|
std::lock_guard<std::mutex> lock(mutexes_[old_shard]);
|
|
|
|
auto it = shards_[old_shard].find(old_id);
|
|
if (it == shards_[old_shard].end()) return old_id;
|
|
|
|
// Generate new ID
|
|
std::string new_id = generate_id();
|
|
size_t new_shard = get_shard(new_id);
|
|
|
|
// Handle cross-shard case
|
|
if (new_shard != old_shard) {
|
|
std::lock_guard<std::mutex> new_lock(mutexes_[new_shard]);
|
|
|
|
// Check session limit on new shard
|
|
if (shards_[new_shard].size() >= MAX_SESSIONS_PER_SHARD) {
|
|
evict_lru_session(new_shard);
|
|
}
|
|
|
|
// Move session data
|
|
SessionData session = std::move(it->second);
|
|
lru_lists_[old_shard].erase(session.lru_iter);
|
|
|
|
// Add to new shard
|
|
lru_lists_[new_shard].push_front(new_id);
|
|
session.lru_iter = lru_lists_[new_shard].begin();
|
|
session.last_access = std::chrono::steady_clock::now();
|
|
session.dirty = true;
|
|
|
|
shards_[new_shard][new_id] = std::move(session);
|
|
shards_[old_shard].erase(it);
|
|
} else {
|
|
// Same shard - just update the key
|
|
SessionData session = std::move(it->second);
|
|
lru_lists_[old_shard].erase(session.lru_iter);
|
|
|
|
lru_lists_[old_shard].push_front(new_id);
|
|
session.lru_iter = lru_lists_[old_shard].begin();
|
|
session.last_access = std::chrono::steady_clock::now();
|
|
session.dirty = true;
|
|
|
|
shards_[old_shard][new_id] = std::move(session);
|
|
shards_[old_shard].erase(it);
|
|
}
|
|
|
|
return new_id;
|
|
}
|
|
|
|
size_t 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) {
|
|
store_.set_file(filename);
|
|
}
|
|
|
|
void load() {
|
|
store_.load();
|
|
}
|
|
|
|
void save() {
|
|
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();
|
|
}
|
|
|
|
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;
|
|
id.reserve(16);
|
|
|
|
for (int i = 0; i < 16; ++i) {
|
|
id += chars[rng_() % (sizeof(chars) - 1)];
|
|
}
|
|
|
|
return id;
|
|
}
|
|
}; |