Defer coverage breakpoints

src/agent/coverage/src/record/windows.rs

@@ -4,12 +4,13 @@
 use std::collections::BTreeMap;
 use std::path::Path;
 
-use anyhow::{anyhow, Error, Result};
+use anyhow::{anyhow, bail, Error, Result};
+use debuggable_module::debuginfo::{DebugInfo, Function};
 use debuggable_module::load_module::LoadModule;
 use debuggable_module::loader::Loader;
 use debuggable_module::path::FilePath;
 use debuggable_module::windows::WindowsModule;
-use debuggable_module::Offset;
+use debuggable_module::{Module, Offset};
 use debugger::{BreakpointId, BreakpointType, DebugEventHandler, Debugger, ModuleLoadInfo};
 
 use crate::allowlist::TargetAllowList;
@@ -18,22 +19,25 @@ use crate::binary::{self, BinaryCoverage};
 pub struct WindowsRecorder<'data> {
     allowlist: TargetAllowList,
     breakpoints: Breakpoints,
+    deferred_breakpoints: BTreeMap<BreakpointId, (Breakpoint, DeferralState)>,
     pub coverage: BinaryCoverage,
     loader: &'data Loader,
-    modules: BTreeMap<FilePath, WindowsModule<'data>>,
+    modules: BTreeMap<FilePath, (WindowsModule<'data>, DebugInfo)>,
     pub stop_error: Option<Error>,
 }
 
 impl<'data> WindowsRecorder<'data> {
     pub fn new(loader: &'data Loader, allowlist: TargetAllowList) -> Self {
         let breakpoints = Breakpoints::default();
+        let deferred_breakpoints = BTreeMap::new();
         let coverage = BinaryCoverage::default();
         let modules = BTreeMap::new();
         let stop_error = None;
 
         Self {
             allowlist,
             breakpoints,
+            deferred_breakpoints,
             coverage,
             loader,
             modules,
@@ -66,11 +70,47 @@ impl<'data> WindowsRecorder<'data> {
         self.insert_module(dbg, module)
     }
 
-    fn try_on_breakpoint(&mut self, _dbg: &mut Debugger, id: BreakpointId) -> Result<()> {
-        let breakpoint = self
-            .breakpoints
-            .remove(id)
-            .ok_or_else(|| anyhow!("stopped on dangling breakpoint"))?;
+    fn try_on_breakpoint(&mut self, dbg: &mut Debugger, id: BreakpointId) -> Result<()> {
+        if let Some((trigger, state)) = self.deferred_breakpoints.remove(&id) {
+            match state {
+                DeferralState::NotEntered => {
+                    // Find the return address.
+                    let frame = dbg.get_current_frame()?;
+                    let ret = frame.return_address();
+                    let id = dbg.new_address_breakpoint(ret, BreakpointType::OneTime)?;
+
+                    // Update the state for this deferral to set module coverage breakpoints on ret.
+                    let thread_id = dbg.get_current_thread_id();
+                    let state = DeferralState::PendingReturn { thread_id };
+                    self.deferred_breakpoints.insert(id, (trigger, state));
+                }
+                DeferralState::PendingReturn { thread_id } => {
+                    if dbg.get_current_thread_id() == thread_id {
+                        // We've returned from the trigger function, and on the same thread.
+                        //
+                        // It's safe to set coverage breakpoints.
+                        self.set_module_breakpoints(dbg, trigger.module)?;
+                    } else {
+                        // Hit a ret breakpoint, but on the wrong thread. Reset it so the correct
+                        // thread has a chance to see it.
+                        //
+                        // We only defer breakpoints in image initialization code, so we don't
+                        // expect to reach this code in practice.
+                        let id = trigger.set(dbg)?;
+                        self.deferred_breakpoints.insert(id, (trigger, state));
+                    }
+                }
+            }
+
+            return Ok(());
+        }
+
+        let breakpoint = self.breakpoints.remove(id);
+
+        let Some(breakpoint) = breakpoint else {
+            let stack = dbg.get_current_stack()?;
+            bail!("stopped on dangling breakpoint, debuggee stack:\n{}", stack);
+        };
 
         let coverage = self
             .coverage
@@ -103,72 +143,121 @@ impl<'data> WindowsRecorder<'data> {
             return Ok(());
         };
 
-        let coverage = binary::find_coverage_sites(&module, &self.allowlist)?;
+        let debuginfo = module.debuginfo()?;
+        self.modules.insert(path.clone(), (module, debuginfo));
+
+        self.set_or_defer_module_breakpoints(dbg, path)?;
+
+        Ok(())
+    }
+
+    fn set_or_defer_module_breakpoints(
+        &mut self,
+        dbg: &mut Debugger,
+        path: FilePath,
+    ) -> Result<()> {
+        let (_module, debuginfo) = &self.modules[&path];
+
+        // For borrowck.
+        let mut trigger = None;
+
+        for function in debuginfo.functions() {
+            // Called on process startup.
+            if function.name.contains("__asan::AsanInitInternal") {
+                trigger = Some(function.clone());
+                break;
+            }
+
+            // Called on shared library load.
+            if function.name.contains("DllMain(") {
+                trigger = Some(function.clone());
+                break;
+            }
+        }
+
+        if let Some(trigger) = trigger {
+            debug!("deferring coverage breakpoints for module {}", path);
+            self.defer_module_breakpoints(dbg, path, trigger)
+        } else {
+            debug!(
+                "immediately setting coverage breakpoints for module {}",
+                path
+            );
+            self.set_module_breakpoints(dbg, path)
+        }
+    }
+
+    fn defer_module_breakpoints(
+        &mut self,
+        dbg: &mut Debugger,
+        path: FilePath,
+        trigger: Function,
+    ) -> Result<()> {
+        let state = DeferralState::NotEntered;
+        let entry_breakpoint = Breakpoint::new(path, trigger.offset);
+        let id = entry_breakpoint.set(dbg)?;
+
+        self.deferred_breakpoints
+            .insert(id, (entry_breakpoint, state));
+
+        Ok(())
+    }
+
+    fn set_module_breakpoints(&mut self, dbg: &mut Debugger, path: FilePath) -> Result<()> {
+        let (module, _) = &self.modules[&path];
+        let coverage = binary::find_coverage_sites(module, &self.allowlist)?;
 
         for offset in coverage.as_ref().keys().copied() {
             let breakpoint = Breakpoint::new(path.clone(), offset);
             self.breakpoints.set(dbg, breakpoint)?;
         }
 
-        self.coverage.modules.insert(path.clone(), coverage);
+        let count = coverage.offsets.len();
+        debug!("set {} breakpoints for module {}", count, path);
 
-        self.modules.insert(path, module);
+        self.coverage.modules.insert(path, coverage);
 
         Ok(())
     }
 }
 
 #[derive(Debug, Default)]
 struct Breakpoints {
-    id_to_offset: BTreeMap<BreakpointId, Offset>,
-    offset_to_breakpoint: BTreeMap<Offset, Breakpoint>,
+    id_to_breakpoint: BTreeMap<BreakpointId, Breakpoint>,
 }
 
 impl Breakpoints {
     pub fn set(&mut self, dbg: &mut Debugger, breakpoint: Breakpoint) -> Result<()> {
-        if self.is_set(&breakpoint) {
-            return Ok(());
-        }
-
-        self.write(dbg, breakpoint)
-    }
-
-    // Unguarded action that ovewrites both the target process address space and our index
-    // of known breakpoints. Callers must use `set()`, which avoids redundant breakpoint
-    // setting.
-    fn write(&mut self, dbg: &mut Debugger, breakpoint: Breakpoint) -> Result<()> {
-        // The `debugger` crates tracks loaded modules by base name. If a path or file
-        // name is used, software breakpoints will not be written.
-        let name = Path::new(breakpoint.module.base_name());
-        let id = dbg.new_rva_breakpoint(name, breakpoint.offset.0, BreakpointType::OneTime)?;
-
-        self.id_to_offset.insert(id, breakpoint.offset);
-        self.offset_to_breakpoint
-            .insert(breakpoint.offset, breakpoint);
-
+        let id = breakpoint.set(dbg)?;
+        self.id_to_breakpoint.insert(id, breakpoint);
         Ok(())
     }
 
-    pub fn is_set(&self, breakpoint: &Breakpoint) -> bool {
-        self.offset_to_breakpoint.contains_key(&breakpoint.offset)
-    }
-
+    /// Remove a registered breakpoint from the index.
+    ///
+    /// This method should be called when a breakpoint is hit to retrieve its associated data.
+    /// It does NOT clear the actual breakpoint in the debugger engine.
     pub fn remove(&mut self, id: BreakpointId) -> Option<Breakpoint> {
-        let offset = self.id_to_offset.remove(&id)?;
-        self.offset_to_breakpoint.remove(&offset)
+        self.id_to_breakpoint.remove(&id)
     }
 }
 
 #[derive(Clone, Debug)]
 struct Breakpoint {
-    module: FilePath,
-    offset: Offset,
+    pub module: FilePath,
+    pub offset: Offset,
 }
 
 impl Breakpoint {
     pub fn new(module: FilePath, offset: Offset) -> Self {
         Self { module, offset }
     }
+
+    pub fn set(&self, dbg: &mut Debugger) -> Result<BreakpointId> {
+        let name = Path::new(self.module.base_name());
+        let id = dbg.new_rva_breakpoint(name, self.offset.0, BreakpointType::OneTime)?;
+        Ok(id)
+    }
 }
 
 impl<'data> DebugEventHandler for WindowsRecorder<'data> {
@@ -193,3 +282,8 @@ impl<'data> DebugEventHandler for WindowsRecorder<'data> {
         }
     }
 }
+
+enum DeferralState {
+    NotEntered,
+    PendingReturn { thread_id: u64 },
+}

src/agent/debuggable-module/src/debuginfo.rs

@@ -40,6 +40,7 @@ impl DebugInfo {
     }
 }
 
+#[derive(Clone, Debug)]
 pub struct Function {
     pub name: String,
     pub offset: Offset,