From a9d018cc7e46224ea528d4d198a8854796841814 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maciej=20W=C3=B3jcik?= Date: Wed, 21 Jan 2026 08:53:32 +0100 Subject: [PATCH 1/6] make truncating optional --- tools/defguard_generator/src/main.rs | 4 ++++ tools/defguard_generator/src/vpn_session_stats.rs | 7 +++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tools/defguard_generator/src/main.rs b/tools/defguard_generator/src/main.rs index e9b5b814c3..322328764f 100644 --- a/tools/defguard_generator/src/main.rs +++ b/tools/defguard_generator/src/main.rs @@ -40,6 +40,8 @@ enum Commands { devices_per_user: u8, #[arg(long)] sessions_per_device: u8, + #[arg(long)] + truncate_sessions_table: bool, }, } @@ -68,12 +70,14 @@ async fn main() -> Result<()> { num_users, devices_per_user, sessions_per_device, + truncate_sessions_table, } => { let config = VpnSessionGeneratorConfig { location_id, num_users, devices_per_user, sessions_per_device, + truncate_sessions_table, }; generate_vpn_session_stats(pool, config).await?; diff --git a/tools/defguard_generator/src/vpn_session_stats.rs b/tools/defguard_generator/src/vpn_session_stats.rs index f2835fa2f1..096c8a23d8 100644 --- a/tools/defguard_generator/src/vpn_session_stats.rs +++ b/tools/defguard_generator/src/vpn_session_stats.rs @@ -26,6 +26,7 @@ pub struct VpnSessionGeneratorConfig { pub num_users: u16, pub devices_per_user: u8, pub sessions_per_device: u8, + pub truncate_sessions_table: bool, } pub async fn generate_vpn_session_stats( @@ -35,8 +36,10 @@ pub async fn generate_vpn_session_stats( let mut rng = rand::thread_rng(); // clear sessions & stats tables - info!("Clearing existing sessions & stats"); - truncate_with_restart(&pool).await?; + if config.truncate_sessions_table { + info!("Clearing existing sessions & stats"); + truncate_with_restart(&pool).await?; + } // fetch specified location let location = WireguardNetwork::find_by_id(&pool, config.location_id) From c666bf4a5d54475f04eb28e074d67bb6cbad8d14 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maciej=20W=C3=B3jcik?= Date: Wed, 21 Jan 2026 08:53:38 +0100 Subject: [PATCH 2/6] add more logs --- .../src/db/models/vpn_client_session.rs | 4 ++-- tools/defguard_generator/src/main.rs | 8 ++++++-- tools/defguard_generator/src/vpn_session_stats.rs | 12 +++++++++++- 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/crates/defguard_common/src/db/models/vpn_client_session.rs b/crates/defguard_common/src/db/models/vpn_client_session.rs index 8f2167e2ef..3ca957d93b 100644 --- a/crates/defguard_common/src/db/models/vpn_client_session.rs +++ b/crates/defguard_common/src/db/models/vpn_client_session.rs @@ -7,7 +7,7 @@ use crate::db::{ models::{WireguardNetwork, vpn_session_stats::VpnSessionStats, wireguard::LocationMfaMode}, }; -#[derive(Default, Type)] +#[derive(Debug, Default, Type)] #[sqlx(type_name = "vpn_client_session_state", rename_all = "lowercase")] pub enum VpnClientSessionState { #[default] @@ -17,7 +17,7 @@ pub enum VpnClientSessionState { } /// Represents a single VPN client session from creation to eventual disconnection -#[derive(Model)] +#[derive(Debug, Model)] #[table(vpn_client_session)] pub struct VpnClientSession { pub id: I, diff --git a/tools/defguard_generator/src/main.rs b/tools/defguard_generator/src/main.rs index 322328764f..be5ea3c697 100644 --- a/tools/defguard_generator/src/main.rs +++ b/tools/defguard_generator/src/main.rs @@ -4,7 +4,7 @@ use defguard_common::db::{Id, init_db}; use defguard_generator::vpn_session_stats::{ VpnSessionGeneratorConfig, generate_vpn_session_stats, }; -use tracing::Level; +use tracing_subscriber::EnvFilter; #[derive(Parser)] #[command(about, long_about = None)] @@ -48,7 +48,11 @@ enum Commands { #[tokio::main] async fn main() -> Result<()> { // Initialize logging - tracing_subscriber::fmt().with_max_level(Level::INFO).init(); + tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")), + ) + .init(); // parse CLI options let cli = Cli::parse(); diff --git a/tools/defguard_generator/src/vpn_session_stats.rs b/tools/defguard_generator/src/vpn_session_stats.rs index 096c8a23d8..0feda22d31 100644 --- a/tools/defguard_generator/src/vpn_session_stats.rs +++ b/tools/defguard_generator/src/vpn_session_stats.rs @@ -21,6 +21,7 @@ use crate::{user_devices::prepare_user_devices, users::prepare_users}; const STATS_COLLECTION_INTERVAL: Duration = Duration::seconds(30); const HANDSHAKE_INTERVAL: Duration = Duration::minutes(2); +#[derive(Debug)] pub struct VpnSessionGeneratorConfig { pub location_id: Id, pub num_users: u16, @@ -33,6 +34,7 @@ pub async fn generate_vpn_session_stats( pool: PgPool, config: VpnSessionGeneratorConfig, ) -> Result<()> { + info!("Running VPN stats generator with config: {config:#?}"); let mut rng = rand::thread_rng(); // clear sessions & stats tables @@ -55,7 +57,10 @@ pub async fn generate_vpn_session_stats( // generate sessions for each user for (i, user) in users.into_iter().enumerate() { - info!("[{i}/{user_count}] Generating VPN sessions for user {user}"); + info!( + "[{}/{user_count}] Generating VPN sessions for user {user}", + i + 1 + ); // begin DB transaction let mut transaction = pool.begin().await?; @@ -65,6 +70,7 @@ pub async fn generate_vpn_session_stats( prepare_user_devices(&pool, &mut rng, &user, config.devices_per_user as usize).await?; for device in devices { + info!("Generating sessions for device {device}"); // generate requested number of sessions for a device // we always start with a session that's currently active // and generate past ones as needed @@ -93,6 +99,8 @@ pub async fn generate_vpn_session_stats( session.save(&mut *transaction).await?; } + info!("Created session {session:?}"); + generate_mock_session_stats( &mut transaction, &mut rng, @@ -103,6 +111,8 @@ pub async fn generate_vpn_session_stats( ) .await?; + info!("Finished generating mock stats for session {session:?}"); + // update end timestamp for next session session_end -= Duration::minutes(rng.gen_range(30..120)); } From 20b17b067fb27dc5569484f2d6ecbbec6880020d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maciej=20W=C3=B3jcik?= Date: Wed, 21 Jan 2026 09:00:57 +0100 Subject: [PATCH 3/6] update readme --- tools/defguard_generator/README.md | 12 ++++++++++++ tools/defguard_generator/src/main.rs | 3 ++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/tools/defguard_generator/README.md b/tools/defguard_generator/README.md index c044980add..de749c9020 100644 --- a/tools/defguard_generator/README.md +++ b/tools/defguard_generator/README.md @@ -2,6 +2,18 @@ This crate contains a simple generator for creating users, devices, stats etc during development. +### Database connection + +The generator uses the same environment variables (or CLI options) for DB connection setup as the core binary: + +- DEFGUARD_DB_HOST +- DEFGUARD_DB_PORT +- DEFGUARD_DB_NAME +- DEFGUARD_DB_USER +- DEFGUARD_DB_PASSWORD + +This means that if you have a development environment set up already it should just work. + ### Usage ```bash diff --git a/tools/defguard_generator/src/main.rs b/tools/defguard_generator/src/main.rs index be5ea3c697..f2d0b274e7 100644 --- a/tools/defguard_generator/src/main.rs +++ b/tools/defguard_generator/src/main.rs @@ -30,7 +30,7 @@ struct Cli { #[derive(Subcommand)] enum Commands { - /// generates VPN session stats + /// generates mock VPN session stats VpnSessionStats { #[arg(long)] location_id: Id, @@ -40,6 +40,7 @@ enum Commands { devices_per_user: u8, #[arg(long)] sessions_per_device: u8, + /// truncate sessions & stats tables before generating stats #[arg(long)] truncate_sessions_table: bool, }, From f423ff49b6a4f71e4c75f7bc2b3564eff74490c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maciej=20W=C3=B3jcik?= Date: Wed, 21 Jan 2026 09:17:00 +0100 Subject: [PATCH 4/6] batch stats inserts --- tools/defguard_generator/src/main.rs | 5 ++ .../src/vpn_session_stats.rs | 64 ++++++++++++++++--- 2 files changed, 61 insertions(+), 8 deletions(-) diff --git a/tools/defguard_generator/src/main.rs b/tools/defguard_generator/src/main.rs index f2d0b274e7..f90113bd45 100644 --- a/tools/defguard_generator/src/main.rs +++ b/tools/defguard_generator/src/main.rs @@ -43,6 +43,9 @@ enum Commands { /// truncate sessions & stats tables before generating stats #[arg(long)] truncate_sessions_table: bool, + /// insert stats records in batches of specified size + #[arg(long, default_value_t = 1000)] + stats_batch_size: u16, }, } @@ -76,6 +79,7 @@ async fn main() -> Result<()> { devices_per_user, sessions_per_device, truncate_sessions_table, + stats_batch_size, } => { let config = VpnSessionGeneratorConfig { location_id, @@ -83,6 +87,7 @@ async fn main() -> Result<()> { devices_per_user, sessions_per_device, truncate_sessions_table, + stats_batch_size, }; generate_vpn_session_stats(pool, config).await?; diff --git a/tools/defguard_generator/src/vpn_session_stats.rs b/tools/defguard_generator/src/vpn_session_stats.rs index 0feda22d31..c8de4e4843 100644 --- a/tools/defguard_generator/src/vpn_session_stats.rs +++ b/tools/defguard_generator/src/vpn_session_stats.rs @@ -13,8 +13,8 @@ use defguard_common::db::{ }, }; use rand::{Rng, rngs::ThreadRng}; -use sqlx::{PgConnection, PgPool}; -use tracing::info; +use sqlx::{PgConnection, PgPool, QueryBuilder}; +use tracing::{debug, info}; use crate::{user_devices::prepare_user_devices, users::prepare_users}; @@ -28,6 +28,7 @@ pub struct VpnSessionGeneratorConfig { pub devices_per_user: u8, pub sessions_per_device: u8, pub truncate_sessions_table: bool, + pub stats_batch_size: u16, } pub async fn generate_vpn_session_stats( @@ -99,7 +100,7 @@ pub async fn generate_vpn_session_stats( session.save(&mut *transaction).await?; } - info!("Created session {session:?}"); + debug!("Created session {session:?}"); generate_mock_session_stats( &mut transaction, @@ -108,10 +109,11 @@ pub async fn generate_vpn_session_stats( gateway.id, session_start, session_end, + config.stats_batch_size, ) .await?; - info!("Finished generating mock stats for session {session:?}"); + debug!("Finished generating mock stats for session {session:?}"); // update end timestamp for next session session_end -= Duration::minutes(rng.gen_range(30..120)); @@ -154,6 +156,7 @@ async fn generate_mock_session_stats( gateway_id: Id, session_start: NaiveDateTime, session_end: NaiveDateTime, + batch_size: u16, ) -> Result<()> { let mut latest_handshake = session_start; let mut next_handshake = latest_handshake + HANDSHAKE_INTERVAL; @@ -164,6 +167,9 @@ async fn generate_mock_session_stats( // assume the IP remains static within a single session let endpoint = random_socket_addr(rng).to_string(); + // Vector to accumulate stats before batch insertion + let mut stats_batch: Vec = Vec::new(); + while collected_at <= session_end { // generate traffic let upload_diff = rng.gen_range(100..100_000); @@ -171,7 +177,7 @@ async fn generate_mock_session_stats( let download_diff = rng.gen_range(100..100_000); total_download += download_diff; - VpnSessionStats::new( + let stats = VpnSessionStats::new( session_id, gateway_id, collected_at, @@ -181,9 +187,15 @@ async fn generate_mock_session_stats( total_download, download_diff, download_diff, - ) - .save(&mut *transaction) - .await?; + ); + + stats_batch.push(stats); + + // If batch is full, insert all at once + if stats_batch.len() >= batch_size.into() { + insert_stats_batch(&mut *transaction, &stats_batch).await?; + stats_batch.clear(); + } // update variables for next sample collected_at += STATS_COLLECTION_INTERVAL; @@ -195,6 +207,42 @@ async fn generate_mock_session_stats( } } + // Insert any remaining stats in the batch + if !stats_batch.is_empty() { + insert_stats_batch(&mut *transaction, &stats_batch).await?; + } + + Ok(()) +} + +/// Insert multiple VpnSessionStats records in a single query +async fn insert_stats_batch( + transaction: &mut PgConnection, + stats_batch: &[VpnSessionStats], +) -> Result<()> { + if stats_batch.is_empty() { + return Ok(()); + } + + let mut query_builder = QueryBuilder::new( + "INSERT INTO vpn_session_stats (session_id, gateway_id, collected_at, latest_handshake, endpoint, total_upload, total_download, upload_diff, download_diff) ", + ); + + query_builder.push_values(stats_batch, |mut b, stats| { + b.push_bind(stats.session_id) + .push_bind(stats.gateway_id) + .push_bind(stats.collected_at) + .push_bind(stats.latest_handshake) + .push_bind(&stats.endpoint) + .push_bind(stats.total_upload) + .push_bind(stats.total_download) + .push_bind(stats.upload_diff) + .push_bind(stats.download_diff); + }); + + let query = query_builder.build(); + query.execute(&mut *transaction).await?; + Ok(()) } From 1ec82c2ba98a5b9ae85d5bfdb9d0f8ceed1d79f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maciej=20W=C3=B3jcik?= Date: Wed, 21 Jan 2026 09:28:40 +0100 Subject: [PATCH 5/6] change truncate option semantics --- tools/defguard_generator/README.md | 10 ++++++++++ tools/defguard_generator/src/main.rs | 8 ++++---- tools/defguard_generator/src/vpn_session_stats.rs | 6 +++--- 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/tools/defguard_generator/README.md b/tools/defguard_generator/README.md index de749c9020..562f45d238 100644 --- a/tools/defguard_generator/README.md +++ b/tools/defguard_generator/README.md @@ -24,3 +24,13 @@ cargo run -p defguard_generator -- vpn-session-stats \ --sessions-per-device 5 ``` +### Session generation logic + +For each device the generator always starts with creating an active (not disconnected) session. +If there are more sessions per device to be generated it goes backwards in time and creates +additional disconnected sessions. +Session duration and gaps between sessions are randomized but there is no logic to verify if +sessions are overlapping so by default the generator runs a `TRUNCATE` query at the start. +To disable this behavior (for example when running it multiple times for separate locations) +use the `--no-truncate` CLI flag. + diff --git a/tools/defguard_generator/src/main.rs b/tools/defguard_generator/src/main.rs index f90113bd45..8c9677d8fe 100644 --- a/tools/defguard_generator/src/main.rs +++ b/tools/defguard_generator/src/main.rs @@ -40,9 +40,9 @@ enum Commands { devices_per_user: u8, #[arg(long)] sessions_per_device: u8, - /// truncate sessions & stats tables before generating stats + /// don't truncate sessions & stats tables before generating stats #[arg(long)] - truncate_sessions_table: bool, + no_truncate: bool, /// insert stats records in batches of specified size #[arg(long, default_value_t = 1000)] stats_batch_size: u16, @@ -78,7 +78,7 @@ async fn main() -> Result<()> { num_users, devices_per_user, sessions_per_device, - truncate_sessions_table, + no_truncate, stats_batch_size, } => { let config = VpnSessionGeneratorConfig { @@ -86,7 +86,7 @@ async fn main() -> Result<()> { num_users, devices_per_user, sessions_per_device, - truncate_sessions_table, + no_truncate, stats_batch_size, }; diff --git a/tools/defguard_generator/src/vpn_session_stats.rs b/tools/defguard_generator/src/vpn_session_stats.rs index c8de4e4843..7e62eb98c9 100644 --- a/tools/defguard_generator/src/vpn_session_stats.rs +++ b/tools/defguard_generator/src/vpn_session_stats.rs @@ -27,7 +27,7 @@ pub struct VpnSessionGeneratorConfig { pub num_users: u16, pub devices_per_user: u8, pub sessions_per_device: u8, - pub truncate_sessions_table: bool, + pub no_truncate: bool, pub stats_batch_size: u16, } @@ -38,8 +38,8 @@ pub async fn generate_vpn_session_stats( info!("Running VPN stats generator with config: {config:#?}"); let mut rng = rand::thread_rng(); - // clear sessions & stats tables - if config.truncate_sessions_table { + // clear sessions & stats tables unless disabled + if !config.no_truncate { info!("Clearing existing sessions & stats"); truncate_with_restart(&pool).await?; } From 58e61dbfd73d7e83318e717132671cbfb13d2198 Mon Sep 17 00:00:00 2001 From: Maciek <19913370+wojcik91@users.noreply.github.com> Date: Wed, 21 Jan 2026 10:10:54 +0100 Subject: [PATCH 6/6] Update tools/defguard_generator/src/main.rs Co-authored-by: Adam --- tools/defguard_generator/src/main.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/defguard_generator/src/main.rs b/tools/defguard_generator/src/main.rs index 8c9677d8fe..3c6feea0e9 100644 --- a/tools/defguard_generator/src/main.rs +++ b/tools/defguard_generator/src/main.rs @@ -30,7 +30,7 @@ struct Cli { #[derive(Subcommand)] enum Commands { - /// generates mock VPN session stats + /// Generates fake VPN session statistics. VpnSessionStats { #[arg(long)] location_id: Id,