reuse session in webdriver

Signed-off-by: Andrei Gherghescu <[email protected]>

Author: andrei-ng

.github/workflows/ci.yml

@@ -129,10 +129,6 @@ jobs:
           
           cargo test --verbose --workspace --features plotly_ndarray,plotly_image,chromedriver --exclude plotly_kaleido
       
-      # Windows-specific cleanup
-      - if: matrix.os == 'windows-latest'
-        run: gci -recurse -filter "*example*"
-
   code-coverage:
     name: Code Coverage
     runs-on: ubuntu-latest

plotly_static/build.rs

@@ -336,13 +336,12 @@ fn main() -> Result<()> {
     } else {
         #[cfg(feature = "chromedriver")]
         {
-            let msg = format!("'webdriver_download' feature disabled. Please install a {CHROMEDRIVER_PATH_ENV} version manually and make the environment variable 'WEBDRIVER_PATH' point to it.");
+            let msg = "'webdriver_download' feature disabled. Please install a 'chromedriver' version manually and make the environment variable 'WEBDRIVER_PATH' point to it.".to_string();
             println!("cargo::warning={msg}");
         }
         #[cfg(feature = "geckodriver")]
         {
-            let msg = format!("'webdriver_download' feature disabled. Please install a {} version manually and make the environment variable 'WEBDRIVER_PATH' point to it.", 
-                GECKODRIVER_PATH_ENV);
+            let msg = format!("'webdriver_download' feature disabled. Please install a 'geckodriver' version manually and make the environment variable 'WEBDRIVER_PATH' point to it.");
             println!("cargo::warning={msg}");
         }
         #[cfg(not(any(feature = "chromedriver", feature = "geckodriver")))]

plotly_static/src/lib.rs

@@ -266,7 +266,7 @@ use std::{println as error, println as warn, println as debug};
 
 use anyhow::{anyhow, Context, Result};
 use base64::{engine::general_purpose, Engine as _};
-use fantoccini::{wd::Capabilities, ClientBuilder};
+use fantoccini::{wd::Capabilities, Client, ClientBuilder};
 #[cfg(not(test))]
 use log::{debug, error, warn};
 use serde::Serialize;
@@ -469,10 +469,15 @@ struct PlotData<'a> {
 /// - Browser capabilities: Default Chrome/Firefox headless options
 /// - Automatic WebDriver detection and connection reuse
 pub struct StaticExporterBuilder {
+    /// WebDriver server port (default: 4444)
     webdriver_port: u32,
+    /// WebDriver server base URL (default: "http://localhost")
     webdriver_url: String,
+    /// Auto-spawn WebDriver if not running (default: true)
     spawn_webdriver: bool,
+    /// Use bundled JS libraries instead of CDN (default: false)
     offline_mode: bool,
+    /// Browser command-line flags (e.g., "--headless", "--no-sandbox")
     webdriver_browser_caps: Vec<String>,
 }
 
@@ -675,6 +680,7 @@ impl StaticExporterBuilder {
             offline_mode: self.offline_mode,
             webdriver_browser_caps: self.webdriver_browser_caps.clone(),
             runtime,
+            webdriver_client: None,
         })
     }
 
@@ -737,12 +743,26 @@ impl StaticExporterBuilder {
 /// - Offline mode support
 /// - Automatic WebDriver management
 pub struct StaticExporter {
+    /// WebDriver server port (default: 4444)
     webdriver_port: u32,
+
+    /// WebDriver server base URL (default: "http://localhost")
     webdriver_url: String,
+
+    /// WebDriver process manager for spawning and cleanup
     webdriver: WebDriver,
+
+    /// Use bundled JS libraries instead of CDN
     offline_mode: bool,
+
+    /// Browser command-line flags (e.g., "--headless", "--no-sandbox")
     webdriver_browser_caps: Vec<String>,
+
+    /// Tokio runtime for async operations
     runtime: std::sync::Arc<tokio::runtime::Runtime>,
+
+    /// Cached WebDriver client for session reuse
+    webdriver_client: Option<Client>,
 }
 
 impl Drop for StaticExporter {
@@ -757,6 +777,16 @@ impl Drop for StaticExporter {
     /// - Leaves externally managed WebDriver sessions running
     /// - Logs errors but doesn't panic if cleanup fails
     fn drop(&mut self) {
+        // Close the WebDriver client if it exists
+        if let Some(client) = self.webdriver_client.take() {
+            let runtime = self.runtime.clone();
+            runtime.block_on(async {
+                if let Err(e) = client.close().await {
+                    error!("Failed to close WebDriver client: {e}");
+                }
+            });
+        }
+
         // Stop the WebDriver process
         if let Err(e) = self.webdriver.stop() {
             error!("Failed to stop WebDriver: {e}");
@@ -941,11 +971,20 @@ impl StaticExporter {
         debug!("Use WebDriver and headless browser to export static plot");
         let webdriver_url = format!("{}:{}", self.webdriver_url, self.webdriver_port,);
 
-        let client = ClientBuilder::native()
-            .capabilities(caps)
-            .connect(&webdriver_url)
-            .await
-            .with_context(|| "WebDriver session errror")?;
+        // Reuse existing client or create new one
+        let client = if let Some(ref client) = self.webdriver_client {
+            debug!("Reusing existing WebDriver session");
+            client.clone()
+        } else {
+            debug!("Creating new WebDriver session");
+            let new_client = ClientBuilder::native()
+                .capabilities(caps)
+                .connect(&webdriver_url)
+                .await
+                .with_context(|| "WebDriver session error")?;
+            self.webdriver_client = Some(new_client.clone());
+            new_client
+        };
 
         // URL-encode the HTML
         let data_uri = format!("data:text/html,{}", encode(data_uri));
@@ -978,7 +1017,8 @@ impl StaticExporter {
 
         let data = client.execute_async(js_script, args).await?;
 
-        client.close().await?;
+        // Don't close the client - keep it for reuse
+        // client.close().await?;
 
         let src = data.as_str().ok_or(anyhow!(
             "Failed to execute Plotly.toImage in browser session"
@@ -1282,46 +1322,50 @@ mod tests {
         init();
         let test_plot = create_test_plot();
 
+        // Use a unique port to test actual WebDriver process reuse
+        let test_port = get_unique_port();
+
         // Create first exporter - this should spawn a new WebDriver
         let mut export1 = StaticExporterBuilder::default()
             .spawn_webdriver(true)
-            .webdriver_port(get_unique_port())
+            .webdriver_port(test_port)
             .build()
             .unwrap();
 
         // Export first image
-        let dst1 = PathBuf::from("session_reuse_1.png");
+        let dst1 = PathBuf::from("process_reuse_1.png");
         export1
             .write_fig(dst1.as_path(), &test_plot, ImageFormat::PNG, 800, 600, 1.0)
             .unwrap();
         assert!(dst1.exists());
         assert!(std::fs::remove_file(dst1.as_path()).is_ok());
 
         // Create second exporter on the same port - this should connect to existing
-        // WebDriver
+        // WebDriver process (but create a new session)
         let mut export2 = StaticExporterBuilder::default()
             .spawn_webdriver(true)
-            .webdriver_port(get_unique_port())
+            .webdriver_port(test_port)
             .build()
             .unwrap();
 
-        // Export second image using the same WebDriver session
-        let dst2 = PathBuf::from("session_reuse_2.png");
+        // Export second image using a new session on the same WebDriver process
+        let dst2 = PathBuf::from("process_reuse_2.png");
         export2
             .write_fig(dst2.as_path(), &test_plot, ImageFormat::PNG, 800, 600, 1.0)
             .unwrap();
         assert!(dst2.exists());
         assert!(std::fs::remove_file(dst2.as_path()).is_ok());
 
         // Create third exporter on the same port - should also connect to existing
+        // WebDriver process
         let mut export3 = StaticExporterBuilder::default()
             .spawn_webdriver(true)
-            .webdriver_port(get_unique_port())
+            .webdriver_port(test_port)
             .build()
             .unwrap();
 
-        // Export third image
-        let dst3 = PathBuf::from("session_reuse_3.png");
+        // Export third image using another new session on the same WebDriver process
+        let dst3 = PathBuf::from("process_reuse_3.png");
         export3
             .write_fig(dst3.as_path(), &test_plot, ImageFormat::PNG, 800, 600, 1.0)
             .unwrap();