WIP: resource sanitizer

Cargo.toml

@@ -43,7 +43,7 @@ path = "src/bin/wasmtime.rs"
 doc = false
 
 [dependencies]
-wasmtime = { workspace = true, features = ['std'] }
+wasmtime = { workspace = true, features = ['std', 'resource-sanitizer'] }
 wasmtime-cache = { workspace = true, optional = true }
 wasmtime-cli-flags = { workspace = true }
 wasmtime-cranelift = { workspace = true, optional = true }

crates/wasi/src/p1.rs

@@ -150,7 +150,9 @@ pub struct WasiP1Ctx {
 impl WasiP1Ctx {
     pub(crate) fn new(wasi: WasiCtx) -> Self {
         Self {
-            table: ResourceTable::new(),
+            table: ResourceTable::builder()
+                .sanitizer(wasmtime::component::ResourceSanitizer::new())
+                .build(),
             wasi,
             adapter: WasiP1Adapter::new(),
         }

crates/wasmtime/Cargo.toml

@@ -408,3 +408,9 @@ component-model-async-bytes = [
   "component-model-async",
   "dep:bytes",
 ]
+
+# Enables support for `component::ResourceSanitizer`
+resource-sanitizer = [
+    "component-model",
+    "std",
+]

crates/wasmtime/src/runtime/component/mod.rs

@@ -137,6 +137,9 @@ pub use self::resources::{Resource, ResourceAny};
 pub use self::types::{ResourceType, Type};
 pub use self::values::Val;
 
+#[cfg(feature = "resource-sanitizer")]
+pub use self::resource_table::{ResourceSanitizer, SanInfo};
+
 pub(crate) use self::instance::RuntimeImport;
 pub(crate) use self::resources::HostResourceData;
 pub(crate) use self::store::ComponentInstanceId;

crates/wasmtime/src/runtime/component/resource_table.rs

@@ -5,6 +5,13 @@ use core::any::Any;
 use core::fmt;
 use core::mem;
 
+#[cfg(feature = "resource-sanitizer")]
+mod sanitizer;
+#[cfg(feature = "resource-sanitizer")]
+pub use sanitizer::{ResourceSanitizer, SanInfo};
+#[cfg(feature = "resource-sanitizer")]
+use std::backtrace::Backtrace;
+
 #[derive(Debug)]
 /// Errors returned by operations on `ResourceTable`
 pub enum ResourceTableError {
@@ -36,6 +43,40 @@ impl core::error::Error for ResourceTableError {}
 pub struct ResourceTable {
     entries: Vec<Entry>,
     free_head: Option<usize>,
+    #[cfg(feature = "resource-sanitizer")]
+    sanitizer: Option<ResourceSanitizer>,
+}
+
+/// Builder for `ResourceTable`
+pub struct ResourceTableBuilder {
+    entries: Option<Vec<Entry>>,
+    #[cfg(feature = "resource-sanitizer")]
+    sanitizer: Option<ResourceSanitizer>,
+}
+
+impl ResourceTableBuilder {
+    /// Allocate at least the specified capacity for table entries
+    pub fn capacity(mut self, capacity: usize) -> Self {
+        self.entries = Some(Vec::with_capacity(capacity));
+        self
+    }
+
+    #[cfg(feature = "resource-sanitizer")]
+    /// Use a `ResourceSanitizer` with the ResourceTable
+    pub fn sanitizer(mut self, sanitizer: ResourceSanitizer) -> Self {
+        self.sanitizer = Some(sanitizer);
+        self
+    }
+
+    /// Construct a `ResourceTable`
+    pub fn build(self) -> ResourceTable {
+        ResourceTable {
+            entries: self.entries.unwrap_or_default(),
+            free_head: None,
+            #[cfg(feature = "resource-sanitizer")]
+            sanitizer: self.sanitizer,
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -101,9 +142,15 @@ impl TableEntry {
 impl ResourceTable {
     /// Create an empty table
     pub fn new() -> Self {
-        ResourceTable {
-            entries: Vec::new(),
-            free_head: None,
+        Self::builder().build()
+    }
+
+    /// Create a builder
+    pub fn builder() -> ResourceTableBuilder {
+        ResourceTableBuilder {
+            entries: None,
+            #[cfg(feature = "resource-sanitizer")]
+            sanitizer: None,
         }
     }
 
@@ -118,21 +165,20 @@ impl ResourceTable {
         })
     }
 
-    /// Create an empty table with at least the specified capacity.
-    pub fn with_capacity(capacity: usize) -> Self {
-        ResourceTable {
-            entries: Vec::with_capacity(capacity),
-            free_head: None,
-        }
-    }
-
     /// Inserts a new value `T` into this table, returning a corresponding
     /// `Resource<T>` which can be used to refer to it after it was inserted.
     pub fn push<T>(&mut self, entry: T) -> Result<Resource<T>, ResourceTableError>
     where
         T: Send + 'static,
     {
         let idx = self.push_(TableEntry::new(Box::new(entry), None))?;
+        #[cfg(feature = "resource-sanitizer")]
+        if let Some(san) = &self.sanitizer {
+            san.log_construction(
+                idx,
+                SanInfo::new(std::any::type_name::<T>(), Backtrace::force_capture()),
+            )
+        }
         Ok(Resource::new_own(idx))
     }
 
@@ -245,6 +291,14 @@ impl ResourceTable {
         self.occupied(parent)?;
         let child = self.push_(TableEntry::new(Box::new(entry), Some(parent)))?;
         self.occupied_mut(parent)?.add_child(child);
+
+        #[cfg(feature = "resource-sanitizer")]
+        if let Some(san) = &self.sanitizer {
+            san.log_construction(
+                child,
+                SanInfo::new(std::any::type_name::<T>(), Backtrace::force_capture()),
+            )
+        }
         Ok(Resource::new_own(child))
     }
 
@@ -279,6 +333,10 @@ impl ResourceTable {
     ///
     /// Multiple shared references can be borrowed at any given time.
     pub fn get<T: Any + Sized>(&self, key: &Resource<T>) -> Result<&T, ResourceTableError> {
+        #[cfg(feature = "resource-sanitizer")]
+        if let Some(san) = &self.sanitizer {
+            san.log_usage(key.rep(), Backtrace::force_capture())
+        }
         self.get_(key.rep())?
             .downcast_ref()
             .ok_or(ResourceTableError::WrongType)
@@ -295,13 +353,21 @@ impl ResourceTable {
         &mut self,
         key: &Resource<T>,
     ) -> Result<&mut T, ResourceTableError> {
+        #[cfg(feature = "resource-sanitizer")]
+        if let Some(san) = &self.sanitizer {
+            san.log_usage(key.rep(), Backtrace::force_capture())
+        }
         self.get_any_mut(key.rep())?
             .downcast_mut()
             .ok_or(ResourceTableError::WrongType)
     }
 
     /// Returns the raw `Any` at the `key` index provided.
     pub fn get_any_mut(&mut self, key: u32) -> Result<&mut dyn Any, ResourceTableError> {
+        #[cfg(feature = "resource-sanitizer")]
+        if let Some(san) = &self.sanitizer {
+            san.log_usage(key, Backtrace::force_capture())
+        }
         let r = self.occupied_mut(key)?;
         Ok(&mut *r.entry)
     }
@@ -311,6 +377,10 @@ impl ResourceTable {
     where
         T: Any,
     {
+        #[cfg(feature = "resource-sanitizer")]
+        if let Some(san) = &self.sanitizer {
+            san.log_delete(resource.rep(), Backtrace::force_capture())
+        }
         self.delete_maybe_debug(resource, cfg!(debug_assertions))
     }
 
@@ -353,6 +423,12 @@ impl ResourceTable {
         &'a mut self,
         map: BTreeMap<u32, T>,
     ) -> impl Iterator<Item = (Result<&'a mut dyn Any, ResourceTableError>, T)> {
+        #[cfg(feature = "resource-sanitizer")]
+        if let Some(san) = &self.sanitizer {
+            for (k, _) in map.iter() {
+                san.log_usage(*k, Backtrace::force_capture())
+            }
+        }
         map.into_iter().map(move |(k, v)| {
             let item = self
                 .occupied_mut(k)
@@ -373,18 +449,39 @@ impl ResourceTable {
     {
         let parent_entry = self.occupied(parent.rep())?;
         Ok(parent_entry.children.iter().map(|child_index| {
+            #[cfg(feature = "resource-sanitizer")]
+            if let Some(san) = &self.sanitizer {
+                san.log_usage(*child_index, Backtrace::force_capture())
+            }
+
             let child = self.occupied(*child_index).expect("missing child");
             child.entry.as_ref()
         }))
     }
 
     /// Iterate over all the entries in this table.
     pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut (dyn Any + Send)> {
+        #[cfg(feature = "resource-sanitizer")]
+        if let Some(san) = &self.sanitizer {
+            for (ix, entry) in self.entries.iter().enumerate() {
+                if matches!(entry, Entry::Occupied { .. }) {
+                    san.log_usage(ix.try_into().unwrap(), Backtrace::force_capture())
+                }
+            }
+        }
+
         self.entries.iter_mut().filter_map(|entry| match entry {
             Entry::Occupied { entry } => Some(&mut *entry.entry),
             Entry::Free { .. } => None,
         })
     }
+
+    /// Get the table's `ResourceSanitizer` if one exists. Some iff
+    /// `ResourceTableBuilder::sanitizer` was used at creation.
+    #[cfg(feature = "resource-sanitizer")]
+    pub fn get_sanitizer(&self) -> Option<&ResourceSanitizer> {
+        self.sanitizer.as_ref()
+    }
 }
 
 impl Default for ResourceTable {

crates/wasmtime/src/runtime/component/resource_table/sanitizer.rs

@@ -0,0 +1,77 @@
+use std::backtrace::Backtrace;
+use std::cell::RefCell;
+use std::io::Write;
+use std::vec::Vec;
+use std::write;
+
+/// Tracks allocation, usage, and deletion of all resources
+#[derive(Debug)]
+pub struct ResourceSanitizer {
+    uses: RefCell<Vec<(u32, SanInfo)>>,
+}
+
+impl ResourceSanitizer {
+    /// XXX
+    pub fn new() -> Self {
+        Self {
+            uses: RefCell::new(Vec::new()),
+        }
+    }
+    pub(super) fn log_construction(&self, index: u32, info: SanInfo) {
+        self.uses.borrow_mut().push((index, info));
+    }
+    pub(super) fn log_usage(&self, index: u32, backtrace: Backtrace) {
+        let mut uses = self.uses.borrow_mut();
+        let (_, info) = uses
+            .iter_mut()
+            .rev()
+            .find(|(ix, _)| index == *ix)
+            .expect("used resource present in sanitizer log");
+        info.last_used = Some(backtrace);
+    }
+    pub(super) fn log_delete(&self, index: u32, backtrace: Backtrace) {
+        let mut uses = self.uses.borrow_mut();
+        let (_, info) = uses
+            .iter_mut()
+            .rev()
+            .find(|(ix, _)| index == *ix)
+            .expect("deleted resource present in sanitizer log");
+        info.deleted = Some(backtrace);
+    }
+
+    /// XXX
+    pub fn report_live_set(&self, w: &mut impl Write) -> Result<(), std::io::Error> {
+        let uses = self.uses.borrow();
+        for (ix, info) in uses.iter() {
+            if info.deleted.is_none() {
+                write!(
+                    w,
+                    "LEAK resource {ix}: {}\nLEAK allocated at {:#?}\nLEAK last used at {:#?}\n",
+                    info.type_name, info.allocated, info.last_used
+                )?;
+            }
+        }
+        Ok(())
+    }
+}
+
+/// sanitizer information for a given resource
+#[derive(Debug)]
+pub struct SanInfo {
+    type_name: &'static str,
+    allocated: std::backtrace::Backtrace,
+    last_used: Option<std::backtrace::Backtrace>,
+    deleted: Option<std::backtrace::Backtrace>,
+}
+
+impl SanInfo {
+    /// XXX
+    pub fn new(type_name: &'static str, allocated: Backtrace) -> Self {
+        SanInfo {
+            type_name,
+            allocated,
+            last_used: None,
+            deleted: None,
+        }
+    }
+}

src/commands/run.rs

@@ -231,6 +231,9 @@ impl RunCommand {
             .await
         });
 
+        if let Some(san) = WasiView::ctx(store.data_mut()).table.get_sanitizer() {
+            san.report_live_set(&mut std::io::stderr())?;
+        }
         // Load the main wasm module.
         match result.unwrap_or_else(|elapsed| {
             Err(anyhow::Error::from(wasmtime::Trap::Interrupt))

src/commands/serve.rs

@@ -190,7 +190,9 @@ impl ServeCommand {
         builder.stderr(LogStream::new(stderr_prefix, Output::Stderr));
 
         let mut host = Host {
-            table: wasmtime::component::ResourceTable::new(),
+            table: wasmtime::component::ResourceTable::builder()
+                .sanitizer(wasmtime::component::ResourceSanitizer::new())
+                .build(),
             ctx: builder.build(),
             http: WasiHttpCtx::new(),
             http_outgoing_body_buffer_chunks: self.run.common.wasi.http_outgoing_body_buffer_chunks,