From 8a827a8a59da4d3c2f2115151e1241cf663c11a1 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Fri, 12 Jul 2024 14:19:51 +0900
Subject: [PATCH v3] Add tap test for pg_signal_autovacuum role

---
 src/backend/postmaster/autovacuum.c           |  7 ++
 .../test_misc/t/006_signal_autovacuum.pl      | 94 +++++++++++++++++++
 2 files changed, 101 insertions(+)
 create mode 100644 src/test/modules/test_misc/t/006_signal_autovacuum.pl

diff --git a/src/backend/postmaster/autovacuum.c b/src/backend/postmaster/autovacuum.c
index 928754b51c..4e4a0ccbef 100644
--- a/src/backend/postmaster/autovacuum.c
+++ b/src/backend/postmaster/autovacuum.c
@@ -100,6 +100,7 @@
 #include "utils/fmgroids.h"
 #include "utils/fmgrprotos.h"
 #include "utils/guc_hooks.h"
+#include "utils/injection_point.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
 #include "utils/ps_status.h"
@@ -1902,6 +1903,12 @@ do_autovacuum(void)
 	/* Start a transaction so our commands have one to play into. */
 	StartTransactionCommand();
 
+	/*
+	 * This injection point is put in a transaction block to work with a wait
+	 * that uses a condition variable.
+	 */
+	INJECTION_POINT("autovacuum-worker-start");
+
 	/*
 	 * Compute the multixact age for which freezing is urgent.  This is
 	 * normally autovacuum_multixact_freeze_max_age, but may be less if we are
diff --git a/src/test/modules/test_misc/t/006_signal_autovacuum.pl b/src/test/modules/test_misc/t/006_signal_autovacuum.pl
new file mode 100644
index 0000000000..ded42b1be9
--- /dev/null
+++ b/src/test/modules/test_misc/t/006_signal_autovacuum.pl
@@ -0,0 +1,94 @@
+# Copyright (c) 2024, PostgreSQL Global Development Group
+
+# Test signaling autovacuum worker with pg_signal_autovacuum_worker.
+#
+# Only roles with privileges of pg_signal_autovacuum_worker are allowed to
+# signal autovacuum workers.  This test uses an injection point located
+# at the beginning of the autovacuum worker startup.
+
+use strict;
+use warnings;
+use PostgreSQL::Test::Cluster;
+use Test::More;
+
+if ($ENV{enable_injection_points} ne 'yes')
+{
+	plan skip_all => 'Injection points not supported by this build';
+}
+
+# Initialize postgres
+my $psql_err = '';
+my $psql_out = '';
+my $node = PostgreSQL::Test::Cluster->new('node');
+$node->init;
+
+# This ensures a quick worker spawn.
+$node->append_conf(
+	'postgresql.conf', 'autovacuum_naptime = 1');
+$node->start;
+$node->safe_psql('postgres', 'CREATE EXTENSION injection_points;');
+
+$node->safe_psql(
+	'postgres', qq(
+    CREATE ROLE regress_regular_role;
+    CREATE ROLE regress_worker_role;
+    GRANT pg_signal_autovacuum_worker TO regress_worker_role;
+));
+
+# From this point, autovacuum worker will wait at startup.
+$node->safe_psql('postgres',
+	"SELECT injection_points_attach('autovacuum-worker-start', 'wait');");
+
+$node->reload();
+
+# Wait until an autovacuum worker starts.
+$node->wait_for_event('autovacuum worker', 'autovacuum-worker-start');
+
+my $av_pid = $node->safe_psql(
+	'postgres', qq(
+    SELECT pid FROM pg_stat_activity WHERE backend_type = 'autovacuum worker';
+));
+
+# Regular role cannot terminate autovacuum worker.
+my $terminate_with_no_pg_signal_av = $node->psql(
+	'postgres', qq(
+    SET ROLE regress_regular_role;
+    SELECT pg_terminate_backend($av_pid);
+),
+	stdout => \$psql_out,
+	stderr => \$psql_err);
+
+like(
+	$psql_err,
+	qr/ERROR:  permission denied to terminate process\nDETAIL:  Only roles with privileges of the "pg_signal_autovacuum_worker" role may terminate autovacuum workers./,
+	"autovacuum worker not signaled with regular role");
+
+my $offset = -s $node->logfile;
+
+# Role with pg_signal_autovacuum can terminate autovacuum worker.
+my $terminate_with_pg_signal_av = $node->psql(
+	'postgres', qq(
+    SET ROLE regress_worker_role;
+    SELECT pg_terminate_backend($av_pid);
+),
+	stdout => \$psql_out,
+	stderr => \$psql_err);
+
+# Wait for the autovacuum worker to exit before scanning the logs.
+$node->poll_query_until('postgres',
+	"SELECT count(*) = 0 FROM pg_stat_activity "
+	. "WHERE pid = $av_pid AND backend_type = 'autovacuum worker';");
+
+# Check that the primary server logs a FATAL indicating that autovacuum
+# is terminated.
+ok( $node->log_contains(
+		qr/FATAL:  terminating autovacuum process due to administrator command/,
+		$offset),
+	"autovacuum worker signaled with pg_signal_autovacuum_worker granted"
+);
+
+# Release injection point.
+$node->safe_psql('postgres',
+	"SELECT injection_points_detach('autovacuum-worker-start');");
+
+done_testing();
-- 
2.45.2