diff --git a/jitexpr/src/compile/compile_fn_builder.rs b/jitexpr/src/compile/compile_fn_builder.rs index f8638aa8c..1df012ee5 100644 --- a/jitexpr/src/compile/compile_fn_builder.rs +++ b/jitexpr/src/compile/compile_fn_builder.rs @@ -316,12 +316,7 @@ impl LoweredFunction { // was built above to exactly match `JitEntry`. The module is retained by // `CompiledFn`, so its executable allocation outlives `entry`. let entry = unsafe { mem::transmute::<*const u8, JitEntry>(code) }; - Ok(CompiledFn { - entry, - _module: module, - inputs: input_vars, - _typed_expr: expression, - }) + Ok(CompiledFn::new(entry, module, input_vars, expression)) } fn into_assembly(mut self) -> Result { diff --git a/jitexpr/src/compile/compiled_fn.rs b/jitexpr/src/compile/compiled_fn.rs index 01e2d9369..8d7fe3818 100644 --- a/jitexpr/src/compile/compiled_fn.rs +++ b/jitexpr/src/compile/compiled_fn.rs @@ -30,11 +30,13 @@ pub(crate) type JitEntry = /// This object owns the JIT module containing its executable memory and every /// resource referenced by the generated code. pub struct CompiledFn { - pub(crate) entry: JitEntry, - pub(crate) _module: JITModule, - pub(super) inputs: Vec, + entry: JitEntry, + // Always `Some` outside of `Drop::drop`, which takes the module to free its + // executable memory. + module: Option, + inputs: Vec, // This AST owns the Arc-backed literals and regexes embedded in generated code. - pub(crate) _typed_expr: Box, + _typed_expr: Box, } // `JITModule` is not `Sync` because it supports lazily looking up symbols through @@ -44,7 +46,33 @@ pub struct CompiledFn { // called concurrently when each caller supplies a distinct `StringArena`. unsafe impl Sync for CompiledFn {} +impl Drop for CompiledFn { + fn drop(&mut self) { + let Some(module) = self.module.take() else { + return; + }; + // SAFETY: `self` is being dropped, so by the invariants documented on + // `CompiledFn`, no function from `module` can be executing and none + // will be called after this point. + unsafe { module.free_memory() }; + } +} + impl CompiledFn { + pub(super) fn new( + entry: JitEntry, + module: JITModule, + inputs: Vec, + typed_expr: Box, + ) -> Self { + CompiledFn { + entry, + module: Some(module), + inputs, + _typed_expr: typed_expr, + } + } + /// Returns the input slots in the exact order expected by [`CompiledFn::call`]. pub fn inputs(&self) -> &[TypedVariable] { &self.inputs diff --git a/jitexpr/src/compile/mod.rs b/jitexpr/src/compile/mod.rs index f96d23105..7a2051e4f 100644 --- a/jitexpr/src/compile/mod.rs +++ b/jitexpr/src/compile/mod.rs @@ -258,6 +258,27 @@ mod tests { assert_send_and_sync::(); } + #[test] + fn test_repeated_compile_and_drop_frees_jit_memory() { + // `CompiledFn::drop` must call `JITModule::free_memory` rather than + // silently leaking the module's executable allocation, since + // `JITModule` itself does not release it on drop. This test does not + // measure process memory directly, but it does drive many compile and + // drop cycles so a use-after-free from freeing the module too early + // (e.g. while `entry` is still reachable) would reliably crash or + // produce wrong results here. + let untyped_expr = UntypedExpr::variable("flag"); + let variable_types = HashMap::from([("flag", VarType::Bool)]); + for _ in 0..10_000 { + let compiled_fn = compile(&untyped_expr, &variable_types).unwrap(); + let mut string_arena = StringArena::new(); + let input = [VariableValue::some(true)]; + let output = unsafe { compiled_fn.call(&input, &mut string_arena) }; + assert_eq!(unsafe { output.as_bool() }, Some(true)); + drop(compiled_fn); + } + } + #[test] fn test_contexts_have_independent_string_arenas() { let untyped_expr = Function::Lower