From 4616732872339c0a39f0647b9455617f4d3ab719 Mon Sep 17 00:00:00 2001 From: "iwata.aya" Date: Thu, 11 Sep 2025 21:16:51 +0900 Subject: [PATCH v0009] 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..6f4fc57e3d9 100644 --- a/doc/src/sgml/bgworker.sgml +++ b/doc/src/sgml/bgworker.sgml @@ -108,6 +108,25 @@ typedef struct BackgroundWorker + + BGWORKER_EXIT_AT_DATABASE_CHANGE + + + BGWORKER_EXIT_AT_DATABASE_CHANGE + Requests termination of the background worker when its connected database is + dropped, renamed, moved to a different tablespace, or used as a template for + CREATE DATABASE. Specifically, the postmaster sends a + termination signal when any of these commands affect the worker's database: + DROP DATABASE, + ALTER DATABASE RENAME TO, + ALTER DATABASE SET TABLESPACE, or + CREATE DATABASE. + Requires both BGWORKER_SHMEM_ACCESS and + BGWORKER_BACKEND_DATABASE_CONNECTION. + + + + diff --git a/src/backend/postmaster/bgworker.c b/src/backend/postmaster/bgworker.c index 1ad65c237c3..462ac749c76 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" @@ -1396,3 +1397,42 @@ GetBackgroundWorkerTypeByPid(pid_t pid) return result; } + +/* + * Terminate all background workers connected to the given database, if they + * had requested it. + */ +void +TerminateBgWorkersByDbOid(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_EXIT_AT_DATABASE_CHANGE)) + { + 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 200f72c6e25..0044f8cbdf6 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" @@ -3689,8 +3691,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. @@ -3715,10 +3718,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; @@ -3768,6 +3780,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_EXIT_AT_DATABASE_CHANGE) + */ + TerminateBgWorkersByDbOid(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..8c69df432a5 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_EXIT_AT_DATABASE_CHANGE 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 TerminateBgWorkersByDbOid(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..9671f7e5770 --- /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_EXIT_AT_DATABASE_CHANGE +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_EXIT_AT_DATABASE_CHANGE 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 bea8339f464..2912abe6cce 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_EXIT_AT_DATABASE_CHANGE; + worker.bgw_start_time = BgWorkerStart_RecoveryFinished; worker.bgw_restart_time = BGW_NEVER_RESTART; sprintf(worker.bgw_library_name, "worker_spi"); -- 2.39.3