0001-Add-support-for-runtime-checks-in-injection-points.patch

From 4767c548b718f30e7f3f427db19fb31146420f3d Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Tue, 23 Jan 2024 10:11:46 +0900
Subject: [PATCH] Add support for runtime checks in injection points

This extends injections points with a 1-argument flavor, so as it
becomes possible for callbacks to perform runtime checks.
---
 src/include/utils/injection_point.h           |  7 +-
 src/backend/utils/misc/injection_point.c      | 81 +++++++++++++++----
 .../expected/injection_points.out             | 25 ++++++
 .../injection_points--1.0.sql                 | 10 +++
 .../injection_points/injection_points.c       | 31 ++++++-
 .../injection_points/sql/injection_points.sql |  7 ++
 doc/src/sgml/xfunc.sgml                       | 10 ++-
 7 files changed, 151 insertions(+), 20 deletions(-)

diff --git a/src/include/utils/injection_point.h b/src/include/utils/injection_point.h
index 55524b568f..8e1527bb4a 100644
--- a/src/include/utils/injection_point.h
+++ b/src/include/utils/injection_point.h
@@ -16,22 +16,27 @@
  */
 #ifdef USE_INJECTION_POINTS
 #define INJECTION_POINT(name) InjectionPointRun(name)
+#define INJECTION_POINT_1ARG(name, arg1) InjectionPointRun1Arg(name, arg1)
 #else
 #define INJECTION_POINT(name) ((void) name)
+#define INJECTION_POINT_1ARG(name) ((void) name, (void) arg1)
 #endif
 
 /*
  * Typedef for callback function launched by an injection point.
  */
 typedef void (*InjectionPointCallback) (const char *name);
+typedef void (*InjectionPointCallback1Arg) (const char *name, void *arg1);
 
 extern Size InjectionPointShmemSize(void);
 extern void InjectionPointShmemInit(void);
 
 extern void InjectionPointAttach(const char *name,
 								 const char *library,
-								 const char *function);
+								 const char *function,
+								 int num_args);
 extern void InjectionPointRun(const char *name);
+extern void InjectionPointRun1Arg(const char *name, void *arg1);
 extern void 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 4f52f19853..082622dde9 100644
--- a/src/backend/utils/misc/injection_point.c
+++ b/src/backend/utils/misc/injection_point.c
@@ -49,6 +49,7 @@ typedef struct InjectionPointEntry
 	char		name[INJ_NAME_MAXLEN];	/* hash key */
 	char		library[INJ_LIB_MAXLEN];	/* library */
 	char		function[INJ_FUNC_MAXLEN];	/* function */
+	int			num_args;	/* number of arguments */
 } InjectionPointEntry;
 
 #define INJECTION_POINT_HASH_INIT_SIZE	16
@@ -61,7 +62,12 @@ typedef struct InjectionPointEntry
 typedef struct InjectionPointCacheEntry
 {
 	char		name[INJ_NAME_MAXLEN];
-	InjectionPointCallback callback;
+	int			num_args;
+	union
+	{
+		InjectionPointCallback callback;
+		InjectionPointCallback1Arg callback_1arg;
+	};
 } InjectionPointCacheEntry;
 
 static HTAB *InjectionPointCache = NULL;
@@ -73,7 +79,8 @@ static HTAB *InjectionPointCache = NULL;
  */
 static void
 injection_point_cache_add(const char *name,
-						  InjectionPointCallback callback)
+						  void *callback,
+						  int num_args)
 {
 	InjectionPointCacheEntry *entry;
 	bool		found;
@@ -99,6 +106,13 @@ injection_point_cache_add(const char *name,
 	Assert(!found);
 	strlcpy(entry->name, name, sizeof(entry->name));
 	entry->callback = callback;
+	entry->num_args = num_args;
+	if (num_args == 0)
+		entry->callback = (InjectionPointCallback) callback;
+	else if (num_args == 1)
+		entry->callback_1arg = (InjectionPointCallback1Arg) callback;
+	else
+		elog(ERROR, "unsupported number of arguments");	/* not reachable */
 }
 
 /*
@@ -123,7 +137,7 @@ injection_point_cache_remove(const char *name)
  *
  * Retrieve an injection point from the local cache, if any.
  */
-static InjectionPointCallback
+static InjectionPointCacheEntry *
 injection_point_cache_get(const char *name)
 {
 	bool		found;
@@ -137,7 +151,7 @@ injection_point_cache_get(const char *name)
 		hash_search(InjectionPointCache, name, HASH_FIND, &found);
 
 	if (found)
-		return entry->callback;
+		return entry;
 
 	return NULL;
 }
@@ -186,7 +200,8 @@ InjectionPointShmemInit(void)
 void
 InjectionPointAttach(const char *name,
 					 const char *library,
-					 const char *function)
+					 const char *function,
+					 int num_args)
 {
 #ifdef USE_INJECTION_POINTS
 	InjectionPointEntry *entry_by_name;
@@ -223,6 +238,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';
+	entry_by_name->num_args = num_args;
 
 	LWLockRelease(InjectionPointLock);
 
@@ -253,18 +269,18 @@ InjectionPointDetach(const char *name)
 }
 
 /*
- * Execute an injection point, if defined.
+ * Workhorse for execution of an injection point, if defined.
  *
  * Check first the shared hash table, and adapt the local cache depending
  * on that as it could be possible that an entry to run has been removed.
  */
-void
-InjectionPointRun(const char *name)
+static inline void
+InjectionPointRunInternal(const char *name, int num_args, void *arg1)
 {
 #ifdef USE_INJECTION_POINTS
 	InjectionPointEntry *entry_by_name;
 	bool		found;
-	InjectionPointCallback injection_callback;
+	InjectionPointCacheEntry *cache_entry;
 
 	LWLockAcquire(InjectionPointLock, LW_SHARED);
 	entry_by_name = (InjectionPointEntry *)
@@ -286,10 +302,11 @@ 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)
+	cache_entry = injection_point_cache_get(name);
+	if (cache_entry == NULL)
 	{
 		char		path[MAXPGPATH];
+		void	   *injection_callback;
 
 		/* not found in local cache, so load and register */
 		snprintf(path, MAXPGPATH, "%s/%s%s", pkglib_path,
@@ -299,19 +316,53 @@ InjectionPointRun(const char *name)
 			elog(ERROR, "could not find library \"%s\" for injection point \"%s\"",
 				 path, name);
 
-		injection_callback = (InjectionPointCallback)
+		injection_callback = (void *)
 			load_external_function(path, entry_by_name->function, true, NULL);
 
 		if (injection_callback == NULL)
 			elog(ERROR, "could not find function \"%s\" in library \"%s\" for injection point \"%s\"",
-				 name, entry_by_name->function, path);
+				 entry_by_name->function, path, name);
+
+		if (entry_by_name->num_args != num_args)
+			elog(ERROR, "incorrect number of arguments in function \"%s\" for injection point \"%s\": defined %d but expected %d",
+				 entry_by_name->function, name,
+				 entry_by_name->num_args, num_args);
 
 		/* add it to the local cache when found */
-		injection_point_cache_add(name, injection_callback);
+		injection_point_cache_add(name, injection_callback, num_args);
+
+		/* and fetch it */
+		cache_entry = injection_point_cache_get(name);
 	}
 
-	injection_callback(name);
+	if (cache_entry->num_args != num_args)
+		elog(ERROR, "unexpected number of arguments for injection point \"%s\": defined %d but expected %d",
+			 name, cache_entry->num_args, num_args);
+
+	if (cache_entry->num_args == 0)
+		cache_entry->callback(name);
+	else if (cache_entry->num_args == 1)
+		cache_entry->callback_1arg(name, arg1);
+	else
+	{
+		Assert(false);	/* cannot be reached */
+	}
 #else
 	elog(ERROR, "Injection points are not supported by this build");
 #endif
 }
+
+/*
+ * Execute an injection point, with no arguments.
+ */
+void
+InjectionPointRun(const char *name)
+{
+	InjectionPointRunInternal(name, 0, NULL);
+}
+/* 1-argument version */
+void
+InjectionPointRun1Arg(const char *name, void *arg1)
+{
+	InjectionPointRunInternal(name, 1, arg1);
+}
diff --git a/src/test/modules/injection_points/expected/injection_points.out b/src/test/modules/injection_points/expected/injection_points.out
index 827456ccc5..b4e70c6831 100644
--- a/src/test/modules/injection_points/expected/injection_points.out
+++ b/src/test/modules/injection_points/expected/injection_points.out
@@ -115,4 +115,29 @@ NOTICE:  notice triggered for injection point TestInjectionLog2
  
 (1 row)
 
+-- 1-argument runs
+SELECT injection_points_run_1arg('TestInjectionLog2', 'blah'); -- error
+ERROR:  unexpected number of arguments for injection point "TestInjectionLog2": defined 0 but expected 1
+SELECT injection_points_attach('TestInjectionLogArg1', 'notice_1arg');
+ injection_points_attach 
+-------------------------
+ 
+(1 row)
+
+SELECT injection_points_run('TestInjectionLogArg1'); -- error
+ERROR:  incorrect number of arguments in function "injection_notice_1arg" for injection point "TestInjectionLogArg1": defined 1 but expected 0
+SELECT injection_points_run_1arg('TestInjectionLogArg1', 'blah'); -- notice
+NOTICE:  notice triggered for injection point TestInjectionLogArg1: blah
+ injection_points_run_1arg 
+---------------------------
+ 
+(1 row)
+
+SELECT injection_points_run_1arg('TestInjectionLogArg1', 'foobar'); -- notice
+NOTICE:  notice triggered for injection point TestInjectionLogArg1: foobar
+ injection_points_run_1arg 
+---------------------------
+ 
+(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 5944c41716..9b99c17199 100644
--- a/src/test/modules/injection_points/injection_points--1.0.sql
+++ b/src/test/modules/injection_points/injection_points--1.0.sql
@@ -24,6 +24,16 @@ RETURNS void
 AS 'MODULE_PATHNAME', 'injection_points_run'
 LANGUAGE C STRICT PARALLEL UNSAFE;
 
+--
+-- injection_points_run_1arg()
+--
+-- Executes the action attached to the injection point.
+--
+CREATE FUNCTION injection_points_run_1arg(IN point_name TEXT, IN arg1 TEXT)
+RETURNS void
+AS 'MODULE_PATHNAME', 'injection_points_run_1arg'
+LANGUAGE C STRICT PARALLEL UNSAFE;
+
 --
 -- injection_points_detach()
 --
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index e843e6594f..862f904a75 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -28,6 +28,7 @@ PG_MODULE_MAGIC;
 
 extern PGDLLEXPORT void injection_error(const char *name);
 extern PGDLLEXPORT void injection_notice(const char *name);
+extern PGDLLEXPORT void injection_notice_1arg(const char *name, void *arg1);
 
 
 /* Set of callbacks available to be attached to an injection point. */
@@ -43,6 +44,13 @@ injection_notice(const char *name)
 	elog(NOTICE, "notice triggered for injection point %s", name);
 }
 
+void
+injection_notice_1arg(const char *name, void *arg1)
+{
+	const char *str = (const char *) arg1;
+	elog(NOTICE, "notice triggered for injection point %s: %s", name, str);
+}
+
 /*
  * SQL function for creating an injection point.
  */
@@ -53,15 +61,21 @@ 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;
+	int			num_args = 0;
 
 	if (strcmp(action, "error") == 0)
 		function = "injection_error";
 	else if (strcmp(action, "notice") == 0)
 		function = "injection_notice";
+	else if (strcmp(action, "notice_1arg") == 0)
+	{
+		function = "injection_notice_1arg";
+		num_args = 1;
+	}
 	else
 		elog(ERROR, "incorrect action \"%s\" for injection point creation", action);
 
-	InjectionPointAttach(name, "injection_points", function);
+	InjectionPointAttach(name, "injection_points", function, num_args);
 
 	PG_RETURN_VOID();
 }
@@ -80,6 +94,21 @@ injection_points_run(PG_FUNCTION_ARGS)
 	PG_RETURN_VOID();
 }
 
+/*
+ * SQL function for triggering an injection point, 1-argument flavor.
+ */
+PG_FUNCTION_INFO_V1(injection_points_run_1arg);
+Datum
+injection_points_run_1arg(PG_FUNCTION_ARGS)
+{
+	char	   *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	char	   *arg1 = text_to_cstring(PG_GETARG_TEXT_PP(1));
+
+	INJECTION_POINT_1ARG(name, (void *) arg1);
+
+	PG_RETURN_VOID();
+}
+
 /*
  * SQL function for dropping 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 23c7e435ad..eca528508b 100644
--- a/src/test/modules/injection_points/sql/injection_points.sql
+++ b/src/test/modules/injection_points/sql/injection_points.sql
@@ -31,4 +31,11 @@ SELECT injection_points_detach('TestInjectionLog'); -- fails
 
 SELECT injection_points_run('TestInjectionLog2'); -- notice
 
+-- 1-argument runs
+SELECT injection_points_run_1arg('TestInjectionLog2', 'blah'); -- error
+SELECT injection_points_attach('TestInjectionLogArg1', 'notice_1arg');
+SELECT injection_points_run('TestInjectionLogArg1'); -- error
+SELECT injection_points_run_1arg('TestInjectionLogArg1', 'blah'); -- notice
+SELECT injection_points_run_1arg('TestInjectionLogArg1', 'foobar'); -- notice
+
 DROP EXTENSION injection_points;
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 8a79ad0943..7659c81c09 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3607,13 +3607,15 @@ uint32 WaitEventExtensionNew(const char *wait_event_name)
      macro:
 <programlisting>
 INJECTION_POINT(name);
+INJECTION_POINT_1ARG(name, arg1);
 </programlisting>
 
      There are a few injection points already declared at strategic points
      within the server code. After adding a new injection point the code needs
      to be compiled in order for that injection point to be available in the
      binary. Add-ins written in C-language can declare injection points in
-     their own code using the same macro.
+     their own code using the same macro. It is possible to pass down arguments
+     to callbacks for runtime checks.
     </para>
 
     <para>
@@ -3622,12 +3624,14 @@ INJECTION_POINT(name);
 <programlisting>
 extern void InjectionPointAttach(const char *name,
                                  const char *library,
-                                 const char *function);
+                                 const char *function,
+                                 int num_args);
 </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>num_args</literal> is
+     the number of arguments that are used by the callback.
     </para>
 
     <para>
-- 
2.43.0