JIT: don't leak shim memory when shutting down the interpreter

Author: diegorusso

Include/internal/pycore_jit.h

@@ -25,6 +25,7 @@ typedef _Py_CODEUNIT *(*jit_func)(
 
 int _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction *trace, size_t length);
 void _PyJIT_Free(_PyExecutorObject *executor);
+void _PyJIT_Fini(void);
 
 #endif  // _Py_JIT
 

Python/jit.c

@@ -151,6 +151,8 @@ typedef struct {
     uintptr_t instruction_starts[UOP_MAX_TRACE_LENGTH];
 } jit_state;
 
+static size_t _Py_jit_shim_size = 0;
+
 // Warning! AArch64 requires you to get your hands dirty. These are your gloves:
 
 // value[value_start : value_start + len]
@@ -676,6 +678,7 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
  * We compile this once only as it effectively a normal
  * function, but we need to use the JIT because it needs
  * to understand the jit-specific calling convention.
+ * Don't forget to call _PyJIT_Fini later!
  */
 static _PyJitEntryFuncPtr
 compile_shim(void)
@@ -717,6 +720,7 @@ compile_shim(void)
         jit_free(memory, total_size);
         return NULL;
     }
+    _Py_jit_shim_size = total_size;
     return (_PyJitEntryFuncPtr)memory;
 }
 
@@ -739,6 +743,7 @@ _Py_LazyJitShim(
     return _Py_jit_entry(executor, frame, stack_pointer, tstate);
 }
 
+// Free executor's memory allocated with _PyJIT_Compile
 void
 _PyJIT_Free(_PyExecutorObject *executor)
 {
@@ -754,4 +759,22 @@ _PyJIT_Free(_PyExecutorObject *executor)
     }
 }
 
+// Free shim memory allocated with compile_shim
+void
+_PyJIT_Fini(void)
+{
+    PyMutex_Lock(&lazy_jit_mutex);
+    unsigned char *memory = (unsigned char *)_Py_jit_entry;
+    size_t size = _Py_jit_shim_size;
+    if (size) {
+        _Py_jit_entry = _Py_LazyJitShim;
+        _Py_jit_shim_size = 0;
+        if (jit_free(memory, size)) {
+            PyErr_FormatUnraisable("Exception ignored while "
+                                   "freeing JIT entry code");
+        }
+    }
+    PyMutex_Unlock(&lazy_jit_mutex);
+}
+
 #endif  // _Py_JIT

Python/pylifecycle.c

@@ -35,6 +35,9 @@
 #include "pycore_uniqueid.h"      // _PyObject_FinalizeUniqueIdPool()
 #include "pycore_warnings.h"      // _PyWarnings_InitState()
 #include "pycore_weakref.h"       // _PyWeakref_GET_REF()
+#ifdef _Py_JIT
+#include "pycore_jit.h"           // _PyJIT_Fini()
+#endif
 
 #include "opcode.h"
 
@@ -2267,6 +2270,7 @@ _Py_Finalize(_PyRuntimeState *runtime)
     /* Print debug stats if any */
     _PyEval_Fini();
 
+
     /* Flush sys.stdout and sys.stderr (again, in case more was printed) */
     if (flush_std_files() < 0) {
         status = -1;
@@ -2346,6 +2350,10 @@ _Py_Finalize(_PyRuntimeState *runtime)
 
     finalize_interp_clear(tstate);
 
+#ifdef _Py_JIT
+    /* Free JIT shim memory */
+    _PyJIT_Fini();
+#endif
 
 #ifdef Py_TRACE_REFS
     /* Display addresses (& refcnts) of all objects still alive.