Skip to content

Commit ff4c6af

Browse files
ayaiwtCommitfest Bot
authored andcommitted
Allow background workers to be terminated at DROP DATABASE
1 parent b4cbc10 commit ff4c6af

File tree

9 files changed

+245
-6
lines changed

9 files changed

+245
-6
lines changed

doc/src/sgml/bgworker.sgml

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,25 @@ typedef struct BackgroundWorker
108108
</listitem>
109109
</varlistentry>
110110

111+
<varlistentry>
112+
<term><literal>BGWORKER_EXIT_AT_DATABASE_CHANGE</literal></term>
113+
<listitem>
114+
<para>
115+
<indexterm><primary>BGWORKER_EXIT_AT_DATABASE_CHANGE</primary></indexterm>
116+
Requests termination of the background worker when its connected database is
117+
dropped, renamed, moved to a different tablespace, or used as a template for
118+
<command>CREATE DATABASE</command>. Specifically, the postmaster sends a
119+
termination signal when any of these commands affect the worker's database:
120+
<command>DROP DATABASE</command>,
121+
<command>ALTER DATABASE RENAME TO</command>,
122+
<command>ALTER DATABASE SET TABLESPACE</command>, or
123+
<command>CREATE DATABASE</command>.
124+
Requires both <literal>BGWORKER_SHMEM_ACCESS</literal> and
125+
<literal>BGWORKER_BACKEND_DATABASE_CONNECTION</literal>.
126+
</para>
127+
</listitem>
128+
</varlistentry>
129+
111130
</variablelist>
112131

113132
</para>

src/backend/postmaster/bgworker.c

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
#include "storage/lwlock.h"
2727
#include "storage/pmsignal.h"
2828
#include "storage/proc.h"
29+
#include "storage/procarray.h"
2930
#include "storage/procsignal.h"
3031
#include "storage/shmem.h"
3132
#include "tcop/tcopprot.h"
@@ -1399,3 +1400,42 @@ GetBackgroundWorkerTypeByPid(pid_t pid)
13991400

14001401
return result;
14011402
}
1403+
1404+
/*
1405+
* Terminate all background workers connected to the given database, if they
1406+
* had requested it.
1407+
*/
1408+
void
1409+
TerminateBgWorkersByDbOid(Oid databaseId)
1410+
{
1411+
bool signal_postmaster = false;
1412+
1413+
LWLockAcquire(BackgroundWorkerLock, LW_EXCLUSIVE);
1414+
1415+
/*
1416+
* Iterate through slots, looking for workers connected to the given
1417+
* database.
1418+
*/
1419+
for (int slotno = 0; slotno < BackgroundWorkerData->total_slots; ++slotno)
1420+
{
1421+
BackgroundWorkerSlot *slot = &BackgroundWorkerData->slot[slotno];
1422+
1423+
if (slot->in_use &&
1424+
(slot->worker.bgw_flags & BGWORKER_EXIT_AT_DATABASE_CHANGE))
1425+
{
1426+
PGPROC *proc = BackendPidGetProc(slot->pid);
1427+
1428+
if (proc && proc->databaseId == databaseId)
1429+
{
1430+
slot->terminate = true;
1431+
signal_postmaster = true;
1432+
}
1433+
}
1434+
}
1435+
1436+
LWLockRelease(BackgroundWorkerLock);
1437+
1438+
/* Make sure the postmaster notices the change to shared memory. */
1439+
if (signal_postmaster)
1440+
SendPostmasterSignal(PMSIGNAL_BACKGROUND_WORKER_CHANGE);
1441+
}

src/backend/storage/ipc/procarray.c

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -56,11 +56,13 @@
5656
#include "catalog/pg_authid.h"
5757
#include "miscadmin.h"
5858
#include "pgstat.h"
59+
#include "postmaster/bgworker.h"
5960
#include "port/pg_lfind.h"
6061
#include "storage/proc.h"
6162
#include "storage/procarray.h"
6263
#include "utils/acl.h"
6364
#include "utils/builtins.h"
65+
#include "utils/injection_point.h"
6466
#include "utils/lsyscache.h"
6567
#include "utils/rel.h"
6668
#include "utils/snapmgr.h"
@@ -3687,8 +3689,9 @@ CountUserBackends(Oid roleid)
36873689
* CountOtherDBBackends -- check for other backends running in the given DB
36883690
*
36893691
* If there are other backends in the DB, we will wait a maximum of 5 seconds
3690-
* for them to exit. Autovacuum backends are encouraged to exit early by
3691-
* sending them SIGTERM, but normal user backends are just waited for.
3692+
* for them to exit. Autovacuum backends and background workers are encouraged
3693+
* to exit early by sending them SIGTERM, but normal user backends are just
3694+
* waited for.
36923695
*
36933696
* The current backend is always ignored; it is caller's responsibility to
36943697
* check whether the current backend uses the given DB, if it's important.
@@ -3713,10 +3716,19 @@ CountOtherDBBackends(Oid databaseId, int *nbackends, int *nprepared)
37133716

37143717
#define MAXAUTOVACPIDS 10 /* max autovacs to SIGTERM per iteration */
37153718
int autovac_pids[MAXAUTOVACPIDS];
3716-
int tries;
37173719

3718-
/* 50 tries with 100ms sleep between tries makes 5 sec total wait */
3719-
for (tries = 0; tries < 50; tries++)
3720+
/*
3721+
* Retry up to 50 times with 100ms between attempts (max 5s total). Can be
3722+
* reduced to 3 attempts (max 0.3s total) to speed up tests.
3723+
*/
3724+
int ntries = 50;
3725+
3726+
#ifdef USE_INJECTION_POINTS
3727+
if (IS_INJECTION_POINT_ATTACHED("reduce-ncounts"))
3728+
ntries = 3;
3729+
#endif
3730+
3731+
for (int tries = 0; tries < ntries; tries++)
37203732
{
37213733
int nautovacs = 0;
37223734
bool found = false;
@@ -3766,6 +3778,12 @@ CountOtherDBBackends(Oid databaseId, int *nbackends, int *nprepared)
37663778
for (index = 0; index < nautovacs; index++)
37673779
(void) kill(autovac_pids[index], SIGTERM); /* ignore any error */
37683780

3781+
/*
3782+
* Terminate all background workers for this database, if they had
3783+
* requested it (BGWORKER_EXIT_AT_DATABASE_CHANGE)
3784+
*/
3785+
TerminateBgWorkersByDbOid(databaseId);
3786+
37693787
/* sleep, then try again */
37703788
pg_usleep(100 * 1000L); /* 100ms */
37713789
}

src/include/postmaster/bgworker.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,13 @@
5959
*/
6060
#define BGWORKER_BACKEND_DATABASE_CONNECTION 0x0002
6161

62+
/*
63+
* Exit the bgworker if its database is involved in a CREATE, ALTER or DROP
64+
* database command.
65+
* Requires BGWORKER_SHMEM_ACCESS and BGWORKER_BACKEND_DATABASE_CONNECTION.
66+
*/
67+
#define BGWORKER_EXIT_AT_DATABASE_CHANGE 0x0004
68+
6269
/*
6370
* This class is used internally for parallel queries, to keep track of the
6471
* number of active parallel workers and make sure we never launch more than
@@ -128,6 +135,7 @@ extern const char *GetBackgroundWorkerTypeByPid(pid_t pid);
128135

129136
/* Terminate a bgworker */
130137
extern void TerminateBackgroundWorker(BackgroundWorkerHandle *handle);
138+
extern void TerminateBgWorkersByDbOid(Oid databaseId);
131139

132140
/* This is valid in a running worker */
133141
extern PGDLLIMPORT BackgroundWorker *MyBgworkerEntry;

src/test/modules/worker_spi/Makefile

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,10 @@ EXTENSION = worker_spi
66
DATA = worker_spi--1.0.sql
77
PGFILEDESC = "worker_spi - background worker example"
88

9+
EXTRA_INSTALL = src/test/modules/injection_points
10+
11+
export enable_injection_points
12+
913
TAP_TESTS = 1
1014

1115
ifdef USE_PGXS

src/test/modules/worker_spi/meson.build

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,12 @@ tests += {
2626
'sd': meson.current_source_dir(),
2727
'bd': meson.current_build_dir(),
2828
'tap': {
29+
'env': {
30+
'enable_injection_points': get_option('injection_points') ? 'yes' : 'no',
31+
},
2932
'tests': [
3033
't/001_worker_spi.pl',
34+
't/002_worker_terminate.pl'
3135
],
3236
},
3337
}
Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
# Copyright (c) 2025, PostgreSQL Global Development Group
2+
3+
# Test background workers can be terminated by db commands
4+
5+
use strict;
6+
use warnings FATAL => 'all';
7+
use PostgreSQL::Test::Cluster;
8+
use PostgreSQL::Test::Utils;
9+
use Test::More;
10+
11+
# This test depends on injection points to detect whether background workers
12+
# remain.
13+
if ($ENV{enable_injection_points} ne 'yes')
14+
{
15+
plan skip_all => 'Injection points not supported by this build';
16+
}
17+
18+
# Ensure the worker_spi dynamic worker is launched on the specified database
19+
sub launch_bgworker
20+
{
21+
my ($node, $database, $testcase, $request_terminate) = @_;
22+
my $offset = -s $node->logfile;
23+
24+
# Launch a background worker on the given database
25+
my $result = $node->safe_psql(
26+
$database, qq(
27+
SELECT worker_spi_launch($testcase, oid, 0, '{}', $request_terminate) IS NOT NULL
28+
FROM pg_database WHERE datname = '$database';
29+
));
30+
is($result, 't', "dynamic bgworker $testcase launched");
31+
32+
# Check the worker is surely initialized
33+
$node->wait_for_log(
34+
qr/LOG: worker_spi dynamic worker $testcase initialized with .*\..*/,
35+
$offset);
36+
}
37+
38+
# Run the given query and verify the background worker can be terminated
39+
sub run_db_command
40+
{
41+
my ($node, $command, $testname) = @_;
42+
my $offset = -s $node->logfile;
43+
44+
$node->safe_psql('postgres', $command);
45+
46+
$node->wait_for_log(
47+
qr/terminating background worker \"worker_spi dynamic\" due to administrator command/,
48+
$offset);
49+
50+
note("background worker can be terminated at $testname");
51+
}
52+
53+
my $node = PostgreSQL::Test::Cluster->new('mynode');
54+
$node->init;
55+
$node->start;
56+
57+
# Check if the extension injection_points is available, as it may be
58+
# possible that this script is run with installcheck, where the module
59+
# would not be installed by default.
60+
if (!$node->check_extension('injection_points'))
61+
{
62+
plan skip_all => 'Extension injection_points not installed';
63+
}
64+
65+
$node->safe_psql('postgres', 'CREATE EXTENSION worker_spi;');
66+
67+
# Launch a background worker without BGWORKER_EXIT_AT_DATABASE_CHANGE
68+
launch_bgworker($node, 'postgres', 0, "false");
69+
70+
# Ensure CREATE DATABASE WITH TEMPLATE fails because background worker retains
71+
72+
# The injection point 'reduce-ncounts' reduces the number of backend
73+
# retries, allowing for shorter test runs. See CountOtherDBBackends().
74+
$node->safe_psql('postgres', "CREATE EXTENSION injection_points;");
75+
$node->safe_psql('postgres',
76+
"SELECT injection_points_attach('reduce-ncounts', 'error');");
77+
78+
my $stderr;
79+
80+
$node->psql(
81+
'postgres',
82+
"CREATE DATABASE testdb WITH TEMPLATE postgres",
83+
stderr => \$stderr);
84+
ok( $stderr =~
85+
"source database \"postgres\" is being accessed by other users",
86+
"background worker blocked the database creation");
87+
88+
# Confirm a background worker is still running
89+
my $result = $node->safe_psql(
90+
"postgres", qq(
91+
SELECT count(1) FROM pg_stat_activity
92+
WHERE backend_type = 'worker_spi dynamic';));
93+
94+
is($result, '1',
95+
"background worker is still running after CREATE DATABASE WITH TEMPLATE");
96+
97+
# Terminate the worker for upcoming tests
98+
$node->safe_psql(
99+
"postgres", qq(
100+
SELECT pg_terminate_backend(pid)
101+
FROM pg_stat_activity WHERE backend_type = 'worker_spi dynamic';));
102+
103+
# The injection point won't be used anymore, release it.
104+
$node->safe_psql('postgres',
105+
"SELECT injection_points_detach('reduce-ncounts');");
106+
107+
# Ensure BGWORKER_EXIT_AT_DATABASE_CHANGE allows background workers to be
108+
# terminated at some database manipulations.
109+
#
110+
# Testcase 1: CREATE DATABASE WITH TEMPLATE
111+
launch_bgworker($node, 'postgres', 1, "true");
112+
run_db_command(
113+
$node,
114+
"CREATE DATABASE testdb WITH TEMPLATE postgres",
115+
"CREATE DATABASE WITH TEMPLATE");
116+
117+
# Testcase 2: ALTER DATABASE RENAME
118+
launch_bgworker($node, 'testdb', 2, "true");
119+
run_db_command(
120+
$node,
121+
"ALTER DATABASE testdb RENAME TO renameddb",
122+
"ALTER DATABASE RENAME");
123+
124+
# Preparation for the next test; create another tablespace
125+
my $tablespace = PostgreSQL::Test::Utils::tempdir;
126+
$node->safe_psql('postgres',
127+
"CREATE TABLESPACE test_tablespace LOCATION '$tablespace'");
128+
129+
# Testcase 3: ALTER DATABASE SET TABLESPACE
130+
launch_bgworker($node, 'renameddb', 3, "true");
131+
run_db_command(
132+
$node,
133+
"ALTER DATABASE renameddb SET TABLESPACE test_tablespace",
134+
"ALTER DATABASE SET TABLESPACE");
135+
136+
# Testcase 4: DROP DATABASE
137+
launch_bgworker($node, 'renameddb', 4, "true");
138+
run_db_command($node, "DROP DATABASE renameddb", "DROP DATABASE");
139+
140+
done_testing();

src/test/modules/worker_spi/worker_spi--1.0.sql

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
CREATE FUNCTION worker_spi_launch(index int4,
88
dboid oid DEFAULT 0,
99
roleoid oid DEFAULT 0,
10-
flags text[] DEFAULT '{}')
10+
flags text[] DEFAULT '{}',
11+
request_termination boolean DEFAULT false)
1112
RETURNS pg_catalog.int4 STRICT
1213
AS 'MODULE_PATHNAME'
1314
LANGUAGE C;

src/test/modules/worker_spi/worker_spi.c

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -404,10 +404,15 @@ worker_spi_launch(PG_FUNCTION_ARGS)
404404
Size ndim;
405405
int nelems;
406406
Datum *datum_flags;
407+
bool request_termination = PG_GETARG_BOOL(4);
407408

408409
memset(&worker, 0, sizeof(worker));
409410
worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
410411
BGWORKER_BACKEND_DATABASE_CONNECTION;
412+
413+
if (request_termination)
414+
worker.bgw_flags |= BGWORKER_EXIT_AT_DATABASE_CHANGE;
415+
411416
worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
412417
worker.bgw_restart_time = BGW_NEVER_RESTART;
413418
sprintf(worker.bgw_library_name, "worker_spi");

0 commit comments

Comments
 (0)