From 42df9177b36a236b0a10cc1c1352b756228ad412 Mon Sep 17 00:00:00 2001
From: Zsolt Parragi <zsolt.parragi@percona.com>
Date: Mon, 20 Jul 2026 22:35:06 +0000
Subject: [PATCH] injection_points: clear waiter slot on error and exit

injection_wait() only clears its slot in the waiter array after the
wait loop finishes. When the waiting query is canceled or the backend
is terminated, the slot leaks. Later wakeups of the same point then
bump the counter of the leaked slot instead of the real waiter, which
sleeps forever. Repeated leaks also exhaust the 8 slots.

Wrap the wait loop in PG_ENSURE_ERROR_CLEANUP, its cleanup callback
covers both ERROR and FATAL.

Add an isolation test: cancel one waiter, terminate another, then
check that a later waiter still receives the wakeup.
---
 src/test/modules/injection_points/Makefile    |  1 +
 .../expected/wait_cleanup.out                 | 87 +++++++++++++++++++
 .../injection_points/injection_points.c       | 37 ++++++--
 src/test/modules/injection_points/meson.build |  1 +
 .../injection_points/specs/wait_cleanup.spec  | 50 +++++++++++
 5 files changed, 168 insertions(+), 8 deletions(-)
 create mode 100644 src/test/modules/injection_points/expected/wait_cleanup.out
 create mode 100644 src/test/modules/injection_points/specs/wait_cleanup.spec

diff --git a/src/test/modules/injection_points/Makefile b/src/test/modules/injection_points/Makefile
index c01d2fb095c..2f77c974e0a 100644
--- a/src/test/modules/injection_points/Makefile
+++ b/src/test/modules/injection_points/Makefile
@@ -13,6 +13,7 @@ REGRESS = injection_points hashagg reindex_conc vacuum
 REGRESS_OPTS = --dlpath=$(top_builddir)/src/test/regress
 
 ISOLATION = basic \
+	    wait_cleanup \
 	    inplace \
 	    repack \
 	    repack_temporal \
diff --git a/src/test/modules/injection_points/expected/wait_cleanup.out b/src/test/modules/injection_points/expected/wait_cleanup.out
new file mode 100644
index 00000000000..c5be17428fc
--- /dev/null
+++ b/src/test/modules/injection_points/expected/wait_cleanup.out
@@ -0,0 +1,87 @@
+Parsed test spec with 3 sessions
+
+starting permutation: wait1 cancel3 noop3 wait2 wakeup3 noop2 detach3
+injection_points_attach
+-----------------------
+                       
+(1 row)
+
+step wait1: SELECT injection_points_run('injection-points-wait'); <waiting ...>
+step cancel3: 
+	SELECT pg_cancel_backend(pid) FROM pg_stat_activity
+	  WHERE wait_event = 'injection-points-wait';
+ <waiting ...>
+step wait1: <... completed>
+ERROR:  canceling statement due to user request
+step cancel3: <... completed>
+pg_cancel_backend
+-----------------
+t                
+(1 row)
+
+step noop3: 
+step wait2: SELECT injection_points_run('injection-points-wait'); <waiting ...>
+step wakeup3: SELECT injection_points_wakeup('injection-points-wait');
+injection_points_wakeup
+-----------------------
+                       
+(1 row)
+
+step wait2: <... completed>
+injection_points_run
+--------------------
+                    
+(1 row)
+
+step noop2: 
+step detach3: SELECT injection_points_detach('injection-points-wait');
+injection_points_detach
+-----------------------
+                       
+(1 row)
+
+
+starting permutation: wait1 terminate3 noop3 wait2 wakeup3 noop2 detach3
+injection_points_attach
+-----------------------
+                       
+(1 row)
+
+step wait1: SELECT injection_points_run('injection-points-wait'); <waiting ...>
+step terminate3: 
+	SELECT pg_terminate_backend(pid) FROM pg_stat_activity
+	  WHERE wait_event = 'injection-points-wait';
+ <waiting ...>
+step wait1: <... completed>
+FATAL:  terminating connection due to administrator command
+server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+
+step terminate3: <... completed>
+pg_terminate_backend
+--------------------
+t                   
+(1 row)
+
+step noop3: 
+step wait2: SELECT injection_points_run('injection-points-wait'); <waiting ...>
+step wakeup3: SELECT injection_points_wakeup('injection-points-wait');
+injection_points_wakeup
+-----------------------
+                       
+(1 row)
+
+step wait2: <... completed>
+injection_points_run
+--------------------
+                    
+(1 row)
+
+step noop2: 
+step detach3: SELECT injection_points_detach('injection-points-wait');
+injection_points_detach
+-----------------------
+                       
+(1 row)
+
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index 2d26ecedd5d..a549720ed75 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -223,6 +223,25 @@ injection_notice(const char *name, const void *private_data, void *arg)
 		elog(NOTICE, "notice triggered for injection point %s", name);
 }
 
+/*
+ * Clear the shared waiter slot of this backend, on error or process exit.
+ *
+ * Without this a canceled or terminated waiter would leak its slot, and
+ * later wakeups of the same injection point would bump the leaked slot's
+ * counter instead of the real waiter's. Registered through
+ * PG_ENSURE_ERROR_CLEANUP so that it also runs on FATAL, which exits
+ * without unwinding the stack.
+ */
+static void
+injection_wait_cleanup(int code, Datum arg)
+{
+	int			index = DatumGetInt32(arg);
+
+	SpinLockAcquire(&inj_state->lock);
+	inj_state->name[index][0] = '\0';
+	SpinLockRelease(&inj_state->lock);
+}
+
 /* Wait until injection_points_wakeup() is called */
 void
 injection_wait(const char *name, const void *private_data, void *arg)
@@ -275,19 +294,21 @@ injection_wait(const char *name, const void *private_data, void *arg)
 	delay_us = INJ_WAIT_INITIAL_US;
 
 	pgstat_report_wait_start(injection_wait_event);
-	while (pg_atomic_read_u32(&inj_state->wait_counts[index]) == old_wait_counts)
+	PG_ENSURE_ERROR_CLEANUP(injection_wait_cleanup, Int32GetDatum(index));
 	{
-		CHECK_FOR_INTERRUPTS();
-		pg_usleep(delay_us);
-		if (delay_us < INJ_WAIT_MAX_US)
-			delay_us = Min(delay_us * 2, INJ_WAIT_MAX_US);
+		while (pg_atomic_read_u32(&inj_state->wait_counts[index]) == old_wait_counts)
+		{
+			CHECK_FOR_INTERRUPTS();
+			pg_usleep(delay_us);
+			if (delay_us < INJ_WAIT_MAX_US)
+				delay_us = Min(delay_us * 2, INJ_WAIT_MAX_US);
+		}
 	}
+	PG_END_ENSURE_ERROR_CLEANUP(injection_wait_cleanup, Int32GetDatum(index));
 	pgstat_report_wait_end();
 
 	/* Remove this injection point from the waiters. */
-	SpinLockAcquire(&inj_state->lock);
-	inj_state->name[index][0] = '\0';
-	SpinLockRelease(&inj_state->lock);
+	injection_wait_cleanup(0, Int32GetDatum(index));
 }
 
 /*
diff --git a/src/test/modules/injection_points/meson.build b/src/test/modules/injection_points/meson.build
index 59dba1cb023..a903bc94a88 100644
--- a/src/test/modules/injection_points/meson.build
+++ b/src/test/modules/injection_points/meson.build
@@ -44,6 +44,7 @@ tests += {
   'isolation': {
     'specs': [
       'basic',
+      'wait_cleanup',
       'inplace',
       'repack',
       'repack_temporal',
diff --git a/src/test/modules/injection_points/specs/wait_cleanup.spec b/src/test/modules/injection_points/specs/wait_cleanup.spec
new file mode 100644
index 00000000000..4ef8c293e83
--- /dev/null
+++ b/src/test/modules/injection_points/specs/wait_cleanup.spec
@@ -0,0 +1,50 @@
+# Check that a canceled or terminated waiter does not leave a stale slot
+# behind in the waiter array. A leaked slot would make later wakeups of
+# the same injection point bump the leaked slot's counter instead of the
+# real waiter's, leaving the real waiter stuck.
+
+setup
+{
+	CREATE EXTENSION injection_points;
+}
+teardown
+{
+	DROP EXTENSION injection_points;
+}
+
+# The first waiter, which gets canceled or terminated. No set_local:
+# the injection point must survive s1's termination so that s3 can
+# still detach it.
+session s1
+setup	{
+	SELECT injection_points_attach('injection-points-wait', 'wait');
+}
+step wait1	{ SELECT injection_points_run('injection-points-wait'); }
+
+# The second waiter, which must still receive the wakeup.
+session s2
+step wait2	{ SELECT injection_points_run('injection-points-wait'); }
+step noop2	{ }
+
+# The control session. The blocker annotations on cancel3/terminate3
+# plus noop3 make the tester wait until wait1 actually finished before
+# starting wait2, otherwise wait2 could register its own waiter slot
+# while s1 still holds the old one.
+session s3
+step cancel3	{
+	SELECT pg_cancel_backend(pid) FROM pg_stat_activity
+	  WHERE wait_event = 'injection-points-wait';
+}
+step terminate3	{
+	SELECT pg_terminate_backend(pid) FROM pg_stat_activity
+	  WHERE wait_event = 'injection-points-wait';
+}
+step wakeup3	{ SELECT injection_points_wakeup('injection-points-wait'); }
+step detach3	{ SELECT injection_points_detach('injection-points-wait'); }
+step noop3	{ }
+
+permutation wait1 cancel3(wait1) noop3 wait2 wakeup3 noop2 detach3
+
+# The terminate permutation has to stay last: s1's connection is dead
+# afterwards, and the tester never reconnects a session.
+permutation wait1 terminate3(wait1) noop3 wait2 wakeup3 noop2 detach3
-- 
2.54.0

