v1-0001-Introduce-pg_signal_autovacuum-role-to-allow-non-sup.patch

From 763f9052f291f4cfe32197379f58b9d9bc35ebb7 Mon Sep 17 00:00:00 2001
From: Anthony Leung <antholeu@amazon.com>
Date: Mon, 1 Apr 2024 04:46:55 +0000
Subject: [PATCH] Introduce pg_signal_autovacuum role to allow non-superuser to
 termiante autovacuum worker

Non-superusers are currently not allowed to termiante autovacuum
backends which can create friction in some scenarios when performing
maintenance. This commit introduce a new pre-defined role
'pg_signal_autovacuum'. Non-superusers granted with this role are
allowed to terminate backends that are running autovacuum.
---
 doc/src/sgml/user-manag.sgml                |  4 ++
 src/backend/storage/ipc/signalfuncs.c       | 31 ++++++++---
 src/backend/utils/activity/backend_status.c | 20 +++++++
 src/include/catalog/pg_authid.dat           |  5 ++
 src/include/utils/backend_status.h          |  2 +-
 src/test/signals/.gitignore                 |  2 +
 src/test/signals/Makefile                   | 23 ++++++++
 src/test/signals/t/001_signal_autovacuum.pl | 59 +++++++++++++++++++++
 8 files changed, 137 insertions(+), 9 deletions(-)
 create mode 100644 src/test/signals/.gitignore
 create mode 100644 src/test/signals/Makefile
 create mode 100644 src/test/signals/t/001_signal_autovacuum.pl

diff --git a/doc/src/sgml/user-manag.sgml b/doc/src/sgml/user-manag.sgml
index 1f7d7e75ce..541a45ab4a 100644
--- a/doc/src/sgml/user-manag.sgml
+++ b/doc/src/sgml/user-manag.sgml
@@ -693,6 +693,10 @@ DROP ROLE doomed_role;
        database to issue
        <link linkend="sql-createsubscription"><command>CREATE SUBSCRIPTION</command></link>.</entry>
       </row>
+      <row>
+       <entry>pg_signal_autovacuum</entry>
+       <entry>Allow terminating backend running autovacuum</entry>
+      </row>
      </tbody>
     </tgroup>
    </table>
diff --git a/src/backend/storage/ipc/signalfuncs.c b/src/backend/storage/ipc/signalfuncs.c
index b595c2d691..a9f0466c46 100644
--- a/src/backend/storage/ipc/signalfuncs.c
+++ b/src/backend/storage/ipc/signalfuncs.c
@@ -79,14 +79,29 @@ pg_signal_backend(int pid, int sig)
 	 * not advertising a role might have the importance of a superuser-owned
 	 * backend, so treat it that way.
 	 */
-	if ((!OidIsValid(proc->roleId) || superuser_arg(proc->roleId)) &&
-		!superuser())
-		return SIGNAL_BACKEND_NOSUPERUSER;
-
-	/* Users can signal backends they have role membership in. */
-	if (!has_privs_of_role(GetUserId(), proc->roleId) &&
-		!has_privs_of_role(GetUserId(), ROLE_PG_SIGNAL_BACKEND))
-		return SIGNAL_BACKEND_NOPERMISSION;
+	if (!superuser())
+	{
+		if (!OidIsValid(proc->roleId))
+		{
+			/*
+			 * We only allow user with pg_signal_autovacuum role to terminate
+			 * autovacuum worker as an exception. 
+			 */
+			if (!(pg_stat_is_backend_autovac_worker(proc->backendId) &&
+				has_privs_of_role(GetUserId(), ROLE_PG_SIGNAL_AUTOVACUUM)))
+				return SIGNAL_BACKEND_NOSUPERUSER;
+		}
+		else
+		{
+			if (superuser_arg(proc->roleId))
+				return SIGNAL_BACKEND_NOSUPERUSER;
+
+			/* Users can signal backends they have role membership in. */
+			if (!has_privs_of_role(GetUserId(), proc->roleId) &&
+				!has_privs_of_role(GetUserId(), ROLE_PG_SIGNAL_BACKEND))
+				return SIGNAL_BACKEND_NOPERMISSION;
+		}
+	}
 
 	/*
 	 * Can the process we just validated above end, followed by the pid being
diff --git a/src/backend/utils/activity/backend_status.c b/src/backend/utils/activity/backend_status.c
index f78ccc204b..b38ba319d2 100644
--- a/src/backend/utils/activity/backend_status.c
+++ b/src/backend/utils/activity/backend_status.c
@@ -1057,6 +1057,26 @@ pgstat_get_my_query_id(void)
 	return MyBEEntry->st_query_id;
 }
 
+/* ----------
+ * pg_stat_is_backend_autovac_worker() -
+ *
+ * Return whether the backend of the given backend id is of type autovacuum worker.
+ */
+bool
+pg_stat_is_backend_autovac_worker(BackendId beid)
+{
+	PgBackendStatus *ret;
+
+	Assert(beid != InvalidBackendId);
+
+	ret = pgstat_get_beentry_by_backend_id(beid);
+
+	if (!ret)
+		return false;
+
+	return ret->st_backendType == B_AUTOVAC_WORKER;
+}
+
 /* ----------
  * cmp_lbestatus
  *
diff --git a/src/include/catalog/pg_authid.dat b/src/include/catalog/pg_authid.dat
index 6b4a0aaaad..556a38f702 100644
--- a/src/include/catalog/pg_authid.dat
+++ b/src/include/catalog/pg_authid.dat
@@ -94,5 +94,10 @@
   rolcreaterole => 'f', rolcreatedb => 'f', rolcanlogin => 'f',
   rolreplication => 'f', rolbypassrls => 'f', rolconnlimit => '-1',
   rolpassword => '_null_', rolvaliduntil => '_null_' },
+{ oid => '6312', oid_symbol => 'ROLE_PG_SIGNAL_AUTOVACUUM',
+  rolname => 'pg_signal_autovacuum', rolsuper => 'f', rolinherit => 't',
+  rolcreaterole => 'f', rolcreatedb => 'f', rolcanlogin => 'f',
+  rolreplication => 'f', rolbypassrls => 'f', rolconnlimit => '-1',
+  rolpassword => '_null_', rolvaliduntil => '_null_' },
 
 ]
diff --git a/src/include/utils/backend_status.h b/src/include/utils/backend_status.h
index d51c840daf..a84aafabd9 100644
--- a/src/include/utils/backend_status.h
+++ b/src/include/utils/backend_status.h
@@ -325,7 +325,7 @@ extern const char *pgstat_get_backend_current_activity(int pid, bool checkUser);
 extern const char *pgstat_get_crashed_backend_activity(int pid, char *buffer,
 													   int buflen);
 extern uint64 pgstat_get_my_query_id(void);
-
+extern bool pg_stat_is_backend_autovac_worker(BackendId beid);
 
 /* ----------
  * Support functions for the SQL-callable functions to
diff --git a/src/test/signals/.gitignore b/src/test/signals/.gitignore
new file mode 100644
index 0000000000..871e943d50
--- /dev/null
+++ b/src/test/signals/.gitignore
@@ -0,0 +1,2 @@
+# Generated by test suite
+/tmp_check/
diff --git a/src/test/signals/Makefile b/src/test/signals/Makefile
new file mode 100644
index 0000000000..688f6a7e9e
--- /dev/null
+++ b/src/test/signals/Makefile
@@ -0,0 +1,23 @@
+#-------------------------------------------------------------------------
+#
+# Makefile for src/test/recovery
+#
+# Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
+# Portions Copyright (c) 1994, Regents of the University of California
+#
+# src/test/recovery/Makefile
+#
+#-------------------------------------------------------------------------
+
+subdir = src/test/signals
+top_builddir = ../../..
+include $(top_builddir)/src/Makefile.global
+
+check:
+	$(prove_check)
+
+installcheck:
+	$(prove_installcheck)
+
+clean distclean maintainer-clean:
+	rm -rf tmp_check
diff --git a/src/test/signals/t/001_signal_autovacuum.pl b/src/test/signals/t/001_signal_autovacuum.pl
new file mode 100644
index 0000000000..93432b3a6b
--- /dev/null
+++ b/src/test/signals/t/001_signal_autovacuum.pl
@@ -0,0 +1,59 @@
+# Copyright (c) 2024-2024, PostgreSQL Global Development Group
+
+# Minimal test testing pg_signal_autovacuum role.
+use strict;
+use warnings;
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+# Initialize postgres
+my $node = PostgreSQL::Test::Cluster->new('node');
+$node->init();
+$node->start;
+
+$node->safe_psql('postgres', qq(
+    CREATE DATABASE regress;
+    CREATE ROLE psa_reg_role_1;
+    CREATE ROLE psa_reg_role_2;
+    GRANT pg_signal_backend TO psa_reg_role_1;
+    GRANT pg_signal_autovacuum TO psa_reg_role_2;
+));
+
+# Create some content in a table and set autovacuum setting such that
+# it would be triggered.
+$node->safe_psql('regress', qq(
+    CREATE TABLE t(i int);
+    INSERT INTO t SELECT * FROM generate_series(1, 1000000);
+    ALTER SYSTEM SET autovacuum_vacuum_cost_limit TO 1;
+    ALTER SYSTEM SET autovacuum_vacuum_cost_delay TO 100;
+    ALTER SYSTEM SET autovacuum_naptime TO 1;    
+));
+
+$node->restart;
+
+# wait for autovacuum to start.
+
+sleep 1;
+
+my $res_pid = $node->safe_psql('regress', qq(
+    SELECT pid FROM pg_stat_activity WHERE backend_type = 'autovacuum worker' and datname = 'regress';;
+));
+
+
+my ($res_reg_psa_1, $stdout_reg_psa_1, $stderr_reg_psa_1)  = $node->psql('regress', qq(
+    SET ROLE psa_reg_role_1;
+    SELECT pg_terminate_backend($res_pid);
+));
+
+ok($res_reg_psa_1 != 0, "should fail for non pg_signal_autovacuum");
+like($stderr_reg_psa_1, qr/Only roles with the SUPERUSER attribute may terminate processes of roles with the SUPERUSER attribute./, "matches");
+
+my ($res_reg_psa_2, $stdout_reg_psa_2, $stderr_reg_psa_2) = $node->psql('regress', qq(
+    SET ROLE psa_reg_role_2;
+    SELECT pg_terminate_backend($res_pid);
+));
+
+ok($res_reg_psa_2 == 0, "should success for pg_signal_autovacuum");
+
+done_testing();
-- 
2.40.1