Support loading of injection points

This can be used to load an injection point and prewarm 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 where no memory allocation can happen
(load_external_function() can do allocations when expanding a library
name).

Tests can use a macro called INJECTION_POINT_LOAD() to load an injection
point.  The test module injection_points gains some tests, and a SQL
function able to load an injection point.

Based on a request from Andrey Borodin, who has implemented a test for
multixacts requiring this facility.

Reviewed-by: Andrey Borodin
Discussion: https://postgr.es/m/ZkrBE1e2q2wGvsoN@paquier.xyz

7 files changed, 168 insertions, 36 deletions

diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index f3a3e4e2f8f..756a9d07fb0 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3619,6 +3619,20 @@ INJECTION_POINT(name);
     </para>
 
     <para>
+     An injection point with a given <literal>name</literal> can be loaded
+     using macro:
+<programlisting>
+INJECTION_POINT_LOAD(name);
+</programlisting>
+
+     This will load the injection point callback into the process cache,
+     doing all memory allocations at this stage without running the callback.
+     This is useful when an injection point is attached in a critical section
+     where no memory can be allocated: load the injection point outside the
+     critical section, then run it in the critical section.
+    </para>
+
+    <para>
      Add-ins can attach callbacks to an already-declared injection point by
      calling:
 <programlisting>
diff --git a/src/backend/utils/misc/injection_point.c b/src/backend/utils/misc/injection_point.c
index afae0dbedf4..48f29e9b60a 100644
--- a/src/backend/utils/misc/injection_point.c
+++ b/src/backend/utils/misc/injection_point.c
@@ -130,19 +130,46 @@ injection_point_cache_remove(const char *name)
 }
 
 /*
+ * 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);
+}
+
+/*
  * injection_point_cache_get
  *
  * Retrieve an injection point from the local cache, if any.
  */
-static InjectionPointCallback
-injection_point_cache_get(const char *name, const void **private_data)
+static InjectionPointCacheEntry *
+injection_point_cache_get(const char *name)
 {
 	bool		found;
 	InjectionPointCacheEntry *entry;
 
-	if (private_data)
-		*private_data = NULL;
-
 	/* no callback if no cache yet */
 	if (InjectionPointCache == NULL)
 		return NULL;
@@ -151,11 +178,7 @@ injection_point_cache_get(const char *name, const void **private_data)
 		hash_search(InjectionPointCache, name, HASH_FIND, &found);
 
 	if (found)
-	{
-		if (private_data)
-			*private_data = entry->private_data;
-		return entry->callback;
-	}
+		return entry;
 
 	return NULL;
 }
@@ -279,6 +302,52 @@ InjectionPointDetach(const char *name)
 }
 
 /*
+ * Load an injection point into the local cache.
+ *
+ * This is useful to be able to load an injection point before running it,
+ * especially if the injection point is called in a code path where memory
+ * allocations cannot happen, like critical sections.
+ */
+void
+InjectionPointLoad(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);
+
+	/*
+	 * If not found, do nothing and remove it from the local cache if it
+	 * existed there.
+	 */
+	if (!found)
+	{
+		injection_point_cache_remove(name);
+		LWLockRelease(InjectionPointLock);
+		return;
+	}
+
+	/* Check first the local cache, and leave if this entry exists. */
+	if (injection_point_cache_get(name) != NULL)
+	{
+		LWLockRelease(InjectionPointLock);
+		return;
+	}
+
+	/* Nothing?  Then load it and leave */
+	injection_point_cache_load(entry_by_name);
+
+	LWLockRelease(InjectionPointLock);
+#else
+	elog(ERROR, "Injection points are not supported by this build");
+#endif
+}
+
+/*
  * Execute an injection point, if defined.
  *
  * Check first the shared hash table, and adapt the local cache depending
@@ -290,8 +359,7 @@ InjectionPointRun(const char *name)
 #ifdef USE_INJECTION_POINTS
 	InjectionPointEntry *entry_by_name;
 	bool		found;
-	InjectionPointCallback injection_callback;
-	const void *private_data;
+	InjectionPointCacheEntry *cache_entry;
 
 	LWLockAcquire(InjectionPointLock, LW_SHARED);
 	entry_by_name = (InjectionPointEntry *)
@@ -313,37 +381,18 @@ InjectionPointRun(const char *name)
 	 * Check if the callback exists in the local cache, to avoid unnecessary
 	 * external loads.
 	 */
-	if (injection_point_cache_get(name, NULL) == NULL)
+	if (injection_point_cache_get(name) == NULL)
 	{
-		char		path[MAXPGPATH];
-		InjectionPointCallback 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 = (InjectionPointCallback)
-			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);
+		/* not found in local cache, so load and register it */
+		injection_point_cache_load(entry_by_name);
 	}
 
 	/* Now loaded, so get it. */
-	injection_callback = injection_point_cache_get(name, &private_data);
+	cache_entry = injection_point_cache_get(name);
 
 	LWLockRelease(InjectionPointLock);
 
-	injection_callback(name, private_data);
+	cache_entry->callback(name, cache_entry->private_data);
 #else
 	elog(ERROR, "Injection points are not supported by this build");
 #endif
diff --git a/src/include/utils/injection_point.h b/src/include/utils/injection_point.h
index a61d5d44391..bd3a62425c3 100644
--- a/src/include/utils/injection_point.h
+++ b/src/include/utils/injection_point.h
@@ -15,8 +15,10 @@
  * Injections points require --enable-injection-points.
  */
 #ifdef USE_INJECTION_POINTS
+#define INJECTION_POINT_LOAD(name) InjectionPointLoad(name)
 #define INJECTION_POINT(name) InjectionPointRun(name)
 #else
+#define INJECTION_POINT_LOAD(name) ((void) name)
 #define INJECTION_POINT(name) ((void) name)
 #endif
 
@@ -34,6 +36,7 @@ extern void InjectionPointAttach(const char *name,
 								 const char *function,
 								 const void *private_data,
 								 int private_data_size);
+extern void InjectionPointLoad(const char *name);
 extern void InjectionPointRun(const char *name);
 extern bool InjectionPointDetach(const char *name);
 
diff --git a/src/test/modules/injection_points/expected/injection_points.out b/src/test/modules/injection_points/expected/injection_points.out
index dd9db06e10b..2f60da900bb 100644
--- a/src/test/modules/injection_points/expected/injection_points.out
+++ b/src/test/modules/injection_points/expected/injection_points.out
@@ -128,6 +128,38 @@ SELECT injection_points_detach('TestInjectionLog2');
  
 (1 row)
 
+-- Loading
+SELECT injection_points_load('TestInjectionLogLoad'); -- nothing
+ injection_points_load 
+-----------------------
+ 
+(1 row)
+
+SELECT injection_points_attach('TestInjectionLogLoad', 'notice');
+ injection_points_attach 
+-------------------------
+ 
+(1 row)
+
+SELECT injection_points_load('TestInjectionLogLoad'); -- nothing happens
+ injection_points_load 
+-----------------------
+ 
+(1 row)
+
+SELECT injection_points_run('TestInjectionLogLoad'); -- runs from cache
+NOTICE:  notice triggered for injection point TestInjectionLogLoad
+ injection_points_run 
+----------------------
+ 
+(1 row)
+
+SELECT injection_points_detach('TestInjectionLogLoad');
+ 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 c16a33b08db..e275c2cf5b6 100644
--- a/src/test/modules/injection_points/injection_points--1.0.sql
+++ b/src/test/modules/injection_points/injection_points--1.0.sql
@@ -15,6 +15,16 @@ AS 'MODULE_PATHNAME', 'injection_points_attach'
 LANGUAGE C STRICT PARALLEL UNSAFE;
 
 --
+-- injection_points_load()
+--
+-- Load an injection point already attached.
+--
+CREATE FUNCTION injection_points_load(IN point_name TEXT)
+RETURNS void
+AS 'MODULE_PATHNAME', 'injection_points_load'
+LANGUAGE C STRICT PARALLEL UNSAFE;
+
+--
 -- injection_points_run()
 --
 -- Executes the action attached to the injection point.
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index 1b695a18203..b6c8e893246 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -303,6 +303,23 @@ injection_points_attach(PG_FUNCTION_ARGS)
 }
 
 /*
+ * SQL function for loading an injection point.
+ */
+PG_FUNCTION_INFO_V1(injection_points_load);
+Datum
+injection_points_load(PG_FUNCTION_ARGS)
+{
+	char	   *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
+
+	if (inj_state == NULL)
+		injection_init_shmem();
+
+	INJECTION_POINT_LOAD(name);
+
+	PG_RETURN_VOID();
+}
+
+/*
  * SQL function for triggering an injection point.
  */
 PG_FUNCTION_INFO_V1(injection_points_run);
diff --git a/src/test/modules/injection_points/sql/injection_points.sql b/src/test/modules/injection_points/sql/injection_points.sql
index 71e2972a7e4..fabf0a8823b 100644
--- a/src/test/modules/injection_points/sql/injection_points.sql
+++ b/src/test/modules/injection_points/sql/injection_points.sql
@@ -41,6 +41,13 @@ SELECT injection_points_detach('TestInjectionLog'); -- fails
 SELECT injection_points_run('TestInjectionLog2'); -- notice
 SELECT injection_points_detach('TestInjectionLog2');
 
+-- Loading
+SELECT injection_points_load('TestInjectionLogLoad'); -- nothing
+SELECT injection_points_attach('TestInjectionLogLoad', 'notice');
+SELECT injection_points_load('TestInjectionLogLoad'); -- nothing happens
+SELECT injection_points_run('TestInjectionLogLoad'); -- runs from cache
+SELECT injection_points_detach('TestInjectionLogLoad');
+
 -- Runtime conditions
 SELECT injection_points_attach('TestConditionError', 'error');
 -- Any follow-up injection point attached will be local to this process.