From f4d30412bd9588051a434a80404b077096a4cc98 Mon Sep 17 00:00:00 2001
From: "iwata.aya" <iwata.aya@fujitsu.com>
Date: Thu, 11 Sep 2025 21:16:51 +0900
Subject: [PATCH v0011 1/2] Allow background workers to be terminated at DROP
 DATABASE

---
 doc/src/sgml/bgworker.sgml                    |  19 +++
 src/backend/postmaster/bgworker.c             |  40 +++++
 src/backend/storage/ipc/procarray.c           |  28 +++-
 src/include/postmaster/bgworker.h             |   8 +
 src/test/modules/worker_spi/Makefile          |   4 +
 src/test/modules/worker_spi/meson.build       |   4 +
 .../worker_spi/t/002_worker_terminate.pl      | 140 ++++++++++++++++++
 .../modules/worker_spi/worker_spi--1.0.sql    |   3 +-
 src/test/modules/worker_spi/worker_spi.c      |   5 +
 9 files changed, 245 insertions(+), 6 deletions(-)
 create mode 100644 src/test/modules/worker_spi/t/002_worker_terminate.pl

diff --git a/doc/src/sgml/bgworker.sgml b/doc/src/sgml/bgworker.sgml
index 2c393385a91..a7c238750fc 100644
--- a/doc/src/sgml/bgworker.sgml
+++ b/doc/src/sgml/bgworker.sgml
@@ -108,6 +108,25 @@ typedef struct BackgroundWorker
      </listitem>
     </varlistentry>
 
+    <varlistentry>
+     <term><literal>BGWORKER_INTERRUPTABLE</literal></term>
+     <listitem>
+     <para>
+      <indexterm><primary>BGWORKER_INTERRUPTABLE</primary></indexterm>
+      Requests termination of the background worker when its connected database is
+      dropped, renamed, moved to a different tablespace, or used as a template for
+      <command>CREATE DATABASE</command>. Specifically, the postmaster sends a
+      termination signal when any of these commands affect the worker's database:
+      <command>DROP DATABASE</command>,
+      <command>ALTER DATABASE RENAME TO</command>,
+      <command>ALTER DATABASE SET TABLESPACE</command>, or
+      <command>CREATE DATABASE</command>.
+      Requires both <literal>BGWORKER_SHMEM_ACCESS</literal> and
+      <literal>BGWORKER_BACKEND_DATABASE_CONNECTION</literal>.
+     </para>
+     </listitem>
+    </varlistentry>
+
    </variablelist>
 
   </para>
diff --git a/src/backend/postmaster/bgworker.c b/src/backend/postmaster/bgworker.c
index 8e1068969ae..8334b75548c 100644
--- a/src/backend/postmaster/bgworker.c
+++ b/src/backend/postmaster/bgworker.c
@@ -26,6 +26,7 @@
 #include "storage/lwlock.h"
 #include "storage/pmsignal.h"
 #include "storage/proc.h"
+#include "storage/procarray.h"
 #include "storage/procsignal.h"
 #include "storage/shmem.h"
 #include "tcop/tcopprot.h"
@@ -1399,3 +1400,42 @@ GetBackgroundWorkerTypeByPid(pid_t pid)
 
 	return result;
 }
+
+/*
+ * Terminate all background workers connected to the given database, if they
+ * had requested it.
+ */
+void
+TerminateInterruptableBgWorkersByDbOid(Oid databaseId)
+{
+	bool		signal_postmaster = false;
+
+	LWLockAcquire(BackgroundWorkerLock, LW_EXCLUSIVE);
+
+	/*
+	 * Iterate through slots, looking for workers connected to the given
+	 * database.
+	 */
+	for (int slotno = 0; slotno < BackgroundWorkerData->total_slots; ++slotno)
+	{
+		BackgroundWorkerSlot *slot = &BackgroundWorkerData->slot[slotno];
+
+		if (slot->in_use &&
+			(slot->worker.bgw_flags & BGWORKER_INTERRUPTABLE))
+		{
+			PGPROC	   *proc = BackendPidGetProc(slot->pid);
+
+			if (proc && proc->databaseId == databaseId)
+			{
+				slot->terminate = true;
+				signal_postmaster = true;
+			}
+		}
+	}
+
+	LWLockRelease(BackgroundWorkerLock);
+
+	/* Make sure the postmaster notices the change to shared memory. */
+	if (signal_postmaster)
+		SendPostmasterSignal(PMSIGNAL_BACKGROUND_WORKER_CHANGE);
+}
diff --git a/src/backend/storage/ipc/procarray.c b/src/backend/storage/ipc/procarray.c
index f3a1603204e..de80422caea 100644
--- a/src/backend/storage/ipc/procarray.c
+++ b/src/backend/storage/ipc/procarray.c
@@ -56,11 +56,13 @@
 #include "catalog/pg_authid.h"
 #include "miscadmin.h"
 #include "pgstat.h"
+#include "postmaster/bgworker.h"
 #include "port/pg_lfind.h"
 #include "storage/proc.h"
 #include "storage/procarray.h"
 #include "utils/acl.h"
 #include "utils/builtins.h"
+#include "utils/injection_point.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
 #include "utils/snapmgr.h"
@@ -3687,8 +3689,9 @@ CountUserBackends(Oid roleid)
  * CountOtherDBBackends -- check for other backends running in the given DB
  *
  * If there are other backends in the DB, we will wait a maximum of 5 seconds
- * for them to exit.  Autovacuum backends are encouraged to exit early by
- * sending them SIGTERM, but normal user backends are just waited for.
+ * for them to exit.  Autovacuum backends and background workers are encouraged
+ * to exit early by sending them SIGTERM, but normal user backends are just
+ * waited for.
  *
  * The current backend is always ignored; it is caller's responsibility to
  * check whether the current backend uses the given DB, if it's important.
@@ -3713,10 +3716,19 @@ CountOtherDBBackends(Oid databaseId, int *nbackends, int *nprepared)
 
 #define MAXAUTOVACPIDS	10		/* max autovacs to SIGTERM per iteration */
 	int			autovac_pids[MAXAUTOVACPIDS];
-	int			tries;
 
-	/* 50 tries with 100ms sleep between tries makes 5 sec total wait */
-	for (tries = 0; tries < 50; tries++)
+	/*
+	 * Retry up to 50 times with 100ms between attempts (max 5s total). Can be
+	 * reduced to 3 attempts (max 0.3s total) to speed up tests.
+	 */
+	int			ntries = 50;
+
+#ifdef USE_INJECTION_POINTS
+	if (IS_INJECTION_POINT_ATTACHED("reduce-ncounts"))
+		ntries = 3;
+#endif
+
+	for (int tries = 0; tries < ntries; tries++)
 	{
 		int			nautovacs = 0;
 		bool		found = false;
@@ -3766,6 +3778,12 @@ CountOtherDBBackends(Oid databaseId, int *nbackends, int *nprepared)
 		for (index = 0; index < nautovacs; index++)
 			(void) kill(autovac_pids[index], SIGTERM);	/* ignore any error */
 
+		/*
+		 * Terminate all background workers for this database, if they had
+		 * requested it (BGWORKER_INTERRUPTABLE)
+		 */
+		TerminateInterruptableBgWorkersByDbOid(databaseId);
+
 		/* sleep, then try again */
 		pg_usleep(100 * 1000L); /* 100ms */
 	}
diff --git a/src/include/postmaster/bgworker.h b/src/include/postmaster/bgworker.h
index 058667a47a0..e43881448d5 100644
--- a/src/include/postmaster/bgworker.h
+++ b/src/include/postmaster/bgworker.h
@@ -59,6 +59,13 @@
  */
 #define BGWORKER_BACKEND_DATABASE_CONNECTION		0x0002
 
+/*
+ * Exit the bgworker if its database is involved in a CREATE, ALTER or DROP
+ * database command.
+ * Requires BGWORKER_SHMEM_ACCESS and BGWORKER_BACKEND_DATABASE_CONNECTION.
+ */
+#define BGWORKER_INTERRUPTABLE			0x0004
+
 /*
  * This class is used internally for parallel queries, to keep track of the
  * number of active parallel workers and make sure we never launch more than
@@ -128,6 +135,7 @@ extern const char *GetBackgroundWorkerTypeByPid(pid_t pid);
 
 /* Terminate a bgworker */
 extern void TerminateBackgroundWorker(BackgroundWorkerHandle *handle);
+extern void TerminateInterruptableBgWorkersByDbOid(Oid databaseId);
 
 /* This is valid in a running worker */
 extern PGDLLIMPORT BackgroundWorker *MyBgworkerEntry;
diff --git a/src/test/modules/worker_spi/Makefile b/src/test/modules/worker_spi/Makefile
index 024b34cdbb3..e7c5c059e32 100644
--- a/src/test/modules/worker_spi/Makefile
+++ b/src/test/modules/worker_spi/Makefile
@@ -6,6 +6,10 @@ EXTENSION = worker_spi
 DATA = worker_spi--1.0.sql
 PGFILEDESC = "worker_spi - background worker example"
 
+EXTRA_INSTALL = src/test/modules/injection_points
+
+export enable_injection_points
+
 TAP_TESTS = 1
 
 ifdef USE_PGXS
diff --git a/src/test/modules/worker_spi/meson.build b/src/test/modules/worker_spi/meson.build
index d673ece48a0..5ba66051396 100644
--- a/src/test/modules/worker_spi/meson.build
+++ b/src/test/modules/worker_spi/meson.build
@@ -26,8 +26,12 @@ tests += {
   'sd': meson.current_source_dir(),
   'bd': meson.current_build_dir(),
   'tap': {
+    'env': {
+      'enable_injection_points': get_option('injection_points') ? 'yes' : 'no',
+    },
     'tests': [
       't/001_worker_spi.pl',
+      't/002_worker_terminate.pl'
     ],
   },
 }
diff --git a/src/test/modules/worker_spi/t/002_worker_terminate.pl b/src/test/modules/worker_spi/t/002_worker_terminate.pl
new file mode 100644
index 00000000000..5467e423fd9
--- /dev/null
+++ b/src/test/modules/worker_spi/t/002_worker_terminate.pl
@@ -0,0 +1,140 @@
+# Copyright (c) 2025, PostgreSQL Global Development Group
+
+# Test background workers can be terminated by db commands
+
+use strict;
+use warnings FATAL => 'all';
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+# This test depends on injection points to detect whether background workers
+# remain.
+if ($ENV{enable_injection_points} ne 'yes')
+{
+	plan skip_all => 'Injection points not supported by this build';
+}
+
+# Ensure the worker_spi dynamic worker is launched on the specified database
+sub launch_bgworker
+{
+	my ($node, $database, $testcase, $request_terminate) = @_;
+	my $offset = -s $node->logfile;
+
+	# Launch a background worker on the given database
+	my $result = $node->safe_psql(
+		$database, qq(
+        SELECT worker_spi_launch($testcase, oid, 0, '{}', $request_terminate) IS NOT NULL
+        FROM pg_database WHERE datname = '$database';
+    ));
+	is($result, 't', "dynamic bgworker $testcase launched");
+
+	# Check the worker is surely initialized
+	$node->wait_for_log(
+		qr/LOG:  worker_spi dynamic worker $testcase initialized with .*\..*/,
+		$offset);
+}
+
+# Run the given query and verify the background worker can be terminated
+sub run_db_command
+{
+	my ($node, $command, $testname) = @_;
+	my $offset = -s $node->logfile;
+
+	$node->safe_psql('postgres', $command);
+
+	$node->wait_for_log(
+		qr/terminating background worker \"worker_spi dynamic\" due to administrator command/,
+		$offset);
+
+	note("background worker can be terminated at $testname");
+}
+
+my $node = PostgreSQL::Test::Cluster->new('mynode');
+$node->init;
+$node->start;
+
+# Check if the extension injection_points is available, as it may be
+# possible that this script is run with installcheck, where the module
+# would not be installed by default.
+if (!$node->check_extension('injection_points'))
+{
+	plan skip_all => 'Extension injection_points not installed';
+}
+
+$node->safe_psql('postgres', 'CREATE EXTENSION worker_spi;');
+
+# Launch a background worker without BGWORKER_INTERRUPTABLE
+launch_bgworker($node, 'postgres', 0, "false");
+
+# Ensure CREATE DATABASE WITH TEMPLATE fails because background worker retains
+
+# The injection point 'reduce-ncounts' reduces the number of backend
+# retries, allowing for shorter test runs. See CountOtherDBBackends().
+$node->safe_psql('postgres', "CREATE EXTENSION injection_points;");
+$node->safe_psql('postgres',
+	"SELECT injection_points_attach('reduce-ncounts', 'error');");
+
+my $stderr;
+
+$node->psql(
+	'postgres',
+	"CREATE DATABASE testdb WITH TEMPLATE postgres",
+	stderr => \$stderr);
+ok( $stderr =~
+	  "source database \"postgres\" is being accessed by other users",
+	"background worker blocked the database creation");
+
+# Confirm a background worker is still running
+my $result = $node->safe_psql(
+	"postgres", qq(
+        SELECT count(1) FROM pg_stat_activity
+		WHERE backend_type = 'worker_spi dynamic';));
+
+is($result, '1',
+	"background worker is still running after CREATE DATABASE WITH TEMPLATE");
+
+# Terminate the worker for upcoming tests
+$node->safe_psql(
+	"postgres", qq(
+        SELECT pg_terminate_backend(pid)
+        FROM pg_stat_activity WHERE backend_type = 'worker_spi dynamic';));
+
+# The injection point won't be used anymore, release it.
+$node->safe_psql('postgres',
+	"SELECT injection_points_detach('reduce-ncounts');");
+
+# Ensure BGWORKER_INTERRUPTABLE allows background workers to be
+# terminated at some database manipulations.
+#
+# Testcase 1: CREATE DATABASE WITH TEMPLATE
+launch_bgworker($node, 'postgres', 1, "true");
+run_db_command(
+	$node,
+	"CREATE DATABASE testdb WITH TEMPLATE postgres",
+	"CREATE DATABASE WITH TEMPLATE");
+
+# Testcase 2: ALTER DATABASE RENAME
+launch_bgworker($node, 'testdb', 2, "true");
+run_db_command(
+	$node,
+	"ALTER DATABASE testdb RENAME TO renameddb",
+	"ALTER DATABASE RENAME");
+
+# Preparation for the next test; create another tablespace
+my $tablespace = PostgreSQL::Test::Utils::tempdir;
+$node->safe_psql('postgres',
+	"CREATE TABLESPACE test_tablespace LOCATION '$tablespace'");
+
+# Testcase 3: ALTER DATABASE SET TABLESPACE
+launch_bgworker($node, 'renameddb', 3, "true");
+run_db_command(
+	$node,
+	"ALTER DATABASE renameddb SET TABLESPACE test_tablespace",
+	"ALTER DATABASE SET TABLESPACE");
+
+# Testcase 4: DROP DATABASE
+launch_bgworker($node, 'renameddb', 4, "true");
+run_db_command($node, "DROP DATABASE renameddb", "DROP DATABASE");
+
+done_testing();
diff --git a/src/test/modules/worker_spi/worker_spi--1.0.sql b/src/test/modules/worker_spi/worker_spi--1.0.sql
index 84deb6199f6..3d12de37bea 100644
--- a/src/test/modules/worker_spi/worker_spi--1.0.sql
+++ b/src/test/modules/worker_spi/worker_spi--1.0.sql
@@ -7,7 +7,8 @@
 CREATE FUNCTION worker_spi_launch(index int4,
   dboid oid DEFAULT 0,
   roleoid oid DEFAULT 0,
-  flags text[] DEFAULT '{}')
+  flags text[] DEFAULT '{}',
+  request_termination boolean DEFAULT false)
 RETURNS pg_catalog.int4 STRICT
 AS 'MODULE_PATHNAME'
 LANGUAGE C;
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index 1a21b8c8876..50d16aaf9a8 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -404,10 +404,15 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	Size		ndim;
 	int			nelems;
 	Datum	   *datum_flags;
+	bool		request_termination = PG_GETARG_BOOL(4);
 
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
 		BGWORKER_BACKEND_DATABASE_CONNECTION;
+
+	if (request_termination)
+		worker.bgw_flags |= BGWORKER_INTERRUPTABLE;
+
 	worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
 	worker.bgw_restart_time = BGW_NEVER_RESTART;
 	sprintf(worker.bgw_library_name, "worker_spi");
-- 
2.39.3

