v12-0002-Tests-for-LWLock-tranche-registration-improvemen.patch

From 45f4762338cdcb76515884e7f05fca787fa4513f Mon Sep 17 00:00:00 2001
From: Sami Imseih <simseih@amazon.com>
Date: Fri, 22 Aug 2025 00:35:45 -0500
Subject: [PATCH v12 2/2] Tests for LWLock tranche registration improvements

---
 src/test/modules/Makefile                     |   3 +-
 src/test/modules/meson.build                  |   1 +
 src/test/modules/test_tranches/Makefile       |  23 +++
 src/test/modules/test_tranches/meson.build    |  33 ++++
 .../test_tranches/t/001_test_tranches.pl      |  60 +++++++
 .../test_tranches/test_tranches--1.0.sql      |  21 +++
 .../modules/test_tranches/test_tranches.c     | 165 ++++++++++++++++++
 .../test_tranches/test_tranches.control       |   6 +
 8 files changed, 311 insertions(+), 1 deletion(-)
 create mode 100644 src/test/modules/test_tranches/Makefile
 create mode 100644 src/test/modules/test_tranches/meson.build
 create mode 100644 src/test/modules/test_tranches/t/001_test_tranches.pl
 create mode 100644 src/test/modules/test_tranches/test_tranches--1.0.sql
 create mode 100644 src/test/modules/test_tranches/test_tranches.c
 create mode 100644 src/test/modules/test_tranches/test_tranches.control

diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 903a8ac151a..fd7aa4675c5 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -44,7 +44,8 @@ SUBDIRS = \
 		  test_tidstore \
 		  unsafe_tests \
 		  worker_spi \
-		  xid_wraparound
+		  xid_wraparound \
+		  test_tranches
 
 
 ifeq ($(enable_injection_points),yes)
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index 93be0f57289..52883dfb496 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -45,3 +45,4 @@ subdir('typcache')
 subdir('unsafe_tests')
 subdir('worker_spi')
 subdir('xid_wraparound')
+subdir('test_tranches')
diff --git a/src/test/modules/test_tranches/Makefile b/src/test/modules/test_tranches/Makefile
new file mode 100644
index 00000000000..5c9e78084da
--- /dev/null
+++ b/src/test/modules/test_tranches/Makefile
@@ -0,0 +1,23 @@
+# src/test/modules/test_tranches/Makefile
+
+MODULE_big = test_tranches
+OBJS = \
+	$(WIN32RES) \
+	test_tranches.o
+PGFILEDESC = "test_tranches - test code LWLock tranche management"
+
+EXTENSION = test_tranches
+DATA = test_tranches--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_tranches
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_tranches/meson.build b/src/test/modules/test_tranches/meson.build
new file mode 100644
index 00000000000..976ce9f0a25
--- /dev/null
+++ b/src/test/modules/test_tranches/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) 2024-2025, PostgreSQL Global Development Group
+
+test_tranches_sources = files(
+  'test_tranches.c',
+)
+
+if host_system == 'windows'
+  test_tranches_sources += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_tranches',
+    '--FILEDESC', 'test_tranches - test code LWLock tranche management',])
+endif
+
+test_tranches = shared_module('test_tranches',
+  test_tranches_sources,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_tranches
+
+test_install_data += files(
+  'test_tranches.control',
+  'test_tranches--1.0.sql',
+)
+
+tests += {
+  'name': 'test_tranches',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'tests': [
+      't/001_test_tranches.pl',
+    ],
+  },
+}
\ No newline at end of file
diff --git a/src/test/modules/test_tranches/t/001_test_tranches.pl b/src/test/modules/test_tranches/t/001_test_tranches.pl
new file mode 100644
index 00000000000..f74fb9a54a1
--- /dev/null
+++ b/src/test/modules/test_tranches/t/001_test_tranches.pl
@@ -0,0 +1,60 @@
+use strict;
+use warnings FATAL => 'all';
+
+use List::Util qw(min);
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node = PostgreSQL::Test::Cluster->new('primary');
+$node->init();
+
+$node->append_conf('postgresql.conf',
+    qq(shared_preload_libraries='test_tranches'));
+$node->append_conf('postgresql.conf',
+    qq(test_tranches.requested_named_tranches=0));
+$node->start();
+
+$node->safe_psql('postgres', q(CREATE EXTENSION test_tranches));
+
+my $log_location = -s $node->logfile;
+
+my ($result, $stdout, $stderr) = $node->psql('postgres',
+    qq{
+        select test_tranches_new(5);
+        select test_tranches_lookup(96);
+        select test_tranches_lookup(97);
+        select test_tranches_lookup(98);
+        select test_tranches_lookup(99);
+        select test_tranches_lookup(100);
+        select test_tranches_lookup(101);
+        select test_tranches_new(5);
+        select test_tranches_lookup(101);
+        select test_tranches_lookup(102);
+        select test_tranches_lookup(103);
+        select test_tranches_lookup(104);
+        select test_tranches_lookup(105);
+        select test_tranches_lookup(106);
+        },
+    on_error_stop => 0);
+like("$stderr", qr/ERROR:  LWLock tranche is not registered\n.*ERROR:  LWLock tranche is not registered/, "unregistered tranche error");
+like("$stdout", qr/test_lock__0.*test_lock__1.*test_lock__2.*test_lock__3.*test_lock__4/s, "matches first 5 tranche names");
+
+my $good_tranche_name = 'A' x 63;
+my $bad_tranche_name = 'B' x 64;
+$node->safe_psql('postgres', qq{select test_tranches_last('$good_tranche_name');});
+($result, $stdout, $stderr) = $node->psql('postgres',
+    qq{
+        select test_tranches_last('$bad_tranche_name');
+        },
+    on_error_stop => 0);
+like("$stderr", qr/ERROR:  tranche name too long/, "tranche name too long");
+
+$node->restart();
+
+$node->safe_psql('postgres', qq{select test_tranches_new(1023)});
+$node->safe_psql('postgres', qq{select test_tranches_new(1)});
+($result, $stdout, $stderr) = $node->psql('postgres', qq{select test_tranches_new(1);}, on_error_stop => 0);
+like("$stderr", qr/ERROR:  too many LWLock tranches registered/, "too many tranches registered");
+
+done_testing();
diff --git a/src/test/modules/test_tranches/test_tranches--1.0.sql b/src/test/modules/test_tranches/test_tranches--1.0.sql
new file mode 100644
index 00000000000..1de8972ab41
--- /dev/null
+++ b/src/test/modules/test_tranches/test_tranches--1.0.sql
@@ -0,0 +1,21 @@
+-- test_tranches--1.0.sql
+
+CREATE FUNCTION test_tranches_new(bigint)
+RETURNS void
+AS 'MODULE_PATHNAME', 'test_tranches_new'
+LANGUAGE C STRICT;
+
+CREATE FUNCTION test_tranches_last(text)
+RETURNS int
+AS 'MODULE_PATHNAME', 'test_tranches_last'
+LANGUAGE C STRICT;
+
+CREATE FUNCTION test_tranches_lookup(int)
+RETURNS text
+AS 'MODULE_PATHNAME', 'test_tranches_lookup'
+LANGUAGE C STRICT;
+
+CREATE FUNCTION test_tranches_get_named_lwlock(text, int)
+RETURNS int
+AS 'MODULE_PATHNAME', 'test_tranches_get_named_lwlock'
+LANGUAGE C STRICT;
diff --git a/src/test/modules/test_tranches/test_tranches.c b/src/test/modules/test_tranches/test_tranches.c
new file mode 100644
index 00000000000..a27455ce5d2
--- /dev/null
+++ b/src/test/modules/test_tranches/test_tranches.c
@@ -0,0 +1,165 @@
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "storage/dsm_registry.h"
+#include "storage/ipc.h"
+#include "storage/lwlock.h"
+#include "storage/shmem.h"
+#include "utils/builtins.h"
+#include "utils/guc.h"
+#include "utils/injection_point.h"
+#include "utils/wait_classes.h"
+
+PG_MODULE_MAGIC;
+
+/* hooks */
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
+static void test_tranches_shmem_request(void);
+static void test_tranches_shmem_startup(void);
+
+/* GUC */
+static int	test_tranches_requested_named_tranches = 0;
+
+typedef struct testTranchesSharedState
+{
+	int			next_index;
+}			testTranchesSharedState;
+
+static testTranchesSharedState * test_lwlock_ss = NULL;
+
+/*
+ * LWLock wait event masks. Copied from src/backend/utils/activity/wait_event.c
+ */
+#define WAIT_EVENT_CLASS_MASK	0xFF000000
+
+/*
+ * Module load callback
+ */
+void
+_PG_init(void)
+{
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = test_tranches_shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = test_tranches_shmem_startup;
+
+	DefineCustomIntVariable("test_tranches.requested_named_tranches",
+							"Sets the number of locks created during shmem request",
+							NULL,
+							&test_tranches_requested_named_tranches,
+							2,
+							0,
+							UINT16_MAX,
+							PGC_POSTMASTER,
+							0,
+							NULL,
+							NULL,
+							NULL);
+}
+
+static void
+test_tranches_shmem_startup(void)
+{
+	bool		found;
+
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	test_lwlock_ss = NULL;
+
+	test_lwlock_ss = ShmemInitStruct("test_tranches",
+									 sizeof(testTranchesSharedState),
+									 &found);
+	if (!found)
+	{
+		test_lwlock_ss->next_index = test_tranches_requested_named_tranches;
+	}
+}
+
+static Size
+test_tranches_memsize(void)
+{
+	Size		size;
+
+	size = MAXALIGN(sizeof(testTranchesSharedState));
+
+	return size;
+}
+
+static void
+test_tranches_shmem_request(void)
+{
+	int			i = 0;
+
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(test_tranches_memsize());
+
+	for (i = 0; i < test_tranches_requested_named_tranches; i++)
+	{
+		char		name[15];
+
+		snprintf(name, sizeof(name), "test_lock_%d", i);
+		RequestNamedLWLockTranche(name, i);
+	}
+}
+
+PG_FUNCTION_INFO_V1(test_tranches_new);
+Datum
+test_tranches_new(PG_FUNCTION_ARGS)
+{
+	int64		num = PG_GETARG_INT64(0);
+	int			i;
+
+	for (i = test_lwlock_ss->next_index; i < num + test_lwlock_ss->next_index; i++)
+	{
+		char		name[50];
+
+		snprintf(name, 50, "test_lock__%d", i);
+
+		LWLockNewTrancheId(name);
+	}
+
+	test_lwlock_ss->next_index = i;
+
+	PG_RETURN_VOID();
+}
+
+PG_FUNCTION_INFO_V1(test_tranches_last);
+Datum
+test_tranches_last(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_INT32(LWLockNewTrancheId(text_to_cstring(PG_GETARG_TEXT_PP(0))));
+}
+
+PG_FUNCTION_INFO_V1(test_tranches_lookup);
+Datum
+test_tranches_lookup(PG_FUNCTION_ARGS)
+{
+	const char *tranche_name = GetLWLockIdentifier(PG_WAIT_LWLOCK & WAIT_EVENT_CLASS_MASK, PG_GETARG_INT32(0));
+
+	PG_RETURN_TEXT_P(cstring_to_text(tranche_name));
+}
+
+PG_FUNCTION_INFO_V1(test_tranches_get_named_lwlock);
+Datum
+test_tranches_get_named_lwlock(PG_FUNCTION_ARGS)
+{
+	int			i;
+	LWLockPadded *locks;
+
+	locks = GetNamedLWLockTranche(text_to_cstring(PG_GETARG_TEXT_PP(0)));
+
+	for (i = 0; i < PG_GETARG_INT32(1); i++)
+	{
+		LWLock	   *lock = &locks[i].lock;
+
+		LWLockAcquire(lock, LW_SHARED);
+		LWLockRelease(lock);
+	}
+
+	PG_RETURN_INT32(i);
+}
diff --git a/src/test/modules/test_tranches/test_tranches.control b/src/test/modules/test_tranches/test_tranches.control
new file mode 100644
index 00000000000..8e6254a7e43
--- /dev/null
+++ b/src/test/modules/test_tranches/test_tranches.control
@@ -0,0 +1,6 @@
+# test_tranches.control
+
+comment = 'Test LWLock tranch names tracking'
+default_version = '1.0'
+relocatable = false
+module_pathname = '$libdir/test_tranches'
\ No newline at end of file
-- 
2.39.5 (Apple Git-154)