injection_points: Introduce runtime conditions

This adds a new SQL function injection_points_set_local() that can be
used to force injection points to be run only in the process where they
are attached.  This is handy for SQL tests to:
- Detach automatically injection points when the process exits.
- Allow tests with injection points to run concurrently with other test
suites, so as such modules do not have to be marked with
NO_INSTALLCHECK.

Currently, the only condition that can be registered is for a PID.
This could be extended to more kinds later, if required, like database
names/OIDs, roles, or more concepts I did not consider.

Using a single function for SQL scripts is an idea from Heikki
Linnakangas.

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

5 files changed, 285 insertions, 0 deletions

diff --git a/src/test/modules/injection_points/expected/injection_points.out b/src/test/modules/injection_points/expected/injection_points.out
index 827456ccc51..3d94016ac9c 100644
--- a/src/test/modules/injection_points/expected/injection_points.out
+++ b/src/test/modules/injection_points/expected/injection_points.out
@@ -115,4 +115,81 @@ NOTICE:  notice triggered for injection point TestInjectionLog2
  
 (1 row)
 
+SELECT injection_points_detach('TestInjectionLog2');
+ injection_points_detach 
+-------------------------
+ 
+(1 row)
+
+-- Runtime conditions
+SELECT injection_points_attach('TestConditionError', 'error');
+ injection_points_attach 
+-------------------------
+ 
+(1 row)
+
+-- Any follow-up injection point attached will be local to this process.
+SELECT injection_points_set_local();
+ injection_points_set_local 
+----------------------------
+ 
+(1 row)
+
+SELECT injection_points_attach('TestConditionLocal1', 'error');
+ injection_points_attach 
+-------------------------
+ 
+(1 row)
+
+SELECT injection_points_attach('TestConditionLocal2', 'notice');
+ injection_points_attach 
+-------------------------
+ 
+(1 row)
+
+SELECT injection_points_run('TestConditionLocal1'); -- error
+ERROR:  error triggered for injection point TestConditionLocal1
+SELECT injection_points_run('TestConditionLocal2'); -- notice
+NOTICE:  notice triggered for injection point TestConditionLocal2
+ injection_points_run 
+----------------------
+ 
+(1 row)
+
+-- reload, local injection points should be gone.
+\c
+SELECT injection_points_run('TestConditionLocal1'); -- nothing
+ injection_points_run 
+----------------------
+ 
+(1 row)
+
+SELECT injection_points_run('TestConditionLocal2'); -- nothing
+ injection_points_run 
+----------------------
+ 
+(1 row)
+
+SELECT injection_points_run('TestConditionError'); -- error
+ERROR:  error triggered for injection point TestConditionError
+SELECT injection_points_detach('TestConditionError');
+ injection_points_detach 
+-------------------------
+ 
+(1 row)
+
+-- Attaching injection points that use the same name as one defined locally
+-- previously should work.
+SELECT injection_points_attach('TestConditionLocal1', 'error');
+ injection_points_attach 
+-------------------------
+ 
+(1 row)
+
+SELECT injection_points_detach('TestConditionLocal1');
+ injection_points_detach 
+-------------------------
+ 
+(1 row)
+
 DROP EXTENSION injection_points;
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 0a2e59aba79..c16a33b08db 100644
--- a/src/test/modules/injection_points/injection_points--1.0.sql
+++ b/src/test/modules/injection_points/injection_points--1.0.sql
@@ -35,6 +35,17 @@ AS 'MODULE_PATHNAME', 'injection_points_wakeup'
 LANGUAGE C STRICT PARALLEL UNSAFE;
 
 --
+-- injection_points_set_local()
+--
+-- Trigger switch to link any future injection points attached to the
+-- current process, useful to make SQL tests concurrently-safe.
+--
+CREATE FUNCTION injection_points_set_local()
+RETURNS void
+AS 'MODULE_PATHNAME', 'injection_points_set_local'
+LANGUAGE C STRICT PARALLEL UNSAFE;
+
+--
 -- injection_points_detach()
 --
 -- Detaches the current action, if any, from the given injection point.
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index 0730254d544..ace386da50b 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -18,8 +18,10 @@
 #include "postgres.h"
 
 #include "fmgr.h"
+#include "miscadmin.h"
 #include "storage/condition_variable.h"
 #include "storage/dsm_registry.h"
+#include "storage/ipc.h"
 #include "storage/lwlock.h"
 #include "storage/shmem.h"
 #include "utils/builtins.h"
@@ -31,6 +33,23 @@ PG_MODULE_MAGIC;
 /* Maximum number of waits usable in injection points at once */
 #define INJ_MAX_WAIT	8
 #define INJ_NAME_MAXLEN	64
+#define	INJ_MAX_CONDITION	4
+
+/*
+ * Conditions related to injection points.  This tracks in shared memory the
+ * runtime conditions under which an injection point is allowed to run.
+ *
+ * If more types of runtime conditions need to be tracked, this structure
+ * should be expanded.
+ */
+typedef struct InjectionPointCondition
+{
+	/* Name of the injection point related to this condition */
+	char		name[INJ_NAME_MAXLEN];
+
+	/* ID of the process where the injection point is allowed to run */
+	int			pid;
+} InjectionPointCondition;
 
 /* Shared state information for injection points. */
 typedef struct InjectionPointSharedState
@@ -46,6 +65,9 @@ typedef struct InjectionPointSharedState
 
 	/* Condition variable used for waits and wakeups */
 	ConditionVariable wait_point;
+
+	/* Conditions to run an injection point */
+	InjectionPointCondition conditions[INJ_MAX_CONDITION];
 } InjectionPointSharedState;
 
 /* Pointer to shared-memory state. */
@@ -55,6 +77,8 @@ extern PGDLLEXPORT void injection_error(const char *name);
 extern PGDLLEXPORT void injection_notice(const char *name);
 extern PGDLLEXPORT void injection_wait(const char *name);
 
+/* track if injection points attached in this process are linked to it */
+static bool injection_point_local = false;
 
 /*
  * Callback for shared memory area initialization.
@@ -67,6 +91,7 @@ injection_point_init_state(void *ptr)
 	SpinLockInit(&state->lock);
 	memset(state->wait_counts, 0, sizeof(state->wait_counts));
 	memset(state->name, 0, sizeof(state->name));
+	memset(state->conditions, 0, sizeof(state->conditions));
 	ConditionVariableInit(&state->wait_point);
 }
 
@@ -87,16 +112,92 @@ injection_init_shmem(void)
 								   &found);
 }
 
+/*
+ * Check runtime conditions associated to an injection point.
+ *
+ * Returns true if the named injection point is allowed to run, and false
+ * otherwise.  Multiple conditions can be associated to a single injection
+ * point, so check them all.
+ */
+static bool
+injection_point_allowed(const char *name)
+{
+	bool		result = true;
+
+	if (inj_state == NULL)
+		injection_init_shmem();
+
+	SpinLockAcquire(&inj_state->lock);
+
+	for (int i = 0; i < INJ_MAX_CONDITION; i++)
+	{
+		InjectionPointCondition *condition = &inj_state->conditions[i];
+
+		if (strcmp(condition->name, name) == 0)
+		{
+			/*
+			 * Check if this injection point is allowed to run in this
+			 * process.
+			 */
+			if (MyProcPid != condition->pid)
+			{
+				result = false;
+				break;
+			}
+		}
+	}
+
+	SpinLockRelease(&inj_state->lock);
+
+	return result;
+}
+
+/*
+ * before_shmem_exit callback to remove injection points linked to a
+ * specific process.
+ */
+static void
+injection_points_cleanup(int code, Datum arg)
+{
+	/* Leave if nothing is tracked locally */
+	if (!injection_point_local)
+		return;
+
+	SpinLockAcquire(&inj_state->lock);
+	for (int i = 0; i < INJ_MAX_CONDITION; i++)
+	{
+		InjectionPointCondition *condition = &inj_state->conditions[i];
+
+		if (condition->name[0] == '\0')
+			continue;
+
+		if (condition->pid != MyProcPid)
+			continue;
+
+		/* Detach the injection point and unregister condition */
+		InjectionPointDetach(condition->name);
+		condition->name[0] = '\0';
+		condition->pid = 0;
+	}
+	SpinLockRelease(&inj_state->lock);
+}
+
 /* Set of callbacks available to be attached to an injection point. */
 void
 injection_error(const char *name)
 {
+	if (!injection_point_allowed(name))
+		return;
+
 	elog(ERROR, "error triggered for injection point %s", name);
 }
 
 void
 injection_notice(const char *name)
 {
+	if (!injection_point_allowed(name))
+		return;
+
 	elog(NOTICE, "notice triggered for injection point %s", name);
 }
 
@@ -111,6 +212,9 @@ injection_wait(const char *name)
 	if (inj_state == NULL)
 		injection_init_shmem();
 
+	if (!injection_point_allowed(name))
+		return;
+
 	/*
 	 * Use the injection point name for this custom wait event.  Note that
 	 * this custom wait event name is not released, but we don't care much for
@@ -182,6 +286,35 @@ injection_points_attach(PG_FUNCTION_ARGS)
 
 	InjectionPointAttach(name, "injection_points", function);
 
+	if (injection_point_local)
+	{
+		int			index = -1;
+
+		/*
+		 * Register runtime condition to link this injection point to the
+		 * current process.
+		 */
+		SpinLockAcquire(&inj_state->lock);
+		for (int i = 0; i < INJ_MAX_CONDITION; i++)
+		{
+			InjectionPointCondition *condition = &inj_state->conditions[i];
+
+			if (condition->name[0] == '\0')
+			{
+				index = i;
+				strlcpy(condition->name, name, INJ_NAME_MAXLEN);
+				condition->pid = MyProcPid;
+				break;
+			}
+		}
+		SpinLockRelease(&inj_state->lock);
+
+		if (index < 0)
+			elog(FATAL,
+				 "could not find free slot for condition of injection point %s",
+				 name);
+	}
+
 	PG_RETURN_VOID();
 }
 
@@ -236,6 +369,32 @@ injection_points_wakeup(PG_FUNCTION_ARGS)
 }
 
 /*
+ * injection_points_set_local
+ *
+ * Track if any injection point created in this process ought to run only
+ * in this process.  Such injection points are detached automatically when
+ * this process exits.  This is useful to make test suites concurrent-safe.
+ */
+PG_FUNCTION_INFO_V1(injection_points_set_local);
+Datum
+injection_points_set_local(PG_FUNCTION_ARGS)
+{
+	/* Enable flag to add a runtime condition based on this process ID */
+	injection_point_local = true;
+
+	if (inj_state == NULL)
+		injection_init_shmem();
+
+	/*
+	 * Register a before_shmem_exit callback to remove any injection points
+	 * linked to this process.
+	 */
+	before_shmem_exit(injection_points_cleanup, (Datum) 0);
+
+	PG_RETURN_VOID();
+}
+
+/*
  * SQL function for dropping an injection point.
  */
 PG_FUNCTION_INFO_V1(injection_points_detach);
@@ -246,5 +405,22 @@ injection_points_detach(PG_FUNCTION_ARGS)
 
 	InjectionPointDetach(name);
 
+	if (inj_state == NULL)
+		injection_init_shmem();
+
+	/* Clean up any conditions associated to this injection point */
+	SpinLockAcquire(&inj_state->lock);
+	for (int i = 0; i < INJ_MAX_CONDITION; i++)
+	{
+		InjectionPointCondition *condition = &inj_state->conditions[i];
+
+		if (strcmp(condition->name, name) == 0)
+		{
+			condition->pid = 0;
+			condition->name[0] = '\0';
+		}
+	}
+	SpinLockRelease(&inj_state->lock);
+
 	PG_RETURN_VOID();
 }
diff --git a/src/test/modules/injection_points/sql/injection_points.sql b/src/test/modules/injection_points/sql/injection_points.sql
index 23c7e435ad2..2aa76a542bb 100644
--- a/src/test/modules/injection_points/sql/injection_points.sql
+++ b/src/test/modules/injection_points/sql/injection_points.sql
@@ -30,5 +30,25 @@ SELECT injection_points_run('TestInjectionLog2'); -- notice
 SELECT injection_points_detach('TestInjectionLog'); -- fails
 
 SELECT injection_points_run('TestInjectionLog2'); -- notice
+SELECT injection_points_detach('TestInjectionLog2');
+
+-- Runtime conditions
+SELECT injection_points_attach('TestConditionError', 'error');
+-- Any follow-up injection point attached will be local to this process.
+SELECT injection_points_set_local();
+SELECT injection_points_attach('TestConditionLocal1', 'error');
+SELECT injection_points_attach('TestConditionLocal2', 'notice');
+SELECT injection_points_run('TestConditionLocal1'); -- error
+SELECT injection_points_run('TestConditionLocal2'); -- notice
+-- reload, local injection points should be gone.
+\c
+SELECT injection_points_run('TestConditionLocal1'); -- nothing
+SELECT injection_points_run('TestConditionLocal2'); -- nothing
+SELECT injection_points_run('TestConditionError'); -- error
+SELECT injection_points_detach('TestConditionError');
+-- Attaching injection points that use the same name as one defined locally
+-- previously should work.
+SELECT injection_points_attach('TestConditionLocal1', 'error');
+SELECT injection_points_detach('TestConditionLocal1');
 
 DROP EXTENSION injection_points;
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index cb48e2f95a8..c9d4ad4a76f 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1219,6 +1219,7 @@ InitSampleScan_function
 InitializeDSMForeignScan_function
 InitializeWorkerForeignScan_function
 InjectionPointCacheEntry
+InjectionPointCondition
 InjectionPointEntry
 InjectionPointSharedState
 InlineCodeBlock