async DB workers, reconnect, and session limit

src/database.cpp

@@ -1,192 +1,225 @@
 #include "database.hpp"
 #include <chrono>
 
-Database::Database(const DatabaseConfig& cfg) 
+Database::Database(const DatabaseConfig& cfg, int workers)
-    : conn(std::make_unique<pqxx::connection>(
+    : connStr(
         "host=" + cfg.host +
         " port=" + std::to_string(cfg.port) +
         " dbname=" + cfg.dbname +
         " user=" + cfg.user +
-        " password=" + cfg.password)),
+        " password=" + cfg.password) {
-      rng(std::random_device{}()) {
+    {
-    ensureTableExists();
+        pqxx::connection probe(connStr);
-}
+        pqxx::work txn(probe);
-
+        txn.exec(
-void Database::ensureTableExists() {
+            "CREATE TABLE IF NOT EXISTS random_data ("
-    pqxx::work txn(*conn);
+            "    id SERIAL PRIMARY KEY,"
-    
+            "    name VARCHAR(100) NOT NULL,"
-    // Создаем таблицу, если её нет
+            "    value INTEGER NOT NULL,"
-    txn.exec(
+            "    description TEXT,"
-        "CREATE TABLE IF NOT EXISTS random_data ("
+            "    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP"
-        "    id SERIAL PRIMARY KEY,"
+            ")"
-        "    name VARCHAR(100) NOT NULL,"
-        "    value INTEGER NOT NULL,"
-        "    description TEXT,"
-        "    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP"
-        ")"
-    );
-    
-    txn.commit();
-}
-
-json Database::getRandomRecord() {
-    std::lock_guard<std::mutex> lock(mtx);
-
-    try {
-        // Получаем общее количество записей
-        int totalCount = getTotalCount();
-
-        if (totalCount == 0) {
-            return json{
-                {"error", "No records found in database"}
-            };
-        }
-
-        // Генерируем случайный ID
-        std::uniform_int_distribution<int> idDist(1, totalCount);
-        int randomId = idDist(rng);
-
-        pqxx::work txn(*conn);
-
-        // Запрашиваем случайную запись
-        auto result = txn.exec_params(
-            "SELECT id, name, value, description, created_at "
-            "FROM random_data "
-            "WHERE id = $1",
-            randomId
         );
+        txn.commit();
+    }
+    for (int i = 0; i < workers; ++i) {
+        threadPool.emplace_back([this, i]() { worker(i); });
+    }
+}
 
-        if (result.empty()) {
+Database::~Database() {
-            // Если не нашли, берем первую (на всякий случай)
+    stop();
-            result = txn.exec("SELECT id, name, value, description, created_at FROM random_data LIMIT 1");
+}
+
+void Database::stop() {
+    {
+        std::lock_guard<std::mutex> lock(mtx);
+        stopFlag = true;
+    }
+    cv.notify_all();
+    for (auto& t : threadPool) {
+        if (t.joinable()) {
+            t.join();
+        }
+    }
+}
+
+void Database::worker(int id) {
+    auto conn = std::make_unique<pqxx::connection>(connStr);
+    std::mt19937 rng(std::random_device{}());
+
+    for (;;) {
+        Task task;
+        {
+            std::unique_lock<std::mutex> lock(mtx);
+            cv.wait(lock, [&] { return stopFlag || !queue.empty(); });
+            if (stopFlag && queue.empty()) return;
+            task = std::move(queue.front());
+            queue.pop();
         }
 
-        if (result.empty()) {
+        json res;
-            txn.abort();
+        try {
-            return json{
+            if (task.op == Op::RandomRecord) {
-                {"error", "No records found in database"}
+                res = getRandomRecord(*conn, rng);
-            };
+            } else {
+                res = getRandomRecords(*conn, rng, task.count, task.med, task.max);
+            }
+        } catch (const std::exception&) {
+            conn = std::make_unique<pqxx::connection>(connStr);
+            res = json{{"error", "Database error"}};
         }
+        task.cb(std::move(res));
+    }
+}
 
-        const auto& row = result[0];
+void Database::enqueue(Task task) {
+    {
+        std::lock_guard<std::mutex> lock(mtx);
+        queue.push(std::move(task));
+    }
+    cv.notify_one();
+}
 
-        json response = {
+void Database::getRandomRecord(Callback cb) {
+    enqueue(Task{Op::RandomRecord, 0, 0, 0, std::move(cb)});
+}
+
+void Database::getRandomRecords(int count, int med, int max, Callback cb) {
+    enqueue(Task{Op::RandomRecords, count, med, max, std::move(cb)});
+}
+
+json Database::getRandomRecord(pqxx::connection& conn, std::mt19937& rng) {
+    int totalCount = getTotalCount(conn);
+
+    if (totalCount == 0) {
+        return json{
+            {"error", "No records found in database"}
+        };
+    }
+
+    std::uniform_int_distribution<int> idDist(1, totalCount);
+    int randomId = idDist(rng);
+
+    pqxx::work txn(conn);
+
+    auto result = txn.exec(
+        "SELECT id, name, value, description, created_at "
+        "FROM random_data "
+        "WHERE id = $1",
+        pqxx::params{randomId}
+    );
+
+    if (result.empty()) {
+        result = txn.exec("SELECT id, name, value, description, created_at FROM random_data LIMIT 1");
+    }
+
+    if (result.empty()) {
+        txn.abort();
+        return json{
+            {"error", "No records found in database"}
+        };
+    }
+
+    const auto& row = result[0];
+
+    json response = {
+        {"id", row[0].as<int>()},
+        {"name", row[1].as<std::string>()},
+        {"value", row[2].as<int>()},
+        {"description", row[3].as<std::string>()},
+        {"created_at", row[4].as<std::string>()},
+        {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
+    };
+
+    txn.commit();
+    return response;
+}
+
+json Database::getRandomRecordsFast(pqxx::connection& conn, const int count) {
+    pqxx::work txn(conn);
+
+    std::string query =
+        "WITH RECURSIVE random_selection AS ("
+        "    WITH bounds AS ("
+        "        SELECT"
+        "            MIN(t.id) AS min_id,"
+        "            MAX(t.id) AS max_id"
+        "        FROM random_data AS t"
+        "    ),"
+        "    first_row AS ("
+        "        SELECT"
+        "            t.id,"
+        "            b.min_id,"
+        "            b.max_id,"
+        "            ARRAY[t.id] AS selected_ids,"
+        "            1 AS depth"
+        "        FROM random_data AS t, bounds AS b"
+        "        WHERE"
+        "            t.id >= b.min_id + ((b.max_id - b.min_id) * random())::INT"
+        "        ORDER BY t.id"
+        "        LIMIT 1"
+        "    ),"
+        "    recursive_rows AS ("
+        "        SELECT"
+        "            t.id,"
+        "            r.min_id,"
+        "            r.max_id,"
+        "            r.selected_ids || t.id AS selected_ids,"
+        "            r.depth + 1 AS depth"
+        "        FROM random_data AS t, first_row AS r"
+        "        WHERE"
+        "            t.id >= r.min_id + ((r.max_id - r.min_id) * random())::INT"
+        "            AND NOT (t.id = ANY(r.selected_ids))"
+        "            AND r.depth + 1 < " + std::to_string(count + 1) +
+        "        ORDER BY t.id"
+        "        LIMIT 1"
+        "    )"
+        "    SELECT * FROM first_row"
+        "    UNION ALL"
+        "    SELECT * FROM recursive_rows"
+        ")"
+        "SELECT"
+        "    rd.id,"
+        "    rd.name,"
+        "    rd.value,"
+        "    rd.description,"
+        "    rd.created_at"
+        "FROM random_selection AS rs"
+        "JOIN random_data AS rd ON rd.id = rs.id"
+        "ORDER BY rs.depth";
+
+    auto result = txn.exec(query);
+
+    if (result.empty()) {
+        return json{
+            {"error", "No records found in database"},
+            {"count", 0}
+        };
+    }
+
+    json response = json::array();
+    for (const auto& row : result) {
+        response.push_back({
             {"id", row[0].as<int>()},
             {"name", row[1].as<std::string>()},
             {"value", row[2].as<int>()},
             {"description", row[3].as<std::string>()},
-            {"created_at", row[4].as<std::string>()},
+            {"created_at", row[4].as<std::string>()}
-            {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
+        });
-        };
-
-        txn.commit();
-        return response;
-
-    } catch (const std::exception& e) {
-        return json{
-            {"error", "Database error in getRandomRecord"},
-        };
     }
+
+    txn.commit();
+
+    return {
+        {"data", response},
+        {"count", static_cast<int>(response.size())},
+        {"strategy", "fast_recursive_cte"},
+        {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
+    };
 }
 
-// БЫСТРЫЙ ВАРИАНТ: Использует рекурсивный CTE
+json Database::getRandomRecordsSlow(pqxx::connection& conn, std::mt19937& rng, const int count, const int totalCount) {
-
-json Database::getRandomRecordsFast(const int count) {
-    pqxx::work txn(*conn);
-    
-    try {
-        // Рекурсивный CTE запрос для получения случайных записей
-        std::string query = 
-            "WITH RECURSIVE random_selection AS ("
-            "    WITH bounds AS ("
-            "        SELECT"
-            "            MIN(t.id) AS min_id,"
-            "            MAX(t.id) AS max_id"
-            "        FROM random_data AS t"
-            "    ),"
-            "    first_row AS ("
-            "        SELECT"
-            "            t.id,"
-            "            b.min_id,"
-            "            b.max_id,"
-            "            ARRAY[t.id] AS selected_ids,"
-            "            1 AS depth"
-            "        FROM random_data AS t, bounds AS b"
-            "        WHERE"
-            "            t.id >= b.min_id + ((b.max_id - b.min_id) * random())::INT"
-            "        ORDER BY t.id"
-            "        LIMIT 1"
-            "    ),"
-            "    recursive_rows AS ("
-            "        SELECT"
-            "            t.id,"
-            "            r.min_id,"
-            "            r.max_id,"
-            "            r.selected_ids || t.id AS selected_ids,"
-            "            r.depth + 1 AS depth"
-            "        FROM random_data AS t, first_row AS r"
-            "        WHERE"
-            "            t.id >= r.min_id + ((r.max_id - r.min_id) * random())::INT"
-            "            AND NOT (t.id = ANY(r.selected_ids))"
-            "            AND r.depth + 1 < " + std::to_string(count + 1) +
-            "        ORDER BY t.id"
-            "        LIMIT 1"
-            "    )"
-            "    SELECT * FROM first_row"
-            "    UNION ALL"
-            "    SELECT * FROM recursive_rows"
-            ")"
-            "SELECT"
-            "    rd.id,"
-            "    rd.name,"
-            "    rd.value,"
-            "    rd.description,"
-            "    rd.created_at"
-            "FROM random_selection AS rs"
-            "JOIN random_data AS rd ON rd.id = rs.id"
-            "ORDER BY rs.depth";
-        
-        auto result = txn.exec(query);
-        
-        if (result.empty()) {
-            return json{
-                {"error", "No records found in database"},
-                {"count", 0}
-            };
-        }
-        
-        json response = json::array();
-        for (const auto& row : result) {
-            response.push_back({
-                {"id", row[0].as<int>()},
-                {"name", row[1].as<std::string>()},
-                {"value", row[2].as<int>()},
-                {"description", row[3].as<std::string>()},
-                {"created_at", row[4].as<std::string>()}
-            });
-        }
-        
-        txn.commit();
-        
-        return {
-            {"data", response},
-            {"count", static_cast<int>(response.size())},
-            {"strategy", "fast_recursive_cte"},
-            {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
-        };
-        
-    } catch (const std::exception& e) {
-        txn.abort();
-        return json{
-            {"error", "Database error in fast recursive CTE strategy"},
-        };
-    }
-}
-
-// МЕДЛЕННЫЙ ВАРИАНТ: Генерирует ID на клиенте, затем запрашивает по IN
-
-json Database::getRandomRecordsSlow(const int count, const int totalCount) {
-    // Генерируем уникальные случайные ID
     std::set<int> ids;
     {
         std::uniform_int_distribution<int> idDist(1, totalCount);
@@ -194,126 +227,104 @@ json Database::getRandomRecordsSlow(const int count, const int totalCount) {
             ids.insert(idDist(rng));
         }
     }
-    
-    // Запрашиваем записи по сгенерированным ID
-    pqxx::work txn(*conn);
-    
-    try {
-        if (ids.empty()) {
-            return json{
-                {"data", json::array()},
-                {"count", 0},
-                {"strategy", "slow"}
-            };
-        }
-        
-        // Формируем запрос с IN
-        std::string query = getQueryRandomIds(static_cast<int>(ids.size()));
 
-        // Подготавливаем параметры для exec_params
+    pqxx::work txn(conn);
-        pqxx::params param_list;
+
-        for (int id : ids) {
+    if (ids.empty()) {
-            param_list.append(id);
-        }
-        
-        auto result = txn.exec_params(query, param_list);
-        
-        json response = json::array();
-        for (const auto& row : result) {
-            response.push_back({
-                {"id", row[0].as<int>()},
-                {"name", row[1].as<std::string>()},
-                {"value", row[2].as<int>()},
-                {"description", row[3].as<std::string>()},
-                {"created_at", row[4].as<std::string>()}
-            });
-        }
-        
-        txn.commit();
-        
-        return {
-            {"data", response},
-            {"count", static_cast<int>(response.size())},
-            {"strategy", "slow"},
-            {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
-        };
-        
-    } catch (const std::exception& e) {
-        txn.abort();
         return json{
-            {"error", "Database error in slow strategy"},
+            {"data", json::array()},
+            {"count", 0},
+            {"strategy", "slow"}
         };
     }
+
+    std::string query = getQueryRandomIds(static_cast<int>(ids.size()));
+
+    pqxx::params param_list;
+    for (int id : ids) {
+        param_list.append(id);
+    }
+
+    auto result = txn.exec(query, param_list);
+
+    json response = json::array();
+    for (const auto& row : result) {
+        response.push_back({
+            {"id", row[0].as<int>()},
+            {"name", row[1].as<std::string>()},
+            {"value", row[2].as<int>()},
+            {"description", row[3].as<std::string>()},
+            {"created_at", row[4].as<std::string>()}
+        });
+    }
+
+    txn.commit();
+
+    return {
+        {"data", response},
+        {"count", static_cast<int>(response.size())},
+        {"strategy", "slow"},
+        {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
+    };
 }
 
-// ГЛАВНАЯ ФУНКЦИЯ: Выбирает стратегию в зависимости от count
+json Database::getRandomRecords(pqxx::connection& conn, std::mt19937& rng, int count, int med, int max) {
-
-json Database::getRandomRecords(int count, int med, int max) {
-    std::lock_guard<std::mutex> lock(mtx);
-    // Валидация входных параметров
     if (count <= 0) {
         return json{
             {"error", "Count must be positive"},
             {"count", 0}
         };
     }
-    //  Получаем общее количество записей с блокировкой
+
-    int totalCount = getTotalCount();
+    int totalCount = getTotalCount(conn);
-    
+
     if (totalCount == 0) {
         return json{
             {"error", "No records found in database"},
             {"count", 0}
         };
     }
-    
+
     if (count > totalCount) {
         count = totalCount;
     }
-    
+
-    // Защита от слишком больших запросов
     if (count > max) {
         count = max;
     }
-    
+
-    // Выбор стратегии:
-    // - Для малых выборок (<= med) используем медленный, но более точный способ
-    // - Для больших выборок (> med) используем быстрый способ с рекурсивным CTE запросом
-    
     json response;
     if (count <= med) {
-        response = getRandomRecordsSlow(count, totalCount);
+        response = getRandomRecordsSlow(conn, rng, count, totalCount);
     } else {
-        response = getRandomRecordsFast(count);
+        response = getRandomRecordsFast(conn, count);
     }
-    if(response.contains("error")){
+    if (response.contains("error")) {
         return response;
     }
-    // Добавляем общую мета-информацию
+
     response["requested_count"] = count;
     response["total_available"] = totalCount;
 
     return response;
 }
 
-// Вспомогательный метод для получения общего количества записей
+int Database::getTotalCount(pqxx::connection& conn) {
-int Database::getTotalCount() {
+    pqxx::work txn(conn);
-    pqxx::work txn(*conn);
     auto result = txn.exec("SELECT COUNT(*) FROM random_data");
     int count = result[0][0].as<int>();
     txn.commit();
     return count;
 }
 
-std::string Database::getQueryRandomIds(const int count){
+std::string Database::getQueryRandomIds(const int count) {
-    std::string query = 
+    std::string query =
-            "SELECT id, name, value, description, created_at "
+        "SELECT id, name, value, description, created_at "
-            "FROM random_data WHERE id IN (";
+        "FROM random_data WHERE id IN (";
-        
+
-    // Подготавливаем параметры
     for (size_t i = 0; i < count; ++i) {
-            if (i > 0) query += ",";
+        if (i > 0) query += ",";
-            query += "$" + std::to_string(i + 1);
+        query += "$" + std::to_string(i + 1);
     }
     query += ") ORDER BY id";
     return query;

src/database.hpp

@@ -5,31 +5,55 @@
 #include <random>
 #include <memory>
 #include <mutex>
+#include <condition_variable>
+#include <queue>
 #include <set>
+#include <thread>
 #include <vector>
+#include <functional>
 
 #include "config.hpp"
 #include "common.hpp"
 
 class Database {
 public:
-    explicit Database(const DatabaseConfig& config);
+    using Callback = std::function<void(json)>;
-    ~Database() = default;
+
-    
+    explicit Database(const DatabaseConfig& config, int workers = 2);
-    json getRandomRecord();
+    ~Database();
-    json getRandomRecords(int count, int med, int max);
+
-    
+    Database(const Database&) = delete;
+    Database& operator=(const Database&) = delete;
+
+    void getRandomRecord(Callback cb);
+    void getRandomRecords(int count, int med, int max, Callback cb);
+
 private:
-    std::unique_ptr<pqxx::connection> conn;
+    enum class Op { RandomRecord, RandomRecords };
-    std::mt19937 rng;
-    std::mutex mtx;
-    // Создание тестовой таблицы, если она не существует
-    void ensureTableExists();
 
-    int getTotalCount();
+    struct Task {
-    // Две стратегии получения случайных записей
+        Op op;
-    json getRandomRecordsFast(const int count);
+        int count = 0;
-    json getRandomRecordsSlow(const int count, const int totalCount);
+        int med = 10;
+        int max = 1000;
+        Callback cb;
+    };
+
+    void worker(int id);
+    void enqueue(Task task);
+    void stop();
+
+    json getRandomRecord(pqxx::connection& conn, std::mt19937& rng);
+    json getRandomRecords(pqxx::connection& conn, std::mt19937& rng, int count, int med, int max);
+    json getRandomRecordsFast(pqxx::connection& conn, const int count);
+    json getRandomRecordsSlow(pqxx::connection& conn, std::mt19937& rng, const int count, const int totalCount);
+    int getTotalCount(pqxx::connection& conn);
     std::string getQueryRandomIds(const int count);
-};
 
+    std::string connStr;
+    std::vector<std::thread> threadPool;
+    std::mutex mtx;
+    std::condition_variable cv;
+    std::queue<Task> queue;
+    bool stopFlag = false;
+};

src/server.cpp

@@ -3,9 +3,20 @@
 #include <iostream>
 #include <chrono>
 
-// Implementation of HttpSession
+HttpSession::HttpSession(net::ip::tcp::socket&& socket,
-HttpSession::HttpSession(net::ip::tcp::socket&& socket, std::shared_ptr<Database> db)
+                         std::shared_ptr<Database> db,
-    : stream(std::move(socket)), db(db) {}
+                         std::shared_ptr<std::atomic<std::size_t>> sessionCounter)
+    : stream(std::move(socket)), db(db), sessionCounter(sessionCounter) {
+    if (sessionCounter) {
+        sessionCounter->fetch_add(1);
+    }
+}
+
+HttpSession::~HttpSession() {
+    if (sessionCounter) {
+        sessionCounter->fetch_sub(1);
+    }
+}
 
 void HttpSession::run() {
     readRequest();
@@ -34,49 +45,70 @@ void HttpSession::onRead(beast::error_code ec, std::size_t) {
 }
 
 void HttpSession::processRequest(http::request<http::string_body>&& request) {
-    http::response<http::string_body> res{http::status::ok, request.version()};
+    auto self = shared_from_this();
-    res.set(http::field::server, "Cards Server");
+
-    res.set(http::field::content_type, "application/json");
+    auto sendJson = [self](http::status status, unsigned version, json body) {
-    
+        http::response<http::string_body> res{status, version};
-    try {
+        res.set(http::field::server, "Cards Server");
-        if (request.method() == http::verb::get && request.target() == "/") {
+        res.set(http::field::content_type, "application/json");
-            // Получаем случайный объект из БД
+        res.body() = body.dump(4);
-            json response = db->getRandomRecord();
+        res.prepare_payload();
-            res.body() = response.dump(4);
+        self->writeResponse(std::move(res));
-        } else if (request.method() == http::verb::get && request.target().starts_with("/cards")) {
+    };
+
+    if (request.method() == http::verb::get && request.target() == "/") {
+        db->getRandomRecord([self, sendJson, version = request.version()](json body) {
+            net::post(
+                self->stream.get_executor(),
+                [self, sendJson, version, body = std::move(body)]() mutable {
+                    sendJson(http::status::ok, version, std::move(body));
+                });
+        });
+        return;
+    }
+
+    if (request.method() == http::verb::get && request.target().starts_with("/cards")) {
+        int count = 0;
+        try {
             std::string number{request.target()};
-            auto start = number.find("?count=") + 7;
+            auto pos = number.find("?count=");
-            auto end = number.find("?", start);
+            if (pos == std::string::npos) {
+                throw std::invalid_argument("missing count");
+            }
+            size_t start = pos + 7;
+            size_t end = number.find('?', start);
             if (end == std::string::npos) {
                 end = number.length();
             }
-            int count = std::stoi(number.substr(start, end - start));
+            count = std::stoi(number.substr(start, end - start));
-            count = std::max(1, std::min(count, 1000)); // валидация
+            count = std::max(1, std::min(count, 1000));
-            json response = db->getRandomRecords(count, 10, 1000);
+        } catch (const std::exception&) {
-            res.body() = response.dump(4);
+            sendJson(http::status::bad_request, request.version(),
-        } else if (request.method() == http::verb::get && request.target() == "/health") {
+                     json{{"error", "Invalid count parameter"}});
-            json response = {
+            return;
-                {"status", "ok"},
-                {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
-            };
-            res.body() = response.dump(4);
-        } else {
-            res.result(http::status::not_found);
-            json error = {
-                {"error", "Not Found"},
-            };
-            res.body() = error.dump(4);
         }
-    } catch (const std::exception& e) {
+        db->getRandomRecords(count, 10, 1000,
-        res.result(http::status::internal_server_error);
+            [self, sendJson, version = request.version()](json body) {
-        json error = {
+                net::post(
-            {"error", "Internal Server Error"},
+                    self->stream.get_executor(),
-        };
+                    [self, sendJson, version, body = std::move(body)]() mutable {
-        res.body() = error.dump(4);
+                        sendJson(http::status::ok, version, std::move(body));
+                    });
+            });
+        return;
     }
-    
+
-    res.prepare_payload();
+    if (request.method() == http::verb::get && request.target() == "/health") {
-    writeResponse(std::move(res));
+        sendJson(http::status::ok, request.version(), json{
+            {"status", "ok"},
+            {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
+        });
+        return;
+    }
+
+    sendJson(http::status::not_found, request.version(), json{
+        {"error", "Not Found"}
+    });
 }
 
 void HttpSession::writeResponse(http::response<http::string_body>&& res) {
@@ -88,7 +120,6 @@ void HttpSession::writeResponse(http::response<http::string_body>&& res) {
             if (ec) {
                 self->handleError(ec, "write");
             }
-            // Close the connection gracefully
             self->stream.socket().shutdown(net::ip::tcp::socket::shutdown_send, ec);
         }
     );
@@ -108,12 +139,12 @@ Server::Server(const std::string& address, int port, int threads, std::shared_pt
       acceptor(ioc, net::ip::tcp::endpoint(net::ip::make_address(address), port), true),
       db(db),
       threads(threads),
-      acceptTimer(ioc)
+      acceptTimer(ioc),
-   {} 
+      activeSessions(std::make_shared<std::atomic<std::size_t>>(0)) {}
 
 void Server::run() {
     accept();
-    
+
     // Run the I/O service on the requested number of threads
     threadPool.reserve(threads);
     for (int i = 0; i < threads; ++i) {
@@ -121,7 +152,7 @@ void Server::run() {
             ioc.run();
         });
     }
-    
+
     // Wait for all threads to finish
     for (auto& t : threadPool) {
         if (t.joinable()) {
@@ -142,18 +173,22 @@ void Server::accept() {
 }
 
 void Server::onAccept(beast::error_code ec, net::ip::tcp::socket socket) {
-    if (!ec) {
+    if (ec) {
-        std::make_shared<HttpSession>(std::move(socket), db)->run();
-        accept();
-    } else {
         std::cerr << "Accept error: " << ec.message() << std::endl;
-        if(!ioc.stopped()) {
+        if (!ioc.stopped()) {
-          acceptTimer.expires_after(std::chrono::milliseconds(100));
+            acceptTimer.expires_after(std::chrono::milliseconds(100));
-          acceptTimer.async_wait([this](beast::error_code ec) {
+            acceptTimer.async_wait([this](beast::error_code ec) {
-              if (!ioc.stopped()) { accept(); }
+                if (!ioc.stopped()) { accept(); }
-              });
+            });
         }
+        return;
     }
-    
-}
 
+    if (activeSessions->load() >= maxSessions) {
+        beast::error_code ignore;
+        socket.close(ignore);
+    } else {
+        std::make_shared<HttpSession>(std::move(socket), db, activeSessions)->run();
+    }
+    accept();
+}

src/server.hpp

@@ -9,6 +9,7 @@
 #include <thread>
 #include <vector>
 #include <string>
+#include <atomic>
 
 #include "common.hpp"
 #include "database.hpp"
@@ -20,19 +21,24 @@ namespace net = boost::asio;
 
 class HttpSession : public std::enable_shared_from_this<HttpSession> {
 public:
-    HttpSession(net::ip::tcp::socket&& socket, std::shared_ptr<Database> db);
+    HttpSession(net::ip::tcp::socket&& socket,
+                std::shared_ptr<Database> db,
+                std::shared_ptr<std::atomic<std::size_t>> sessionCounter);
+    ~HttpSession();
+
     void run();
-    
+
 private:
     void readRequest();
     void onRead(beast::error_code ec, std::size_t bytes_transferred);
     void processRequest(http::request<http::string_body>&& req);
     void writeResponse(http::response<http::string_body>&& res);
     void handleError(beast::error_code ec, const char* what);
-    
+
     beast::tcp_stream stream;
     beast::flat_buffer buffer;
     std::shared_ptr<Database> db;
+    std::shared_ptr<std::atomic<std::size_t>> sessionCounter;
     http::request<http::string_body> req;
 };
 
@@ -41,15 +47,17 @@ public:
     Server(const std::string& address, int port, int threads, std::shared_ptr<Database> db);
     void run();
     void stop();
-    
+
 private:
     void accept();
     void onAccept(beast::error_code ec, net::ip::tcp::socket socket);
-    
+
     net::io_context ioc;
     net::ip::tcp::acceptor acceptor;
     std::shared_ptr<Database> db;
     int threads;
     std::vector<std::thread> threadPool;
     net::steady_timer acceptTimer;
+    std::shared_ptr<std::atomic<std::size_t>> activeSessions;
+    static constexpr std::size_t maxSessions = 512;
 };