v4-0001-Add-chunk-area-for-injection-points.patch

From 233e710c8fc2242bdd4e40d5a20f8de2828765b7 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Fri, 10 May 2024 09:40:42 +0900
Subject: [PATCH v4] Add chunk area for injection points

---
 src/include/utils/injection_point.h           |   9 +-
 src/backend/utils/misc/injection_point.c      |  53 ++++-
 .../expected/injection_points.out             |   2 +-
 .../injection_points/injection_points.c       | 191 +++++++-----------
 doc/src/sgml/xfunc.sgml                       |  14 +-
 src/tools/pgindent/typedefs.list              |   1 +
 6 files changed, 131 insertions(+), 139 deletions(-)

diff --git a/src/include/utils/injection_point.h b/src/include/utils/injection_point.h
index 55524b568f..a61d5d4439 100644
--- a/src/include/utils/injection_point.h
+++ b/src/include/utils/injection_point.h
@@ -23,15 +23,18 @@
 /*
  * Typedef for callback function launched by an injection point.
  */
-typedef void (*InjectionPointCallback) (const char *name);
+typedef void (*InjectionPointCallback) (const char *name,
+										const void *private_data);
 
 extern Size InjectionPointShmemSize(void);
 extern void InjectionPointShmemInit(void);
 
 extern void InjectionPointAttach(const char *name,
 								 const char *library,
-								 const char *function);
+								 const char *function,
+								 const void *private_data,
+								 int private_data_size);
 extern void InjectionPointRun(const char *name);
-extern void InjectionPointDetach(const char *name);
+extern bool InjectionPointDetach(const char *name);
 
 #endif							/* INJECTION_POINT_H */
diff --git a/src/backend/utils/misc/injection_point.c b/src/backend/utils/misc/injection_point.c
index 0cf4d51cac..d5a8726644 100644
--- a/src/backend/utils/misc/injection_point.c
+++ b/src/backend/utils/misc/injection_point.c
@@ -42,6 +42,7 @@ static HTAB *InjectionPointHash;	/* find points from names */
 #define INJ_NAME_MAXLEN		64
 #define INJ_LIB_MAXLEN		128
 #define INJ_FUNC_MAXLEN		128
+#define INJ_PRIVATE_MAXLEN	1024
 
 /* Single injection point stored in InjectionPointHash */
 typedef struct InjectionPointEntry
@@ -49,6 +50,12 @@ typedef struct InjectionPointEntry
 	char		name[INJ_NAME_MAXLEN];	/* hash key */
 	char		library[INJ_LIB_MAXLEN];	/* library */
 	char		function[INJ_FUNC_MAXLEN];	/* function */
+
+	/*
+	 * Opaque data area that modules can use to pass some custom data to
+	 * callbacks, registered when attached.
+	 */
+	char		private_data[INJ_PRIVATE_MAXLEN];
 } InjectionPointEntry;
 
 #define INJECTION_POINT_HASH_INIT_SIZE	16
@@ -61,6 +68,7 @@ typedef struct InjectionPointEntry
 typedef struct InjectionPointCacheEntry
 {
 	char		name[INJ_NAME_MAXLEN];
+	char		private_data[INJ_PRIVATE_MAXLEN];
 	InjectionPointCallback callback;
 } InjectionPointCacheEntry;
 
@@ -73,7 +81,8 @@ static HTAB *InjectionPointCache = NULL;
  */
 static void
 injection_point_cache_add(const char *name,
-						  InjectionPointCallback callback)
+						  InjectionPointCallback callback,
+						  const void *private_data)
 {
 	InjectionPointCacheEntry *entry;
 	bool		found;
@@ -99,6 +108,7 @@ injection_point_cache_add(const char *name,
 	Assert(!found);
 	strlcpy(entry->name, name, sizeof(entry->name));
 	entry->callback = callback;
+	memcpy(entry->private_data, private_data, INJ_PRIVATE_MAXLEN);
 }
 
 /*
@@ -124,11 +134,14 @@ injection_point_cache_remove(const char *name)
  * Retrieve an injection point from the local cache, if any.
  */
 static InjectionPointCallback
-injection_point_cache_get(const char *name)
+injection_point_cache_get(const char *name, const void **private_data)
 {
 	bool		found;
 	InjectionPointCacheEntry *entry;
 
+	if (private_data)
+		*private_data = NULL;
+
 	/* no callback if no cache yet */
 	if (InjectionPointCache == NULL)
 		return NULL;
@@ -137,7 +150,11 @@ injection_point_cache_get(const char *name)
 		hash_search(InjectionPointCache, name, HASH_FIND, &found);
 
 	if (found)
+	{
+		if (private_data)
+			*private_data = entry->private_data;
 		return entry->callback;
+	}
 
 	return NULL;
 }
@@ -186,7 +203,9 @@ InjectionPointShmemInit(void)
 void
 InjectionPointAttach(const char *name,
 					 const char *library,
-					 const char *function)
+					 const char *function,
+					 const void *private_data,
+					 int private_data_size)
 {
 #ifdef USE_INJECTION_POINTS
 	InjectionPointEntry *entry_by_name;
@@ -201,6 +220,9 @@ InjectionPointAttach(const char *name,
 	if (strlen(function) >= INJ_FUNC_MAXLEN)
 		elog(ERROR, "injection point function %s too long (maximum of %u)",
 			 function, INJ_FUNC_MAXLEN);
+	if (private_data_size >= INJ_PRIVATE_MAXLEN)
+		elog(ERROR, "injection point data too long (maximum of %u)",
+			 INJ_PRIVATE_MAXLEN);
 
 	/*
 	 * Allocate and register a new injection point.  A new point should not
@@ -223,6 +245,7 @@ InjectionPointAttach(const char *name,
 	entry_by_name->library[INJ_LIB_MAXLEN - 1] = '\0';
 	strlcpy(entry_by_name->function, function, sizeof(entry_by_name->function));
 	entry_by_name->function[INJ_FUNC_MAXLEN - 1] = '\0';
+	memcpy(entry_by_name->private_data, private_data, private_data_size);
 
 	LWLockRelease(InjectionPointLock);
 
@@ -233,8 +256,10 @@ InjectionPointAttach(const char *name,
 
 /*
  * Detach an existing injection point.
+ *
+ * Returns true if the injection point was detached, false otherwise.
  */
-void
+bool
 InjectionPointDetach(const char *name)
 {
 #ifdef USE_INJECTION_POINTS
@@ -245,10 +270,12 @@ InjectionPointDetach(const char *name)
 	LWLockRelease(InjectionPointLock);
 
 	if (!found)
-		elog(ERROR, "injection point \"%s\" not found", name);
+		return false;
 
+	return true;
 #else
 	elog(ERROR, "Injection points are not supported by this build");
+	return true;				/* silence compiler */
 #endif
 }
 
@@ -265,6 +292,7 @@ InjectionPointRun(const char *name)
 	InjectionPointEntry *entry_by_name;
 	bool		found;
 	InjectionPointCallback injection_callback;
+	const void *private_data;
 
 	LWLockAcquire(InjectionPointLock, LW_SHARED);
 	entry_by_name = (InjectionPointEntry *)
@@ -286,10 +314,10 @@ InjectionPointRun(const char *name)
 	 * Check if the callback exists in the local cache, to avoid unnecessary
 	 * external loads.
 	 */
-	injection_callback = injection_point_cache_get(name);
-	if (injection_callback == NULL)
+	if (injection_point_cache_get(name, NULL) == 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,
@@ -299,18 +327,21 @@ InjectionPointRun(const char *name)
 			elog(ERROR, "could not find library \"%s\" for injection point \"%s\"",
 				 path, name);
 
-		injection_callback = (InjectionPointCallback)
+		injection_callback_local = (InjectionPointCallback)
 			load_external_function(path, entry_by_name->function, false, NULL);
 
-		if (injection_callback == 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);
+		injection_point_cache_add(name, injection_callback_local,
+								  entry_by_name->private_data);
 	}
 
-	injection_callback(name);
+	/* Now loaded, so get it. */
+	injection_callback = injection_point_cache_get(name, &private_data);
+	injection_callback(name, private_data);
 #else
 	elog(ERROR, "Injection points are not supported by this build");
 #endif
diff --git a/src/test/modules/injection_points/expected/injection_points.out b/src/test/modules/injection_points/expected/injection_points.out
index 1341d66c92..dd9db06e10 100644
--- a/src/test/modules/injection_points/expected/injection_points.out
+++ b/src/test/modules/injection_points/expected/injection_points.out
@@ -114,7 +114,7 @@ NOTICE:  notice triggered for injection point TestInjectionLog2
 (1 row)
 
 SELECT injection_points_detach('TestInjectionLog'); -- fails
-ERROR:  injection point "TestInjectionLog" not found
+ERROR:  could not detach injection point "TestInjectionLog"
 SELECT injection_points_run('TestInjectionLog2'); -- notice
 NOTICE:  notice triggered for injection point TestInjectionLog2
  injection_points_run 
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index a74a4a28af..5c44625d1d 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -19,6 +19,8 @@
 
 #include "fmgr.h"
 #include "miscadmin.h"
+#include "nodes/pg_list.h"
+#include "nodes/value.h"
 #include "storage/condition_variable.h"
 #include "storage/dsm_registry.h"
 #include "storage/ipc.h"
@@ -26,6 +28,7 @@
 #include "storage/shmem.h"
 #include "utils/builtins.h"
 #include "utils/injection_point.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 
 PG_MODULE_MAGIC;
@@ -33,24 +36,37 @@ 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.
+ * runtime conditions under which an injection point is allowed to run,
+ * stored as private_data when an injection point is attached, and passed as
+ * argument to the callback.
  *
  * If more types of runtime conditions need to be tracked, this structure
  * should be expanded.
  */
+typedef enum InjectionPointConditionType
+{
+	INJ_CONDITION_ALWAYS = 0,	/* always run */
+	INJ_CONDITION_PID,			/* PID restriction */
+} InjectionPointConditionType;
+
 typedef struct InjectionPointCondition
 {
-	/* Name of the injection point related to this condition */
-	char		name[INJ_NAME_MAXLEN];
+	/* Type of the condition */
+	InjectionPointConditionType type;
 
 	/* ID of the process where the injection point is allowed to run */
 	int			pid;
 } InjectionPointCondition;
 
+/*
+ * List of injection points stored in TopMemoryContext attached
+ * locally to this process.
+ */
+static List *inj_list_local = NIL;
+
 /* Shared state information for injection points. */
 typedef struct InjectionPointSharedState
 {
@@ -65,17 +81,17 @@ 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. */
 static InjectionPointSharedState *inj_state = NULL;
 
-extern PGDLLEXPORT void injection_error(const char *name);
-extern PGDLLEXPORT void injection_notice(const char *name);
-extern PGDLLEXPORT void injection_wait(const char *name);
+extern PGDLLEXPORT void injection_error(const char *name,
+										const void *private_data);
+extern PGDLLEXPORT void injection_notice(const char *name,
+										 const void *private_data);
+extern PGDLLEXPORT void injection_wait(const char *name,
+									   const void *private_data);
 
 /* track if injection points attached in this process are linked to it */
 static bool injection_point_local = false;
@@ -91,7 +107,6 @@ 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);
 }
 
@@ -116,39 +131,23 @@ injection_init_shmem(void)
  * 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.
+ * otherwise.
  */
 static bool
-injection_point_allowed(const char *name)
+injection_point_allowed(InjectionPointCondition *condition)
 {
 	bool		result = true;
 
-	if (inj_state == NULL)
-		injection_init_shmem();
-
-	SpinLockAcquire(&inj_state->lock);
-
-	for (int i = 0; i < INJ_MAX_CONDITION; i++)
+	switch (condition->type)
 	{
-		InjectionPointCondition *condition = &inj_state->conditions[i];
-
-		if (strcmp(condition->name, name) == 0)
-		{
-			/*
-			 * Check if this injection point is allowed to run in this
-			 * process.
-			 */
+		case INJ_CONDITION_PID:
 			if (MyProcPid != condition->pid)
-			{
 				result = false;
-				break;
-			}
-		}
+			break;
+		case INJ_CONDITION_ALWAYS:
+			break;
 	}
 
-	SpinLockRelease(&inj_state->lock);
-
 	return result;
 }
 
@@ -159,70 +158,39 @@ injection_point_allowed(const char *name)
 static void
 injection_points_cleanup(int code, Datum arg)
 {
-	char		names[INJ_MAX_CONDITION][INJ_NAME_MAXLEN] = {0};
-	int			count = 0;
+	ListCell   *lc;
 
 	/* Leave if nothing is tracked locally */
 	if (!injection_point_local)
 		return;
 
-	/*
-	 * This is done in three steps: detect the points to detach, detach them
-	 * and release their conditions.
-	 */
-	SpinLockAcquire(&inj_state->lock);
-	for (int i = 0; i < INJ_MAX_CONDITION; i++)
+	/* Detach all the local points */
+	foreach(lc, inj_list_local)
 	{
-		InjectionPointCondition *condition = &inj_state->conditions[i];
+		char	   *name = strVal(lfirst(lc));
 
-		if (condition->name[0] == '\0')
-			continue;
-
-		if (condition->pid != MyProcPid)
-			continue;
-
-		/* Extract the point name to detach */
-		strlcpy(names[count], condition->name, INJ_NAME_MAXLEN);
-		count++;
+		(void) InjectionPointDetach(name);
 	}
-	SpinLockRelease(&inj_state->lock);
-
-	/* Detach, without holding the spinlock */
-	for (int i = 0; i < count; i++)
-		InjectionPointDetach(names[i]);
-
-	/* Clear all the conditions */
-	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;
-
-		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)
+injection_error(const char *name, const void *private_data)
 {
-	if (!injection_point_allowed(name))
+	InjectionPointCondition *condition = (InjectionPointCondition *) private_data;
+
+	if (!injection_point_allowed(condition))
 		return;
 
 	elog(ERROR, "error triggered for injection point %s", name);
 }
 
 void
-injection_notice(const char *name)
+injection_notice(const char *name, const void *private_data)
 {
-	if (!injection_point_allowed(name))
+	InjectionPointCondition *condition = (InjectionPointCondition *) private_data;
+
+	if (!injection_point_allowed(condition))
 		return;
 
 	elog(NOTICE, "notice triggered for injection point %s", name);
@@ -230,16 +198,17 @@ injection_notice(const char *name)
 
 /* Wait on a condition variable, awaken by injection_points_wakeup() */
 void
-injection_wait(const char *name)
+injection_wait(const char *name, const void *private_data)
 {
 	uint32		old_wait_counts = 0;
 	int			index = -1;
 	uint32		injection_wait_event = 0;
+	InjectionPointCondition *condition = (InjectionPointCondition *) private_data;
 
 	if (inj_state == NULL)
 		injection_init_shmem();
 
-	if (!injection_point_allowed(name))
+	if (!injection_point_allowed(condition))
 		return;
 
 	/*
@@ -301,6 +270,7 @@ injection_points_attach(PG_FUNCTION_ARGS)
 	char	   *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
 	char	   *action = text_to_cstring(PG_GETARG_TEXT_PP(1));
 	char	   *function;
+	InjectionPointCondition condition = {0};
 
 	if (strcmp(action, "error") == 0)
 		function = "injection_error";
@@ -311,37 +281,24 @@ injection_points_attach(PG_FUNCTION_ARGS)
 	else
 		elog(ERROR, "incorrect action \"%s\" for injection point creation", action);
 
-	InjectionPointAttach(name, "injection_points", function);
+	if (injection_point_local)
+	{
+		condition.type = INJ_CONDITION_PID;
+		condition.pid = MyProcPid;
+	}
+
+	InjectionPointAttach(name, "injection_points", function, &condition,
+						 sizeof(InjectionPointCondition));
 
 	if (injection_point_local)
 	{
-		int			index = -1;
+		MemoryContext oldctx;
 
-		/*
-		 * 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);
+		/* Local injection point, so track it for automated cleanup */
+		oldctx = MemoryContextSwitchTo(TopMemoryContext);
+		inj_list_local = lappend(inj_list_local, makeString(pstrdup(name)));
+		MemoryContextSwitchTo(oldctx);
 	}
-
 	PG_RETURN_VOID();
 }
 
@@ -430,24 +387,18 @@ injection_points_detach(PG_FUNCTION_ARGS)
 {
 	char	   *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
 
-	InjectionPointDetach(name);
+	if (!InjectionPointDetach(name))
+		elog(ERROR, "could not detach injection point \"%s\"", 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++)
+	/* Remove point from local list, if required */
+	if (inj_list_local != NIL)
 	{
-		InjectionPointCondition *condition = &inj_state->conditions[i];
+		MemoryContext oldctx;
 
-		if (strcmp(condition->name, name) == 0)
-		{
-			condition->pid = 0;
-			condition->name[0] = '\0';
-		}
+		oldctx = MemoryContextSwitchTo(TopMemoryContext);
+		inj_list_local = list_delete(inj_list_local, makeString(name));
+		MemoryContextSwitchTo(oldctx);
 	}
-	SpinLockRelease(&inj_state->lock);
 
 	PG_RETURN_VOID();
 }
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 7d053698a2..f5cb4b2212 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3624,12 +3624,16 @@ INJECTION_POINT(name);
 <programlisting>
 extern void InjectionPointAttach(const char *name,
                                  const char *library,
-                                 const char *function);
+                                 const char *function,
+                                 const void *private_data,
+                                 int private_data_size);
 </programlisting>
 
      <literal>name</literal> is the name of the injection point, which when
      reached during execution will execute the <literal>function</literal>
-     loaded from <literal>library</literal>.
+     loaded from <literal>library</literal>. <literal>private_data</literal>
+     is a private area of data of size <literal>private_data_size</literal>
+     given as argument to the callback when executed,
     </para>
 
     <para>
@@ -3637,7 +3641,7 @@ extern void InjectionPointAttach(const char *name,
      <literal>InjectionPointCallback</literal>:
 <programlisting>
 static void
-custom_injection_callback(const char *name)
+custom_injection_callback(const char *name, const void *private_data)
 {
     elog(NOTICE, "%s: executed custom callback", name);
 }
@@ -3650,8 +3654,10 @@ custom_injection_callback(const char *name)
     <para>
      Optionally, it is possible to detach an injection point by calling:
 <programlisting>
-extern void InjectionPointDetach(const char *name);
+extern bool InjectionPointDetach(const char *name);
 </programlisting>
+     On success, <literal>true</literal> is returned, <literal>false</literal>
+     otherwise.
     </para>
 
     <para>
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 2311f82d81..34ec87a85e 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1219,6 +1219,7 @@ InitializeDSMForeignScan_function
 InitializeWorkerForeignScan_function
 InjectionPointCacheEntry
 InjectionPointCondition
+InjectionPointConditionType
 InjectionPointEntry
 InjectionPointSharedState
 InlineCodeBlock
-- 
2.43.0