llvmjit: Use explicit LLVMContextRef for inlining

When performing inlining LLVM unfortunately "leaks" types (the
types survive and are usable, but a new round of inlining will
recreate new structurally equivalent types). This accumulation
will over time amount to a memory leak which for some queries
can be large enough to trigger the OOM process killer.

To avoid accumulation of types, all IR related data is stored
in an LLVMContextRef which is dropped and recreated in order
to release all types.  Dropping and recreating incurs overhead,
so it will be done only after 100 queries. This is a heuristic
which might be revisited, but until we can get the size of the
context from LLVM we are flying a bit blind.

This issue has been reported several times, there may be more
references to it in the archives on top of the threads linked
below.

This is a backpatch of 9dce220 to all supported branches.

Reported-By: Justin Pryzby <[email protected]>
Reported-By: Kurt Roeckx <[email protected]>
Reported-By: Jaime Casanova <[email protected]>
Reported-By: Lauri Laanmets <[email protected]>
Author: Andres Freund and Daniel Gustafsson
Discussion: https://postgr.es/m/[email protected]
Discussion: https://postgr.es/m/[email protected]
Discussion: https://postgr.es/m/CAPH-tTxLf44s3CvUUtQpkDr1D8Hxqc2NGDzGXS1ODsfiJ6WSqA@mail.gmail.com
Backpatch-through: v12

Author: danielgustafsson

src/backend/jit/llvm/llvmjit.c

@@ -47,6 +47,8 @@
 #include "utils/memutils.h"
 #include "utils/resowner_private.h"
 
+#define LLVMJIT_LLVM_CONTEXT_REUSE_MAX 100
+
 /* Handle of a module emitted via ORC JIT */
 typedef struct LLVMJitHandle
 {
@@ -100,8 +102,15 @@ LLVMModuleRef llvm_types_module = NULL;
 
 static bool llvm_session_initialized = false;
 static size_t llvm_generation = 0;
+
+/* number of LLVMJitContexts that currently are in use */
+static size_t llvm_jit_context_in_use_count = 0;
+
+/* how many times has the current LLVMContextRef been used */
+static size_t llvm_llvm_context_reuse_count = 0;
 static const char *llvm_triple = NULL;
 static const char *llvm_layout = NULL;
+static LLVMContextRef llvm_context;
 
 
 static LLVMTargetRef llvm_targetref;
@@ -122,6 +131,8 @@ static void llvm_compile_module(LLVMJitContext *context);
 static void llvm_optimize_module(LLVMJitContext *context, LLVMModuleRef module);
 
 static void llvm_create_types(void);
+static void llvm_set_target(void);
+static void llvm_recreate_llvm_context(void);
 static uint64_t llvm_resolve_symbol(const char *name, void *ctx);
 
 #if LLVM_VERSION_MAJOR > 11
@@ -143,6 +154,63 @@ _PG_jit_provider_init(JitProviderCallbacks *cb)
 	cb->compile_expr = llvm_compile_expr;
 }
 
+
+/*
+ * Every now and then create a new LLVMContextRef. Unfortunately, during every
+ * round of inlining, types may "leak" (they can still be found/used via the
+ * context, but new types will be created the next time in inlining is
+ * performed). To prevent that from slowly accumulating problematic amounts of
+ * memory, recreate the LLVMContextRef we use. We don't want to do so too
+ * often, as that implies some overhead (particularly re-loading the module
+ * summaries / modules is fairly expensive). A future TODO would be to make
+ * this more finegrained and only drop/recreate the LLVMContextRef when we know
+ * there has been inlining. If we can get the size of the context from LLVM
+ * then that might be a better way to determine when to drop/recreate rather
+ * then the usagecount heuristic currently employed.
+ */
+static void
+llvm_recreate_llvm_context(void)
+{
+	if (!llvm_context)
+		elog(ERROR, "Trying to recreate a non-existing context");
+
+	/*
+	 * We can only safely recreate the LLVM context if no other code is being
+	 * JITed, otherwise we'd release the types in use for that.
+	 */
+	if (llvm_jit_context_in_use_count > 0)
+	{
+		llvm_llvm_context_reuse_count++;
+		return;
+	}
+
+	if (llvm_llvm_context_reuse_count <= LLVMJIT_LLVM_CONTEXT_REUSE_MAX)
+	{
+		llvm_llvm_context_reuse_count++;
+		return;
+	}
+
+	/*
+	 * Need to reset the modules that the inlining code caches before
+	 * disposing of the context. LLVM modules exist within a specific LLVM
+	 * context, therefore disposing of the context before resetting the cache
+	 * would lead to dangling pointers to modules.
+	 */
+	llvm_inline_reset_caches();
+
+	LLVMContextDispose(llvm_context);
+	llvm_context = LLVMContextCreate();
+	llvm_llvm_context_reuse_count = 0;
+
+	/*
+	 * Re-build cached type information, so code generation code can rely on
+	 * that information to be present (also prevents the variables to be
+	 * dangling references).
+	 */
+	llvm_create_types();
+}
+
+
 /*
  * Create a context for JITing work.
  *
@@ -159,6 +227,8 @@ llvm_create_context(int jitFlags)
 
 	llvm_session_initialize();
 
+	llvm_recreate_llvm_context();
+
 	ResourceOwnerEnlargeJIT(CurrentResourceOwner);
 
 	context = MemoryContextAllocZero(TopMemoryContext,
@@ -169,6 +239,8 @@ llvm_create_context(int jitFlags)
 	context->base.resowner = CurrentResourceOwner;
 	ResourceOwnerRememberJIT(CurrentResourceOwner, PointerGetDatum(context));
 
+	llvm_jit_context_in_use_count++;
+
 	return context;
 }
 
@@ -178,9 +250,15 @@ llvm_create_context(int jitFlags)
 static void
 llvm_release_context(JitContext *context)
 {
-	LLVMJitContext *llvm_context = (LLVMJitContext *) context;
+	LLVMJitContext *llvm_jit_context = (LLVMJitContext *) context;
 	ListCell   *lc;
 
+	/*
+	 * Consider as cleaned up even if we skip doing so below, that way we can
+	 * verify the tracking is correct (see llvm_shutdown()).
+	 */
+	llvm_jit_context_in_use_count--;
+
 	/*
 	 * When this backend is exiting, don't clean up LLVM. As an error might
 	 * have occurred from within LLVM, we do not want to risk reentering. All
@@ -191,13 +269,13 @@ llvm_release_context(JitContext *context)
 
 	llvm_enter_fatal_on_oom();
 
-	if (llvm_context->module)
+	if (llvm_jit_context->module)
 	{
-		LLVMDisposeModule(llvm_context->module);
-		llvm_context->module = NULL;
+		LLVMDisposeModule(llvm_jit_context->module);
+		llvm_jit_context->module = NULL;
 	}
 
-	foreach(lc, llvm_context->handles)
+	foreach(lc, llvm_jit_context->handles)
 	{
 		LLVMJitHandle *jit_handle = (LLVMJitHandle *) lfirst(lc);
 
@@ -227,8 +305,8 @@ llvm_release_context(JitContext *context)
 
 		pfree(jit_handle);
 	}
-	list_free(llvm_context->handles);
-	llvm_context->handles = NIL;
+	list_free(llvm_jit_context->handles);
+	llvm_jit_context->handles = NIL;
 
 	llvm_leave_fatal_on_oom();
 }
@@ -248,7 +326,7 @@ llvm_mutable_module(LLVMJitContext *context)
 	{
 		context->compiled = false;
 		context->module_generation = llvm_generation++;
-		context->module = LLVMModuleCreateWithName("pg");
+		context->module = LLVMModuleCreateWithNameInContext("pg", llvm_context);
 		LLVMSetTarget(context->module, llvm_triple);
 		LLVMSetDataLayout(context->module, llvm_layout);
 	}
@@ -832,6 +910,14 @@ llvm_session_initialize(void)
 	LLVMInitializeNativeAsmPrinter();
 	LLVMInitializeNativeAsmParser();
 
+	if (llvm_context == NULL)
+	{
+		llvm_context = LLVMContextCreate();
+
+		llvm_jit_context_in_use_count = 0;
+		llvm_llvm_context_reuse_count = 0;
+	}
+
 	/*
 	 * When targeting LLVM 15, turn off opaque pointers for the context we
 	 * build our code in.  We don't need to do so for other contexts (e.g.
@@ -851,6 +937,11 @@ llvm_session_initialize(void)
 	 */
 	llvm_create_types();
 
+	/*
+	 * Extract target information from loaded module.
+	 */
+	llvm_set_target();
+
 	if (LLVMGetTargetFromTriple(llvm_triple, &llvm_targetref, &error) != 0)
 	{
 		elog(FATAL, "failed to query triple %s", error);
@@ -946,6 +1037,10 @@ llvm_shutdown(int code, Datum arg)
 		return;
 	}
 
+	if (llvm_jit_context_in_use_count != 0)
+		elog(PANIC, "LLVMJitContext in use count not 0 at exit (is %zu)",
+			 llvm_jit_context_in_use_count);
+
 #if LLVM_VERSION_MAJOR > 11
 	{
 		if (llvm_opt3_orc)
@@ -1008,6 +1103,23 @@ load_return_type(LLVMModuleRef mod, const char *name)
 	return typ;
 }
 
+/*
+ * Load triple & layout from clang emitted file so we're guaranteed to be
+ * compatible.
+ */
+static void
+llvm_set_target(void)
+{
+	if (!llvm_types_module)
+		elog(ERROR, "failed to extract target information, llvmjit_types.c not loaded");
+
+	if (llvm_triple == NULL)
+		llvm_triple = pstrdup(LLVMGetTarget(llvm_types_module));
+
+	if (llvm_layout == NULL)
+		llvm_layout = pstrdup(LLVMGetDataLayoutStr(llvm_types_module));
+}
+
 /*
  * Load required information, types, function signatures from llvmjit_types.c
  * and make them available in global variables.
@@ -1031,19 +1143,12 @@ llvm_create_types(void)
 	}
 
 	/* eagerly load contents, going to need it all */
-	if (LLVMParseBitcode2(buf, &llvm_types_module))
+	if (LLVMParseBitcodeInContext2(llvm_context, buf, &llvm_types_module))
 	{
-		elog(ERROR, "LLVMParseBitcode2 of %s failed", path);
+		elog(ERROR, "LLVMParseBitcodeInContext2 of %s failed", path);
 	}
 	LLVMDisposeMemoryBuffer(buf);
 
-	/*
-	 * Load triple & layout from clang emitted file so we're guaranteed to be
-	 * compatible.
-	 */
-	llvm_triple = pstrdup(LLVMGetTarget(llvm_types_module));
-	llvm_layout = pstrdup(LLVMGetDataLayoutStr(llvm_types_module));
-
 	TypeSizeT = llvm_pg_var_type("TypeSizeT");
 	TypeParamBool = load_return_type(llvm_types_module, "FunctionReturningBool");
 	TypeStorageBool = llvm_pg_var_type("TypeStorageBool");