Optionally reuse existing debuggee handle

src/agent/coverage/src/block/windows.rs

@@ -171,7 +171,7 @@ impl<'c> Recorder<'c> {
             return Ok(());
         }
 
-        match self.cache.fetch(&path) {
+        match self.cache.fetch(&path, dbg.target().process_handle()) {
             Ok(Some(info)) => {
                 let new = self.coverage.insert(&path, info.blocks.iter().copied());
 

src/agent/coverage/src/cache.rs

@@ -6,6 +6,9 @@ use std::collections::{BTreeSet, HashMap};
 use anyhow::Result;
 use serde::{Deserialize, Serialize};
 
+#[cfg(target_os = "windows")]
+use winapi::um::winnt::HANDLE;
+
 use crate::code::{ModuleIndex, ModulePath};
 
 #[derive(Clone, Debug, Default, Deserialize, Serialize)]
@@ -20,6 +23,7 @@ impl ModuleCache {
         Self { cached }
     }
 
+    #[cfg(target_os = "linux")]
     pub fn fetch(&mut self, path: &ModulePath) -> Result<Option<&ModuleInfo>> {
         if !self.cached.contains_key(path) {
             self.insert(path)?;
@@ -28,6 +32,19 @@ impl ModuleCache {
         Ok(self.cached.get(path))
     }
 
+    #[cfg(target_os = "windows")]
+    pub fn fetch(
+        &mut self,
+        path: &ModulePath,
+        handle: impl Into<Option<HANDLE>>,
+    ) -> Result<Option<&ModuleInfo>> {
+        if !self.cached.contains_key(path) {
+            self.insert(path, handle)?;
+        }
+
+        Ok(self.cached.get(path))
+    }
+
     #[cfg(target_os = "linux")]
     pub fn insert(&mut self, path: &ModulePath) -> Result<()> {
         let entry = ModuleInfo::new_elf(path)?;
@@ -36,8 +53,8 @@ impl ModuleCache {
     }
 
     #[cfg(target_os = "windows")]
-    pub fn insert(&mut self, path: &ModulePath) -> Result<()> {
-        let entry = ModuleInfo::new_pe(path)?;
+    pub fn insert(&mut self, path: &ModulePath, handle: impl Into<Option<HANDLE>>) -> Result<()> {
+        let entry = ModuleInfo::new_pe(path, handle)?;
         self.cached.insert(path.clone(), entry);
         Ok(())
     }
@@ -65,13 +82,13 @@ impl ModuleInfo {
     }
 
     #[cfg(target_os = "windows")]
-    pub fn new_pe(path: &ModulePath) -> Result<Self> {
+    pub fn new_pe(path: &ModulePath, handle: impl Into<Option<HANDLE>>) -> Result<Self> {
         let file = std::fs::File::open(path)?;
         let data = unsafe { memmap2::Mmap::map(&file)? };
 
         let pe = goblin::pe::PE::parse(&data)?;
         let module = ModuleIndex::index_pe(path.clone(), &pe);
-        let offsets = crate::pe::process_module(path, &data, &pe, false)?;
+        let offsets = crate::pe::process_module(path, &data, &pe, false, handle.into())?;
         let blocks = offsets.ones().map(|off| off as u64).collect();
 
         Ok(Self { module, blocks })

src/agent/coverage/src/pe.rs

@@ -233,13 +233,14 @@ pub fn process_module(
     data: &[u8],
     pe: &PE,
     functions_only: bool,
+    target_handle: Option<HANDLE>,
 ) -> Result<FixedBitSet> {
     if let Some(DebugData {
         image_debug_directory: _,
         codeview_pdb70_debug_info: Some(cv),
     }) = pe.debug_data
     {
-        let pdb_path = find_pdb_path(pe_path.as_ref(), &cv)
+        let pdb_path = find_pdb_path(pe_path.as_ref(), &cv, target_handle.into())
             .with_context(|| format!("searching for PDB for PE: {}", pe_path.as_ref().display()))?;
         log::info!("found PDB: {}", pdb_path.display());
 
@@ -302,7 +303,11 @@ fn process_pdb(data: &[u8], pe: &PE, functions_only: bool, pdb_path: &Path) -> R
 // See: https://docs.microsoft.com/en-us/windows/win32/api/dbghelp/nf-dbghelp-syminitializew
 const PSEUDO_HANDLE: HANDLE = -2i64 as _;
 
-fn find_pdb_path(pe_path: &Path, cv: &CodeviewPDB70DebugInfo) -> Result<PathBuf> {
+fn find_pdb_path(
+    pe_path: &Path,
+    cv: &CodeviewPDB70DebugInfo,
+    target_handle: Option<HANDLE>,
+) -> Result<PathBuf> {
     let cv_filename = CStr::from_bytes_with_nul(cv.filename)?.to_str()?;
 
     // This field is named `filename`, but it may be an absolute path.
@@ -314,11 +319,22 @@ fn find_pdb_path(pe_path: &Path, cv: &CodeviewPDB70DebugInfo) -> Result<PathBuf>
         return Ok(cv_filename.to_owned());
     }
 
-    let handle = PSEUDO_HANDLE;
+    // If we have one, use the the process handle for an existing debug
+    let handle = target_handle.unwrap_or(PSEUDO_HANDLE);
 
     let dbghelp = debugger::dbghelp::lock()?;
-    dbghelp.sym_initialize(handle)?;
-    let _cleanup = DbgHelpCleanupGuard::new(&dbghelp, handle);
+
+    // If a target handle was provided, we assume the caller initialized the
+    // dbghelp symbol handler, and will clean up after itself.
+    //
+    // Otherwise, initialize a symbol handler with our own pseudo-path, and use
+    // a drop guard to ensure we don't leak resources.
+    let _cleanup = if target_handle.is_some() {
+        None
+    } else {
+        dbghelp.sym_initialize(handle)?;
+        Some(DbgHelpCleanupGuard::new(&dbghelp, handle))
+    };
 
     // Enable signature and age checking.
     let options = dbghelp.sym_get_options();
@@ -371,10 +387,14 @@ impl<'d> Drop for DbgHelpCleanupGuard<'d> {
     }
 }
 
-pub fn process_image(path: impl AsRef<Path>, functions_only: bool) -> Result<FixedBitSet> {
+pub fn process_image(
+    path: impl AsRef<Path>,
+    functions_only: bool,
+    handle: Option<HANDLE>,
+) -> Result<FixedBitSet> {
     let file = File::open(path.as_ref())?;
     let mmap = unsafe { Mmap::map(&file)? };
     let pe = PE::parse(&mmap)?;
 
-    process_module(path, &mmap, &pe, functions_only)
+    process_module(path, &mmap, &pe, functions_only, handle)
 }