diff --git a/src/python_spy.rs b/src/python_spy.rs index 0ee24abd..797c5600 100644 --- a/src/python_spy.rs +++ b/src/python_spy.rs @@ -36,7 +36,7 @@ pub struct PythonSpy { pub native: Option, pub short_filenames: HashMap>, pub python_thread_ids: HashMap, - pub python_thread_names: HashMap, + pub python_thread_names: HashMap>, #[cfg(target_os = "linux")] pub dockerized: bool, } @@ -229,6 +229,8 @@ impl PythonSpy { let gil_thread_id = get_gil_threadid::(self.threadstate_address, &self.process) .context("failed to get gil_thread_id")?; + let mut did_thread_name_lookup = false; + let mut traces = Vec::new(); let mut threads = threads_head; while !threads.is_null() { @@ -283,7 +285,8 @@ impl PythonSpy { trace.os_thread_id = os_thread_id.map(|id| id as u64); } - trace.thread_name = self._get_python_thread_name(python_thread_id); + trace.thread_name = + self._get_python_thread_name(&mut did_thread_name_lookup, python_thread_id); trace.owns_gil = owns_gil; trace.pid = self.process.pid; @@ -536,12 +539,29 @@ impl PythonSpy { Ok(None) } - fn _get_python_thread_name(&mut self, python_thread_id: u64) -> Option { + fn _get_python_thread_name( + &mut self, + did_thread_name_lookup: &mut bool, + python_thread_id: u64, + ) -> Option { match self.python_thread_names.get(&python_thread_id) { - Some(thread_name) => Some(thread_name.clone()), + Some(thread_name) => thread_name.clone(), None => { - self.python_thread_names = thread_name_lookup(self).unwrap_or_default(); - self.python_thread_names.get(&python_thread_id).cloned() + // avoid triggering a lookup if we already did one during this sample + if !*did_thread_name_lookup { + info!("looking up thread names"); + self.python_thread_names = thread_name_lookup(self) + .unwrap_or_default() + .into_iter() + .map(|(k, v)| (k, Some(v))) + .collect(); + *did_thread_name_lookup = true; + } + // avoid triggering a lookup next time if the thread has no name + self.python_thread_names + .entry(python_thread_id) + .or_insert(None) + .clone() } } }