May 20, 2022 · May 14, 2022 · May 14, 2022 · May 14, 2022 · May 14, 2022
diff --git a/collector.c b/collector.c
 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));
 * 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
 boolwrite_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
 PG_MODULE_MAGIC;

 void_PG_init(void);
 void_PG_fini(void);

 /* Global variables */
 boolshmem_initialized = false;
 shm_mq_handle  *recv_mqh = NULL;
 LOCKTAGqueueTag;

 #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,
 * 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;
 */
 #if PG_VERSION_NUM >= 120000
 count += max_wal_senders;
 #endif
 #endif/* pg 12+ */
 #endif/* pg 15- */
 /* End of MaxBackends calculation. */

 /* AuxiliaryProcs */
 /*AddAuxiliaryProcs */
 count += NUM_AUXILIARY_PROCS;

 return count;
 }
 }

 #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.
 */
 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;
 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.
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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));
Expand DownExpand Up		@@ -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
Expand DownExpand Up		@@ -379,6 +387,9 @@ collector_main(Datum main_arg)
		boolwrite_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();

Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -39,7 +39,6 @@
		PG_MODULE_MAGIC;

		void_PG_init(void);
		void_PG_fini(void);

		/* Global variables */
		boolshmem_initialized = false;
Expand All		@@ -59,6 +58,9 @@ shm_mq *recv_mq = NULL;
		shm_mq_handle *recv_mqh = NULL;
		LOCKTAGqueueTag;

		#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,
Expand All		@@ -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
maksm90 marked this conversation as resolved. Show resolvedHide resolved
		* 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;
Expand All		@@ -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 */
		/AddAuxiliaryProcs /
		count += NUM_AUXILIARY_PROCS;

		return count;
Expand DownExpand Up		@@ -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.
		*/
Expand DownExpand Up		@@ -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;
Expand All		@@ -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.
Expand Down