Add cache for coverage locations and debuginfo

src/agent/coverage/examples/record.rs

@@ -6,7 +6,7 @@ use anyhow::{bail, Result};
 use clap::Parser;
 use cobertura::CoberturaCoverage;
 use coverage::allowlist::{AllowList, TargetAllowList};
-use coverage::binary::BinaryCoverage;
+use coverage::binary::{BinaryCoverage, DebugInfoCache};
 use coverage::record::{CoverageRecorder, Recorded};
 use debuggable_module::loader::Loader;
 
@@ -66,6 +66,7 @@ fn main() -> Result<()> {
 
     let mut coverage = BinaryCoverage::default();
     let loader = Arc::new(Loader::new());
+    let cache = Arc::new(DebugInfoCache::new(allowlist.source_files.clone()));
 
     if let Some(dir) = args.input_dir {
         check_for_input_marker(&args.command)?;
@@ -77,6 +78,7 @@ fn main() -> Result<()> {
             let recorded = CoverageRecorder::new(cmd)
                 .allowlist(allowlist.clone())
                 .loader(loader.clone())
+                .debuginfo_cache(cache.clone())
                 .timeout(timeout)
                 .record()?;
 
@@ -91,6 +93,7 @@ fn main() -> Result<()> {
         let recorded = CoverageRecorder::new(cmd)
             .allowlist(allowlist.clone())
             .loader(loader)
+            .debuginfo_cache(cache)
             .timeout(timeout)
             .record()?;
 

src/agent/coverage/src/binary.rs

@@ -1,15 +1,17 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::BTreeMap;
+use std::sync::{Arc, Mutex};
 
-use anyhow::Result;
+use anyhow::{bail, Result};
+use debuggable_module::block::Blocks;
 use debuggable_module::Module;
 pub use debuggable_module::{block, path::FilePath, Offset};
 use symbolic::debuginfo::Object;
 use symbolic::symcache::{SymCache, SymCacheConverter};
 
-use crate::allowlist::TargetAllowList;
+use crate::allowlist::{AllowList, TargetAllowList};
 
 #[derive(Clone, Debug, Default, Eq, PartialEq)]
 pub struct BinaryCoverage {
@@ -102,6 +104,114 @@ impl std::ops::AddAssign for Count {
     }
 }
 
+/// Cache of analyzed binary metadata for a set of modules, relative to a common
+/// source allowlist.
+pub struct DebugInfoCache {
+    allowlist: Arc<AllowList>,
+    modules: Arc<Mutex<BTreeMap<FilePath, CachedDebugInfo>>>,
+}
+
+impl DebugInfoCache {
+    pub fn new(allowlist: AllowList) -> Self {
+        let allowlist = Arc::new(allowlist);
+        let modules = Arc::new(Mutex::new(BTreeMap::new()));
+
+        Self { allowlist, modules }
+    }
+
+    pub fn get_or_insert(&self, module: &dyn Module) -> Result<CachedDebugInfo> {
+        if !self.is_cached(module) {
+            self.insert(module)?;
+        }
+
+        if let Some(cached) = self.get(module.executable_path()) {
+            Ok(cached)
+        } else {
+            // Unreachable.
+            bail!("module should be cached but data is missing")
+        }
+    }
+
+    fn get(&self, path: &FilePath) -> Option<CachedDebugInfo> {
+        self.modules.lock().unwrap().get(path).cloned()
+    }
+
+    fn insert(&self, module: &dyn Module) -> Result<()> {
+        let debuginfo = module.debuginfo()?;
+
+        let mut symcache = vec![];
+        let mut converter = SymCacheConverter::new();
+        let exe = Object::parse(module.executable_data())?;
+        converter.process_object(&exe)?;
+        let di = Object::parse(module.debuginfo_data())?;
+        converter.process_object(&di)?;
+        converter.serialize(&mut std::io::Cursor::new(&mut symcache))?;
+        let symcache = SymCache::parse(&symcache)?;
+
+        let mut blocks = Blocks::new();
+
+        for function in debuginfo.functions() {
+            if let Some(location) = symcache.lookup(function.offset.0).next() {
+                if let Some(file) = location.file() {
+                    if !self.allowlist.is_allowed(file.full_path()) {
+                        debug!(
+                            "skipping sweep of `{}:{}` due to excluded source path `{}`",
+                            module.executable_path(),
+                            function.name,
+                            file.full_path(),
+                        );
+                        continue;
+                    }
+                }
+            }
+
+            let fn_blocks =
+                block::sweep_region(module, &debuginfo, function.offset, function.size)?;
+
+            for block in &fn_blocks {
+                if let Some(location) = symcache.lookup(block.offset.0).next() {
+                    if let Some(file) = location.file() {
+                        let path = file.full_path();
+
+                        // Apply allowlists per block, to account for inlining. The `location` values
+                        // here describe the top of the inline-inclusive call stack.
+                        if !self.allowlist.is_allowed(path) {
+                            continue;
+                        }
+
+                        blocks.map.insert(block.offset, *block);
+                    }
+                }
+            }
+        }
+
+        let coverage = ModuleBinaryCoverage::from((&blocks).into_iter().map(|b| b.offset));
+        let cached = CachedDebugInfo::new(blocks, coverage);
+        self.modules
+            .lock()
+            .unwrap()
+            .insert(module.executable_path().clone(), cached);
+
+        Ok(())
+    }
+
+    fn is_cached(&self, module: &dyn Module) -> bool {
+        self.get(module.executable_path()).is_some()
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct CachedDebugInfo {
+    pub blocks: Blocks,
+    pub coverage: ModuleBinaryCoverage,
+}
+
+impl CachedDebugInfo {
+    pub fn new(blocks: Blocks, coverage: ModuleBinaryCoverage) -> Self {
+        Self { blocks, coverage }
+    }
+}
+
 pub fn find_coverage_sites(
     module: &dyn Module,
     allowlist: &TargetAllowList,
@@ -117,7 +227,7 @@ pub fn find_coverage_sites(
     converter.serialize(&mut std::io::Cursor::new(&mut symcache))?;
     let symcache = SymCache::parse(&symcache)?;
 
-    let mut offsets = BTreeSet::new();
+    let mut blocks = Blocks::new();
 
     for function in debuginfo.functions() {
         if let Some(location) = symcache.lookup(function.offset.0).next() {
@@ -134,9 +244,9 @@ pub fn find_coverage_sites(
             }
         }
 
-        let blocks = block::sweep_region(module, &debuginfo, function.offset, function.size)?;
+        let fn_blocks = block::sweep_region(module, &debuginfo, function.offset, function.size)?;
 
-        for block in &blocks {
+        for block in &fn_blocks {
             if let Some(location) = symcache.lookup(block.offset.0).next() {
                 if let Some(file) = location.file() {
                     let path = file.full_path();
@@ -147,13 +257,13 @@ pub fn find_coverage_sites(
                         continue;
                     }
 
-                    offsets.insert(block.offset);
+                    blocks.map.insert(block.offset, *block);
                 }
             }
         }
     }
 
-    let coverage = ModuleBinaryCoverage::from(offsets.into_iter());
+    let coverage = ModuleBinaryCoverage::from((&blocks).into_iter().map(|b| b.offset));
 
     Ok(coverage)
 }

src/agent/coverage/src/record.rs

@@ -9,7 +9,7 @@ use anyhow::Result;
 use debuggable_module::loader::Loader;
 
 use crate::allowlist::TargetAllowList;
-use crate::binary::BinaryCoverage;
+use crate::binary::{BinaryCoverage, DebugInfoCache};
 
 #[cfg(target_os = "linux")]
 pub mod linux;
@@ -19,6 +19,7 @@ pub mod windows;
 
 pub struct CoverageRecorder {
     allowlist: TargetAllowList,
+    cache: Arc<DebugInfoCache>,
     cmd: Command,
     loader: Arc<Loader>,
     timeout: Duration,
@@ -30,11 +31,13 @@ impl CoverageRecorder {
         cmd.stderr(Stdio::piped());
 
         let allowlist = TargetAllowList::default();
+        let cache = Arc::new(DebugInfoCache::new(allowlist.source_files.clone()));
         let loader = Arc::new(Loader::new());
         let timeout = Duration::from_secs(5);
 
         Self {
             allowlist,
+            cache,
             cmd,
             loader,
             timeout,
@@ -51,6 +54,11 @@ impl CoverageRecorder {
         self
     }
 
+    pub fn debuginfo_cache(mut self, cache: impl Into<Arc<DebugInfoCache>>) -> Self {
+        self.cache = cache.into();
+        self
+    }
+
     pub fn timeout(mut self, timeout: Duration) -> Self {
         self.timeout = timeout;
         self
@@ -74,7 +82,7 @@ impl CoverageRecorder {
             let child_pid = child_pid.clone();
 
             timer::timed(self.timeout, move || {
-                let mut recorder = LinuxRecorder::new(&loader, self.allowlist);
+                let mut recorder = LinuxRecorder::new(&loader, self.allowlist, &self.cache);
                 let mut dbg = Debugger::new(&mut recorder);
                 let child = dbg.spawn(self.cmd)?;
 
@@ -120,7 +128,7 @@ impl CoverageRecorder {
         let loader = self.loader.clone();
 
         crate::timer::timed(self.timeout, move || {
-            let mut recorder = WindowsRecorder::new(&loader, self.allowlist);
+            let mut recorder = WindowsRecorder::new(&loader, self.allowlist, &self.cache);
             let (mut dbg, child) = Debugger::init(self.cmd, &mut recorder)?;
             dbg.run(&mut recorder)?;
 

src/agent/coverage/src/record/linux.rs

@@ -15,22 +15,28 @@ pub mod debugger;
 use debugger::{DebugEventHandler, DebuggerContext, ModuleImage};
 
 use crate::allowlist::TargetAllowList;
-use crate::binary::{self, BinaryCoverage};
+use crate::binary::{BinaryCoverage, DebugInfoCache};
 
-pub struct LinuxRecorder<'data> {
+pub struct LinuxRecorder<'cache, 'data> {
     allowlist: TargetAllowList,
+    cache: &'cache DebugInfoCache,
     pub coverage: BinaryCoverage,
     loader: &'data Loader,
     modules: BTreeMap<FilePath, LinuxModule<'data>>,
 }
 
-impl<'data> LinuxRecorder<'data> {
-    pub fn new(loader: &'data Loader, allowlist: TargetAllowList) -> Self {
+impl<'cache, 'data> LinuxRecorder<'cache, 'data> {
+    pub fn new(
+        loader: &'data Loader,
+        allowlist: TargetAllowList,
+        cache: &'cache DebugInfoCache,
+    ) -> Self {
         let coverage = BinaryCoverage::default();
         let modules = BTreeMap::new();
 
         Self {
             allowlist,
+            cache,
             coverage,
             loader,
             modules,
@@ -86,7 +92,7 @@ impl<'data> LinuxRecorder<'data> {
             return Ok(());
         };
 
-        let coverage = binary::find_coverage_sites(&module, &self.allowlist)?;
+        let coverage = self.cache.get_or_insert(&module)?.coverage;
 
         for offset in coverage.as_ref().keys().copied() {
             let addr = image.base().offset_by(offset)?;
@@ -101,7 +107,7 @@ impl<'data> LinuxRecorder<'data> {
     }
 }
 
-impl<'data> DebugEventHandler for LinuxRecorder<'data> {
+impl<'cache, 'data> DebugEventHandler for LinuxRecorder<'cache, 'data> {
     fn on_breakpoint(&mut self, context: &mut DebuggerContext, tracee: &mut Tracee) -> Result<()> {
         self.do_on_breakpoint(context, tracee)
     }

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

@@ -14,7 +14,7 @@ use debuggable_module::{Module, Offset};
 use debugger::{BreakpointId, BreakpointType, DebugEventHandler, Debugger, ModuleLoadInfo};
 
 use crate::allowlist::TargetAllowList;
-use crate::binary::{self, BinaryCoverage};
+use crate::binary::{BinaryCoverage, DebugInfoCache};
 
 // For a new module image, we defer setting coverage breakpoints until exit from one of these
 // functions (when present). This avoids breaking hotpatching routines in the ASan interceptor
@@ -25,18 +25,23 @@ use crate::binary::{self, BinaryCoverage};
 const PROCESS_IMAGE_DEFERRAL_TRIGGER: &str = "__asan::AsanInitInternal(";
 const LIBRARY_IMAGE_DEFERRAL_TRIGGER: &str = "DllMain(";
 
-pub struct WindowsRecorder<'data> {
+pub struct WindowsRecorder<'cache, 'data> {
     allowlist: TargetAllowList,
     breakpoints: Breakpoints,
+    cache: &'cache DebugInfoCache,
     deferred_breakpoints: BTreeMap<BreakpointId, (Breakpoint, DeferralState)>,
     pub coverage: BinaryCoverage,
     loader: &'data Loader,
     modules: BTreeMap<FilePath, (WindowsModule<'data>, DebugInfo)>,
     pub stop_error: Option<Error>,
 }
 
-impl<'data> WindowsRecorder<'data> {
-    pub fn new(loader: &'data Loader, allowlist: TargetAllowList) -> Self {
+impl<'cache, 'data> WindowsRecorder<'cache, 'data> {
+    pub fn new(
+        loader: &'data Loader,
+        allowlist: TargetAllowList,
+        cache: &'cache DebugInfoCache,
+    ) -> Self {
         let breakpoints = Breakpoints::default();
         let deferred_breakpoints = BTreeMap::new();
         let coverage = BinaryCoverage::default();
@@ -46,6 +51,7 @@ impl<'data> WindowsRecorder<'data> {
         Self {
             allowlist,
             breakpoints,
+            cache,
             deferred_breakpoints,
             coverage,
             loader,
@@ -231,7 +237,7 @@ impl<'data> WindowsRecorder<'data> {
 
     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)?;
+        let coverage = self.cache.get_or_insert(module)?.coverage;
 
         for offset in coverage.as_ref().keys().copied() {
             let breakpoint = Breakpoint::new(path.clone(), offset);
@@ -288,7 +294,7 @@ impl Breakpoint {
     }
 }
 
-impl<'data> DebugEventHandler for WindowsRecorder<'data> {
+impl<'cache, 'data> DebugEventHandler for WindowsRecorder<'cache, 'data> {
     fn on_create_process(&mut self, dbg: &mut Debugger, module: &ModuleLoadInfo) {
         if let Err(err) = self.try_on_create_process(dbg, module) {
             warn!("{err}");