Add test module for Custom WAL Resource Manager feature

From baadff89aeb4ae0dded760b2b6e8eda8c9a07c84 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddyforpostgres@gmail.com>
Date: Mon, 14 Nov 2022 03:59:35 +0000
Subject: [PATCH v2] Add test module for Custom WAL Resource Manager feature

---
 doc/src/sgml/custom-rmgr.sgml                 |   7 +
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 src/test/modules/test_custom_rmgrs/.gitignore |   4 +
 src/test/modules/test_custom_rmgrs/Makefile   |  24 +++
 .../modules/test_custom_rmgrs/meson.build     |  34 ++++
 .../modules/test_custom_rmgrs/t/001_basic.pl  |  63 +++++++
 .../test_custom_rmgrs--1.0.sql                |  16 ++
 .../test_custom_rmgrs/test_custom_rmgrs.c     | 163 ++++++++++++++++++
 .../test_custom_rmgrs.control                 |   4 +
 10 files changed, 317 insertions(+)
 create mode 100644 src/test/modules/test_custom_rmgrs/.gitignore
 create mode 100644 src/test/modules/test_custom_rmgrs/Makefile
 create mode 100644 src/test/modules/test_custom_rmgrs/meson.build
 create mode 100644 src/test/modules/test_custom_rmgrs/t/001_basic.pl
 create mode 100644 src/test/modules/test_custom_rmgrs/test_custom_rmgrs--1.0.sql
 create mode 100644 src/test/modules/test_custom_rmgrs/test_custom_rmgrs.c
 create mode 100644 src/test/modules/test_custom_rmgrs/test_custom_rmgrs.control

diff --git a/doc/src/sgml/custom-rmgr.sgml b/doc/src/sgml/custom-rmgr.sgml
index 2893016cef..6d6909fc12 100644
--- a/doc/src/sgml/custom-rmgr.sgml
+++ b/doc/src/sgml/custom-rmgr.sgml
@@ -57,6 +57,13 @@ typedef struct RmgrData
 } RmgrData;
 </programlisting>
  </para>
+
+  <para>
+   The <filename>src/test/modules/test_custom_rmgrs</filename> module
+   contains a working example, which demonstrates usage of custom WAL
+   resource managers.
+  </para>
+
  <para>
   Then, register your new resource
   manager.
diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 7b3f292965..548469f7c1 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -16,6 +16,7 @@ SUBDIRS = \
 		  spgist_name_ops \
 		  test_bloomfilter \
 		  test_copy_callbacks \
+		  test_custom_rmgrs \
 		  test_ddl_deparse \
 		  test_extensions \
 		  test_ginpostinglist \
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index c2e5f5ffd5..f2df05b1bc 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -10,6 +10,7 @@ subdir('spgist_name_ops')
 subdir('ssl_passphrase_callback')
 subdir('test_bloomfilter')
 subdir('test_copy_callbacks')
+subdir('test_custom_rmgrs')
 subdir('test_ddl_deparse')
 subdir('test_extensions')
 subdir('test_ginpostinglist')
diff --git a/src/test/modules/test_custom_rmgrs/.gitignore b/src/test/modules/test_custom_rmgrs/.gitignore
new file mode 100644
index 0000000000..5dcb3ff972
--- /dev/null
+++ b/src/test/modules/test_custom_rmgrs/.gitignore
@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/
diff --git a/src/test/modules/test_custom_rmgrs/Makefile b/src/test/modules/test_custom_rmgrs/Makefile
new file mode 100644
index 0000000000..b557e5888e
--- /dev/null
+++ b/src/test/modules/test_custom_rmgrs/Makefile
@@ -0,0 +1,24 @@
+# src/test/modules/test_custom_rmgrs/Makefile
+
+MODULE_big = test_custom_rmgrs
+OBJS = \
+	$(WIN32RES) \
+	test_custom_rmgrs.o
+PGFILEDESC = "test_custom_rmgrs - test custom WAL resource managers"
+
+EXTENSION = test_custom_rmgrs
+DATA = test_custom_rmgrs--1.0.sql
+
+EXTRA_INSTALL = contrib/pg_walinspect
+TAP_TESTS = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_custom_rmgrs
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_custom_rmgrs/meson.build b/src/test/modules/test_custom_rmgrs/meson.build
new file mode 100644
index 0000000000..05ec06d6d5
--- /dev/null
+++ b/src/test/modules/test_custom_rmgrs/meson.build
@@ -0,0 +1,34 @@
+# FIXME: prevent install during main install, but not during test :/
+
+test_custom_rmgrs_sources = files(
+  'test_custom_rmgrs.c',
+)
+
+if host_system == 'windows'
+  test_custom_rmgrs_sources += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_custom_rmgrs',
+    '--FILEDESC', 'test_custom_rmgrs - test custom WAL resource managers',])
+endif
+
+test_custom_rmgrs = shared_module('test_custom_rmgrs',
+  test_custom_rmgrs_sources,
+  kwargs: pg_mod_args,
+)
+testprep_targets += test_custom_rmgrs
+
+install_data(
+  'test_custom_rmgrs.control',
+  'test_custom_rmgrs--1.0.sql',
+  kwargs: contrib_data_args,
+)
+
+tests += {
+  'name': 'test_custom_rmgrs',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'tests': [
+      't/001_basic.pl',
+    ],
+  },
+}
diff --git a/src/test/modules/test_custom_rmgrs/t/001_basic.pl b/src/test/modules/test_custom_rmgrs/t/001_basic.pl
new file mode 100644
index 0000000000..6210c3538b
--- /dev/null
+++ b/src/test/modules/test_custom_rmgrs/t/001_basic.pl
@@ -0,0 +1,63 @@
+# Copyright (c) 2021-2022, 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',
+	qq{shared_preload_libraries = 'test_custom_rmgrs'});
+$node->append_conf(
+	'postgresql.conf', q{
+wal_level = 'replica'
+max_wal_senders = 4
+shared_preload_libraries = 'test_custom_rmgrs'
+});
+$node->start;
+
+# setup
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_rmgrs');
+
+# pg_walinspect is required only for verifying test_custom_rmgrs output.
+# test_custom_rmgrs doesn't use/depend on it internally.
+$node->safe_psql('postgres', 'CREATE EXTENSION pg_walinspect');
+
+# make sure checkpoints don't interfere with the test.
+my $start_lsn = $node->safe_psql('postgres',
+	qq[SELECT lsn FROM pg_create_physical_replication_slot('regress_test_slot1', true, false);]);
+
+# write and save the WAL record's returned end LSN for verifying it later
+my $record_end_lsn = $node->safe_psql('postgres',
+	'SELECT * FROM test_custom_rmgrs_synthesize_wal_record(16)');
+
+# ensure the WAL is written and flushed to disk
+$node->safe_psql('postgres', 'SELECT pg_switch_wal()');
+
+my $end_lsn = $node->safe_psql('postgres', 'SELECT pg_current_wal_flush_lsn()');
+
+# check if our custom WAL resource manager has successfully registered with the server
+my $row_count =
+  $node->safe_psql('postgres',
+	qq[SELECT count(*) FROM pg_get_wal_resource_managers()
+		WHERE rm_name = 'test_custom_rmgrs';]);
+is($row_count, '1',
+	'custom WAL resource manager has successfully registered with the server'
+);
+
+# check if our custom WAL resource manager has successfully written a WAL record
+my $result =
+  $node->safe_psql('postgres',
+	qq[SELECT end_lsn FROM pg_get_wal_records_info('$start_lsn', '$end_lsn')
+		WHERE resource_manager = 'test_custom_rmgrs';]);
+is($result, $record_end_lsn,
+	'custom WAL resource manager has successfully written a WAL record'
+);
+
+$node->stop;
+done_testing();
diff --git a/src/test/modules/test_custom_rmgrs/test_custom_rmgrs--1.0.sql b/src/test/modules/test_custom_rmgrs/test_custom_rmgrs--1.0.sql
new file mode 100644
index 0000000000..4a72ca3f73
--- /dev/null
+++ b/src/test/modules/test_custom_rmgrs/test_custom_rmgrs--1.0.sql
@@ -0,0 +1,16 @@
+/* src/test/modules/test_custom_rmgrs/test_custom_rmgrs--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_custom_rmgrs" to load this file. \quit
+
+--
+-- test_custom_rmgrs_synthesize_wal_record()
+--
+-- Writes a synthesized message into WAL with the help of custom WAL resource
+-- manager.
+--
+CREATE FUNCTION test_custom_rmgrs_synthesize_wal_record(IN size int8,
+    OUT lsn pg_lsn
+)
+AS 'MODULE_PATHNAME', 'test_custom_rmgrs_synthesize_wal_record'
+LANGUAGE C STRICT PARALLEL UNSAFE;
diff --git a/src/test/modules/test_custom_rmgrs/test_custom_rmgrs.c b/src/test/modules/test_custom_rmgrs/test_custom_rmgrs.c
new file mode 100644
index 0000000000..64e485c2fb
--- /dev/null
+++ b/src/test/modules/test_custom_rmgrs/test_custom_rmgrs.c
@@ -0,0 +1,163 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_custom_rmgrs.c
+ *		Code for testing custom WAL resource managers.
+ *
+ * Portions Copyright (c) 1996-2022, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_custom_rmgrs/test_custom_rmgrs.c
+ *
+ * -------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "access/xlog.h"
+#include "access/xlog_internal.h"
+#include "access/xloginsert.h"
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "utils/pg_lsn.h"
+
+PG_MODULE_MAGIC;
+
+/*
+ * test_custom_rmgrs WAL record message.
+ */
+typedef struct xl_testcustomrmgrs_message
+{
+	Size		message_size;   /* size of the message */
+	char		message[FLEXIBLE_ARRAY_MEMBER]; /* payload */
+} xl_testcustomrmgrs_message;
+
+#define SizeOfTestCustomRmgrsMessage	(offsetof(xl_testcustomrmgrs_message, message))
+#define XLOG_TEST_CUSTOM_RMGRS_MESSAGE	0x00
+
+/*
+ * Typically one needs to reserve unique RmgrId at
+ * https://wiki.postgresql.org/wiki/CustomWALResourceManagers. However, we use
+ * experimental RmgrId in this test module.
+ */
+static RmgrId rmid = RM_EXPERIMENTAL_ID;
+static RmgrData rmgr;
+
+/*
+ * Define these RMGR APIs to the taste, see xlog_internal.h and rmgrlist.h for
+ * reference.
+ */
+void		testcustomrmgrs_redo(XLogReaderState *record);
+void		testcustomrmgrs_desc(StringInfo buf, XLogReaderState *record);
+const char *testcustomrmgrs_identify(uint8 info);
+
+PG_FUNCTION_INFO_V1(test_custom_rmgrs_synthesize_wal_record);
+
+/*
+ * Module load callback
+ */
+void
+_PG_init(void)
+{
+	/*
+	 * In order to create our own custom resource manager, we have to be loaded
+	 * via shared_preload_libraries.
+	 */
+	if (!process_shared_preload_libraries_in_progress)
+		return;
+
+	MemSet(&rmgr, 0, sizeof(rmgr));
+	rmgr.rm_name = "test_custom_rmgrs";
+	rmgr.rm_redo = testcustomrmgrs_redo;
+	rmgr.rm_desc = testcustomrmgrs_desc;
+	rmgr.rm_identify = testcustomrmgrs_identify;
+
+	RegisterCustomRmgr(rmid, &rmgr);
+}
+
+/* RMGR API implementation */
+
+/*
+ * Redo is basically just a noop for this module.
+ */
+void
+testcustomrmgrs_redo(XLogReaderState *record)
+{
+	uint8		info = XLogRecGetInfo(record) & ~XLR_INFO_MASK;
+
+	if (info != XLOG_TEST_CUSTOM_RMGRS_MESSAGE)
+		elog(PANIC, "testcustomrmgrs__redo: unknown op code %u", info);
+}
+
+void
+testcustomrmgrs_desc(StringInfo buf, XLogReaderState *record)
+{
+	char	   *rec = XLogRecGetData(record);
+	uint8		info = XLogRecGetInfo(record) & ~XLR_INFO_MASK;
+
+	if (info == XLOG_TEST_CUSTOM_RMGRS_MESSAGE)
+	{
+		xl_testcustomrmgrs_message *xlrec = (xl_testcustomrmgrs_message *) rec;
+
+		appendStringInfo(buf, "payload (%zu bytes): %s", xlrec->message_size, xlrec->message);
+	}
+}
+
+const char *
+testcustomrmgrs_identify(uint8 info)
+{
+	if ((info & ~XLR_INFO_MASK) == XLOG_TEST_CUSTOM_RMGRS_MESSAGE)
+		return "TEST_CUSTOM_RMGRS_MESSAGE";
+
+	return NULL;
+}
+
+/*
+ * SQL function for writing a synthesized message into WAL with the help of custom WAL
+ * resource manager.
+ */
+Datum
+test_custom_rmgrs_synthesize_wal_record(PG_FUNCTION_ARGS)
+{
+	int64		size = PG_GETARG_INT64(0);
+	XLogRecPtr	lsn;
+	char 		*message;
+	static char charset[] = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
+	xl_testcustomrmgrs_message xlrec;
+
+	if (size <= 0)
+		ereport(ERROR,
+				(errmsg("size %lld must be greater than zero to synthesize WAL record",
+						(long long) size)));
+
+	/*
+	 * Since this a test module, let's limit maximum size of WAL record to be
+	 * 1MB to not over-allocate memory.
+	 */
+	if (size > 1024*1024)
+		size = 1024*1024;
+
+	message = (char *) palloc0(size);
+
+	/* Prepare a random message */
+	for (int i = 0; i < size; i++)
+	{
+		int key = rand() % (int)(sizeof(charset) -1);
+		message[i] = charset[key];
+	}
+
+	xlrec.message_size = size;
+
+	XLogBeginInsert();
+	XLogRegisterData((char *) &xlrec, SizeOfTestCustomRmgrsMessage);
+	XLogRegisterData((char *) &message, size);
+
+	/* Let's mark this record as unimportant, just in case. */
+	XLogSetRecordFlags(XLOG_MARK_UNIMPORTANT);
+
+	lsn = XLogInsert(rmid, XLOG_TEST_CUSTOM_RMGRS_MESSAGE);
+
+	pfree(message);
+
+	PG_RETURN_LSN(lsn);
+}
diff --git a/src/test/modules/test_custom_rmgrs/test_custom_rmgrs.control b/src/test/modules/test_custom_rmgrs/test_custom_rmgrs.control
new file mode 100644
index 0000000000..8113e0007d
--- /dev/null
+++ b/src/test/modules/test_custom_rmgrs/test_custom_rmgrs.control
@@ -0,0 +1,4 @@
+comment = 'Test code for custom WAL resource managers'
+default_version = '1.0'
+module_pathname = '$libdir/test_custom_rmgrs'
+relocatable = true
-- 
2.34.1