v3-0002-Add-test-codes-for-custom-wait-events.patch

From 901b75d31070ef0029557db6981c98e06f5c16c3 Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda.us@hco.ntt.co.jp>
Date: Tue, 20 Jun 2023 17:59:34 +0900
Subject: [PATCH 2/2] Add test codes for custom wait events

---
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 .../test_custom_wait_events/.gitignore        |   2 +
 .../modules/test_custom_wait_events/Makefile  |  23 +++
 .../test_custom_wait_events/meson.build       |  33 ++++
 .../test_custom_wait_events/t/001_basic.pl    |  34 ++++
 .../test_custom_wait_events--1.0.sql          |  14 ++
 .../test_custom_wait_events.c                 | 182 ++++++++++++++++++
 .../test_custom_wait_events.control           |   3 +
 9 files changed, 293 insertions(+)
 create mode 100644 src/test/modules/test_custom_wait_events/.gitignore
 create mode 100644 src/test/modules/test_custom_wait_events/Makefile
 create mode 100644 src/test/modules/test_custom_wait_events/meson.build
 create mode 100644 src/test/modules/test_custom_wait_events/t/001_basic.pl
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.c
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.control

diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 6331c976dc..84312b7e57 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -17,6 +17,7 @@ SUBDIRS = \
 		  test_bloomfilter \
 		  test_copy_callbacks \
 		  test_custom_rmgrs \
+		  test_custom_wait_events \
 		  test_ddl_deparse \
 		  test_extensions \
 		  test_ginpostinglist \
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index 17d369e378..83a1d8fd19 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -14,6 +14,7 @@ subdir('ssl_passphrase_callback')
 subdir('test_bloomfilter')
 subdir('test_copy_callbacks')
 subdir('test_custom_rmgrs')
+subdir('test_custom_wait_events')
 subdir('test_ddl_deparse')
 subdir('test_extensions')
 subdir('test_ginpostinglist')
diff --git a/src/test/modules/test_custom_wait_events/.gitignore b/src/test/modules/test_custom_wait_events/.gitignore
new file mode 100644
index 0000000000..716e17f5a2
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/.gitignore
@@ -0,0 +1,2 @@
+# Generated subdirectories
+/tmp_check/
diff --git a/src/test/modules/test_custom_wait_events/Makefile b/src/test/modules/test_custom_wait_events/Makefile
new file mode 100644
index 0000000000..dda365d523
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/Makefile
@@ -0,0 +1,23 @@
+# src/test/modules/test_custom_wait_events/Makefile
+
+MODULE_big = test_custom_wait_events
+OBJS = \
+	$(WIN32RES) \
+	test_custom_wait_events.o
+PGFILEDESC = "test_custom_wait_events - test custom wait events"
+
+EXTENSION = test_custom_wait_events
+DATA = test_custom_wait_events--1.0.sql
+
+TAP_TESTS = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_custom_wait_events
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_custom_wait_events/meson.build b/src/test/modules/test_custom_wait_events/meson.build
new file mode 100644
index 0000000000..8ad073e577
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) PostgreSQL Global Development Group
+
+test_custom_wait_events = files(
+  'test_custom_wait_events.c',
+)
+
+if host_system == 'windows'
+  test_custom_wait_events += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_custom_wait_events',
+    '--FILEDESC', 'test_custom_wait_events - test custom wait events',])
+endif
+
+test_custom_wait_events = shared_module('test_custom_wait_events',
+  test_custom_wait_events,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_custom_wait_events
+
+test_install_data += files(
+  'test_custom_wait_events.control',
+  'test_custom_wait_events--1.0.sql',
+)
+
+tests += {
+  'name': 'test_custom_wait_events',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'tests': [
+      't/001_basic.pl',
+    ],
+  },
+}
diff --git a/src/test/modules/test_custom_wait_events/t/001_basic.pl b/src/test/modules/test_custom_wait_events/t/001_basic.pl
new file mode 100644
index 0000000000..2809afbd47
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/t/001_basic.pl
@@ -0,0 +1,34 @@
+# Copyright (c) 2023, PostgreSQL Global Development Group
+
+use strict;
+use warnings;
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node = PostgreSQL::Test::Cluster->new('main');
+
+$node->init;
+$node->append_conf(
+	'postgresql.conf', 
+	"shared_preload_libraries = 'test_custom_wait_events'"
+);
+$node->start;
+
+# setup
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_wait_events;');
+
+$node->safe_psql('postgres', 'SELECT wait_worker_launch();');
+
+ok($node->poll_query_until('postgres',
+	qq[SELECT 
+		(SELECT count(*) FROM pg_stat_activity
+	 		WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND wait_event = 'custom_wait_event'
+		) > 0;])
+) or die "Timed out waiting the custom wait event to be showed up";
+
+$node->stop;
+done_testing();
\ No newline at end of file
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
new file mode 100644
index 0000000000..1fc2130f84
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
@@ -0,0 +1,14 @@
+/* src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_custom_wait_events" to load this file. \quit
+
+--
+-- wait_worker_launch()
+--
+-- Launch a background worker to wait forever with a custom wait event.
+--
+CREATE FUNCTION wait_worker_launch()
+RETURNS int
+AS 'MODULE_PATHNAME', 'wait_worker_launch'
+LANGUAGE C STRICT PARALLEL SAFE;
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.c b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
new file mode 100644
index 0000000000..ba3fe45ad6
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
@@ -0,0 +1,182 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_custom_wait_events.c
+ *		Code for testing custom wait events
+ *
+ * This code initializes a custom wait event in shmem_request_hook() and
+ * provide a function to launch a background worker waiting forever
+ * with the custom wait event.
+ * 
+ * Portions Copyright (c) 1996-2023, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_custom_wait_events/test_custom_wait_events.c
+ *
+ * -------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "postmaster/bgworker.h"
+#include "postmaster/interrupt.h"
+#include "storage/ipc.h"
+#include "storage/latch.h"
+#include "storage/lwlock.h"
+#include "storage/proc.h"
+#include "storage/shmem.h"
+#include "tcop/utility.h"
+#include "utils/wait_event.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(wait_worker_launch);
+
+PGDLLEXPORT void worker_main(Datum main_arg);
+
+/* GUC variables */
+typedef struct state
+{
+	uint32			wait_event_info;  /* a allocated wait event */
+}			state;
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static char * worker_database = "postgres";   /* connected by a background worker */
+static state * ss = NULL;
+
+static void shmem_request(void);
+static void shmem_startup(void);
+static Size memsize(void);
+
+void
+_PG_init(void)
+{
+	if (!process_shared_preload_libraries_in_progress)
+		return;
+
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = shmem_startup;
+}
+
+static void
+shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(memsize());
+}
+
+static Size
+memsize(void)
+{
+	return sizeof(state);
+}
+
+/*
+ * Allocate and register a new wait event.
+ *
+ * This routine allocate a wait event id when it's called first time.
+ * After that, the name of the wait event is associated with the id. It's
+ * registered in the lookup table of each processes dynamically.
+ */
+static void
+shmem_startup(void)
+{
+	bool		found;
+
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	ss = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	ss = ShmemInitStruct("custom_wait_event",
+						 sizeof(state),
+						 &found);
+
+	/* Allocate a new wait event. */
+	if (!found)
+		ss->wait_event_info = ExtensionWaitEventNewTranche();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	ExtensionWaitEventRegisterTranche(ss->wait_event_info, "custom_wait_event");
+
+	return;
+}
+
+/*
+ * Connect to the database and wait forever with the custom wait event
+ */
+void
+worker_main(Datum main_arg)
+{
+	pqsignal(SIGTERM, SignalHandlerForShutdownRequest);
+	BackgroundWorkerUnblockSignals();
+
+	/* Connect to our database to show up a wait event in pg_stat_activity */
+	BackgroundWorkerInitializeConnection(worker_database, NULL, 0);
+
+	/* Wait forever with the custom wait event */
+	while (!ShutdownRequestPending)
+	{
+		(void) WaitLatch(MyLatch,
+						 WL_LATCH_SET | WL_EXIT_ON_PM_DEATH,
+						 -1L,
+						 ss->wait_event_info);  /* the custom wait event */
+		ResetLatch(MyLatch);
+	}
+}
+
+/*
+ * Dynamically launch a background worker waiting forever
+ */
+Datum
+wait_worker_launch(PG_FUNCTION_ARGS)
+{
+	BackgroundWorker worker;
+	BackgroundWorkerHandle *handle;
+	BgwHandleStatus status;
+	pid_t		pid;
+
+	memset(&worker, 0, sizeof(worker));
+	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
+		BGWORKER_BACKEND_DATABASE_CONNECTION;
+	worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
+	worker.bgw_restart_time = BGW_NEVER_RESTART;
+	sprintf(worker.bgw_function_name, "worker_main");
+	sprintf(worker.bgw_library_name, "test_custom_wait_events");
+	snprintf(worker.bgw_name, BGW_MAXLEN, "test_custom_wait_events worker");
+	snprintf(worker.bgw_type, BGW_MAXLEN, "test_custom_wait_events");
+	worker.bgw_notify_pid = MyProcPid;
+
+	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
+		PG_RETURN_NULL();
+
+	status = WaitForBackgroundWorkerStartup(handle, &pid);
+
+	if (status == BGWH_STOPPED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("could not start background process"),
+				 errhint("More details may be available in the server log.")));
+	if (status == BGWH_POSTMASTER_DIED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("cannot start background processes without postmaster"),
+				 errhint("Kill all remaining database processes and restart the database.")));
+	Assert(status == BGWH_STARTED);
+
+	PG_RETURN_INT32(pid);
+}
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.control b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
new file mode 100644
index 0000000000..8dd875a252
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
@@ -0,0 +1,3 @@
+comment = 'Test code for custom wait events'
+default_version = '1.0'
+module_pathname = '$libdir/test_custom_wait_events'
-- 
2.25.1