Free JITModule executable memory when CompiledFn is dropped (#3090)

JITModule never releases its executable allocation on drop; the only
way to reclaim it is the consuming, unsafe JITModule::free_memory.
CompiledFn previously stored the module in a plain field, so every
compiled expression's executable memory leaked for the remainder of
the process once the CompiledFn was dropped.

Wrap the module in ManuallyDrop and add an explicit Drop impl for
CompiledFn that calls free_memory. This is safe because CompiledFn is
only ever constructed behind an Arc and never exposes its entry point
or other raw pointers into the module outside &self-bounded calls, so
drop only runs once no call into the module can be in flight or happen
afterward.

Verified empirically: 200k compile+drop cycles peaked at ~3.3 GB RSS
before this fix and ~6 MB after.

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<String, CompileError> {

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<TypedVariable>,
+    entry: JitEntry,
+    // Always `Some` outside of `Drop::drop`, which takes the module to free its
+    // executable memory.
+    module: Option<JITModule>,
+    inputs: Vec<TypedVariable>,
     // This AST owns the Arc-backed literals and regexes embedded in generated code.
-    pub(crate) _typed_expr: Box<TypedExpr>,
+    _typed_expr: Box<TypedExpr>,
 }
 
 // `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<TypedVariable>,
+        typed_expr: Box<TypedExpr>,
+    ) -> 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

jitexpr/src/compile/mod.rs

@@ -258,6 +258,27 @@ mod tests {
         assert_send_and_sync::<CompiledFn>();
     }
 
+    #[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