Review Rust browser automation architecture and code

Cargo.toml

@@ -36,6 +36,9 @@ clap = { version = "4.0", features = ["derive"] }
 axum = { version = "0.8", features = ["macros"] }
 tower-http = { version = "0.6", features = ["cors"] }
 
+# Lazy initialization
+once_cell = "1.19"
+
 # OAuth dependencies
 oauth2 = { version = "4.4" }
 openidconnect = { version = "3.5" }

src/client.rs

@@ -3,13 +3,23 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
 use fantoccini::{Client, ClientBuilder, Locator, elements::Element};
 use futures::lock::Mutex;
 
-use crate::{config::Config, driver::DriverManager, error::Result};
+use crate::{config::Config, driver::DriverManager, error::Result, pool::ConnectionPool};
+
+/// Metadata about a session for pool management
+#[derive(Clone, Debug)]
+struct SessionMetadata {
+    driver_type: crate::driver::DriverType,
+}
 
 #[derive(Clone)]
 pub struct ClientManager {
     clients: Arc<Mutex<HashMap<String, Client>>>,
+    /// Metadata for each session (driver type, etc.)
+    session_metadata: Arc<Mutex<HashMap<String, SessionMetadata>>>,
     config: Config,
     driver_manager: DriverManager,
+    /// Connection pool for reusing sessions
+    pool: Arc<ConnectionPool>,
 }
 
 impl ClientManager {
@@ -18,10 +28,14 @@ impl ClientManager {
             .validate()
             .map_err(|e| anyhow::anyhow!("Configuration error: {}", e))?;
 
+        let pool = Arc::new(ConnectionPool::new(&config));
+
         Ok(Self {
             clients: Arc::new(Mutex::new(HashMap::new())),
+            session_metadata: Arc::new(Mutex::new(HashMap::new())),
             config,
             driver_manager: DriverManager::new(),
+            pool,
         })
     }
 
@@ -68,22 +82,56 @@ impl ClientManager {
 
     /// Full multi-tenant client creation for HTTP mode
     async fn get_or_create_client_http(&self, session_id: Option<String>) -> Result<(String, Client)> {
-        let mut clients = self.clients.lock().await;
         let session = session_id.unwrap_or_else(|| "default".to_string());
 
-        if let Some(client) = clients.get(&session) {
-            match client.current_url().await {
-                Ok(_) => return Ok((session, client.clone())),
-                Err(_) => {
-                    clients.remove(&session);
+        // Check active clients first
+        {
+            let mut clients = self.clients.lock().await;
+            if let Some(client) = clients.get(&session) {
+                match client.current_url().await {
+                    Ok(_) => return Ok((session, client.clone())),
+                    Err(_) => {
+                        clients.remove(&session);
+                        // Also remove from metadata
+                        let mut metadata = self.session_metadata.lock().await;
+                        metadata.remove(&session);
+                    }
                 }
             }
         }
 
-        // Determine the actual endpoint to use based on session preferences
+        // Determine driver type for this session
+        let driver_type = self.extract_browser_preference_from_session(&session)
+            .unwrap_or(crate::driver::DriverType::Chrome);
+
+        // Try to acquire from pool
+        if let Ok(Some((pooled_session, client))) = self.pool.acquire(&driver_type).await {
+            tracing::debug!(
+                "Reusing pooled {} connection for session '{}'",
+                driver_type.browser_name(),
+                session
+            );
+
+            // Store in active clients
+            let mut clients = self.clients.lock().await;
+            clients.insert(session.clone(), client.clone());
+
+            // Store metadata
+            let mut metadata = self.session_metadata.lock().await;
+            metadata.insert(session.clone(), SessionMetadata {
+                driver_type: driver_type.clone(),
+            });
+
+            // Update pool to track with new session id
+            self.pool.release(&driver_type, &pooled_session).await;
+            self.pool.add(driver_type, client.clone(), session.clone()).await;
+
+            return Ok((session, client));
+        }
+
+        // No pooled connection available, create a new one
         let endpoint = self.resolve_webdriver_endpoint_for_session(&session).await?;
 
-        // Create client with proper browser configuration
         let client = self
             .create_configured_client(&endpoint, &session)
             .await
@@ -96,7 +144,22 @@ impl ClientManager {
                 )
             })?;
 
+        // Add to pool
+        let added_to_pool = self.pool.add(driver_type.clone(), client.clone(), session.clone()).await;
+        if added_to_pool {
+            tracing::debug!("Added new {} connection to pool: {}", driver_type.browser_name(), session);
+        }
+
+        // Store in active clients
+        let mut clients = self.clients.lock().await;
         clients.insert(session.clone(), client.clone());
+
+        // Store metadata
+        let mut metadata = self.session_metadata.lock().await;
+        metadata.insert(session.clone(), SessionMetadata {
+            driver_type,
+        });
+
         Ok((session, client))
     }
 
@@ -198,7 +261,7 @@ impl ClientManager {
         // If endpoint is "auto", try to use pre-started drivers first
         if self.config.webdriver_endpoint == "auto" {
             // Check for healthy pre-started drivers
-            let healthy_endpoints = self.driver_manager.get_healthy_endpoints();
+            let healthy_endpoints = self.driver_manager.get_healthy_endpoints().await;
             tracing::debug!("Available healthy endpoints: {:?}", healthy_endpoints);
 
             if !healthy_endpoints.is_empty() {
@@ -345,14 +408,53 @@ impl ClientManager {
         &self.config
     }
 
+    /// Get access to the connection pool
+    pub fn get_pool(&self) -> &ConnectionPool {
+        &self.pool
+    }
+
+    /// Release a session back to the pool (marks it as idle for reuse)
+    pub async fn release_session(&self, session_id: &str) {
+        // Get the driver type for this session
+        let driver_type = {
+            let metadata = self.session_metadata.lock().await;
+            metadata.get(session_id).map(|m| m.driver_type.clone())
+        };
+
+        if let Some(driver_type) = driver_type {
+            self.pool.release(&driver_type, session_id).await;
+            tracing::debug!("Released session '{}' back to pool", session_id);
+        }
+    }
+
+    /// Get pool statistics
+    pub async fn get_pool_stats(&self) -> std::collections::HashMap<crate::driver::DriverType, crate::pool::PoolStats> {
+        self.pool.get_stats().await
+    }
+
     /// Close all active WebDriver sessions
     pub async fn close_all_sessions(&self) -> Result<()> {
         tracing::info!("Closing all active WebDriver sessions...");
+
+        // Close all pooled connections first
+        if self.pool.is_enabled() {
+            tracing::debug!("Closing pooled connections...");
+            if let Err(e) = self.pool.close_all().await {
+                tracing::warn!("Error closing pool connections: {}", e);
+            }
+        }
+
+        // Clear session metadata
+        {
+            let mut metadata = self.session_metadata.lock().await;
+            metadata.clear();
+        }
+
         let mut clients = self.clients.lock().await;
-        
+
         for (session_id, client) in clients.drain() {
             tracing::debug!("Closing session: {}", session_id);
-            
+
             // Add timeout to individual session close operations
             let close_timeout = Duration::from_secs(2);
             match tokio::time::timeout(close_timeout, client.close()).await {
@@ -364,13 +466,13 @@ impl ClientManager {
                 }
             }
         }
-        
+
         // CRITICAL FIX: Force cleanup of orphaned browser processes
         self.force_cleanup_orphaned_processes().await?;
-        
+
         // Also cleanup all managed processes through driver manager
         self.driver_manager.force_cleanup_all_processes().await?;
-        
+
         tracing::info!("All WebDriver sessions closed and orphaned processes cleaned");
         Ok(())
     }
@@ -541,8 +643,6 @@ impl ClientManager {
     }
 }
 
-impl Default for ClientManager {
-    fn default() -> Self {
-        Self::new(Config::from_env()).expect("Failed to create ClientManager with default config")
-    }
-}
+// Note: Default is intentionally not implemented for ClientManager
+// because ClientManager::new() can fail if configuration validation fails.
+// Use ClientManager::new(Config::from_env()) with proper error handling instead.

src/config.rs

@@ -13,6 +13,16 @@ pub struct Config {
     pub driver_startup_timeout_ms: u64,
     /// Enable Chrome performance memory APIs
     pub enable_performance_memory: bool,
+
+    // Connection pool settings
+    /// Maximum number of connections per driver type in the pool
+    pub pool_max_connections_per_driver: usize,
+    /// Idle timeout in seconds before closing unused connections
+    pub pool_idle_timeout_secs: u64,
+    /// Maximum time to wait for a connection in milliseconds
+    pub pool_acquire_timeout_ms: u64,
+    /// Enable connection pooling (true by default)
+    pub pool_enabled: bool,
 }
 
 impl Config {
@@ -41,6 +51,23 @@ impl Config {
             enable_performance_memory: env::var("WEBDRIVER_ENABLE_PERFORMANCE_MEMORY")
                 .map(|v| v.to_lowercase() == "true" || v == "1")
                 .unwrap_or(false), // Default to false for compatibility
+
+            // Connection pool settings
+            pool_max_connections_per_driver: env::var("WEBDRIVER_POOL_MAX_CONNECTIONS")
+                .ok()
+                .and_then(|s| s.parse().ok())
+                .unwrap_or(3), // Default to 3 connections per driver
+            pool_idle_timeout_secs: env::var("WEBDRIVER_POOL_IDLE_TIMEOUT_SECS")
+                .ok()
+                .and_then(|s| s.parse().ok())
+                .unwrap_or(300), // Default to 5 minutes
+            pool_acquire_timeout_ms: env::var("WEBDRIVER_POOL_ACQUIRE_TIMEOUT_MS")
+                .ok()
+                .and_then(|s| s.parse().ok())
+                .unwrap_or(30000), // Default to 30 seconds
+            pool_enabled: env::var("WEBDRIVER_POOL_ENABLED")
+                .map(|v| v.to_lowercase() == "true" || v == "1")
+                .unwrap_or(true), // Default to enabled
         }
     }
 
@@ -76,6 +103,19 @@ impl Config {
             }
         }
 
+        // Validate pool settings
+        if self.pool_max_connections_per_driver == 0 {
+            return Err("Pool max connections must be greater than 0".to_string());
+        }
+
+        if self.pool_idle_timeout_secs == 0 {
+            return Err("Pool idle timeout must be greater than 0".to_string());
+        }
+
+        if self.pool_acquire_timeout_ms == 0 {
+            return Err("Pool acquire timeout must be greater than 0".to_string());
+        }
+
         Ok(())
     }
 
@@ -114,6 +154,12 @@ WebDriver MCP Server Setup:
    - WEBDRIVER_STARTUP_TIMEOUT_MS: Driver startup timeout (default: 10000)
    - WEBDRIVER_ENABLE_PERFORMANCE_MEMORY: true or false (default: false) - enables Chrome memory APIs
 
+   Connection Pool Settings:
+   - WEBDRIVER_POOL_ENABLED: true (default) or false - enable connection pooling
+   - WEBDRIVER_POOL_MAX_CONNECTIONS: max connections per driver (default: 3)
+   - WEBDRIVER_POOL_IDLE_TIMEOUT_SECS: idle timeout before closing (default: 300)
+   - WEBDRIVER_POOL_ACQUIRE_TIMEOUT_MS: timeout to acquire connection (default: 30000)
+
 3. Manual Setup (if auto-start disabled):
    - Chrome: chromedriver --port=9515
    - Firefox: geckodriver --port=4444