0002-Support-preloading-of-injection-points.patch

From 4dd79183e070069287c403c18a72768acf88f316 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Mon, 20 May 2024 11:48:06 +0900
Subject: [PATCH 2/2] Support preloading of injection points

This can be used to load an injection point in the backend-level cache
before running it, to avoid issues if the point cannot be loaded due to
restrictions in the code path where it would be run, like a critical
section.
---
 src/include/utils/injection_point.h           |  3 +
 src/backend/utils/misc/injection_point.c      | 96 ++++++++++++++-----
 .../expected/injection_points.out             | 32 +++++++
 .../injection_points--1.0.sql                 | 10 ++
 .../injection_points/injection_points.c       | 14 +++
 .../injection_points/sql/injection_points.sql |  7 ++
 6 files changed, 140 insertions(+), 22 deletions(-)

diff --git a/src/include/utils/injection_point.h b/src/include/utils/injection_point.h
index c2c0840706..9ef411cf29 100644
--- a/src/include/utils/injection_point.h
+++ b/src/include/utils/injection_point.h
@@ -15,9 +15,11 @@
  * Injections points require --enable-injection-points.
  */
 #ifdef USE_INJECTION_POINTS
+#define INJECTION_POINT_PRELOAD(name) InjectionPointPreload(name)
 #define INJECTION_POINT(name) InjectionPointRun(name)
 #define INJECTION_POINT_1ARG(name, arg1) InjectionPointRun1Arg(name, arg1)
 #else
+#define INJECTION_POINT_PRELOAD(name) ((void) name)
 #define INJECTION_POINT(name) ((void) name)
 #define INJECTION_POINT_1ARG(name) ((void) name, (void) arg1)
 #endif
@@ -40,6 +42,7 @@ extern void InjectionPointAttach(const char *name,
 								 const void *private_data,
 								 int private_data_size,
 								 int num_args);
+extern void InjectionPointPreload(const char *name);
 extern void InjectionPointRun(const char *name);
 extern void InjectionPointRun1Arg(const char *name, void *arg1);
 extern bool InjectionPointDetach(const char *name);
diff --git a/src/backend/utils/misc/injection_point.c b/src/backend/utils/misc/injection_point.c
index 2bcdb2708c..e4d317a286 100644
--- a/src/backend/utils/misc/injection_point.c
+++ b/src/backend/utils/misc/injection_point.c
@@ -145,6 +145,37 @@ injection_point_cache_remove(const char *name)
 	(void) hash_search(InjectionPointCache, name, HASH_REMOVE, NULL);
 }
 
+/*
+ * injection_point_cache_load
+ *
+ * Load an injection point into the local cache.
+ */
+static void
+injection_point_cache_load(InjectionPointEntry *entry_by_name)
+{
+	char		path[MAXPGPATH];
+	void	   *injection_callback_local;
+
+	snprintf(path, MAXPGPATH, "%s/%s%s", pkglib_path,
+			 entry_by_name->library, DLSUFFIX);
+
+	if (!pg_file_exists(path))
+		elog(ERROR, "could not find library \"%s\" for injection point \"%s\"",
+			 path, entry_by_name->name);
+
+	injection_callback_local = (void *)
+		load_external_function(path, entry_by_name->function, false, NULL);
+
+	if (injection_callback_local == NULL)
+		elog(ERROR, "could not find function \"%s\" in library \"%s\" for injection point \"%s\"",
+			 entry_by_name->function, path, entry_by_name->name);
+
+	/* add it to the local cache when found */
+	injection_point_cache_add(entry_by_name->name, injection_callback_local,
+							  entry_by_name->private_data,
+							  entry_by_name->num_args);
+}
+
 /*
  * injection_point_cache_get
  *
@@ -292,6 +323,47 @@ InjectionPointDetach(const char *name)
 #endif
 }
 
+/*
+ * Preload an injection point into the local cache.
+ *
+ * This is useful to be able to load a function pointer at a stage earlier
+ * than it would be run, especially if the injection point is called in a code
+ * path where memory allocations cannot happen, like critical sections.
+ */
+void
+InjectionPointPreload(const char *name)
+{
+#ifdef USE_INJECTION_POINTS
+	InjectionPointEntry *entry_by_name;
+	bool		found;
+
+	LWLockAcquire(InjectionPointLock, LW_SHARED);
+	entry_by_name = (InjectionPointEntry *)
+		hash_search(InjectionPointHash, name,
+					HASH_FIND, &found);
+	LWLockRelease(InjectionPointLock);
+
+	/*
+	 * If not found, do nothing and remove it from the local cache if it
+	 * existed there.
+	 */
+	if (!found)
+	{
+		injection_point_cache_remove(name);
+		return;
+	}
+
+	/* Check first the local cache, and leave if this entry exists. */
+	if (injection_point_cache_get(name) != NULL)
+		return;
+
+	/* Nothing?  Then load it and leave */
+	injection_point_cache_load(entry_by_name);
+#else
+	elog(ERROR, "Injection points are not supported by this build");
+#endif
+}
+
 /*
  * Workhorse for execution of an injection point, if defined.
  *
@@ -328,28 +400,8 @@ InjectionPointRunInternal(const char *name, int num_args, void *arg1)
 	 */
 	if (injection_point_cache_get(name) == NULL)
 	{
-		char		path[MAXPGPATH];
-		void	   *injection_callback_local;
-
-		/* not found in local cache, so load and register */
-		snprintf(path, MAXPGPATH, "%s/%s%s", pkglib_path,
-				 entry_by_name->library, DLSUFFIX);
-
-		if (!pg_file_exists(path))
-			elog(ERROR, "could not find library \"%s\" for injection point \"%s\"",
-				 path, name);
-
-		injection_callback_local = (void *)
-			load_external_function(path, entry_by_name->function, false, NULL);
-
-		if (injection_callback_local == NULL)
-			elog(ERROR, "could not find function \"%s\" in library \"%s\" for injection point \"%s\"",
-				 entry_by_name->function, path, name);
-
-		/* add it to the local cache when found */
-		injection_point_cache_add(name, injection_callback_local,
-								  entry_by_name->private_data,
-								  entry_by_name->num_args);
+		/* not found in local cache, so load and register it */
+		injection_point_cache_load(entry_by_name);
 	}
 
 	/* Now loaded, so get it. */
diff --git a/src/test/modules/injection_points/expected/injection_points.out b/src/test/modules/injection_points/expected/injection_points.out
index 740bdc8cfd..392dd85234 100644
--- a/src/test/modules/injection_points/expected/injection_points.out
+++ b/src/test/modules/injection_points/expected/injection_points.out
@@ -159,6 +159,38 @@ SELECT injection_points_detach('TestInjectionLogArg1');
  
 (1 row)
 
+-- Preloading
+SELECT injection_points_preload('TestInjectionLogPreload'); -- nothing
+ injection_points_preload 
+--------------------------
+ 
+(1 row)
+
+SELECT injection_points_attach('TestInjectionLogPreload', 'notice');
+ injection_points_attach 
+-------------------------
+ 
+(1 row)
+
+SELECT injection_points_preload('TestInjectionLogPreload'); -- nothing happens
+ injection_points_preload 
+--------------------------
+ 
+(1 row)
+
+SELECT injection_points_run('TestInjectionLogPreload'); -- runs from cache
+NOTICE:  notice triggered for injection point TestInjectionLogPreload
+ injection_points_run 
+----------------------
+ 
+(1 row)
+
+SELECT injection_points_detach('TestInjectionLogPreload');
+ injection_points_detach 
+-------------------------
+ 
+(1 row)
+
 -- Runtime conditions
 SELECT injection_points_attach('TestConditionError', 'error');
  injection_points_attach 
diff --git a/src/test/modules/injection_points/injection_points--1.0.sql b/src/test/modules/injection_points/injection_points--1.0.sql
index 074b7b2ea7..29d43cc556 100644
--- a/src/test/modules/injection_points/injection_points--1.0.sql
+++ b/src/test/modules/injection_points/injection_points--1.0.sql
@@ -14,6 +14,16 @@ RETURNS void
 AS 'MODULE_PATHNAME', 'injection_points_attach'
 LANGUAGE C STRICT PARALLEL UNSAFE;
 
+--
+-- injection_points_preload()
+--
+-- Preload an injection point already attached.
+--
+CREATE FUNCTION injection_points_preload(IN point_name TEXT)
+RETURNS void
+AS 'MODULE_PATHNAME', 'injection_points_preload'
+LANGUAGE C STRICT PARALLEL UNSAFE;
+
 --
 -- injection_points_run()
 --
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index 7ef1c471a0..086b8a6a5e 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -324,6 +324,20 @@ injection_points_attach(PG_FUNCTION_ARGS)
 	PG_RETURN_VOID();
 }
 
+/*
+ * SQL function for preloading an injection point.
+ */
+PG_FUNCTION_INFO_V1(injection_points_preload);
+Datum
+injection_points_preload(PG_FUNCTION_ARGS)
+{
+	char	   *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
+
+	INJECTION_POINT_PRELOAD(name);
+
+	PG_RETURN_VOID();
+}
+
 /*
  * SQL function for triggering an injection point.
  */
diff --git a/src/test/modules/injection_points/sql/injection_points.sql b/src/test/modules/injection_points/sql/injection_points.sql
index 5b612a75c3..69822be647 100644
--- a/src/test/modules/injection_points/sql/injection_points.sql
+++ b/src/test/modules/injection_points/sql/injection_points.sql
@@ -50,6 +50,13 @@ SELECT injection_points_run_1arg('TestInjectionLogArg1', 'foobar'); -- notice
 SELECT injection_points_detach('TestInjectionLog2');
 SELECT injection_points_detach('TestInjectionLogArg1');
 
+-- Preloading
+SELECT injection_points_preload('TestInjectionLogPreload'); -- nothing
+SELECT injection_points_attach('TestInjectionLogPreload', 'notice');
+SELECT injection_points_preload('TestInjectionLogPreload'); -- nothing happens
+SELECT injection_points_run('TestInjectionLogPreload'); -- runs from cache
+SELECT injection_points_detach('TestInjectionLogPreload');
+
 -- Runtime conditions
 SELECT injection_points_attach('TestConditionError', 'error');
 -- Any follow-up injection point attached will be local to this process.
-- 
2.43.0