capabable spawning multiple openvpn instances

src/main.rs

@@ -8,23 +8,28 @@ mod scraper;
 use anyhow::Result;
 use config::Config;
 use scraper::webdriver::ChromeDriverPool;
+use scraper::vpn_manager::VpnPool;
 use util::directories::DataPaths;
 use util::{logger, opnv};
 use std::sync::Arc;
 
 /// The entry point of the application.
 ///
-/// This function loads the configuration, initializes a shared ChromeDriver pool,
-/// fetches the latest VPNBook OpenVPN configurations if VPN rotation is enabled,
+/// This function loads the configuration, optionally initializes a VPN pool,
+/// initializes a shared ChromeDriver pool bound to the VPN pool (if enabled),
 /// and sequentially runs the full updates for corporate and economic data.
-/// Sequential execution helps prevent resource exhaustion from concurrent
-/// chromedriver instances and avoids spamming the target websites with too many requests.
+/// 
+/// If VPN rotation is enabled:
+/// 1. Fetches latest VPNBook OpenVPN configurations
+/// 2. Creates a VPN pool and connects all VPN instances
+/// 3. Binds each ChromeDriver instance to a different VPN for IP rotation
+/// 4. Performs periodic health checks to reconnect unhealthy VPN instances
 ///
 /// # Errors
 ///
 /// Returns an error if configuration loading fails, pool initialization fails,
 /// VPN fetching fails (if enabled), or if either update function encounters an issue
-/// (e.g., network errors, scraping failures, or chromedriver spawn failures like "program not found").
+/// (e.g., network errors, scraping failures, or chromedriver spawn failures).
 #[tokio::main]
 async fn main() -> Result<()> {
     let config = Config::load().map_err(|err| {
@@ -41,27 +46,100 @@ async fn main() -> Result<()> {
     })?;
 
     logger::log_info("=== Application started ===").await;
-    logger::log_info(&format!("Config: economic_start_date={}, corporate_start_date={}, lookahead_months={}, max_parallel_instances={}, enable_vpn_rotation={}", 
-        config.economic_start_date, config.corporate_start_date, config.economic_lookahead_months, config.max_parallel_instances, config.enable_vpn_rotation)).await;
+    logger::log_info(&format!("Config: economic_start_date={}, corporate_start_date={}, lookahead_months={}, max_parallel_instances={}, enable_vpn_rotation={}, max_tasks_per_instance={}", 
+        config.economic_start_date, config.corporate_start_date, config.economic_lookahead_months, config.max_parallel_instances, config.enable_vpn_rotation, config.max_tasks_per_instance)).await;
 
-    // Initialize the shared ChromeDriver pool once
+    // Initialize VPN pool if enabled
+    let vpn_pool = if config.enable_vpn_rotation {
+        logger::log_info("=== VPN Rotation Enabled ===").await;
+        logger::log_info("--- Fetching latest VPNBook OpenVPN configurations ---").await;
+        
+        let (username, password, _files) = 
+            util::opnv::fetch_vpnbook_configs(&Arc::new(ChromeDriverPool::new(1).await?), paths.cache_dir()).await?;
+
+        let amount_of_openvpn_servers = _files.len();
+        
+        logger::log_info(&format!("✓ Fetched VPN credentials - Username: {}", username)).await;
+
+        // Create VPN pool
+        let openvpn_dir = paths.cache_dir().join("openvpn");
+        logger::log_info("--- Initializing VPN Pool ---").await;
+        let vp = Arc::new(VpnPool::new(
+            &openvpn_dir,
+            username,
+            password,
+            true,  // enable rotation
+            config.tasks_per_vpn_session,
+            amount_of_openvpn_servers,
+        ).await?);
+
+        // Connect all VPN instances (gracefully handles failures)
+        logger::log_info("--- Connecting to VPN servers ---").await;
+        match vp.connect_all().await {
+            Ok(()) => {
+                logger::log_info("✓ VPN initialization complete").await;
+                Some(vp)
+            }
+            Err(e) => {
+                logger::log_warn(&format!(
+                    "⚠ VPN initialization failed: {}. Continuing without VPN.",
+                    e
+                )).await;
+                None
+            }
+        }
+    } else {
+        None
+    };
+
+    // Initialize the shared ChromeDriver pool with VPN pool
     let pool_size = config.max_parallel_instances;
-    logger::log_info(&format!("Initializing ChromeDriver pool with size: {}", pool_size)).await;
+    let max_tasks_per_instance = config.max_tasks_per_instance;
+    
+    logger::log_info(&format!(
+        "Initializing ChromeDriver pool with size: {}{}",
+        pool_size,
+        if max_tasks_per_instance > 0 { &format!(" (max {} tasks/instance)", max_tasks_per_instance) } else { "" }
+    )).await;
+    
+    let pool = Arc::new(
+        if max_tasks_per_instance > 0 {
+            ChromeDriverPool::new_with_vpn_and_task_limit(pool_size, vpn_pool.clone(), max_tasks_per_instance).await?
+        } else if vpn_pool.is_some() {
+            ChromeDriverPool::new_with_vpn(pool_size, vpn_pool.clone()).await?
+        } else {
+            ChromeDriverPool::new(pool_size).await?
+        }
+    );
     
-    let pool = Arc::new(ChromeDriverPool::new(pool_size).await?);
     logger::log_info("✓ ChromeDriver pool initialized successfully").await;
 
-    // Fetch VPNBook configs if VPN rotation is enabled
-    if config.enable_vpn_rotation {
-        logger::log_info("--- Fetching latest VPNBook OpenVPN configurations ---").await;
-        let (username, password, files) = 
-            util::opnv::fetch_vpnbook_configs(&pool, paths.cache_dir()).await?;
-        logger::log_info(&format!("Fetched VPN username: {}, password: {}", username, password)).await;
-        for file in &files {
-            logger::log_info(&format!("Extracted OVPN: {:?}", file)).await;
-        }
-        // Optionally, store username/password for rotation use (e.g., in a file or global state)
-        // For now, just log them; extend as needed for rotation integration
+    // Spawn background Ctrl-C handler to gracefully shutdown pool and VPNs
+    {
+        let pool_for_signal = Arc::clone(&pool);
+        let vpn_for_signal = vpn_pool.clone();
+        tokio::spawn(async move {
+            if let Err(e) = tokio::signal::ctrl_c().await {
+                let _ = util::logger::log_error(&format!("Ctrl-C handler failed to install: {}", e)).await;
+                return;
+            }
+
+            let _ = util::logger::log_info("Ctrl-C received — initiating graceful shutdown").await;
+
+            if let Err(e) = pool_for_signal.shutdown().await {
+                let _ = util::logger::log_warn(&format!("Error shutting down ChromeDriver pool: {}", e)).await;
+            }
+
+            if let Some(vp) = vpn_for_signal {
+                if let Err(e) = vp.disconnect_all().await {
+                    let _ = util::logger::log_warn(&format!("Error disconnecting VPNs: {}", e)).await;
+                }
+            }
+
+            let _ = util::logger::log_info("Graceful shutdown complete (from Ctrl-C)").await;
+            // Exit the process now that cleanup is done
+            std::process::exit(0);
+        });
     }
 
     // Run economic update first, passing the shared pool
@@ -74,6 +152,18 @@ async fn main() -> Result<()> {
     corporate::run_full_update(&config, &pool).await?;
     logger::log_info("✓ Corporate data update completed").await;
 
+    // Shutdown ChromeDriver pool before disconnecting VPNs so instances can
+    // cleanly terminate any network-bound processes.
+    logger::log_info("--- Shutting down ChromeDriver pool ---").await;
+    pool.shutdown().await?;
+    logger::log_info("✓ ChromeDriver pool shutdown complete").await;
+
+    // Disconnect all VPN instances if enabled
+    if let Some(vp) = vpn_pool {
+        logger::log_info("--- Disconnecting VPN instances ---").await;
+        vp.disconnect_all().await?;
+    }
+
     logger::log_info("=== Application completed successfully ===").await;
     Ok(())
 }