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

From 3f7ad775a9fc82d2f183a370c2e0cfaf761b1393 Mon Sep 17 00:00:00 2001
From: Sami Imseih <simseih@amazon.com>
Date: Fri, 22 Aug 2025 00:35:45 -0500
Subject: [PATCH v11 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      | 115 ++++++++++++
 .../test_tranches/test_tranches--1.0.sql      |  21 +++
 .../modules/test_tranches/test_tranches.c     | 165 ++++++++++++++++++
 .../test_tranches/test_tranches.control       |   6 +
 8 files changed, 366 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..9a56b25e204
--- /dev/null
+++ b/src/test/modules/test_tranches/t/001_test_tranches.pl
@@ -0,0 +1,115 @@
+use strict;
+use warnings FATAL => 'all';
+
+use List::Util qw(min);
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $UINT16_MAX = 65535;
+my $MAX_DYNAMIC_TRANCHES = 1024;
+
+# Helper function: wait for one or more logs
+sub maybe_wait_for_log {
+    my ($node, $logs, $log_loc) = @_;
+
+    # Normalize single regex to array
+    $logs = [$logs] unless ref($logs) eq 'ARRAY';
+
+    foreach my $regex (@$logs) {
+        $log_loc = $node->wait_for_log($regex, $log_loc);
+    }
+    return $log_loc;
+}
+
+my $node = PostgreSQL::Test::Cluster->new('primary');
+$node->init();
+
+my $requested_named_tranches = 4;
+
+$node->append_conf('postgresql.conf',
+    qq(shared_preload_libraries='test_tranches'));
+$node->append_conf('postgresql.conf',
+    qq(test_tranches.requested_named_tranches=$requested_named_tranches));
+$node->append_conf('postgresql.conf',
+    qq(log_min_messages=DEBUG3));
+$node->start();
+
+$node->safe_psql('postgres', q(CREATE EXTENSION test_tranches));
+
+my $next_index = 0;
+my $lookup_tranche_id = 0;
+
+my $log_location = -s $node->logfile;
+
+#
+# Lookup requested tranches
+#
+my $last_tranche_id = $node->safe_psql('postgres',
+    qq{select test_tranches_last('fake')});
+
+my $first_tranche_id = $last_tranche_id - $requested_named_tranches;
+my $second_tranche_id = $first_tranche_id + 1;
+
+my $requested_tranches = $node->safe_psql('postgres',
+    qq{select test_tranches_lookup($first_tranche_id);
+       select test_tranches_lookup($second_tranche_id);});
+
+ok ("test_lock_0\ntest_lock_1" eq $requested_tranches, "looked up requested tranches");
+
+#
+# Create tranches
+#
+$node->safe_psql('postgres', qq{select test_tranches_new(100)});
+my $next_tranche_id = $second_tranche_id + 1;
+$requested_tranches = $node->safe_psql('postgres',
+    qq{select test_tranches_lookup($next_tranche_id);});
+ok ("test_lock_2" eq $requested_tranches, "looked up first dynamic tranches");
+
+#
+# get LWLocks
+#
+$node->safe_psql('postgres',
+    qq{select test_tranches_get_named_lwlock('test_lock_3', 3);});
+
+#
+# restart node without requested tranches to test limits
+#
+$node->append_conf('postgresql.conf',
+    qq(test_tranches.requested_named_tranches=0));
+$node->restart;
+
+# long 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');});
+$node->psql('postgres', qq{select test_tranches_last('$bad_tranche_name');});
+$log_location = maybe_wait_for_log($node, qr/ ERROR:  tranche name too long/, $log_location);
+ok(1, "check for error on tranche name too long");
+
+#
+# Check for errors with not registered tranches
+#
+my $beyond_limit_id = $MAX_DYNAMIC_TRANCHES + 100;
+$node->psql('postgres',
+    qq{select test_tranches_lookup($beyond_limit_id)});
+$log_location = maybe_wait_for_log($node, qr/ ERROR:  LWLock tranche is not registered/, $log_location);
+
+my $within_limit_id = $MAX_DYNAMIC_TRANCHES - 100;
+$node->psql('postgres',
+    qq{select test_tranches_lookup($within_limit_id)});
+$log_location = maybe_wait_for_log($node, qr/ ERROR:  LWLock tranche is not registered/, $log_location);
+
+ok(1, "check for error with not registered tranches");
+
+
+# Check for errors with too many tranches registered
+my $dynamic_tranches_count = $MAX_DYNAMIC_TRANCHES;
+$node->safe_psql('postgres',
+    qq{select test_tranches_new($dynamic_tranches_count)});
+$node->psql('postgres',
+    qq{select test_tranches_new(1)});
+$log_location = maybe_wait_for_log($node, qr/ ERROR:  too many LWLock tranches registered/, $log_location);
+ok(1, "check for error with too many dynamic tranches");
+
+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)