From 87fd8c92ee5584e8a0703f03d052a6161f4b746a Mon Sep 17 00:00:00 2001
From: Pavel Borisov <pashkin.elfe@gmail.com>
Date: Thu, 3 Sep 2020 15:06:52 +0400
Subject: [PATCH v3] Custom signal handler for extensions

The patch allows to define custom signals and their handlers from extension
side. A fixed set of signals can be reserved on postgres initialization
stage (in _PG_init function of shared preloaded module) through specific
interface functions. The relationship between custom signals and assigned
handlers is replicated from postmaster to child processes, including working
backends. The handler on a remote backend called by this signals (via
SendProcSignal function) comes into action at CHECK_FOR_INTERRUPTS point.

This feature is used by pg_query_state extension
https://github.com/postgrespro/pg_query_state and this API can also be usedful
for profiling extensions etc.

Authors: Maksim Milyutin, Alexey Kondratov
Discussion: https://postgrespro.ru/list/thread-id/2362085#590e3bc2-0eb3-322c-3c1f-4e79ff562bfe@gmail.com
---
 src/backend/storage/ipc/procsignal.c | 93 ++++++++++++++++++++++++++++
 src/backend/tcop/postgres.c          |  2 +
 src/include/storage/procsignal.h     | 20 ++++++
 3 files changed, 115 insertions(+)

diff --git a/src/backend/storage/ipc/procsignal.c b/src/backend/storage/ipc/procsignal.c
index 4fa385b0ec..3e8328d341 100644
--- a/src/backend/storage/ipc/procsignal.c
+++ b/src/backend/storage/ipc/procsignal.c
@@ -83,6 +83,12 @@ typedef struct
  */
 #define NumProcSignalSlots	(MaxBackends + NUM_AUXPROCTYPES)
 
+#define IsCustomProcSignalReason(reason) \
+	((reason) >= PROCSIG_CUSTOM_1 && (reason) <= PROCSIG_CUSTOM_N)
+
+static bool CustomSignalPendings[NUM_CUSTOM_PROCSIGNALS];
+static ProcSignalHandler_type CustomInterruptHandlers[NUM_CUSTOM_PROCSIGNALS];
+
 /* Check whether the relevant type bit is set in the flags. */
 #define BARRIER_SHOULD_CHECK(flags, type) \
 	(((flags) & (((uint32) 1) << (uint32) (type))) != 0)
@@ -92,6 +98,9 @@ static volatile ProcSignalSlot *MyProcSignalSlot = NULL;
 
 static bool CheckProcSignal(ProcSignalReason reason);
 static void CleanupProcSignalState(int status, Datum arg);
+
+static void CheckAndSetCustomSignalInterrupts(void);
+
 static void ProcessBarrierPlaceholder(void);
 
 /*
@@ -235,6 +244,36 @@ CleanupProcSignalState(int status, Datum arg)
 	slot->pss_pid = 0;
 }
 
+/*
+ * RegisterCustomProcSignalHandler
+ *		Assign specific handler of custom process signal with new
+ *		ProcSignalReason key.
+ *
+ * This function has to be called in _PG_init function of extensions at the
+ * stage of loading shared preloaded libraries. Otherwise it throws fatal error.
+ *
+ * Return INVALID_PROCSIGNAL if all slots for custom signals are occupied.
+ */
+ProcSignalReason
+RegisterCustomProcSignalHandler(ProcSignalHandler_type handler)
+{
+	ProcSignalReason reason;
+
+	if (!process_shared_preload_libraries_in_progress)
+		ereport(FATAL, (errcode(ERRCODE_INTERNAL_ERROR),
+						errmsg("cannot register custom signal after startup")));
+
+	/* Iterate through custom signal slots to find a free one */
+	for (reason = PROCSIG_CUSTOM_1; reason <= PROCSIG_CUSTOM_N; reason++)
+		if (!CustomInterruptHandlers[reason - PROCSIG_CUSTOM_1])
+		{
+			CustomInterruptHandlers[reason - PROCSIG_CUSTOM_1] = handler;
+			return reason;
+		}
+
+	return INVALID_PROCSIGNAL;
+}
+
 /*
  * SendProcSignal
  *		Send a signal to a Postgres process
@@ -585,9 +624,63 @@ procsignal_sigusr1_handler(SIGNAL_ARGS)
 	if (CheckProcSignal(PROCSIG_RECOVERY_CONFLICT_BUFFERPIN))
 		RecoveryConflictInterrupt(PROCSIG_RECOVERY_CONFLICT_BUFFERPIN);
 
+	CheckAndSetCustomSignalInterrupts();
+
 	SetLatch(MyLatch);
 
 	latch_sigusr1_handler();
 
 	errno = save_errno;
 }
+
+/*
+ * Handle receipt of an interrupt indicating any of custom process signals.
+ */
+static void
+CheckAndSetCustomSignalInterrupts()
+{
+	ProcSignalReason	reason;
+
+	for (reason = PROCSIG_CUSTOM_1; reason <= PROCSIG_CUSTOM_N; reason++)
+	{
+		if (CheckProcSignal(reason))
+		{
+
+			/* set interrupt flags */
+			InterruptPending = true;
+			CustomSignalPendings[reason - PROCSIG_CUSTOM_1] = true;
+		}
+	}
+
+	SetLatch(MyLatch);
+}
+
+/*
+ * CheckAndHandleCustomSignals
+ *		Check custom signal flags and call handler assigned to that signal
+ *		if it is not NULL
+ *
+ * This function is called within CHECK_FOR_INTERRUPTS if interrupt occurred.
+ */
+void
+CheckAndHandleCustomSignals(void)
+{
+	int i;
+
+	/* Disable interrupts to avoid recursive calls */
+	HOLD_INTERRUPTS();
+
+	/* Check on expiring of custom signals and call its handlers if exist */
+	for (i = 0; i < NUM_CUSTOM_PROCSIGNALS; i++)
+		if (CustomSignalPendings[i])
+		{
+			ProcSignalHandler_type  handler;
+
+			CustomSignalPendings[i] = false;
+			handler = CustomInterruptHandlers[i];
+			if (handler != NULL)
+				handler();
+		}
+
+	RESUME_INTERRUPTS();
+}
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index c9424f167c..ea330c9045 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -3209,6 +3209,8 @@ ProcessInterrupts(void)
 
 	if (ParallelMessagePending)
 		HandleParallelMessages();
+
+	CheckAndHandleCustomSignals();
 }
 
 
diff --git a/src/include/storage/procsignal.h b/src/include/storage/procsignal.h
index 5cb39697f3..14cddb99de 100644
--- a/src/include/storage/procsignal.h
+++ b/src/include/storage/procsignal.h
@@ -16,6 +16,7 @@
 
 #include "storage/backendid.h"
 
+#define NUM_CUSTOM_PROCSIGNALS 64
 
 /*
  * Reasons for signaling a Postgres child process (a backend or an auxiliary
@@ -29,6 +30,8 @@
  */
 typedef enum
 {
+	INVALID_PROCSIGNAL = -1,	/* Must be first */
+
 	PROCSIG_CATCHUP_INTERRUPT,	/* sinval catchup interrupt */
 	PROCSIG_NOTIFY_INTERRUPT,	/* listen/notify interrupt */
 	PROCSIG_PARALLEL_MESSAGE,	/* message from cooperating parallel backend */
@@ -43,9 +46,20 @@ typedef enum
 	PROCSIG_RECOVERY_CONFLICT_BUFFERPIN,
 	PROCSIG_RECOVERY_CONFLICT_STARTUP_DEADLOCK,
 
+	PROCSIG_CUSTOM_1,
+	/*
+	 * PROCSIG_CUSTOM_2,
+	 * ...,
+	 * PROCSIG_CUSTOM_N-1,
+	 */
+	PROCSIG_CUSTOM_N = PROCSIG_CUSTOM_1 + NUM_CUSTOM_PROCSIGNALS - 1,
+
 	NUM_PROCSIGNALS				/* Must be last! */
 } ProcSignalReason;
 
+/* Handler of custom process signal */
+typedef void (*ProcSignalHandler_type) (void);
+
 typedef enum
 {
 	/*
@@ -63,6 +77,10 @@ extern Size ProcSignalShmemSize(void);
 extern void ProcSignalShmemInit(void);
 
 extern void ProcSignalInit(int pss_idx);
+
+extern ProcSignalReason
+	RegisterCustomProcSignalHandler(ProcSignalHandler_type handler);
+
 extern int	SendProcSignal(pid_t pid, ProcSignalReason reason,
 						   BackendId backendId);
 
@@ -70,6 +88,8 @@ extern uint64 EmitProcSignalBarrier(ProcSignalBarrierType type);
 extern void WaitForProcSignalBarrier(uint64 generation);
 extern void ProcessProcSignalBarrier(void);
 
+extern void CheckAndHandleCustomSignals(void);
+
 extern void procsignal_sigusr1_handler(SIGNAL_ARGS);
 
 #endif							/* PROCSIGNAL_H */
-- 
2.28.0

