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 (wait look has a CHECK_FOR_INTERRUPS), the slot leaks.
Later wakeups of the same point then bump the counter of the leaked slot
instead of the real waiter, that sleeps forever.  Repeated leaks can
exhaust all the slots.

The code is changed so as the waiting loop is wrapped with
PG_ENSURE_ERROR_CLEANUP, so as the injection point slots, that are
shared resources, can be cleaned up on ERROR as much as a FATAL.

An isolation test is added: cancel one waiter, terminate another waiter,
then check that a later waiter still receives a wakeup.  Without the
fixed code, the test would fail on timeout.

Author: Zsolt Parragi <zsolt.parragi@percona.com>
Discussion: https://postgr.es/m/CAN4CZFO+KF=cc0-iEg28RhqRBp_fTs6D4b8b7D7DB-pGYP3Ccg@mail.gmail.com
Backpatch-through: 17

src/test/modules/injection_points/Makefile

@@ -17,6 +17,7 @@ REGRESS_OPTS = --dlpath=$(top_builddir)/src/test/regress
 ISOLATION = basic \
 	    inplace \
 	    syscache-update-pruned \
+	    wait_cleanup \
 	    heap_lock_update
 
 TAP_TESTS = 1

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)
+

src/test/modules/injection_points/injection_points.c

@@ -254,6 +254,19 @@ injection_notice(const char *name, const void *private_data, void *arg)
 		elog(NOTICE, "notice triggered for injection point %s", name);
 }
 
+/*
+ * Error cleanup callback for injection point waits.
+ */
+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 on a condition variable, awaken by injection_points_wakeup() */
 void
 injection_wait(const char *name, const void *private_data, void *arg)
@@ -300,24 +313,26 @@ injection_wait(const char *name, const void *private_data, void *arg)
 
 	/* And sleep.. */
 	ConditionVariablePrepareToSleep(&inj_state->wait_point);
-	for (;;)
+	PG_ENSURE_ERROR_CLEANUP(injection_wait_cleanup, Int32GetDatum(index));
 	{
-		uint32		new_wait_counts;
+		for (;;)
+		{
+			uint32		new_wait_counts;
 
-		SpinLockAcquire(&inj_state->lock);
-		new_wait_counts = inj_state->wait_counts[index];
-		SpinLockRelease(&inj_state->lock);
+			SpinLockAcquire(&inj_state->lock);
+			new_wait_counts = inj_state->wait_counts[index];
+			SpinLockRelease(&inj_state->lock);
 
-		if (old_wait_counts != new_wait_counts)
-			break;
-		ConditionVariableSleep(&inj_state->wait_point, injection_wait_event);
+			if (old_wait_counts != new_wait_counts)
+				break;
+			ConditionVariableSleep(&inj_state->wait_point, injection_wait_event);
+		}
+		ConditionVariableCancelSleep();
 	}
-	ConditionVariableCancelSleep();
+	PG_END_ENSURE_ERROR_CLEANUP(injection_wait_cleanup, Int32GetDatum(index));
 
 	/* 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));
 }
 
 /*

src/test/modules/injection_points/meson.build

@@ -48,6 +48,7 @@ tests += {
       'basic',
       'inplace',
       'syscache-update-pruned',
+      'wait_cleanup',
       'heap_lock_update',
     ],
     'runningcheck': false, # see syscache-update-pruned

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, that gets canceled or terminated.  This does not
+# use injection_points_set_local() on purpose: 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, that receives a wakeup.
+session s2
+step wait2	{ SELECT injection_points_run('injection-points-wait'); }
+step noop2	{ }
+
+# Control session.  The blocker annotations on cancel3/terminate3,
+# together with noop3, make the tester wait until wait1 has fully
+# completed before starting wait2.  Otherwise, wait2 could register a
+# new waiter slot while s1 still owns the previous 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