From 0f6f4b1bdc69a10b1048731d5b1d744567978cce Mon Sep 17 00:00:00 2001
From: Thomas Munro <thomas.munro@gmail.com>
Date: Tue, 29 Aug 2023 18:12:07 +1200
Subject: [PATCH] XXX cache psql processes


diff --git a/contrib/amcheck/t/001_verify_heapam.pl b/contrib/amcheck/t/001_verify_heapam.pl
index 46d5b53181..316e876c5d 100644
--- a/contrib/amcheck/t/001_verify_heapam.pl
+++ b/contrib/amcheck/t/001_verify_heapam.pl
@@ -18,6 +18,7 @@ $node = PostgreSQL::Test::Cluster->new('test');
 $node->init;
 $node->append_conf('postgresql.conf', 'autovacuum=off');
 $node->start;
+$node->enable_auto_background_psql(1);
 $node->safe_psql('postgres', q(CREATE EXTENSION amcheck));
 
 #
diff --git a/src/test/perl/PostgreSQL/Test/BackgroundPsql.pm b/src/test/perl/PostgreSQL/Test/BackgroundPsql.pm
index 924b57ab21..53ee69c801 100644
--- a/src/test/perl/PostgreSQL/Test/BackgroundPsql.pm
+++ b/src/test/perl/PostgreSQL/Test/BackgroundPsql.pm
@@ -221,7 +221,7 @@ sub query
 	$output = $self->{stdout};
 
 	# remove banner again, our caller doesn't care
-	$output =~ s/\n$banner$//s;
+	$output =~ s/\n?$banner$//s;
 
 	# clear out output for the next query
 	$self->{stdout} = '';
diff --git a/src/test/perl/PostgreSQL/Test/Cluster.pm b/src/test/perl/PostgreSQL/Test/Cluster.pm
index 227c34ab4d..f53844e25e 100644
--- a/src/test/perl/PostgreSQL/Test/Cluster.pm
+++ b/src/test/perl/PostgreSQL/Test/Cluster.pm
@@ -986,6 +986,8 @@ sub stop
 
 	local %ENV = $self->_get_env();
 
+	$self->_close_auto_background_psqls;
+
 	$mode = 'fast' unless defined $mode;
 	return 1 unless defined $self->{_pid};
 
@@ -1025,6 +1027,8 @@ sub reload
 
 	local %ENV = $self->_get_env();
 
+	$self->_close_auto_background_psqls;
+
 	print "### Reloading node \"$name\"\n";
 	PostgreSQL::Test::Utils::system_or_bail('pg_ctl', '-D', $pgdata,
 		'reload');
@@ -1049,6 +1053,8 @@ sub restart
 
 	local %ENV = $self->_get_env(PGAPPNAME => undef);
 
+	$self->_close_auto_background_psqls;
+
 	print "### Restarting node \"$name\"\n";
 
 	# -w is now the default but having it here does no harm and helps
@@ -1349,7 +1355,9 @@ sub new
 		_logfile_base =>
 		  "$PostgreSQL::Test::Utils::log_path/${testname}_${name}",
 		_logfile =>
-		  "$PostgreSQL::Test::Utils::log_path/${testname}_${name}.log"
+		  "$PostgreSQL::Test::Utils::log_path/${testname}_${name}.log",
+		_auto_background_psqls => {},
+		_use_auto_background_psqls => 0
 	};
 
 	if ($params{install_path})
@@ -1508,6 +1516,39 @@ sub _get_env
 	return (%inst_env);
 }
 
+# Private routine to close and forget any psql processes we may be
+# holding onto.
+sub _close_auto_background_psqls
+{
+	my ($self) = @_;
+	foreach my $dbname (keys %{$self->{_auto_background_psqls}})
+	{
+		my $psql = $self->{_auto_background_psqls}{$dbname};
+		delete $self->{_auto_background_psqls}{$dbname};
+		$psql->quit;
+	}
+}
+
+=pod
+
+=item enable_auto_background_psql()
+
+Enable or disable the automatic reuse of long-lived psql sessions.
+
+=cut
+
+sub enable_auto_background_psql
+{
+	my $self = shift;
+	my $value = shift;
+
+	$self->{_use_auto_background_psqls} = $value;
+	if (!$value)
+	{
+		$self->_close_auto_background_psqls;
+	}
+}
+
 # Private routine to get an installation path qualified command.
 #
 # IPC::Run maintains a cache, %cmd_cache, mapping commands to paths.  Tests
@@ -1744,13 +1785,32 @@ sub safe_psql
 
 	my ($stdout, $stderr);
 
-	my $ret = $self->psql(
-		$dbname, $sql,
-		%params,
-		stdout => \$stdout,
-		stderr => \$stderr,
-		on_error_die => 1,
-		on_error_stop => 1);
+
+	# If there are no special parameters, and re-use of sessions
+	# hasn't been disabled, create or re-use a long-lived psql
+	# process.
+	if (!%params && $self->{_use_auto_background_psqls})
+	{
+		if (!exists($self->{_auto_background_psqls}{$dbname}))
+		{
+			$self->{_auto_background_psqls}{$dbname} =
+			  $self->background_psql($dbname);
+		}
+		my $psql = $self->{_auto_background_psqls}{$dbname};
+		$stdout = $psql->query($sql);
+		chomp($stdout);
+		$stderr = $psql->{stderr};
+	}
+	else
+	{
+		my $ret = $self->psql(
+			$dbname, $sql,
+			%params,
+			stdout        => \$stdout,
+			stderr        => \$stderr,
+			on_error_die  => 1,
+			on_error_stop => 1);
+	}
 
 	# psql can emit stderr from NOTICEs etc
 	if ($stderr ne "")
-- 
2.39.2