From c789f0f6cc5da084f75f2556bada572a1a05e39a Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Thu, 5 Oct 2023 10:23:22 +0900
Subject: [PATCH] dblink: Replace WAIT_EVENT_EXTENSION with custom wait events

Two custom wait events are added here:
- "DblinkConnect", when waiting to establish a connection to a remote
server.
- "DblinkGetConnect", when waiting to establish a connection to a remote
server but it could not be found in the list of already-opened ones.

Author: Masahiro Ikeda
Discussion: https://postgr.es/m/197bce267fa691a0ac62c86c4ab904c4@oss.nttdata.com
---
 contrib/dblink/dblink.c  | 16 ++++++++++++++--
 doc/src/sgml/dblink.sgml | 26 ++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/contrib/dblink/dblink.c b/contrib/dblink/dblink.c
index 41e1f6c91d..195b278f55 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;
 
+/* custom wait event values, retrieved from shared memory */
+static uint32 dblink_we_connect = 0;
+static uint32 dblink_we_get_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_connect == 0)
+		dblink_we_connect = WaitEventExtensionNew("DblinkConnect");
+
 	/* OK to make connection */
-	conn = libpqsrv_connect(connstr, WAIT_EVENT_EXTENSION);
+	conn = libpqsrv_connect(connstr, dblink_we_connect);
 
 	if (PQstatus(conn) == CONNECTION_BAD)
 	{
diff --git a/doc/src/sgml/dblink.sgml b/doc/src/sgml/dblink.sgml
index 7d25f24f49..e8de5a65cb 100644
--- a/doc/src/sgml/dblink.sgml
+++ b/doc/src/sgml/dblink.sgml
@@ -13,6 +13,32 @@
   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.
+    </para>
+   </listitem>
+  </varlistentry>
+
+  <varlistentry>
+   <term><literal>DblinkGetConnect</literal></term>
+   <listitem>
+    <para>
+     Waiting to establish a connection to a remote server when it could not
+     be found in the list of already-opened connections.
+    </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.