From e7e56f0f7b36ea973fa472aadd16dbfe6db7a6ec Mon Sep 17 00:00:00 2001
From: Daniel Gustafsson <daniel@yesql.se>
Date: Tue, 2 May 2023 14:58:44 +0200
Subject: [PATCH v4 5/5] Address review comments

---
 src/backend/jit/llvm/llvmjit.c | 155 +++++++++++++++++++--------------
 1 file changed, 89 insertions(+), 66 deletions(-)

diff --git a/src/backend/jit/llvm/llvmjit.c b/src/backend/jit/llvm/llvmjit.c
index f0c70d68ab..8fe68a64ba 100644
--- a/src/backend/jit/llvm/llvmjit.c
+++ b/src/backend/jit/llvm/llvmjit.c
@@ -112,6 +112,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
@@ -133,6 +135,61 @@ _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).
+ */
+static void
+llvm_recreate_llvm_context(void)
+{
+	Assert(llvm_context != NULL);
+
+	if (llvm_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();
+
+	/*
+	 * 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.
  *
@@ -149,54 +206,7 @@ llvm_create_context(int jitFlags)
 
 	llvm_session_initialize();
 
-	/*
-	 * 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).
-	 *
-	 * 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.
-	 *
-	 * FIXME: Extract into helper function.
-	 */
-	if (llvm_jit_context_in_use_count == 0 &&
-		llvm_llvm_context_reuse_count > LLVMJIT_LLVM_CONTEXT_REUSE_MAX)
-	{
-		llvm_llvm_context_reuse_count = 0;
-
-		Assert(llvm_context != NULL);
-
-		/*
-		 * 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();
-
-		if (llvm_context != NULL)
-			LLVMContextDispose(llvm_context);
-		llvm_context = LLVMContextCreate();
-
-		/*
-		 * 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).
-		 *
-		 * FIXME: should split the handling of llvm_triple / llvm_layout out
-		 * of llvm_create_types() - that doesn't need to be redone.
-		 */
-		llvm_create_types();
-	}
-	else
-	{
-		llvm_llvm_context_reuse_count++;
-	}
+	llvm_recreate_llvm_context();
 
 	ResourceOwnerEnlargeJIT(CurrentResourceOwner);
 
@@ -219,7 +229,7 @@ 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;
 
 	/*
@@ -238,13 +248,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);
 
@@ -274,8 +284,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;
 }
 
 /*
@@ -853,7 +863,8 @@ llvm_session_initialize(void)
 	LLVMInitializeNativeAsmPrinter();
 	LLVMInitializeNativeAsmParser();
 
-	llvm_context = LLVMContextCreate();
+	if (llvm_context == NULL)
+		llvm_context = LLVMContextCreate();
 
 	/*
 	 * When targeting an LLVM version with opaque pointers enabled by
@@ -871,6 +882,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);
@@ -1040,6 +1056,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.
@@ -1069,16 +1102,6 @@ llvm_create_types(void)
 	}
 	LLVMDisposeMemoryBuffer(buf);
 
-	/*
-	 * Load triple & layout from clang emitted file so we're guaranteed to be
-	 * compatible.
-	 */
-	if (llvm_triple == NULL)
-	{
-		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");
-- 
2.32.1 (Apple Git-133)

