Hexor/khm
1
0
Fork 0
mirror of https://github.com/house-of-vanity/khm.git synced 2025-08-21 14:27:14 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Improved server performance

Browse Source
This commit is contained in:
A B
2025-04-29 21:17:40 +00:00
parent 3f852e50f4
commit 32adb309ee
3 changed files with 100 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
.github/workflows/main.yml vendored
View File

@@ -57,10 +57,6 @@ jobs:
- uses: dtolnay/rust-toolchain@stable - uses: dtolnay/rust-toolchain@stable
- uses: mbrobbel/rustfmt-check@master
with:
token: ${{ secrets.GITHUB_TOKEN }}
- name: Install rust targets - name: Install rust targets
run: rustup target add ${{ matrix.build_target }} run: rustup target add ${{ matrix.build_target }}

136
src/db.rs
View File

@@ -1,21 +1,21 @@
use log::{info}; use crate::server::SshKey;
use log::info;
use std::collections::HashMap; use std::collections::HashMap;
use std::collections::HashSet; use std::collections::HashSet;
use tokio_postgres::Client; use tokio_postgres::Client;
use crate::server::SshKey;
// Структура для хранения статистики обработки ключей // Структура для хранения статистики обработки ключей
pub struct KeyInsertStats { pub struct KeyInsertStats {
pub total: usize, // Общее количество полученных ключей pub total: usize, // Общее количество полученных ключей
pub inserted: usize, // Количество новых ключей pub inserted: usize, // Количество новых ключей
pub updated: usize, // Количество обновленных ключей pub updated: usize, // Количество обновленных ключей
pub unchanged: usize, // Количество неизмененных ключей pub unchanged: usize, // Количество неизмененных ключей
pub key_id_map: Vec<(SshKey, i32)>, // Связь ключей с их ID в базе pub key_id_map: Vec<(SshKey, i32)>, // Связь ключей с их ID в базе
} }
pub async fn initialize_db_schema(client: &Client) -> Result<(), tokio_postgres::Error> { pub async fn initialize_db_schema(client: &Client) -> Result<(), tokio_postgres::Error> {
info!("Checking and initializing database schema if needed"); info!("Checking and initializing database schema if needed");
// Check if tables exist by querying information_schema // Check if tables exist by querying information_schema
let tables_exist = client let tables_exist = client
.query( .query(
@@ -37,7 +37,7 @@ pub async fn initialize_db_schema(client: &Client) -> Result<(), tokio_postgres:
if !tables_exist { if !tables_exist {
info!("Database schema doesn't exist. Creating tables..."); info!("Database schema doesn't exist. Creating tables...");
// Create the keys table // Create the keys table
client client
.execute( .execute(
@@ -51,7 +51,7 @@ pub async fn initialize_db_schema(client: &Client) -> Result<(), tokio_postgres:
&[], &[],
) )
.await?; .await?;
// Create the flows table // Create the flows table
client client
.execute( .execute(
@@ -68,7 +68,7 @@ pub async fn initialize_db_schema(client: &Client) -> Result<(), tokio_postgres:
&[], &[],
) )
.await?; .await?;
// Create an index for faster lookups // Create an index for faster lookups
client client
.execute( .execute(
@@ -76,12 +76,12 @@ pub async fn initialize_db_schema(client: &Client) -> Result<(), tokio_postgres:
&[], &[],
) )
.await?; .await?;
info!("Database schema created successfully"); info!("Database schema created successfully");
} else { } else {
info!("Database schema already exists"); info!("Database schema already exists");
} }
Ok(()) Ok(())
} }
@@ -102,42 +102,43 @@ pub async fn batch_insert_keys(
// Подготавливаем массивы для пакетной вставки // Подготавливаем массивы для пакетной вставки
let mut host_values: Vec<&str> = Vec::with_capacity(keys.len()); let mut host_values: Vec<&str> = Vec::with_capacity(keys.len());
let mut key_values: Vec<&str> = Vec::with_capacity(keys.len()); let mut key_values: Vec<&str> = Vec::with_capacity(keys.len());
for key in keys { for key in keys {
host_values.push(&key.server); host_values.push(&key.server);
key_values.push(&key.public_key); key_values.push(&key.public_key);
} }
// Сначала проверяем, какие ключи уже существуют в базе // Сначала проверяем, какие ключи уже существуют в базе
let mut existing_keys = HashMap::new(); let mut existing_keys = HashMap::new();
let mut key_query = String::from("SELECT host, key, key_id FROM public.keys WHERE "); let mut key_query = String::from("SELECT host, key, key_id FROM public.keys WHERE ");
for i in 0..keys.len() { for i in 0..keys.len() {
if i > 0 { if i > 0 {
key_query.push_str(" OR "); key_query.push_str(" OR ");
} }
key_query.push_str(&format!("(host = ${} AND key = ${})", i * 2 + 1, i * 2 + 2)); key_query.push_str(&format!("(host = ${} AND key = ${})", i * 2 + 1, i * 2 + 2));
} }
let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::with_capacity(keys.len() * 2); let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> =
Vec::with_capacity(keys.len() * 2);
for i in 0..keys.len() { for i in 0..keys.len() {
params.push(&host_values[i]); params.push(&host_values[i]);
params.push(&key_values[i]); params.push(&key_values[i]);
} }
let rows = client.query(&key_query, &params[..]).await?; let rows = client.query(&key_query, &params[..]).await?;
for row in rows { for row in rows {
let host: String = row.get(0); let host: String = row.get(0);
let key: String = row.get(1); let key: String = row.get(1);
let key_id: i32 = row.get(2); let key_id: i32 = row.get(2);
existing_keys.insert((host, key), key_id); existing_keys.insert((host, key), key_id);
} }
// Определяем, какие ключи нужно вставить, а какие уже существуют // Определяем, какие ключи нужно вставить, а какие уже существуют
let mut keys_to_insert = Vec::new(); let mut keys_to_insert = Vec::new();
let mut unchanged_keys = Vec::new(); let mut unchanged_keys = Vec::new();
for key in keys { for key in keys {
let key_tuple = (key.server.clone(), key.public_key.clone()); let key_tuple = (key.server.clone(), key.public_key.clone());
if existing_keys.contains_key(&key_tuple) { if existing_keys.contains_key(&key_tuple) {
@@ -146,18 +147,16 @@ pub async fn batch_insert_keys(
keys_to_insert.push(key.clone()); keys_to_insert.push(key.clone());
} }
} }
let mut inserted_keys = Vec::new(); let mut inserted_keys = Vec::new();
// Если есть ключи для вставки, выполняем вставку // Если есть ключи для вставки, выполняем вставку
if !keys_to_insert.is_empty() { if !keys_to_insert.is_empty() {
let mut insert_sql = String::from( let mut insert_sql = String::from("INSERT INTO public.keys (host, key, updated) VALUES ");
"INSERT INTO public.keys (host, key, updated) VALUES "
);
let mut insert_params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); let mut insert_params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new();
let mut param_count = 1; let mut param_count = 1;
for (i, key) in keys_to_insert.iter().enumerate() { for (i, key) in keys_to_insert.iter().enumerate() {
if i > 0 { if i > 0 {
insert_sql.push_str(", "); insert_sql.push_str(", ");
@@ -167,31 +166,34 @@ pub async fn batch_insert_keys(
insert_params.push(&key.public_key); insert_params.push(&key.public_key);
param_count += 2; param_count += 2;
} }
insert_sql.push_str(" RETURNING key_id, host, key"); insert_sql.push_str(" RETURNING key_id, host, key");
let inserted_rows = client.query(&insert_sql, &insert_params[..]).await?; let inserted_rows = client.query(&insert_sql, &insert_params[..]).await?;
for row in inserted_rows { for row in inserted_rows {
let host: String = row.get(1); let host: String = row.get(1);
let key_text: String = row.get(2); let key_text: String = row.get(2);
let key_id: i32 = row.get(0); let key_id: i32 = row.get(0);
if let Some(orig_key) = keys_to_insert.iter().find(|k| k.server == host && k.public_key == key_text) { if let Some(orig_key) = keys_to_insert
.iter()
.find(|k| k.server == host && k.public_key == key_text)
{
inserted_keys.push((orig_key.clone(), key_id)); inserted_keys.push((orig_key.clone(), key_id));
} }
} }
} }
// Сохраняем количество элементов перед объединением // Сохраняем количество элементов перед объединением
let inserted_count = inserted_keys.len(); let inserted_count = inserted_keys.len();
let unchanged_count = unchanged_keys.len(); let unchanged_count = unchanged_keys.len();
// Комбинируем результаты и формируем статистику // Комбинируем результаты и формируем статистику
let mut key_id_map = Vec::with_capacity(unchanged_count + inserted_count); let mut key_id_map = Vec::with_capacity(unchanged_count + inserted_count);
key_id_map.extend(unchanged_keys); key_id_map.extend(unchanged_keys);
key_id_map.extend(inserted_keys); key_id_map.extend(inserted_keys);
let stats = KeyInsertStats { let stats = KeyInsertStats {
total: keys.len(), total: keys.len(),
inserted: inserted_count, inserted: inserted_count,
@@ -199,10 +201,12 @@ pub async fn batch_insert_keys(
unchanged: unchanged_count, unchanged: unchanged_count,
key_id_map, key_id_map,
}; };
info!("Keys stats: received={}, new={}, unchanged={}", info!(
stats.total, stats.inserted, stats.unchanged); "Keys stats: received={}, new={}, unchanged={}",
stats.total, stats.inserted, stats.unchanged
);
Ok(stats) Ok(stats)
} }
@@ -217,10 +221,9 @@ pub async fn batch_insert_flow_keys(
} }
// Сначала проверим, какие связи уже существуют // Сначала проверим, какие связи уже существуют
let mut existing_query = String::from( let mut existing_query =
"SELECT key_id FROM public.flows WHERE name = $1 AND key_id IN (" String::from("SELECT key_id FROM public.flows WHERE name = $1 AND key_id IN (");
);
for i in 0..key_ids.len() { for i in 0..key_ids.len() {
if i > 0 { if i > 0 {
existing_query.push_str(", "); existing_query.push_str(", ");
@@ -228,35 +231,39 @@ pub async fn batch_insert_flow_keys(
existing_query.push_str(&format!("${}", i + 2)); existing_query.push_str(&format!("${}", i + 2));
} }
existing_query.push_str(")"); existing_query.push_str(")");
let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::with_capacity(key_ids.len() + 1); let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> =
Vec::with_capacity(key_ids.len() + 1);
params.push(&flow_name); params.push(&flow_name);
for key_id in key_ids { for key_id in key_ids {
params.push(key_id); params.push(key_id);
} }
let rows = client.query(&existing_query, &params[..]).await?; let rows = client.query(&existing_query, &params[..]).await?;
let mut existing_associations = HashSet::new(); let mut existing_associations = HashSet::new();
for row in rows { for row in rows {
let key_id: i32 = row.get(0); let key_id: i32 = row.get(0);
existing_associations.insert(key_id); existing_associations.insert(key_id);
} }
// Фильтруем только те ключи, которые еще не связаны с потоком // Фильтруем только те ключи, которые еще не связаны с потоком
let new_key_ids: Vec<&i32> = key_ids.iter() let new_key_ids: Vec<&i32> = key_ids
.iter()
.filter(|&id| !existing_associations.contains(id)) .filter(|&id| !existing_associations.contains(id))
.collect(); .collect();
if new_key_ids.is_empty() { if new_key_ids.is_empty() {
info!("All {} keys are already associated with flow '{}'", key_ids.len(), flow_name); info!(
"All {} keys are already associated with flow '{}'",
key_ids.len(),
flow_name
);
return Ok(0); return Ok(0);
} }
// Строим SQL запрос с множественными значениями только для новых связей // Строим SQL запрос с множественными значениями только для новых связей
let mut sql = String::from( let mut sql = String::from("INSERT INTO public.flows (name, key_id) VALUES ");
"INSERT INTO public.flows (name, key_id) VALUES "
);
for i in 0..new_key_ids.len() { for i in 0..new_key_ids.len() {
if i > 0 { if i > 0 {
@@ -268,7 +275,8 @@ pub async fn batch_insert_flow_keys(
sql.push_str(" ON CONFLICT (name, key_id) DO NOTHING"); sql.push_str(" ON CONFLICT (name, key_id) DO NOTHING");
// Подготавливаем параметры для запроса // Подготавливаем параметры для запроса
let mut insert_params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::with_capacity(new_key_ids.len() + 1); let mut insert_params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> =
Vec::with_capacity(new_key_ids.len() + 1);
insert_params.push(&flow_name); insert_params.push(&flow_name);
for key_id in &new_key_ids { for key_id in &new_key_ids {
insert_params.push(*key_id); insert_params.push(*key_id);
@@ -276,11 +284,15 @@ pub async fn batch_insert_flow_keys(
// Выполняем запрос // Выполняем запрос
let affected = client.execute(&sql, &insert_params[..]).await?; let affected = client.execute(&sql, &insert_params[..]).await?;
let affected_usize = affected as usize; let affected_usize = affected as usize;
info!("Added {} new key-flow associations for flow '{}' (skipped {} existing)", info!(
affected_usize, flow_name, existing_associations.len()); "Added {} new key-flow associations for flow '{}' (skipped {} existing)",
affected_usize,
flow_name,
existing_associations.len()
);
Ok(affected_usize) Ok(affected_usize)
} }

40
src/server.rs
View File

@@ -163,8 +163,12 @@ pub async fn add_keys(
valid_keys.push(new_key.clone()); valid_keys.push(new_key.clone());
} }
info!("Processing batch of {} keys for flow: {}", valid_keys.len(), flow_id_str); info!(
"Processing batch of {} keys for flow: {}",
valid_keys.len(),
flow_id_str
);
// Батчевая вставка ключей с получением статистики // Батчевая вставка ключей с получением статистики
let key_stats = match crate::db::batch_insert_keys(&db_client, &valid_keys).await { let key_stats = match crate::db::batch_insert_keys(&db_client, &valid_keys).await {
Ok(stats) => stats, Ok(stats) => stats,
@@ -178,19 +182,21 @@ pub async fn add_keys(
// Если нет новых ключей, нет необходимости обновлять связи с flow // Если нет новых ключей, нет необходимости обновлять связи с flow
if key_stats.inserted > 0 { if key_stats.inserted > 0 {
// Извлекаем только ID ключей из статистики // Извлекаем только ID ключей из статистики
let key_ids: Vec<i32> = key_stats.key_id_map let key_ids: Vec<i32> = key_stats.key_id_map.iter().map(|(_, id)| *id).collect();
.iter()
.map(|(_, id)| *id)
.collect();
// Батчевая вставка связей ключей с flow // Батчевая вставка связей ключей с flow
if let Err(e) = crate::db::batch_insert_flow_keys(&db_client, &flow_id_str, &key_ids).await { if let Err(e) = crate::db::batch_insert_flow_keys(&db_client, &flow_id_str, &key_ids).await
{
error!("Failed to batch insert flow keys into database: {}", e); error!("Failed to batch insert flow keys into database: {}", e);
return HttpResponse::InternalServerError() return HttpResponse::InternalServerError()
.body("Failed to batch insert flow keys into database"); .body("Failed to batch insert flow keys into database");
} }
info!("Added flow associations for {} keys in flow '{}'", key_ids.len(), flow_id_str); info!(
"Added flow associations for {} keys in flow '{}'",
key_ids.len(),
flow_id_str
);
} else { } else {
info!("No new keys to associate with flow '{}'", flow_id_str); info!("No new keys to associate with flow '{}'", flow_id_str);
} }
@@ -211,15 +217,21 @@ pub async fn add_keys(
let updated_flow = flows_guard.iter().find(|flow| flow.name == flow_id_str); let updated_flow = flows_guard.iter().find(|flow| flow.name == flow_id_str);
if let Some(flow) = updated_flow { if let Some(flow) = updated_flow {
let servers: Vec<&SshKey> = flow.servers.iter().collect(); let servers: Vec<&SshKey> = flow.servers.iter().collect();
info!("Keys summary for flow '{}': total received={}, new={}, unchanged={}, total in flow={}", info!(
flow_id_str, key_stats.total, key_stats.inserted, key_stats.unchanged, servers.len()); "Keys summary for flow '{}': total received={}, new={}, unchanged={}, total in flow={}",
flow_id_str,
key_stats.total,
key_stats.inserted,
key_stats.unchanged,
servers.len()
);
// Добавляем статистику в HTTP заголовки ответа // Добавляем статистику в HTTP заголовки ответа
let mut response = HttpResponse::Ok(); let mut response = HttpResponse::Ok();
response.append_header(("X-Keys-Total", key_stats.total.to_string())); response.append_header(("X-Keys-Total", key_stats.total.to_string()));
response.append_header(("X-Keys-New", key_stats.inserted.to_string())); response.append_header(("X-Keys-New", key_stats.inserted.to_string()));
response.append_header(("X-Keys-Unchanged", key_stats.unchanged.to_string())); response.append_header(("X-Keys-Unchanged", key_stats.unchanged.to_string()));
response.json(servers) response.json(servers)
} else { } else {
error!("Flow ID not found after update: {}", flow_id_str); error!("Flow ID not found after update: {}", flow_id_str);