From 0c0530eeb4ae57ab8d764c8d3cd4a54ff79ec3ba Mon Sep 17 00:00:00 2001
From: Paul Masurel <paul.masurel@datadoghq.com>
Date: Thu, 10 Sep 2026 21:22:01 +0200
Subject: [PATCH] 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 |  7 +----
 jitexpr/src/compile/compiled_fn.rs        | 36 ++++++++++++++++++++---
 jitexpr/src/compile/mod.rs                | 21 +++++++++++++
 3 files changed, 54 insertions(+), 10 deletions(-)

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