T0b1/wasmtime
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix a memory leak with command modules (#2017)

Browse Source 
This commit fixes a memory leak that can happen with `Linker::module`
when the provided module is a command. This function creates a closure
but the closure closed over a strong reference to `Store` (and
transitively through any imports provided). Unfortunately a `Store`
keeps everything alive, including `Func`, so this meant that `Store` was
inserted into a cycle which caused the leak.

The cycle here is manually broken by closing over the raw value of each
external value rather than the external value itself (which has a
strong reference to `Store`).
This commit is contained in:
Alex Crichton
2021-02-01 13:04:27 -06:00
committed by GitHub
parent 23b8c6be79
commit cb7b1aadcc
4 changed files with 120 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
crates/wasmtime/src/externals.rs
View File

@@ -160,6 +160,17 @@ impl Extern {
Extern::Module(_) => "module",
}
}
pub(crate) fn wasmtime_export(&self) -> wasmtime_runtime::Export {
match self {
Extern::Func(f) => f.wasmtime_export().clone().into(),
Extern::Global(f) => f.wasmtime_export().clone().into(),
Extern::Table(f) => f.wasmtime_export().clone().into(),
Extern::Memory(f) => f.wasmtime_export().clone().into(),
Extern::Instance(f) => wasmtime_runtime::Export::Instance(f.wasmtime_export().clone()),
Extern::Module(f) => wasmtime_runtime::Export::Module(Box::new(f.clone())),
}
}
}
impl From<Func> for Extern {

4
crates/wasmtime/src/instance.rs
View File

@@ -121,6 +121,10 @@ impl Instance {
ty
}
pub(crate) fn wasmtime_export(&self) -> &RuntimeInstance {
&self.items
}
/// Returns the associated [`Store`] that this `Instance` is compiled into.
///
/// This is the [`Store`] that generally serves as a sort of global cache

65
crates/wasmtime/src/linker.rs
View File

@@ -402,34 +402,53 @@ impl Linker {
fn command(&mut self, module_name: &str, module: &Module) -> Result<&mut Self> {
for export in module.exports() {
if let Some(func_ty) = export.ty().func() {
let imports = self.compute_imports(module)?;
let store = self.store.clone();
let imports = self
.compute_imports(module)?
.into_iter()
.map(|e| e.wasmtime_export())
.collect::<Vec<_>>();
let module = module.clone();
let export_name = export.name().to_owned();
let func = Func::new(&self.store, func_ty.clone(), move |_, params, results| {
// Create a new instance for this command execution.
let instance = Instance::new(&store, &module, &imports)?;
let func = Func::new(
&self.store,
func_ty.clone(),
move |caller, params, results| {
let store = caller.store();
// `unwrap()` everything here because we know the instance contains a
// function export with the given name and signature because we're
// iterating over the module it was instantiated from.
let command_results = instance
.get_export(&export_name)
.unwrap()
.into_func()
.unwrap()
.call(params)
.map_err(|error| error.downcast::<Trap>().unwrap())?;
// Note that the unsafety here is due to the validity of
// `i` and the validity of `i` within `store`. For our
// case though these items all come from `imports` above
// so they're all valid. They're also all kept alive by
// the store itself used here so this should be safe.
let imports = imports
.iter()
.map(|i| unsafe { Extern::from_wasmtime_export(&i, &store) })
.collect::<Vec<_>>();
// Copy the return values into the output slice.
for (result, command_result) in
results.iter_mut().zip(command_results.into_vec())
{
*result = command_result;
}
// Create a new instance for this command execution.
let instance = Instance::new(&store, &module, &imports)?;
Ok(())
});
// `unwrap()` everything here because we know the instance contains a
// function export with the given name and signature because we're
// iterating over the module it was instantiated from.
let command_results = instance
.get_export(&export_name)
.unwrap()
.into_func()
.unwrap()
.call(params)
.map_err(|error| error.downcast::<Trap>().unwrap())?;
// Copy the return values into the output slice.
for (result, command_result) in
results.iter_mut().zip(command_results.into_vec())
{
*result = command_result;
}
Ok(())
},
);
self.insert(module_name, export.name(), func.into())?;
} else if export.name() == "memory" && export.ty().memory().is_some() {
// Allow an exported "memory" memory for now.

63
tests/all/linker.rs
View File

@@ -1,4 +1,6 @@
use anyhow::Result;
use std::cell::Cell;
use std::rc::Rc;
use wasmtime::*;
#[test]
@@ -160,3 +162,64 @@ fn module_interposition() -> Result<()> {
assert_eq!(func()?, 112);
Ok(())
}
#[test]
fn no_leak() -> Result<()> {
struct DropMe(Rc<Cell<bool>>);
impl Drop for DropMe {
fn drop(&mut self) {
self.0.set(true);
}
}
let flag = Rc::new(Cell::new(false));
{
let store = Store::default();
let mut linker = Linker::new(&store);
let drop_me = DropMe(flag.clone());
linker.func("", "", move || drop(&drop_me))?;
let module = Module::new(
store.engine(),
r#"
(module
(func (export "_start"))
)
"#,
)?;
linker.module("a", &module)?;
}
assert!(flag.get(), "store was leaked");
Ok(())
}
#[test]
fn no_leak_with_imports() -> Result<()> {
struct DropMe(Rc<Cell<bool>>);
impl Drop for DropMe {
fn drop(&mut self) {
self.0.set(true);
}
}
let flag = Rc::new(Cell::new(false));
{
let store = Store::default();
let mut linker = Linker::new(&store);
let drop_me = DropMe(flag.clone());
linker.func("", "", move || drop(&drop_me))?;
let module = Module::new(
store.engine(),
r#"
(module
(import "" "" (func))
(func (export "_start"))
)
"#,
)?;
linker.module("a", &module)?;
}
assert!(flag.get(), "store was leaked");
Ok(())
}