Merge pull request #49 from rjuju/shmem_request_hook

shinderuk · web-flow · commit af532056ced2 · 2022-05-20T15:46:17.000+03:00
Various fixes for compatibility with pg 15
diff --git a/collector.c b/collector.c
@@ -44,7 +44,7 @@ register_wait_collector(void)
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS;
 	worker.bgw_start_time = BgWorkerStart_ConsistentState;
-	worker.bgw_restart_time = 0;
+	worker.bgw_restart_time = 1;
 	worker.bgw_notify_pid = 0;
 	snprintf(worker.bgw_library_name, BGW_MAXLEN, "pg_wait_sampling");
 	snprintf(worker.bgw_function_name, BGW_MAXLEN, CppAsString(collector_main));
@@ -339,8 +339,16 @@ collector_main(Datum main_arg)
 	 * any equivalent of the backend's command-read loop, where interrupts can
 	 * be processed immediately, so make sure ImmediateInterruptOK is turned
 	 * off.
+	 *
+	 * We also want to respond to the ProcSignal notifications.  This is done
+	 * in the upstream provided procsignal_sigusr1_handler, which is
+	 * automatically used if a bgworker connects to a database.  But since our
+	 * worker doesn't connect to any database even though it calls
+	 * InitPostgres, which will still initializze a new backend and thus
+	 * partitipate to the ProcSignal infrastructure.
 	 */
 	pqsignal(SIGTERM, handle_sigterm);
+	pqsignal(SIGUSR1, procsignal_sigusr1_handler);
 	BackgroundWorkerUnblockSignals();
 
 #if PG_VERSION_NUM >= 110000
@@ -379,6 +387,9 @@ collector_main(Datum main_arg)
 		bool			write_history,
 						write_profile;
 
+		/* We need an explicit call for at least ProcSignal notifications. */
+		CHECK_FOR_INTERRUPTS();
+
 		/* Wait calculate time to next sample for history or profile */
 		current_ts = GetCurrentTimestamp();
 
diff --git a/pg_wait_sampling.c b/pg_wait_sampling.c
@@ -39,7 +39,6 @@
 PG_MODULE_MAGIC;
 
 void		_PG_init(void);
-void		_PG_fini(void);
 
 /* Global variables */
 bool					shmem_initialized = false;
@@ -59,6 +58,9 @@ shm_mq		   *recv_mq = NULL;
 shm_mq_handle  *recv_mqh = NULL;
 LOCKTAG			queueTag;
 
+#if PG_VERSION_NUM >= 150000
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+#endif
 static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
 static PGPROC * search_proc(int backendPid);
 static PlannedStmt *pgws_planner_hook(Query *parse,
@@ -74,28 +76,40 @@ static void pgws_ExecutorEnd(QueryDesc *queryDesc);
  * The value has to be in sync with ProcGlobal->allProcCount, initialized in
  * InitProcGlobal() (proc.c).
  *
- * We calculate the value here as it won't initialized when we need it during
- * _PG_init().
- *
- * Note that the value returned during _PG_init() might be different from the
- * value returned later if some third-party modules change one of the
- * underlying GUC.  This isn't ideal but can't lead to a crash, as the value
- * returned during _PG_init() is only used to ask for additional shmem with
- * RequestAddinShmemSpace(), and postgres has an extra 100kB of shmem to
- * compensate some small unaccounted usage.  So if the value later changes, we
- * will allocate and initialize the new (and correct) memory size, which
- * will either work thanks for the extra 100kB of shmem, of fail (and prevent
- * postgres startup) due to an out of shared memory error.
  */
 static int
 get_max_procs_count(void)
 {
 	int count = 0;
 
+	/* First, add the maximum number of backends (MaxBackends). */
+#if PG_VERSION_NUM >= 150000
 	/*
-	 * MaxBackends: bgworkers, autovacuum workers and launcher.
+	 * On pg15+, we can directly access the MaxBackends variable, as it will
+	 * have already been initialized in shmem_request_hook.
+	 */
+	Assert(MaxBackends > 0);
+	count += MaxBackends;
+#else
+	/*
+	 * On older versions, we need to compute MaxBackends: bgworkers, autovacuum
+	 * workers and launcher.
 	 * This has to be in sync with the value computed in
 	 * InitializeMaxBackends() (postinit.c)
+	 *
+	 * Note that we need to calculate the value as it won't initialized when we
+	 * need it during _PG_init().
+	 *
+	 * Note also that the value returned during _PG_init() might be different
+	 * from the value returned later if some third-party modules change one of
+	 * the underlying GUC.  This isn't ideal but can't lead to a crash, as the
+	 * value returned during _PG_init() is only used to ask for additional
+	 * shmem with RequestAddinShmemSpace(), and postgres has an extra 100kB of
+	 * shmem to compensate some small unaccounted usage.  So if the value later
+	 * changes, we will allocate and initialize the new (and correct) memory
+	 * size, which will either work thanks for the extra 100kB of shmem, of
+	 * fail (and prevent postgres startup) due to an out of shared memory
+	 * error.
 	 */
 	count += MaxConnections + autovacuum_max_workers + 1
 			+ max_worker_processes;
@@ -106,9 +120,11 @@ get_max_procs_count(void)
 	 */
 #if PG_VERSION_NUM >= 120000
 	count += max_wal_senders;
-#endif
+#endif		/* pg 12+ */
+#endif		/* pg 15- */
+	/* End of MaxBackends calculation. */
 
-	/* AuxiliaryProcs */
+	/* Add AuxiliaryProcs */
 	count += NUM_AUXILIARY_PROCS;
 
 	return count;
@@ -266,6 +282,23 @@ setup_gucs()
 	}
 }
 
+#if PG_VERSION_NUM >= 150000
+/*
+ * shmem_request hook: request additional shared memory resources.
+ *
+ * If you change code here, don't forget to also report the modifications in
+ * _PG_init() for pg14 and below.
+ */
+static void
+pgws_shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(pgws_shmem_size());
+}
+#endif
+
 /*
  * Distribute shared memory.
  */
@@ -345,18 +378,27 @@ _PG_init(void)
 	if (!process_shared_preload_libraries_in_progress)
 		return;
 
+#if PG_VERSION_NUM < 150000
 	/*
 	 * Request additional shared resources.  (These are no-ops if we're not in
 	 * the postmaster process.)  We'll allocate or attach to the shared
 	 * resources in pgws_shmem_startup().
+	 *
+	 * If you change code here, don't forget to also report the modifications
+	 * in pgsp_shmem_request() for pg15 and later.
 	 */
 	RequestAddinShmemSpace(pgws_shmem_size());
+#endif
 
 	register_wait_collector();
 
 	/*
 	 * Install hooks.
 	 */
+#if PG_VERSION_NUM >= 150000
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook		= pgws_shmem_request;
+#endif
 	prev_shmem_startup_hook = shmem_startup_hook;
 	shmem_startup_hook		= pgws_shmem_startup;
 	planner_hook_next		= planner_hook;
@@ -365,16 +407,6 @@ _PG_init(void)
 	ExecutorEnd_hook		= pgws_ExecutorEnd;
 }
 
-/*
- * Module unload callback
- */
-void
-_PG_fini(void)
-{
-	/* Uninstall hooks. */
-	shmem_startup_hook = prev_shmem_startup_hook;
-}
-
 /*
  * Find PGPROC entry responsible for given pid assuming ProcArrayLock was
  * already taken.
