Add flag to limit analysis time per file

Author: Hendrik van Antwerpen

tree-sitter-stack-graphs/src/cli/analyze.rs

@@ -16,31 +16,52 @@ use stack_graphs::stitching::Database;
 use std::collections::HashMap;
 use std::path::Path;
 use std::path::PathBuf;
+use std::sync::Arc;
+use std::time::Duration;
 use tree_sitter_graph::Variables;
 use walkdir::WalkDir;
 
+use crate::cli::util::duration_from_seconds_str;
 use crate::cli::util::map_parse_errors;
 use crate::cli::util::path_exists;
 use crate::loader::Loader;
+use crate::CancelAfterDuration;
+use crate::CancellationFlag;
 use crate::LoadError;
 use crate::NoCancellation;
 
 /// Analyze sources
 #[derive(Args)]
 pub struct AnalyzeArgs {
     /// Source file or directory paths.
-    #[clap(value_name = "SOURCE_PATH", required = true, value_hint = ValueHint::AnyPath, parse(from_os_str), validator_os = path_exists)]
+    #[clap(
+        value_name = "SOURCE_PATH",
+        required = true,
+        value_hint = ValueHint::AnyPath,
+        parse(from_os_str),
+        validator_os = path_exists,
+    )]
     pub source_paths: Vec<PathBuf>,
 
-    #[clap(short = 'v')]
+    #[clap(long, short = 'v')]
     pub verbose: bool,
+
+    /// Maximum runtime per file in seconds.
+    #[clap(
+        long,
+        value_name = "SECONDS",
+        parse(try_from_str = duration_from_seconds_str),
+        require_equals = true,
+    )]
+    pub max_file_time: Option<Duration>,
 }
 
 impl AnalyzeArgs {
     pub fn new(source_paths: Vec<PathBuf>) -> Self {
         Self {
             source_paths,
             verbose: false,
+            max_file_time: None,
         }
     }
 
@@ -76,36 +97,41 @@ impl AnalyzeArgs {
             .with_context(|| format!("Error analyzing file {}", source_path.display()))
     }
 
-    /// Run test file.
     fn analyze_file(
         &self,
         source_root: &Path,
         source_path: &Path,
         loader: &mut Loader,
     ) -> anyhow::Result<()> {
         if self.verbose {
-            eprint!("{} ", source_path.display());
+            eprint!("{}: ", source_path.display());
+        }
+
+        let mut cancellation_flag: Arc<dyn CancellationFlag> = Arc::new(NoCancellation);
+        if let Some(max_file_time) = self.max_file_time {
+            cancellation_flag = CancelAfterDuration::new(max_file_time);
         }
 
         let source = std::fs::read_to_string(source_path)?;
-        let lc = match loader.load_for_file(source_path, Some(&source), &NoCancellation)? {
-            Some(sgl) => sgl,
-            None => {
-                if self.verbose {
-                    eprintln!("{}", "⦵".dimmed());
+        let lc = match loader.load_for_file(source_path, Some(&source), cancellation_flag.as_ref())
+        {
+            Ok(Some(sgl)) => sgl,
+            Ok(None) => return Ok(()),
+            Err(crate::loader::LoadError::Cancelled(_)) => {
+                if !self.verbose {
+                    eprint!("{}: ", source_path.display());
                 }
+                eprintln!("{}", "language loading timed out".yellow());
                 return Ok(());
             }
+            Err(e) => return Err(e.into()),
         };
 
         let mut graph = StackGraph::new();
         let file = graph
             .add_file(&source_path.to_string_lossy())
             .map_err(|_| anyhow!("Duplicate file {}", source_path.display()))?;
 
-        if self.verbose {
-            eprint!("{} ", "🕸".yellow());
-        }
         let relative_source_path = source_path.strip_prefix(source_root).unwrap();
         let result = if let Some(fa) = source_path
             .file_name()
@@ -118,45 +144,63 @@ impl AnalyzeArgs {
                 &source,
                 &mut std::iter::empty(),
                 &HashMap::new(),
-                &NoCancellation,
+                cancellation_flag.as_ref(),
             )
         } else {
             let globals = Variables::new();
-            lc.sgl
-                .build_stack_graph_into(&mut graph, file, &source, &globals, &NoCancellation)
+            lc.sgl.build_stack_graph_into(
+                &mut graph,
+                file,
+                &source,
+                &globals,
+                cancellation_flag.as_ref(),
+            )
         };
         match result {
             Err(LoadError::ParseErrors(parse_errors)) => {
-                let parse_error = map_parse_errors(source_path, &parse_errors, &source);
+                let parse_error = map_parse_errors(source_path, &parse_errors, &source, "");
                 if !self.verbose {
-                    eprint!("{}", source_path.display());
+                    eprint!("{}: ", source_path.display());
                 }
-                eprintln!("{}", "✗".red());
+                eprintln!("{}", "parsing failed".red());
                 eprintln!("{}", parse_error);
                 return Ok(());
             }
+            Err(LoadError::Cancelled(_)) => {
+                if !self.verbose {
+                    eprint!("{}: ", source_path.display());
+                }
+                eprintln!("{}", "parsing timed out".yellow());
+                return Ok(());
+            }
             Err(e) => return Err(e.into()),
             Ok(_) => {}
         };
 
-        if self.verbose {
-            eprint!("{} ", "🦶".yellow());
-        }
         let mut partials = PartialPaths::new();
         let mut db = Database::new();
-        partials.find_all_partial_paths_in_file(
+        match partials.find_all_partial_paths_in_file(
             &graph,
             file,
-            &stack_graphs::NoCancellation,
+            &cancellation_flag.as_ref(),
             |g, ps, p| {
                 if p.is_complete_as_possible(g) {
                     db.add_partial_path(g, ps, p);
                 }
             },
-        )?;
+        ) {
+            Ok(_) => {}
+            Err(_) => {
+                if !self.verbose {
+                    eprint!("{}: ", source_path.display());
+                }
+                eprintln!("{}", "path computation timed out".yellow());
+                return Ok(());
+            }
+        }
 
         if self.verbose {
-            eprintln!("{}", "✓".green());
+            eprintln!("{}", "success".green());
         }
         Ok(())
     }

tree-sitter-stack-graphs/src/cli/test.rs

@@ -281,7 +281,7 @@ impl TestArgs {
     ) -> anyhow::Result<()> {
         match sgl.build_stack_graph_into(graph, file, source, globals, &NoCancellation) {
             Err(LoadError::ParseErrors(parse_errors)) => {
-                Err(map_parse_errors(test_path, &parse_errors, source))
+                Err(map_parse_errors(test_path, &parse_errors, source, "  "))
             }
             Err(e) => Err(e.into()),
             Ok(_) => Ok(()),

tree-sitter-stack-graphs/src/cli/util.rs

@@ -10,6 +10,7 @@ use std::ffi::OsStr;
 use std::ffi::OsString;
 use std::path::Path;
 use std::path::PathBuf;
+use std::time::Duration;
 use tree_sitter_graph::parse_error::TreeWithParseErrorVec;
 
 use crate::cli::MAX_PARSE_ERRORS;
@@ -126,14 +127,16 @@ pub fn map_parse_errors(
     test_path: &Path,
     parse_errors: &TreeWithParseErrorVec,
     source: &str,
+    prefix: &str,
 ) -> anyhow::Error {
     let mut error = String::new();
     let parse_errors = parse_errors.errors();
     for parse_error in parse_errors.iter().take(MAX_PARSE_ERRORS) {
         let line = parse_error.node().start_position().row;
         let column = parse_error.node().start_position().column;
         error.push_str(&format!(
-            "  {}:{}:{}: {}\n",
+            "{}{}:{}:{}: {}\n",
+            prefix,
             test_path.display(),
             line + 1,
             column + 1,
@@ -150,3 +153,7 @@ pub fn map_parse_errors(
     }
     anyhow!(error)
 }
+
+pub fn duration_from_seconds_str(s: &str) -> Result<Duration, anyhow::Error> {
+    Ok(Duration::new(s.parse()?, 0))
+}

tree-sitter-stack-graphs/src/lib.rs

@@ -322,6 +322,9 @@ use std::mem::transmute;
 use std::path::Path;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
+use std::sync::Arc;
+use std::time::Duration;
+use std::time::Instant;
 use thiserror::Error;
 use tree_sitter::Parser;
 use tree_sitter_graph::functions::Functions;
@@ -612,6 +615,43 @@ impl CancellationFlag for NoCancellation {
     }
 }
 
+pub struct CancelAfterDuration {
+    flag: AtomicUsize,
+}
+
+impl CancelAfterDuration {
+    pub fn new(limit: Duration) -> Arc<Self> {
+        let result = Arc::new(Self {
+            flag: AtomicUsize::new(0),
+        });
+        let start = Instant::now();
+        let timer = Arc::downgrade(&result);
+        std::thread::spawn(move || {
+            loop {
+                std::thread::sleep(Duration::from_millis(10));
+                if let Some(timer) = timer.upgrade() {
+                    // the flag is still in use
+                    if start.elapsed().ge(&limit) {
+                        // set flag and stop polling
+                        timer.flag.store(1, Ordering::Relaxed);
+                        return;
+                    }
+                } else {
+                    // the flag is not in use anymore, stop polling
+                    return;
+                }
+            }
+        });
+        result
+    }
+}
+
+impl CancellationFlag for CancelAfterDuration {
+    fn flag(&self) -> Option<&AtomicUsize> {
+        Some(&self.flag)
+    }
+}
+
 /// An error that can occur while loading a stack graph from a TSG file
 #[derive(Debug, Error)]
 pub enum LoadError {