Introduce custom executable id (#46)

Before this commit the only unique identifier for an executable file or
object file was the build id. As the GNU's build ID might not be present,
we fall back to the sha256 hash of the `.text` section. This works fine but
takes quite a bit of space as we refer to it in every mapping.

By using a 64 bit id derived from the first 8 bytes of the sha256 hash
of the code, we can save some memory, make comparisons and hash the IDs faster,
and helps simplify the code. For example, there's no need for a separate ID in
the native unwind state. There might be a bit more CPU and IO needed to hash
the code, but this might be something we can cache in the future.

This is also useful for future profile formats that require unique 64 bit
IDs for object files.

Also fixes other minor things, such as using the lower hexadecimal representation.

Test Plan
=========

Lots of manual tests, everything seems good.

src/bpf/profiler.bpf.c

@@ -52,7 +52,7 @@ struct {
 struct {
   __uint(type, BPF_MAP_TYPE_HASH);
   __uint(max_entries, 5 * 1000);
-  __type(key, u32);
+  __type(key, u64);
   __type(value, unwind_info_chunks_t);
 } unwind_info_chunks SEC(".maps");
 
@@ -116,9 +116,9 @@ static __always_inline u64 find_offset_for_pc(stack_unwind_table_t *table, u64 p
 // address.
 static __always_inline chunk_info_t*
 find_chunk(mapping_t *mapping, u64 object_relative_pc) {
-  u32 executable_id = mapping->executable_id;
+  u64 executable_id = mapping->executable_id;
 
-  LOG("~about to check chunks, executable_id=%d", executable_id);
+  LOG("~about to check chunks, executable_id=%lld", executable_id);
   // Find the chunk where this unwind table lives.
   // Each chunk maps to exactly one shard.
   unwind_info_chunks_t *chunks =

src/bpf/profiler.h

@@ -125,11 +125,11 @@ typedef struct {
 
 // Represents an executable mapping.
 typedef struct {
-  u32 executable_id;
-  u32 type;
+  u64 executable_id;
   u64 load_address;
   u64 begin;
   u64 end;
+  u32 type;
 } mapping_t;
 
 // Key for the longest prefix matching. This is defined

src/collector.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
 use tracing::{debug, span, Level};
 
-use crate::object::BuildId;
+use crate::object::ExecutableId;
 use crate::profile::symbolize_profile;
 use crate::profiler::ObjectFileInfo;
 use crate::profiler::ProcessInfo;
@@ -12,7 +12,7 @@ use crate::profiler::SymbolizedAggregatedProfile;
 pub struct Collector {
     profiles: Vec<RawAggregatedProfile>,
     procs: HashMap<i32, ProcessInfo>,
-    objs: HashMap<BuildId, ObjectFileInfo>,
+    objs: HashMap<ExecutableId, ObjectFileInfo>,
 }
 
 type ThreadSafeCollector = Arc<Mutex<Collector>>;
@@ -30,7 +30,7 @@ impl Collector {
         &mut self,
         profile: RawAggregatedProfile,
         procs: &HashMap<i32, ProcessInfo>,
-        objs: &HashMap<BuildId, ObjectFileInfo>,
+        objs: &HashMap<ExecutableId, ObjectFileInfo>,
     ) {
         self.profiles.push(profile);
 
@@ -40,7 +40,7 @@ impl Collector {
 
         for (k, v) in objs {
             self.objs.insert(
-                k.clone(),
+                *k,
                 ObjectFileInfo {
                     file: std::fs::File::open(v.path.clone()).unwrap(),
                     path: v.path.clone(),

src/object.rs

@@ -3,7 +3,7 @@ use std::io::Read;
 use std::path::PathBuf;
 
 use anyhow::{anyhow, Result};
-use data_encoding::HEXUPPER;
+use data_encoding::HEXLOWER;
 use memmap2;
 use ring::digest::{Context, Digest, SHA256};
 
@@ -16,6 +16,8 @@ use object::Object;
 use object::ObjectKind;
 use object::ObjectSection;
 
+pub type ExecutableId = u64;
+
 #[derive(Hash, Eq, PartialEq, Clone, Debug)]
 pub enum BuildId {
     Gnu(String),
@@ -32,6 +34,7 @@ pub struct ElfLoad {
 pub struct ObjectFile<'a> {
     leaked_mmap_ptr: *const memmap2::Mmap,
     object: object::File<'a>,
+    code_hash: Digest,
 }
 
 impl Drop for ObjectFile<'_> {
@@ -49,13 +52,26 @@ impl ObjectFile<'_> {
         let mmap = Box::new(mmap);
         let leaked = Box::leak(mmap);
         let object = object::File::parse(&**leaked)?;
+        let Some(code_hash) = code_hash(&object) else {
+            return Err(anyhow!("code hash is None"));
+        };
 
         Ok(ObjectFile {
             leaked_mmap_ptr: leaked as *const memmap2::Mmap,
             object,
+            code_hash,
         })
     }
 
+    /// Returns an identifier for the executable using the first 8 bytes of the Sha256 of the code section.
+    pub fn id(&self) -> Result<ExecutableId> {
+        let mut buffer = [0; 8];
+        let _ = self.code_hash.as_ref().read(&mut buffer)?;
+        Ok(u64::from_ne_bytes(buffer))
+    }
+
+    /// Returns the executable build ID if present. If no GNU build ID and no Go build ID
+    /// are found it returns the hash of the text section.
     pub fn build_id(&self) -> anyhow::Result<BuildId> {
         let object = &self.object;
         let build_id = object.build_id()?;
@@ -85,18 +101,8 @@ impl ObjectFile<'_> {
             }
         }
 
-        // No build id (rust, some other libraries).
-        for section in object.sections() {
-            if section.name().unwrap() == ".text" {
-                if let Ok(section) = section.data() {
-                    return Ok(BuildId::Sha256(
-                        HEXUPPER.encode(sha256_digest(section).as_ref()),
-                    ));
-                }
-            }
-        }
-
-        unreachable!("A build id should always be returned");
+        // No build id (Rust, some compilers and Linux distributions).
+        return Ok(BuildId::Sha256(HEXLOWER.encode(self.code_hash.as_ref())));
     }
 
     pub fn is_dynamic(&self) -> bool {
@@ -160,6 +166,18 @@ impl ObjectFile<'_> {
     }
 }
 
+pub fn code_hash(object: &object::File) -> Option<Digest> {
+    for section in object.sections() {
+        if section.name().unwrap() == ".text" {
+            if let Ok(section) = section.data() {
+                return Some(sha256_digest(section));
+            }
+        }
+    }
+
+    None
+}
+
 fn sha256_digest<R: Read>(mut reader: R) -> Digest {
     let mut context = Context::new(&SHA256);
     let mut buffer = [0; 1024];

src/profile.rs

@@ -1,10 +1,10 @@
 use std::collections::HashMap;
 use std::path::PathBuf;
-use tracing::{debug, span, Level};
+use tracing::{debug, error, span, Level};
 
 use crate::bpf::profiler_bindings::native_stack_t;
 use crate::ksym::KsymIter;
-use crate::object::BuildId;
+use crate::object::ExecutableId;
 use crate::profiler::ExecutableMapping;
 use crate::profiler::ObjectFileInfo;
 use crate::profiler::ProcessInfo;
@@ -16,7 +16,7 @@ use crate::usym::symbolize_native_stack_blaze;
 pub fn symbolize_profile(
     profile: &RawAggregatedProfile,
     procs: &HashMap<i32, ProcessInfo>,
-    objs: &HashMap<BuildId, ObjectFileInfo>,
+    objs: &HashMap<ExecutableId, ObjectFileInfo>,
 ) -> SymbolizedAggregatedProfile {
     let _span = span!(Level::DEBUG, "symbolize_profile").entered();
     let mut r = SymbolizedAggregatedProfile::new();
@@ -61,8 +61,6 @@ pub fn symbolize_profile(
                 symbolized_sample.kstack.push(le_symbol.symbol_name);
             }
         };
-        // debug!("--- symbolized sample: {}", symbolized_sample);
-
         r.push(symbolized_sample);
     }
 
@@ -82,7 +80,7 @@ fn find_mapping(mappings: &[ExecutableMapping], addr: u64) -> Option<ExecutableM
 fn fetch_symbols_for_profile(
     profile: &RawAggregatedProfile,
     procs: &HashMap<i32, ProcessInfo>,
-    objs: &HashMap<BuildId, ObjectFileInfo>,
+    objs: &HashMap<ExecutableId, ObjectFileInfo>,
 ) -> HashMap<PathBuf, HashMap<u64, Vec<String>>> {
     let mut addresses_per_sample: HashMap<PathBuf, HashMap<u64, Vec<String>>> = HashMap::new();
 
@@ -104,28 +102,21 @@ fn fetch_symbols_for_profile(
                 continue; //return Err(anyhow!("could not find mapping"));
             };
 
-            match &mapping.build_id {
-                Some(build_id) => {
-                    match objs.get(build_id) {
-                        Some(obj) => {
-                            // We need the normalized address for normal object files
-                            // and might need the absolute addresses for JITs
-                            let normalized_addr = addr - mapping.start_addr + mapping.offset
-                                - obj.load_offset
-                                + obj.load_vaddr;
-
-                            let key = obj.path.clone();
-                            let addrs = addresses_per_sample.entry(key).or_default();
-                            addrs.insert(normalized_addr, vec!["<default>".to_string()]);
-                            // <- default value is a bit janky
-                        }
-                        None => {
-                            println!("\t\t - [no build id found]");
-                        }
-                    }
+            match objs.get(&mapping.executable_id) {
+                Some(obj) => {
+                    // We need the normalized address for normal object files
+                    // and might need the absolute addresses for JITs
+                    let normalized_addr = addr - mapping.start_addr + mapping.offset
+                        - obj.load_offset
+                        + obj.load_vaddr;
+
+                    let key = obj.path.clone();
+                    let addrs = addresses_per_sample.entry(key).or_default();
+                    addrs.insert(normalized_addr, vec!["<default>".to_string()]);
+                    // <- default value is a bit janky
                 }
                 None => {
-                    println!("\t\t - mapping is not backed by a file, could be a JIT segment");
+                    error!("executable with id {} not found", mapping.executable_id);
                 }
             }
         }
@@ -146,7 +137,7 @@ fn fetch_symbols_for_profile(
 fn symbolize_native_stack(
     addresses_per_sample: &HashMap<PathBuf, HashMap<u64, Vec<String>>>,
     procs: &HashMap<i32, ProcessInfo>,
-    objs: &HashMap<BuildId, ObjectFileInfo>,
+    objs: &HashMap<ExecutableId, ObjectFileInfo>,
     pid: i32,
     native_stack: &native_stack_t,
 ) -> Vec<String> {
@@ -168,39 +159,31 @@ fn symbolize_native_stack(
         };
 
         // finally
-        match &mapping.build_id {
-            Some(build_id) => match objs.get(build_id) {
-                Some(obj) => {
-                    let failed_to_fetch_symbol =
-                        vec!["<failed to fetch symbol for addr>".to_string()];
-                    let failed_to_symbolize = vec!["<failed to symbolize>".to_string()];
-
-                    let normalized_addr = addr - mapping.start_addr + mapping.offset
-                        - obj.load_offset
-                        + obj.load_vaddr;
-
-                    let func_names = match addresses_per_sample.get(&obj.path) {
-                        Some(value) => match value.get(&normalized_addr) {
-                            Some(v) => v,
-                            None => &failed_to_fetch_symbol,
-                        },
-                        None => &failed_to_symbolize,
-                    };
+        match objs.get(&mapping.executable_id) {
+            Some(obj) => {
+                let failed_to_fetch_symbol = vec!["<failed to fetch symbol for addr>".to_string()];
+                let failed_to_symbolize = vec!["<failed to symbolize>".to_string()];
+
+                let normalized_addr =
+                    addr - mapping.start_addr + mapping.offset - obj.load_offset + obj.load_vaddr;
+
+                let func_names = match addresses_per_sample.get(&obj.path) {
+                    Some(value) => match value.get(&normalized_addr) {
+                        Some(v) => v,
+                        None => &failed_to_fetch_symbol,
+                    },
+                    None => &failed_to_symbolize,
+                };
 
-                    for func_name in func_names {
-                        r.push(func_name.clone());
-                    }
+                for func_name in func_names {
+                    r.push(func_name.clone());
                 }
-                None => {
-                    debug!("\t\t - [no build id found]");
-                }
-            },
+            }
             None => {
-                debug!("\t\t - mapping is not backed by a file, could be a JIT segment");
+                debug!("build id not found");
             }
         }
     }
 
     r
-    // Ok(())
 }