Rethink the wait event names for postgres_fdw, dblink and etc

From 025525eae164e33117d94f2a90a4219808072b3c Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <mshr.ikeda@ntt.com>
Date: Mon, 21 Aug 2023 10:36:10 +0900
Subject: [PATCH] Make to use custom wait events for modules

---
 contrib/dblink/dblink.c                       | 16 +++++++-
 contrib/dblink/expected/dblink.out            |  9 +++++
 contrib/dblink/sql/dblink.sql                 |  4 ++
 contrib/pg_prewarm/autoprewarm.c              | 22 ++++++++++-
 contrib/postgres_fdw/connection.c             | 23 +++++++++--
 .../postgres_fdw/expected/postgres_fdw.out    | 12 ++++++
 contrib/postgres_fdw/sql/postgres_fdw.sql     |  7 ++++
 doc/src/sgml/dblink.sgml                      | 28 +++++++++++++
 doc/src/sgml/postgres-fdw.sgml                | 39 +++++++++++++++++++
 doc/src/sgml/xfunc.sgml                       |  6 +--
 src/test/modules/test_shm_mq/setup.c          |  9 ++++-
 src/test/modules/test_shm_mq/test.c           |  9 ++++-
 .../modules/worker_spi/t/001_worker_spi.pl    | 12 +++---
 src/test/modules/worker_spi/worker_spi.c      |  2 +-
 14 files changed, 179 insertions(+), 19 deletions(-)

diff --git a/contrib/dblink/dblink.c b/contrib/dblink/dblink.c
index 41e1f6c91d..104d6e9f37 100644
--- a/contrib/dblink/dblink.c
+++ b/contrib/dblink/dblink.c
@@ -130,6 +130,10 @@ static void restoreLocalGucs(int nestlevel);
 static remoteConn *pconn = NULL;
 static HTAB *remoteConnHash = NULL;
 
+/* value cached, fetched from shared memory */
+static uint32 dblink_we_get_conn = 0;
+static uint32 dblink_we_conn = 0;
+
 /*
  *	Following is list that holds multiple remote connections.
  *	Calling convention of each dblink function changes to accept
@@ -202,8 +206,12 @@ dblink_get_conn(char *conname_or_str,
 			connstr = conname_or_str;
 		dblink_connstr_check(connstr);
 
+		/* first time, allocate or get the custom wait event */
+		if (dblink_we_get_conn == 0)
+			dblink_we_get_conn = WaitEventExtensionNew("DblinkGetConnect");
+
 		/* OK to make connection */
-		conn = libpqsrv_connect(connstr, WAIT_EVENT_EXTENSION);
+		conn = libpqsrv_connect(connstr, dblink_we_get_conn);
 
 		if (PQstatus(conn) == CONNECTION_BAD)
 		{
@@ -292,8 +300,12 @@ dblink_connect(PG_FUNCTION_ARGS)
 	/* check password in connection string if not superuser */
 	dblink_connstr_check(connstr);
 
+	/* first time, allocate or get the custom wait event */
+	if (dblink_we_conn == 0)
+		dblink_we_conn = WaitEventExtensionNew("DblinkConnect");
+
 	/* OK to make connection */
-	conn = libpqsrv_connect(connstr, WAIT_EVENT_EXTENSION);
+	conn = libpqsrv_connect(connstr, dblink_we_conn);
 
 	if (PQstatus(conn) == CONNECTION_BAD)
 	{
diff --git a/contrib/dblink/expected/dblink.out b/contrib/dblink/expected/dblink.out
index 7809f58d96..c17c7b1361 100644
--- a/contrib/dblink/expected/dblink.out
+++ b/contrib/dblink/expected/dblink.out
@@ -1209,6 +1209,15 @@ SHOW intervalstyle;
  postgres
 (1 row)
 
+-- Check custom wait events are registered
+SELECT name FROM pg_wait_events WHERE name ~ '^Dblink'
+ORDER BY name COLLATE "C";
+       name       
+------------------
+ DblinkConnect
+ DblinkGetConnect
+(2 rows)
+
 -- Clean up GUC-setting tests
 SELECT dblink_disconnect('myconn');
  dblink_disconnect 
diff --git a/contrib/dblink/sql/dblink.sql b/contrib/dblink/sql/dblink.sql
index 7870ce5d5a..519b3f5266 100644
--- a/contrib/dblink/sql/dblink.sql
+++ b/contrib/dblink/sql/dblink.sql
@@ -627,6 +627,10 @@ FROM dblink_fetch('myconn','error_cursor', 1) AS t(i int);
 SHOW datestyle;
 SHOW intervalstyle;
 
+-- Check custom wait events are registered
+SELECT name FROM pg_wait_events WHERE name ~ '^Dblink'
+ORDER BY name COLLATE "C";
+
 -- Clean up GUC-setting tests
 SELECT dblink_disconnect('myconn');
 RESET datestyle;
diff --git a/contrib/pg_prewarm/autoprewarm.c b/contrib/pg_prewarm/autoprewarm.c
index d0efc9e524..2160d40705 100644
--- a/contrib/pg_prewarm/autoprewarm.c
+++ b/contrib/pg_prewarm/autoprewarm.c
@@ -105,6 +105,16 @@ static AutoPrewarmSharedState *apw_state = NULL;
 static bool autoprewarm = true; /* start worker? */
 static int	autoprewarm_interval = 300; /* dump interval */
 
+/*
+ * Cached custom wait events, fetched from shared memory.
+ *
+ * Note that they are not reported on pg_stat_activity. pgstat_bestart() isn't
+ * called because the autoprewarm worker isn't an auxiliary process and it
+ * doesn't connect to a database.
+ */
+static uint32 autoprewarm_we_shutdown = 0;
+static uint32 autoprewarm_we_delay = 0;
+
 /*
  * Module load callback.
  */
@@ -233,11 +243,15 @@ autoprewarm_main(Datum main_arg)
 
 		if (autoprewarm_interval <= 0)
 		{
+			/* first time, allocate or get the custom wait event */
+			if (autoprewarm_we_shutdown == 0)
+				autoprewarm_we_shutdown = WaitEventExtensionNew("PgPrewarmDumpShutdown");
+
 			/* We're only dumping at shutdown, so just wait forever. */
 			(void) WaitLatch(MyLatch,
 							 WL_LATCH_SET | WL_EXIT_ON_PM_DEATH,
 							 -1L,
-							 WAIT_EVENT_EXTENSION);
+							autoprewarm_we_shutdown);
 		}
 		else
 		{
@@ -260,11 +274,15 @@ autoprewarm_main(Datum main_arg)
 				continue;
 			}
 
+			/* first time, allocate or get the custom wait event */
+			if (autoprewarm_we_delay == 0)
+				autoprewarm_we_delay = WaitEventExtensionNew("PgPrewarmDumpDelay");
+
 			/* Sleep until the next dump time. */
 			(void) WaitLatch(MyLatch,
 							 WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
 							 delay_in_ms,
-							 WAIT_EVENT_EXTENSION);
+							 autoprewarm_we_delay);
 		}
 
 		/* Reset the latch, loop. */
diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 7e12b722ec..7d6fbc56a1 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -83,6 +83,11 @@ static unsigned int prep_stmt_number = 0;
 /* tracks whether any work is needed in callback functions */
 static bool xact_got_connection = false;
 
+/* value cached, fetched from shared memory */
+static uint32 pgfdw_we_connect = 0;
+static uint32 pgfdw_we_receive = 0;
+static uint32 pgfdw_we_cleanup_receive = 0;
+
 /*
  * Milliseconds to wait to cancel an in-progress query or execute a cleanup
  * query; if it takes longer than 30 seconds to do these, we assume the
@@ -527,10 +532,14 @@ connect_pg_server(ForeignServer *server, UserMapping *user)
 		/* verify the set of connection parameters */
 		check_conn_params(keywords, values, user);
 
+		/* first time, allocate or get the custom wait event */
+		if (pgfdw_we_connect == 0)
+			pgfdw_we_connect = WaitEventExtensionNew("PostgresFdwConnect");
+
 		/* OK to make connection */
 		conn = libpqsrv_connect_params(keywords, values,
 									   false,	/* expand_dbname */
-									   WAIT_EVENT_EXTENSION);
+									   pgfdw_we_connect);
 
 		if (!conn || PQstatus(conn) != CONNECTION_OK)
 			ereport(ERROR,
@@ -858,12 +867,16 @@ pgfdw_get_result(PGconn *conn, const char *query)
 			{
 				int			wc;
 
+				/* first time, allocate or get the custom wait event */
+				if (pgfdw_we_receive == 0)
+					pgfdw_we_receive = WaitEventExtensionNew("PostgresFdwReceive");
+
 				/* Sleep until there's something to do */
 				wc = WaitLatchOrSocket(MyLatch,
 									   WL_LATCH_SET | WL_SOCKET_READABLE |
 									   WL_EXIT_ON_PM_DEATH,
 									   PQsocket(conn),
-									   -1L, WAIT_EVENT_EXTENSION);
+									   -1L, pgfdw_we_receive);
 				ResetLatch(MyLatch);
 
 				CHECK_FOR_INTERRUPTS();
@@ -1562,12 +1575,16 @@ pgfdw_get_cleanup_result(PGconn *conn, TimestampTz endtime, PGresult **result,
 					goto exit;
 				}
 
+				/* first time, allocate or get the custom wait event */
+				if (pgfdw_we_cleanup_receive == 0)
+					pgfdw_we_cleanup_receive = WaitEventExtensionNew("PostgresFdwCleanupReceive");
+
 				/* Sleep until there's something to do */
 				wc = WaitLatchOrSocket(MyLatch,
 									   WL_LATCH_SET | WL_SOCKET_READABLE |
 									   WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
 									   PQsocket(conn),
-									   cur_timeout, WAIT_EVENT_EXTENSION);
+									   cur_timeout, pgfdw_we_cleanup_receive);
 				ResetLatch(MyLatch);
 
 				CHECK_FOR_INTERRUPTS();
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 77df7eb8e4..616fb55d8b 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -11901,6 +11901,18 @@ ALTER SERVER loopback OPTIONS (SET analyze_sampling 'random');
 ANALYZE analyze_table;
 ALTER SERVER loopback OPTIONS (SET analyze_sampling 'off');
 ANALYZE analyze_table;
+-- ===================================================================
+-- test for custom wait events
+-- ===================================================================
+SELECT name FROM pg_wait_events WHERE name ~ '^PostgresFdw'
+ORDER BY name COLLATE "C";
+           name            
+---------------------------
+ PostgresFdwCleanupReceive
+ PostgresFdwConnect
+ PostgresFdwReceive
+(3 rows)
+
 -- cleanup
 DROP FOREIGN TABLE analyze_ftable;
 DROP TABLE analyze_table;
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index cfb1b57e33..bab922aeae 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -4039,6 +4039,13 @@ ANALYZE analyze_table;
 ALTER SERVER loopback OPTIONS (SET analyze_sampling 'off');
 ANALYZE analyze_table;
 
+-- ===================================================================
+-- test for custom wait events
+-- ===================================================================
+
+SELECT name FROM pg_wait_events WHERE name ~ '^PostgresFdw'
+ORDER BY name COLLATE "C";
+
 -- cleanup
 DROP FOREIGN TABLE analyze_ftable;
 DROP TABLE analyze_table;
diff --git a/doc/src/sgml/dblink.sgml b/doc/src/sgml/dblink.sgml
index 7d25f24f49..7800b80bdb 100644
--- a/doc/src/sgml/dblink.sgml
+++ b/doc/src/sgml/dblink.sgml
@@ -13,6 +13,34 @@
   session.
  </para>
 
+ <para>
+  <filename>dblink</filename> can report the following wait events under the wait
+  event type <literal>Extension</literal>.
+ </para>
+
+ <variablelist>
+  <varlistentry>
+   <term><literal>DblinkConnect</literal></term>
+   <listitem>
+    <para>
+     Waiting to establish a connection to a remote server. It can be reported when calling
+     <function>dblink_connect</function>.
+    </para>
+   </listitem>
+  </varlistentry>
+
+  <varlistentry>
+   <term><literal>DblinkGetConnect</literal></term>
+   <listitem>
+    <para>
+     Waiting to establish a connection to a remote server when it can't look up the
+     specified persistent connection's name. It can be reported when calling
+     <function>dblink</function> and <function>dblink_exec</function>.
+    </para>
+   </listitem>
+  </varlistentry>
+ </variablelist>
+
  <para>
   See also <xref linkend="postgres-fdw"/>, which provides roughly the same
   functionality using a more modern and standards-compliant infrastructure.
diff --git a/doc/src/sgml/postgres-fdw.sgml b/doc/src/sgml/postgres-fdw.sgml
index 5062d712e7..f38136e94f 100644
--- a/doc/src/sgml/postgres-fdw.sgml
+++ b/doc/src/sgml/postgres-fdw.sgml
@@ -1042,6 +1042,45 @@ postgres=# SELECT postgres_fdw_disconnect_all();
   </para>
  </sect2>
 
+ <sect2 id="postgres-fdw-wait-events">
+  <title>Wait Events</title>
+
+  <para>
+   <filename>postgres_fdw</filename> can report the following wait events under the wait
+   event type <literal>Extension</literal>.
+  </para>
+
+  <variablelist>
+   <varlistentry>
+    <term><literal>PostgresFdwConnect</literal></term>
+    <listitem>
+     <para>
+      Waiting to establish a connection to a remote server.
+     </para>
+    </listitem>
+   </varlistentry>
+
+   <varlistentry>
+    <term><literal>PostgresFdwReceive</literal></term>
+    <listitem>
+     <para>
+      Waiting to receive the results of a query from a remote server.
+     </para>
+    </listitem>
+   </varlistentry>
+
+   <varlistentry>
+    <term><literal>PostgresFdwCleanupReceive</literal></term>
+    <listitem>
+     <para>
+      Waiting for transaction abort on remote server.
+     </para>
+    </listitem>
+   </varlistentry>
+  </variablelist>
+
+ </sect2>
+
  <sect2 id="postgres-fdw-configuration-parameters">
   <title>Configuration Parameters</title>
 
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 281c178b0e..96ba95818c 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3472,9 +3472,9 @@ uint32 WaitEventExtensionNew(const char *wait_event_name)
 <screen>
 =# SELECT wait_event_type, wait_event FROM pg_stat_activity
      WHERE backend_type ~ 'worker_spi';
- wait_event_type |   wait_event
------------------+-----------------
- Extension       | worker_spi_main
+ wait_event_type |  wait_event
+-----------------+---------------
+ Extension       | WorkerSpiMain
 (1 row)
 </screen>
     </para>
diff --git a/src/test/modules/test_shm_mq/setup.c b/src/test/modules/test_shm_mq/setup.c
index 192e5cc2ab..abc79352b9 100644
--- a/src/test/modules/test_shm_mq/setup.c
+++ b/src/test/modules/test_shm_mq/setup.c
@@ -40,6 +40,9 @@ static void wait_for_workers_to_become_ready(worker_state *wstate,
 											 volatile test_shm_mq_header *hdr);
 static bool check_worker_status(worker_state *wstate);
 
+/* value cached, fetched from shared memory */
+static uint32 we_bgworker_startup = 0;
+
 /*
  * Set up a dynamic shared memory segment and zero or more background workers
  * for a test run.
@@ -278,9 +281,13 @@ wait_for_workers_to_become_ready(worker_state *wstate,
 			break;
 		}
 
+		/* first time, allocate or get the custom wait event */
+		if (we_bgworker_startup == 0)
+			we_bgworker_startup = WaitEventExtensionNew("TestShmMqBgWorkerStartup");
+
 		/* Wait to be signaled. */
 		(void) WaitLatch(MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0,
-						 WAIT_EVENT_EXTENSION);
+						 we_bgworker_startup);
 
 		/* Reset the latch so we don't spin. */
 		ResetLatch(MyLatch);
diff --git a/src/test/modules/test_shm_mq/test.c b/src/test/modules/test_shm_mq/test.c
index d9be703350..cb52a680a5 100644
--- a/src/test/modules/test_shm_mq/test.c
+++ b/src/test/modules/test_shm_mq/test.c
@@ -28,6 +28,9 @@ PG_FUNCTION_INFO_V1(test_shm_mq_pipelined);
 static void verify_message(Size origlen, char *origdata, Size newlen,
 						   char *newdata);
 
+/* value cached, fetched from shared memory */
+static uint32 we_message_queue = 0;
+
 /*
  * Simple test of the shared memory message queue infrastructure.
  *
@@ -225,6 +228,10 @@ test_shm_mq_pipelined(PG_FUNCTION_ARGS)
 
 		if (wait)
 		{
+			/* first time, allocate or get the custom wait event */
+			if (we_message_queue == 0)
+				we_message_queue = WaitEventExtensionNew("TestShmMqMessageQueue");
+
 			/*
 			 * If we made no progress, wait for one of the other processes to
 			 * which we are connected to set our latch, indicating that they
@@ -232,7 +239,7 @@ test_shm_mq_pipelined(PG_FUNCTION_ARGS)
 			 * for us to do.
 			 */
 			(void) WaitLatch(MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0,
-							 WAIT_EVENT_EXTENSION);
+							 we_message_queue);
 			ResetLatch(MyLatch);
 			CHECK_FOR_INTERRUPTS();
 		}
diff --git a/src/test/modules/worker_spi/t/001_worker_spi.pl b/src/test/modules/worker_spi/t/001_worker_spi.pl
index 2965acd789..f3ef395554 100644
--- a/src/test/modules/worker_spi/t/001_worker_spi.pl
+++ b/src/test/modules/worker_spi/t/001_worker_spi.pl
@@ -43,15 +43,15 @@ is($result, qq(total|1), 'dynamic bgworker correctly consumed tuple data');
 $result = $node->poll_query_until(
 	'postgres',
 	qq[SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
-	qq[worker_spi_main]);
+	qq[WorkerSpiMain]);
 is($result, 1,
-	'dynamic bgworker has reported "worker_spi_main" as wait event');
+	'dynamic bgworker has reported "WorkerSpiMain" as wait event');
 
 # Check the wait event used by the dynamic bgworker appears in pg_wait_events
 $result = $node->safe_psql('postgres',
-	q[SELECT count(*) > 0 from pg_wait_events where type = 'Extension' and name = 'worker_spi_main';]
+	q[SELECT count(*) > 0 from pg_wait_events where type = 'Extension' and name = 'WorkerSpiMain';]
 );
-is($result, 't', '"worker_spi_main" is reported in pg_wait_events');
+is($result, 't', '"WorkerSpiMain" is reported in pg_wait_events');
 
 note "testing bgworkers loaded with shared_preload_libraries";
 
@@ -74,7 +74,7 @@ ok( $node->poll_query_until(
 		'mydb',
 		qq[SELECT datname, count(datname), wait_event FROM pg_stat_activity
             WHERE backend_type = 'worker_spi' GROUP BY datname, wait_event;],
-		'mydb|3|worker_spi_main'),
+		'mydb|3|WorkerSpiMain'),
 	'bgworkers all launched'
 ) or die "Timed out while waiting for bgworkers to be launched";
 
@@ -89,7 +89,7 @@ ok( $node->poll_query_until(
 		qq[SELECT datname, count(datname), wait_event FROM pg_stat_activity
             WHERE backend_type = 'worker_spi dynamic' AND
             pid IN ($worker1_pid, $worker2_pid) GROUP BY datname, wait_event;],
-		'mydb|2|worker_spi_main'),
+		'mydb|2|WorkerSpiMain'),
 	'dynamic bgworkers all launched'
 ) or die "Timed out while waiting for dynamic bgworkers to be launched";
 
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index 98f8d4194b..2e3114990e 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -194,7 +194,7 @@ worker_spi_main(Datum main_arg)
 
 		/* First time, allocate or get the custom wait event */
 		if (worker_spi_wait_event_main == 0)
-			worker_spi_wait_event_main = WaitEventExtensionNew("worker_spi_main");
+			worker_spi_wait_event_main = WaitEventExtensionNew("WorkerSpiMain");
 
 		/*
 		 * Background workers mustn't call usleep() or any direct equivalent:
-- 
2.25.1