From 60b0225d8e82d412237297710a7f007b006a7773 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Thu, 31 Aug 2023 11:08:01 -0500
Subject: [PATCH] Add PG_MODULE_MAGIC_REENTRANT

Extensions should use this if and only if they support reentrancy. If
the multithreaded GUC is set to true, then we can only load extensions
which support reentrancy.
---
 src/backend/utils/fmgr/dfmgr.c | 14 +++++++++++---
 src/include/fmgr.h             | 23 +++++++++++++++++++++++
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/backend/utils/fmgr/dfmgr.c b/src/backend/utils/fmgr/dfmgr.c
index 6ee9053044917..c2e6c22127067 100644
--- a/src/backend/utils/fmgr/dfmgr.c
+++ b/src/backend/utils/fmgr/dfmgr.c
@@ -85,7 +85,7 @@ static char *substitute_libpath_macro(const char *name);
 static char *find_in_dynamic_libpath(const char *basename);
 
 /* Magic structure that module needs to match to be accepted */
-static static_singleton const Pg_magic_struct magic_data = PG_MODULE_MAGIC_DATA;
+static static_singleton const Pg_magic_struct magic_data = PG_MODULE_MAGIC_REENTRANT_DATA;
 
 
 /*
@@ -255,8 +255,9 @@ internal_load_library(const char *libname)
 		{
 			const Pg_magic_struct *magic_data_ptr = (*magic_func) ();
 
+			/* We will check reentrancy later. */
 			if (magic_data_ptr->len != magic_data.len ||
-				memcmp(magic_data_ptr, &magic_data, magic_data.len) != 0)
+				memcmp(magic_data_ptr, &magic_data, offsetof(Pg_magic_struct, reentrant)) != 0)
 			{
 				/* copy data block before unlinking library */
 				Pg_magic_struct module_magic_data = *magic_data_ptr;
@@ -268,6 +269,13 @@ internal_load_library(const char *libname)
 				/* issue suitable complaint */
 				incompatible_module_error(libname, &module_magic_data);
 			}
+
+			if (IsMultiThreaded && !magic_data_ptr->reentrant)
+				ereport(ERROR,
+						(errmsg("incompatible library \"%s\": not reentrant",
+								libname),
+								errhint("Extension libraries must use the PG_MODULE_MAGIC_REENTRANT macro to be loaded "
+										"when multithreading is turned on.")));
 		}
 		else
 		{
@@ -278,7 +286,7 @@ internal_load_library(const char *libname)
 			ereport(ERROR,
 					(errmsg("incompatible library \"%s\": missing magic block",
 							libname),
-					 errhint("Extension libraries are required to use the PG_MODULE_MAGIC macro.")));
+					 errhint("Extension libraries are required to use the PG_MODULE_MAGIC or PG_MODULE_MAGIC_REENTRANT macros.")));
 		}
 
 		/*
diff --git a/src/include/fmgr.h b/src/include/fmgr.h
index 3ecb98a907d33..b49901eacae10 100644
--- a/src/include/fmgr.h
+++ b/src/include/fmgr.h
@@ -470,6 +470,7 @@ typedef struct
 	int			namedatalen;	/* NAMEDATALEN */
 	int			float8byval;	/* FLOAT8PASSBYVAL */
 	char		abi_extra[32];	/* see pg_config_manual.h */
+	bool		reentrant;		/* Whether the extension supports reentrancy */
 } Pg_magic_struct;
 
 /* The actual data block contents */
@@ -482,6 +483,19 @@ typedef struct
 	NAMEDATALEN, \
 	FLOAT8PASSBYVAL, \
 	FMGR_ABI_EXTRA, \
+	false, \
+}
+
+#define PG_MODULE_MAGIC_REENTRANT_DATA \
+{ \
+	sizeof(Pg_magic_struct), \
+	PG_VERSION_NUM / 100, \
+	FUNC_MAX_ARGS, \
+	INDEX_MAX_KEYS, \
+	NAMEDATALEN, \
+	FLOAT8PASSBYVAL, \
+	FMGR_ABI_EXTRA, \
+	true, \
 }
 
 StaticAssertDecl(sizeof(FMGR_ABI_EXTRA) <= sizeof(((Pg_magic_struct *) 0)->abi_extra),
@@ -506,6 +520,15 @@ PG_MAGIC_FUNCTION_NAME(void) \
 } \
 extern int no_such_variable
 
+#define PG_MODULE_MAGIC_REENTRANT \
+extern PGDLLEXPORT const Pg_magic_struct *PG_MAGIC_FUNCTION_NAME(void); \
+const Pg_magic_struct * \
+PG_MAGIC_FUNCTION_NAME(void) \
+{ \
+	static const Pg_magic_struct Pg_magic_data = PG_MODULE_MAGIC_REENTRANT_DATA; \
+	return &Pg_magic_data; \
+} \
+extern int no_such_variable
 
 /*-------------------------------------------------------------------------
  *		Support routines and macros for callers of fmgr-compatible functions