Merge pull request #375 from pythonspeed/374-rust-line-cache

Rust-based line cache, to prevent reentrancy deadlocks and GIL deadlocks

.changelog/347.bugfix

@@ -0,0 +1 @@
+Give more accurate message when running in no-browser mode (thanks to Paul-Louis NECH).

.changelog/374.bugfix

@@ -0,0 +1 @@
+Fix a number of potential deadlock scenarios when writing out reports.

memapi/Cargo.toml

@@ -38,6 +38,7 @@ cgroups-rs = "0.2.9"
 [dev-dependencies]
 proptest = "1.0"
 proc-maps = "0.2.1"
+tempfile = "3.3.0"
 
 [features]
 default = []

memapi/src/lib.rs

@@ -1,6 +1,7 @@
 #![deny(unsafe_op_in_unsafe_fn)]
 pub mod ffi;
 pub mod flamegraph;
+mod linecache;
 pub mod memorytracking;
 pub mod mmap;
 pub mod oom;

memapi/src/linecache.rs

@@ -0,0 +1,72 @@
+//! A line caching library, a bit like Python's linecache.
+
+use std::{
+    collections::HashMap,
+    fs::File,
+    io::{BufRead, BufReader},
+};
+
+use ahash::RandomState;
+
+/// Store a cache of files in memory, allow easy reading of lines.
+#[derive(Default)]
+pub struct LineCacher {
+    file_lines: HashMap<String, Vec<String>, RandomState>,
+}
+
+static EMPTY_STRING: String = String::new();
+
+impl LineCacher {
+    /// Get the source code line for the given file. If the file doesn't exist
+    /// or line number is too big, an empty string is returned. Line endings are
+    /// stripped.
+    pub fn get_source_line<'a>(&'a mut self, filename: &str, line_number: usize) -> &'a str {
+        if line_number == 0 {
+            return &EMPTY_STRING;
+        }
+        let entry =
+            self.file_lines
+                .entry(filename.to_string())
+                .or_insert_with(|| -> Vec<String> {
+                    File::open(filename)
+                        .map(|f| {
+                            BufReader::new(f)
+                                .lines()
+                                .map(|l| l.unwrap_or_else(|_| EMPTY_STRING.clone()))
+                                .collect()
+                        })
+                        .unwrap_or_else(|_| vec![])
+                });
+        entry.get(line_number - 1).unwrap_or(&EMPTY_STRING)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io::Write;
+
+    use super::LineCacher;
+
+    #[test]
+    fn linecache() {
+        let mut cache = LineCacher::default();
+
+        // Non-existent file
+        assert_eq!(cache.get_source_line("/no/such/file", 1), "");
+
+        let mut f = tempfile::NamedTempFile::new().unwrap();
+        f.as_file_mut().write_all(b"abc\ndef\r\nghijk").unwrap();
+        let path = f.path().as_os_str().to_str().unwrap();
+
+        // 0 line number
+        assert_eq!(cache.get_source_line(path, 0), "");
+
+        // Too high line number
+        assert_eq!(cache.get_source_line(path, 4), "");
+
+        // Present line numbers
+        assert_eq!(cache.get_source_line(path, 1), "abc");
+        assert_eq!(cache.get_source_line(path, 2), "def");
+        assert_eq!(cache.get_source_line(path, 3), "ghijk");
+    }
+}

memapi/src/memorytracking.rs

@@ -1,5 +1,6 @@
 use crate::flamegraph::filter_to_useful_callstacks;
 use crate::flamegraph::write_flamegraphs;
+use crate::linecache::LineCacher;
 use crate::python::get_runpy_path;
 
 use super::rangemap::RangeMap;
@@ -168,6 +169,7 @@ impl Callstack {
         to_be_post_processed: bool,
         functions: &dyn FunctionLocations,
         separator: &'static str,
+        linecache: &mut LineCacher,
     ) -> String {
         if self.calls.is_empty() {
             return "[No Python stack]".to_string();
@@ -190,8 +192,7 @@ impl Callstack {
             .map(|(id, (function, filename))| {
                 if to_be_post_processed {
                     // Get Python code.
-                    let code = crate::python::get_source_line(filename, id.line_number)
-                        .unwrap_or_else(|_| "".to_string());
+                    let code = linecache.get_source_line(filename, id.line_number as usize);
                     // Leading whitespace is dropped by SVG, so we'd like to
                     // replace it with non-breaking space. However, inferno
                     // trims whitespace
@@ -386,7 +387,7 @@ impl<FL: FunctionLocations> AllocationTracker<FL> {
         eprintln!("=fil-profile= {}", message);
         eprintln!(
             "=| {}",
-            callstack.as_string(false, &self.functions, "\n=| ")
+            callstack.as_string(false, &self.functions, "\n=| ", &mut LineCacher::default())
         );
     }
 
@@ -611,13 +612,15 @@ impl<FL: FunctionLocations> AllocationTracker<FL> {
     ) -> impl ExactSizeIterator<Item = String> + '_ {
         let by_call = self.combine_callstacks(peak).into_iter();
         let id_to_callstack = self.interner.get_reverse_map();
+        let mut linecache = LineCacher::default();
         by_call.map(move |(callstack_id, size)| {
             format!(
                 "{} {}",
                 id_to_callstack.get(&callstack_id).unwrap().as_string(
                     to_be_post_processed,
                     &self.functions,
-                    ";"
+                    ";",
+                    &mut linecache,
                 ),
                 size,
             )

memapi/src/python.rs

@@ -3,18 +3,6 @@ use once_cell::sync::Lazy;
 use pyo3::prelude::*;
 use pyo3::types::PyModule;
 
-// Get the source code line from a given filename.
-pub fn get_source_line(filename: &str, line_number: u16) -> PyResult<String> {
-    Python::with_gil(|py| {
-        let linecache = PyModule::import(py, "linecache")?;
-        let result: String = linecache
-            .getattr("getline")?
-            .call1((filename, line_number))?
-            .to_string();
-        Ok(result)
-    })
-}
-
 // Return the filesystem path of the stdlib's runpy module.
 pub fn get_runpy_path() -> &'static str {
     static PATH: Lazy<String> = Lazy::new(|| {