Batch SQL operations

This commit is contained in:
dynco-nym
2025-06-27 13:00:52 +02:00
parent 61361d84cf
commit 31df4c127c
7 changed files with 175 additions and 88 deletions
@@ -3,7 +3,7 @@
set -eu
export ENVIRONMENT=${ENVIRONMENT:-"mainnet"}
probe_git_ref="nym-vpn-core-v1.10.0"
probe_git_ref="nym-vpn-core-v1.4.0"
crate_root=$(dirname $(realpath "$0"))
monorepo_root=$(realpath "${crate_root}/../..")
@@ -14,7 +14,7 @@ use serde::{Deserialize, Serialize};
use sqlx::FromRow;
use std::str::FromStr;
use strum_macros::{EnumString, FromRepr};
use time::{Date, OffsetDateTime};
use time::{Date, OffsetDateTime, UtcDateTime};
use utoipa::ToSchema;
macro_rules! serialize_opt_to_value {
@@ -362,7 +362,7 @@ impl TryFrom<GatewaySessionsRecord> for http::models::SessionStats {
}
}
#[derive(strum_macros::Display)]
#[derive(strum_macros::Display, Clone)]
pub(crate) enum ScrapeNodeKind {
LegacyMixnode { mix_id: i64 },
MixingNymNode { node_id: i64 },
@@ -520,3 +520,10 @@ pub struct NodeStats {
pub packets_sent: i64,
pub packets_dropped: i64,
}
pub struct InsertStatsRecord {
pub node_kind: ScrapeNodeKind,
pub timestamp_utc: UtcDateTime,
pub unix_timestamp: i64,
pub stats: NodeStats,
}
@@ -19,7 +19,7 @@ pub(crate) use nym_nodes::{
get_described_node_bond_info, get_node_self_description, update_nym_nodes,
};
pub(crate) use packet_stats::{
get_raw_node_stats, insert_daily_node_stats, insert_node_packet_stats,
batch_store_packet_stats, get_raw_node_stats, insert_daily_node_stats_uncommitted,
};
pub(crate) use scraper::{get_nodes_for_scraping, insert_scraped_node_description};
pub(crate) use summary::{get_summary, get_summary_history};
@@ -1,17 +1,70 @@
use crate::db::{
models::{NodeStats, ScrapeNodeKind, ScraperNodeInfo},
DbPool,
use crate::{
db::{
models::{InsertStatsRecord, NodeStats, ScrapeNodeKind},
DbPool,
},
node_scraper::helpers::update_daily_stats_uncommitted,
utils::now_utc,
};
use anyhow::Result;
use sqlx::Transaction;
use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::{info, instrument};
pub(crate) async fn insert_node_packet_stats(
#[instrument(level = "info", skip_all)]
pub(crate) async fn batch_store_packet_stats(
pool: &DbPool,
results: Arc<Mutex<Vec<InsertStatsRecord>>>,
) -> anyhow::Result<()> {
let results_iter = results.lock().await;
info!(
"📊 ⏳ Storing {} packet stats into the DB",
results_iter.len()
);
let started_at = now_utc();
let mut tx = pool
.begin()
.await
.map_err(|err| anyhow::anyhow!("Failed to begin transaction: {err}"))?;
for stats_record in &(*results_iter) {
insert_node_packet_stats_uncommitted(
&mut tx,
&stats_record.node_kind,
&stats_record.stats,
stats_record.unix_timestamp,
)
.await?;
update_daily_stats_uncommitted(
&mut tx,
&stats_record.node_kind,
stats_record.timestamp_utc,
&stats_record.stats,
)
.await?;
}
tx.commit()
.await
.inspect(|_| {
let elapsed = now_utc() - started_at;
info!(
"📊 ☑️ Packet stats successfully committed to DB (took {}s)",
elapsed.as_seconds_f32()
);
})
.map_err(|err| anyhow::anyhow!("Failed to commit: {err}"))
}
async fn insert_node_packet_stats_uncommitted(
tx: &mut Transaction<'static, sqlx::Sqlite>,
node_kind: &ScrapeNodeKind,
stats: &NodeStats,
timestamp_utc: i64,
) -> Result<()> {
let mut conn = pool.acquire().await?;
match node_kind {
ScrapeNodeKind::LegacyMixnode { mix_id } => {
sqlx::query!(
@@ -26,7 +79,7 @@ pub(crate) async fn insert_node_packet_stats(
stats.packets_sent,
stats.packets_dropped,
)
.execute(&mut *conn)
.execute(tx.as_mut())
.await?;
}
ScrapeNodeKind::MixingNymNode { node_id }
@@ -43,7 +96,7 @@ pub(crate) async fn insert_node_packet_stats(
stats.packets_sent,
stats.packets_dropped,
)
.execute(&mut *conn)
.execute(tx.as_mut())
.await?;
}
}
@@ -52,12 +105,10 @@ pub(crate) async fn insert_node_packet_stats(
}
pub(crate) async fn get_raw_node_stats(
pool: &DbPool,
node: &ScraperNodeInfo,
tx: &mut Transaction<'static, sqlx::Sqlite>,
node_kind: &ScrapeNodeKind,
) -> Result<Option<NodeStats>> {
let mut conn = pool.acquire().await?;
let packets = match node.node_kind {
let packets = match node_kind {
// if no packets are found, it's fine to assume 0 because that's also
// SQL default value if none provided
ScrapeNodeKind::LegacyMixnode { mix_id } => {
@@ -75,7 +126,7 @@ pub(crate) async fn get_raw_node_stats(
"#,
mix_id
)
.fetch_optional(&mut *conn)
.fetch_optional(tx.as_mut())
.await?
}
ScrapeNodeKind::MixingNymNode { node_id }
@@ -94,7 +145,7 @@ pub(crate) async fn get_raw_node_stats(
"#,
node_id
)
.fetch_optional(&mut *conn)
.fetch_optional(tx.as_mut())
.await?
}
};
@@ -102,15 +153,13 @@ pub(crate) async fn get_raw_node_stats(
Ok(packets)
}
pub(crate) async fn insert_daily_node_stats(
pool: &DbPool,
node: &ScraperNodeInfo,
pub(crate) async fn insert_daily_node_stats_uncommitted(
tx: &mut Transaction<'static, sqlx::Sqlite>,
node_kind: &ScrapeNodeKind,
date_utc: &str,
packets: NodeStats,
) -> Result<()> {
let mut conn = pool.acquire().await?;
match node.node_kind {
match node_kind {
ScrapeNodeKind::LegacyMixnode { mix_id } => {
let total_stake = sqlx::query_scalar!(
r#"
@@ -121,7 +170,7 @@ pub(crate) async fn insert_daily_node_stats(
"#,
mix_id
)
.fetch_one(&mut *conn)
.fetch_one(tx.as_mut())
.await?;
sqlx::query!(
@@ -144,7 +193,7 @@ pub(crate) async fn insert_daily_node_stats(
packets.packets_sent,
packets.packets_dropped,
)
.execute(&mut *conn)
.execute(tx.as_mut())
.await?;
}
ScrapeNodeKind::MixingNymNode { node_id }
@@ -158,7 +207,7 @@ pub(crate) async fn insert_daily_node_stats(
"#,
node_id
)
.fetch_one(&mut *conn)
.fetch_one(tx.as_mut())
.await?;
sqlx::query!(
@@ -181,7 +230,7 @@ pub(crate) async fn insert_daily_node_stats(
packets.packets_sent,
packets.packets_dropped,
)
.execute(&mut *conn)
.execute(tx.as_mut())
.await?;
}
}
@@ -160,11 +160,11 @@ async fn submit_testrun(
.map(unix_timestamp_to_utc_rfc3339)
.unwrap_or_else(|| String::from("never"));
tracing::info!(
"✅ Testrun row_id {} for gateway {} complete (last assigned {}, created at {})",
gateway_id = gw_identity,
last_assigned = last_assigned,
created_at = created_at,
"✅ Testrun row_id {} for gateway complete",
assigned_testrun.id,
gw_identity,
last_assigned,
created_at
);
Ok(StatusCode::CREATED)
@@ -1,8 +1,8 @@
use crate::{
db::{
models::{NodeStats, ScraperNodeInfo},
models::{InsertStatsRecord, NodeStats, ScrapeNodeKind, ScraperNodeInfo},
queries::{
get_raw_node_stats, insert_daily_node_stats, insert_node_packet_stats,
get_raw_node_stats, insert_daily_node_stats_uncommitted,
insert_scraped_node_description,
},
},
@@ -12,7 +12,7 @@ use ammonia::Builder;
use anyhow::{anyhow, Result};
use reqwest;
use serde::{Deserialize, Serialize};
use sqlx::SqlitePool;
use sqlx::{SqlitePool, Transaction};
use std::time::Duration;
use time::UtcDateTime;
@@ -156,10 +156,7 @@ pub async fn scrape_and_store_description(pool: &SqlitePool, node: &ScraperNodeI
Ok(())
}
pub async fn scrape_and_store_packet_stats(
pool: &SqlitePool,
node: &ScraperNodeInfo,
) -> Result<()> {
pub async fn scrape_packet_stats(node: &ScraperNodeInfo) -> Result<InsertStatsRecord> {
let client = build_client()?;
let urls = node.contact_addresses();
@@ -187,19 +184,21 @@ pub async fn scrape_and_store_packet_stats(
anyhow::anyhow!("Failed to fetch description from any URL: {}", err_msg)
})?;
let timestamp = now_utc();
let timestamp_utc = timestamp.unix_timestamp();
insert_node_packet_stats(pool, &node.node_kind, &stats, timestamp_utc).await?;
let timestamp_utc = now_utc();
let unix_timestamp = timestamp_utc.unix_timestamp();
let result = InsertStatsRecord {
node_kind: node.node_kind.to_owned(),
timestamp_utc,
unix_timestamp,
stats,
};
// TODO dz does this need to run every time?
update_daily_stats(pool, node, timestamp, &stats).await?;
Ok(())
Ok(result)
}
pub async fn update_daily_stats(
pool: &SqlitePool,
node: &ScraperNodeInfo,
pub async fn update_daily_stats_uncommitted(
tx: &mut Transaction<'static, sqlx::Sqlite>,
node_kind: &ScrapeNodeKind,
timestamp: UtcDateTime,
current_stats: &NodeStats,
) -> Result<()> {
@@ -211,7 +210,7 @@ pub async fn update_daily_stats(
);
// Get previous stats
let previous_stats = get_raw_node_stats(pool, node).await?;
let previous_stats = get_raw_node_stats(tx, &node_kind).await?;
let (diff_received, diff_sent, diff_dropped) = if let Some(prev) = previous_stats {
(
@@ -223,9 +222,9 @@ pub async fn update_daily_stats(
(0, 0, 0) // No previous stats available
};
insert_daily_node_stats(
pool,
node,
insert_daily_node_stats_uncommitted(
tx,
node_kind,
&date_utc,
NodeStats {
packets_received: diff_received,
@@ -1,17 +1,19 @@
use super::helpers::scrape_and_store_packet_stats;
use anyhow::Result;
use super::helpers::scrape_packet_stats;
use sqlx::SqlitePool;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::sync::Arc;
use std::time::Duration;
use tracing::{debug, error, instrument, warn};
use tokio::sync::Mutex;
use tokio::task::JoinSet;
use tracing::{debug, error, info, instrument, warn};
use crate::db::models::ScraperNodeInfo;
use crate::db::queries::get_nodes_for_scraping;
use crate::db::models::{InsertStatsRecord, ScraperNodeInfo};
use crate::db::queries;
const PACKET_SCRAPE_INTERVAL: Duration = Duration::from_secs(60 * 60);
const QUEUE_CHECK_INTERVAL: Duration = Duration::from_millis(250);
const MAX_CONCURRENT_TASKS: usize = 5;
// TODO dz should be env configurable
const MAX_CONCURRENT_TASKS: usize = 25;
static TASK_COUNTER: AtomicUsize = AtomicUsize::new(0);
static TASK_ID_COUNTER: AtomicUsize = AtomicUsize::new(0);
@@ -53,49 +55,59 @@ impl PacketScraper {
async fn run_packet_scraper(
pool: &SqlitePool,
queue: Arc<Mutex<Vec<ScraperNodeInfo>>>,
) -> Result<()> {
let nodes = get_nodes_for_scraping(pool).await?;
tracing::info!("Querying {} mixing nodes", nodes.len());
if let Ok(mut queue_lock) = queue.lock() {
) -> anyhow::Result<()> {
let nodes = queries::get_nodes_for_scraping(pool).await?;
{
// TODO dz why do we use mut queue instead of initializing a new queue for each run?
let mut queue_lock = queue.lock().await;
tracing::info!(
"Adding {} nodes to the queue (queue total={})",
nodes.len(),
queue_lock.len()
);
queue_lock.extend(nodes);
} else {
warn!("Failed to acquire packet queue lock");
return Ok(());
}
Self::process_packet_queue(pool, queue).await;
Ok(())
let results = Self::process_packet_queue(queue).await;
queries::batch_store_packet_stats(pool, results)
.await
.map_err(|err| anyhow::anyhow!("Failed to store packet stats to DB: {err}"))
}
async fn process_packet_queue(pool: &SqlitePool, queue: Arc<Mutex<Vec<ScraperNodeInfo>>>) {
async fn process_packet_queue(
queue: Arc<Mutex<Vec<ScraperNodeInfo>>>,
) -> Arc<Mutex<Vec<InsertStatsRecord>>> {
let results = Arc::new(Mutex::new(Vec::new()));
let mut task_set = JoinSet::new();
loop {
let running_tasks = TASK_COUNTER.load(Ordering::Relaxed);
if running_tasks < MAX_CONCURRENT_TASKS {
let node = {
if let Ok(mut queue_lock) = queue.lock() {
if queue_lock.is_empty() {
TASK_ID_COUNTER.store(0, Ordering::Relaxed);
break;
}
queue_lock.remove(0)
} else {
warn!("Failed to acquire packet queue lock");
let mut queue_lock = queue.lock().await;
if queue_lock.is_empty() {
TASK_ID_COUNTER.store(0, Ordering::Relaxed);
break;
}
queue_lock.remove(0)
};
TASK_COUNTER.fetch_add(1, Ordering::Relaxed);
let task_id = TASK_ID_COUNTER.fetch_add(1, Ordering::Relaxed);
let pool = pool.clone();
let results_clone = Arc::clone(&results);
tokio::spawn(async move {
match scrape_and_store_packet_stats(&pool, &node).await {
Ok(_) => debug!(
"📊 ✅ Packet stats task #{} for node {} complete",
task_id,
node.node_id()
),
task_set.spawn(async move {
match scrape_packet_stats(&node).await {
Ok(result) => {
// each task contributes their result to a shared vec
results_clone.lock().await.push(result);
debug!(
"📊 ✅ Packet stats task #{} for node {} complete",
task_id,
node.node_id()
)
}
Err(e) => debug!(
"📊 ❌ Packet stats task #{} for {} {} failed: {}",
task_id,
@@ -111,6 +123,26 @@ impl PacketScraper {
}
}
// TODO After all tasks complete, write results to the DB
// wait for all the tasks to complete before returning their results
let total_count = task_set.len();
let mut success_count = 0;
while let Some(res) = task_set.join_next().await {
if let Err(err) = res {
warn!("Packet stats task panicked: {err}");
} else {
success_count += 1;
}
}
let msg = format!(
"Successfully completed {}/{} tasks ",
success_count, total_count
);
if success_count != total_count {
warn!(msg);
} else {
info!(msg);
}
return results;
}
}