375 lines
11 KiB
C++
375 lines
11 KiB
C++
#include "database.hpp"
|
|
#include <chrono>
|
|
#include <iostream>
|
|
#include <syncstream>
|
|
|
|
Database::Database(const DatabaseConfig& cfg, int workers, bool debug)
|
|
: connStr(
|
|
"host=" + cfg.host +
|
|
" port=" + std::to_string(cfg.port) +
|
|
" dbname=" + cfg.dbname +
|
|
" user=" + cfg.user +
|
|
" password=" + cfg.password),
|
|
debug(debug) {
|
|
logDebug("Creating schema for random_data");
|
|
{
|
|
pqxx::connection probe(connStr);
|
|
pqxx::work txn(probe);
|
|
txn.exec(
|
|
"CREATE TABLE IF NOT EXISTS random_data ("
|
|
" id SERIAL PRIMARY KEY,"
|
|
" name VARCHAR(100) NOT NULL,"
|
|
" text TEXT NOT NULL,"
|
|
" description TEXT,"
|
|
" created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP"
|
|
")"
|
|
);
|
|
txn.commit();
|
|
}
|
|
logDebug("Schema is ready; starting " + std::to_string(workers) + " database worker(s)");
|
|
for (int i = 0; i < workers; ++i) {
|
|
threadPool.emplace_back([this, i]() { worker(i); });
|
|
}
|
|
}
|
|
|
|
Database::~Database() {
|
|
stop();
|
|
}
|
|
|
|
void Database::stop() {
|
|
{
|
|
std::lock_guard<std::mutex> lock(mtx);
|
|
stopFlag = true;
|
|
}
|
|
logDebug("Stopping database workers");
|
|
cv.notify_all();
|
|
for (auto& t : threadPool) {
|
|
if (t.joinable()) {
|
|
t.join();
|
|
}
|
|
}
|
|
logDebug("All database workers stopped");
|
|
}
|
|
|
|
void Database::worker(int id) {
|
|
logDebug("Worker " + std::to_string(id) + " connecting to database");
|
|
auto conn = std::make_unique<pqxx::connection>(connStr);
|
|
logDebug("Worker " + std::to_string(id) + " connected to database");
|
|
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()) {
|
|
logDebug("Worker " + std::to_string(id) + " stopped");
|
|
return;
|
|
}
|
|
task = std::move(queue.front());
|
|
queue.pop();
|
|
}
|
|
|
|
json res;
|
|
try {
|
|
if (task.op == Op::RandomRecord) {
|
|
logDebug("Worker " + std::to_string(id) + " fetching one random record");
|
|
res = getRandomRecord(*conn, rng);
|
|
} else {
|
|
logDebug("Worker " + std::to_string(id) + " fetching " + std::to_string(task.count)
|
|
+ " random record(s)");
|
|
res = getRandomRecords(*conn, rng, task.count, task.med, task.max);
|
|
}
|
|
} catch (const std::exception& e) {
|
|
logDebug("Worker " + std::to_string(id) + " database error: " + e.what()
|
|
+ "; reconnecting");
|
|
conn = std::make_unique<pqxx::connection>(connStr);
|
|
logDebug("Worker " + std::to_string(id) + " reconnected to database");
|
|
res = json{{"error", "Database error"}};
|
|
}
|
|
logDebug("Worker " + std::to_string(id) + " completed request"
|
|
+ (res.contains("error") ? ": " + res["error"].get<std::string>() : ""));
|
|
task.cb(std::move(res));
|
|
}
|
|
}
|
|
|
|
void Database::enqueue(Task task) {
|
|
std::size_t queueSize;
|
|
const char* operation = task.op == Op::RandomRecord ? "one random record" : "random records";
|
|
{
|
|
std::lock_guard<std::mutex> lock(mtx);
|
|
queue.push(std::move(task));
|
|
queueSize = queue.size();
|
|
}
|
|
logDebug("Queued request for " + std::string(operation)
|
|
+ "; pending requests: " + std::to_string(queueSize));
|
|
cv.notify_one();
|
|
}
|
|
|
|
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);
|
|
logDebug("random_data contains " + std::to_string(totalCount) + " record(s)");
|
|
|
|
if (totalCount == 0) {
|
|
return json{
|
|
{"error", "No records found in database"}
|
|
};
|
|
}
|
|
|
|
std::uniform_int_distribution<int> idDist(1, totalCount);
|
|
int randomId = idDist(rng);
|
|
logDebug("Selected random record ID " + std::to_string(randomId));
|
|
|
|
pqxx::work txn(conn);
|
|
|
|
auto result = txn.exec(
|
|
"SELECT id, name, text, description, created_at "
|
|
"FROM random_data "
|
|
"WHERE id = $1",
|
|
pqxx::params{randomId}
|
|
);
|
|
|
|
if (result.empty()) {
|
|
logDebug("Selected record no longer exists; falling back to the first available record");
|
|
result = txn.exec("SELECT id, name, text, 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>()},
|
|
{"text", row[2].as<std::string>()},
|
|
{"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) {
|
|
logDebug("Using recursive CTE strategy for " + std::to_string(count) + " record(s)");
|
|
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.text,"
|
|
" 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>()},
|
|
{"text", row[2].as<std::string>()},
|
|
{"description", row[3].as<std::string>()},
|
|
{"created_at", row[4].as<std::string>()}
|
|
});
|
|
}
|
|
|
|
txn.commit();
|
|
logDebug("Recursive CTE returned " + std::to_string(response.size()) + " record(s)");
|
|
|
|
return {
|
|
{"data", response},
|
|
{"count", static_cast<int>(response.size())},
|
|
{"strategy", "fast_recursive_cte"},
|
|
{"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
|
|
};
|
|
}
|
|
|
|
json Database::getRandomRecordsSlow(pqxx::connection& conn, std::mt19937& rng, const int count, const int totalCount) {
|
|
logDebug("Using individual-ID strategy for " + std::to_string(count) + " of "
|
|
+ std::to_string(totalCount) + " record(s)");
|
|
std::set<int> ids;
|
|
{
|
|
std::uniform_int_distribution<int> idDist(1, totalCount);
|
|
while (ids.size() < static_cast<size_t>(count)) {
|
|
ids.insert(idDist(rng));
|
|
}
|
|
}
|
|
|
|
pqxx::work txn(conn);
|
|
|
|
if (ids.empty()) {
|
|
return json{
|
|
{"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>()},
|
|
{"text", row[2].as<std::string>()},
|
|
{"description", row[3].as<std::string>()},
|
|
{"created_at", row[4].as<std::string>()}
|
|
});
|
|
}
|
|
|
|
txn.commit();
|
|
logDebug("Individual-ID query returned " + std::to_string(response.size()) + " record(s)");
|
|
|
|
return {
|
|
{"data", response},
|
|
{"count", static_cast<int>(response.size())},
|
|
{"strategy", "slow"},
|
|
{"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
|
|
};
|
|
}
|
|
|
|
json Database::getRandomRecords(pqxx::connection& conn, std::mt19937& rng, int count, int med, int max) {
|
|
if (count <= 0) {
|
|
return json{
|
|
{"error", "Count must be positive"},
|
|
{"count", 0}
|
|
};
|
|
}
|
|
|
|
int totalCount = getTotalCount(conn);
|
|
logDebug("Requested " + std::to_string(count) + " random record(s); "
|
|
+ std::to_string(totalCount) + " available");
|
|
|
|
if (totalCount == 0) {
|
|
return json{
|
|
{"error", "No records found in database"},
|
|
{"count", 0}
|
|
};
|
|
}
|
|
|
|
if (count > totalCount) {
|
|
logDebug("Reducing requested count to available record count");
|
|
count = totalCount;
|
|
}
|
|
|
|
if (count > max) {
|
|
logDebug("Reducing requested count to configured maximum of " + std::to_string(max));
|
|
count = max;
|
|
}
|
|
|
|
json response;
|
|
if (count <= med) {
|
|
response = getRandomRecordsSlow(conn, rng, count, totalCount);
|
|
} else {
|
|
response = getRandomRecordsFast(conn, count);
|
|
}
|
|
if (response.contains("error")) {
|
|
return response;
|
|
}
|
|
|
|
response["requested_count"] = count;
|
|
response["total_available"] = totalCount;
|
|
|
|
return response;
|
|
}
|
|
|
|
int Database::getTotalCount(pqxx::connection& 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 query =
|
|
"SELECT id, name, text, description, created_at "
|
|
"FROM random_data WHERE id IN (";
|
|
|
|
for (size_t i = 0; i < count; ++i) {
|
|
if (i > 0) query += ",";
|
|
query += "$" + std::to_string(i + 1);
|
|
}
|
|
query += ") ORDER BY id";
|
|
return query;
|
|
}
|
|
|
|
void Database::logDebug(const std::string& message) const {
|
|
if (debug) {
|
|
std::osyncstream(std::cout) << "[debug][database] " << message << '\n';
|
|
}
|
|
}
|