From 8caea1d434aa1cbd6f1da777b89dd4895a88b44f Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Thu, 6 Mar 2025 13:06:05 +0900
Subject: [PATCH v2] Fix race condition in pre-auth test

---
 src/test/authentication/t/007_pre_auth.pl     | 36 ++++++---------
 .../perl/PostgreSQL/Test/BackgroundPsql.pm    | 44 +++++++++++++++++++
 2 files changed, 57 insertions(+), 23 deletions(-)

diff --git a/src/test/authentication/t/007_pre_auth.pl b/src/test/authentication/t/007_pre_auth.pl
index a638226dbaf1..fb4241cac4ab 100644
--- a/src/test/authentication/t/007_pre_auth.pl
+++ b/src/test/authentication/t/007_pre_auth.pl
@@ -43,36 +43,26 @@ $psql->query_safe("SELECT injection_points_attach('init-pre-auth', 'wait')");
 # authentication. Use the $psql connection handle for server interaction.
 my $conn = $node->background_psql('postgres', wait => 0);
 
-# Wait for the connection to show up.
-my $pid;
-while (1)
-{
-	$pid = $psql->query(
-		"SELECT pid FROM pg_stat_activity WHERE state = 'starting';");
-	last if $pid ne "";
+# Wait for the connection to show up in pg_stat_activity, with the wait_event
+# of the injection point.
+my $res = $psql->poll_query_until(
+	qq{SELECT count(pid) > 0 FROM pg_stat_activity
+  WHERE state = 'starting' and wait_event = 'init-pre-auth';});
+ok($res, 'authenticating connections are recorded in pg_stat_activity');
 
-	usleep(100_000);
-}
-
-note "backend $pid is authenticating";
-ok(1, 'authenticating connections are recorded in pg_stat_activity');
+# Get the PID of the backend waiting, for the next checks.
+my $pid = $psql->query(
+	qq{SELECT pid FROM pg_stat_activity
+  WHERE state = 'starting' and wait_event = 'init-pre-auth';});
 
 # Detach the waitpoint and wait for the connection to complete.
 $psql->query_safe("SELECT injection_points_wakeup('init-pre-auth');");
 $conn->wait_connect();
 
 # Make sure the pgstat entry is updated eventually.
-while (1)
-{
-	my $state =
-	  $psql->query("SELECT state FROM pg_stat_activity WHERE pid = $pid;");
-	last if $state eq "idle";
-
-	note "state for backend $pid is '$state'; waiting for 'idle'...";
-	usleep(100_000);
-}
-
-ok(1, 'authenticated connections reach idle state in pg_stat_activity');
+$res = $psql->poll_query_until(
+	qq{SELECT state FROM pg_stat_activity WHERE pid = $pid;}, 'idle');
+ok($res, 'authenticated connections reach idle state in pg_stat_activity');
 
 $psql->query_safe("SELECT injection_points_detach('init-pre-auth');");
 $psql->quit();
diff --git a/src/test/perl/PostgreSQL/Test/BackgroundPsql.pm b/src/test/perl/PostgreSQL/Test/BackgroundPsql.pm
index c611a61cf4e6..83ecf8b3e720 100644
--- a/src/test/perl/PostgreSQL/Test/BackgroundPsql.pm
+++ b/src/test/perl/PostgreSQL/Test/BackgroundPsql.pm
@@ -61,6 +61,7 @@ use Config;
 use IPC::Run;
 use PostgreSQL::Test::Utils qw(pump_until);
 use Test::More;
+use Time::HiRes qw(usleep);
 
 =pod
 
@@ -172,6 +173,49 @@ sub wait_connect
 	die "psql startup timed out" if $self->{timeout}->is_expired;
 }
 
+
+=pod
+
+=item $session->poll_query_until($query [, $expected ])
+
+Run B<$query> repeatedly, until it returns the B<$expected> result
+('t', or SQL boolean true, by default).
+Continues polling if B<query> returns an error result.
+
+Times out after $PostgreSQL::Test::Utils::timeout_default seconds.
+
+Returns 1 if successful, 0 if timed out.
+
+=cut
+
+sub poll_query_until
+{
+	my ($self, $query, $expected) = @_;
+
+	$expected = 't' unless defined($expected);    # default value
+
+	my $max_attempts = 10 * $PostgreSQL::Test::Utils::timeout_default;
+	my $attempts = 0;
+
+	while ($attempts < $max_attempts)
+	{
+		my $ret = $self->query($query);
+
+		if ($ret eq $expected)
+		{
+			return 1;
+		}
+
+		# Wait 0.1 second before retrying.
+		usleep(100_000);
+		$attempts++;
+	}
+
+	# Give up.  The output of the last attempt is logged by query(),
+	# so no need to do anything here.
+	return 0;
+}
+
 =pod
 
 =item $session->quit
-- 
2.47.2