From 5b15c5938e8641d0103d225246dbe03911e0ea3f Mon Sep 17 00:00:00 2001 From: roman khapov Date: Sat, 13 Dec 2025 07:06:00 +0000 Subject: [PATCH 2/2] pg_term_reason: POC for pg_terminate_backend_msg This commits adds proof-of-concept implementation for extensions, that defines some function to terminate connections with additional message from admin. Ex.: postgres=# create extension pg_term_reason; CREATE EXTENSION postgres=# select pg_terminate_backend_msg(pg_backend_pid(), 0, 'Have you seen my coffee cup?'); FATAL: terminating connection due to administrator command: Have you seen my coffee cup? Signed-off-by: roman khapov --- contrib/Makefile | 1 + contrib/meson.build | 1 + contrib/pg_term_reason/.gitignore | 4 + contrib/pg_term_reason/Makefile | 22 +++++ contrib/pg_term_reason/meson.build | 23 +++++ .../pg_term_reason/pg_term_reason--1.0.sql | 14 +++ contrib/pg_term_reason/pg_term_reason.c | 89 +++++++++++++++++++ contrib/pg_term_reason/pg_term_reason.control | 5 ++ contrib/pg_term_reason/t/001_basic.pl | 28 ++++++ 9 files changed, 187 insertions(+) create mode 100644 contrib/pg_term_reason/.gitignore create mode 100644 contrib/pg_term_reason/Makefile create mode 100644 contrib/pg_term_reason/meson.build create mode 100644 contrib/pg_term_reason/pg_term_reason--1.0.sql create mode 100644 contrib/pg_term_reason/pg_term_reason.c create mode 100644 contrib/pg_term_reason/pg_term_reason.control create mode 100644 contrib/pg_term_reason/t/001_basic.pl diff --git a/contrib/Makefile b/contrib/Makefile index 2f0a88d3f7..82eb2c0d12 100644 --- a/contrib/Makefile +++ b/contrib/Makefile @@ -37,6 +37,7 @@ SUBDIRS = \ pg_prewarm \ pg_stat_statements \ pg_surgery \ + pg_term_reason \ pg_trgm \ pgrowlocks \ pgstattuple \ diff --git a/contrib/meson.build b/contrib/meson.build index ed30ee7d63..398d58d71e 100644 --- a/contrib/meson.build +++ b/contrib/meson.build @@ -53,6 +53,7 @@ subdir('pgrowlocks') subdir('pg_stat_statements') subdir('pgstattuple') subdir('pg_surgery') +subdir('pg_term_reason') subdir('pg_trgm') subdir('pg_visibility') subdir('pg_walinspect') diff --git a/contrib/pg_term_reason/.gitignore b/contrib/pg_term_reason/.gitignore new file mode 100644 index 0000000000..5dcb3ff972 --- /dev/null +++ b/contrib/pg_term_reason/.gitignore @@ -0,0 +1,4 @@ +# Generated subdirectories +/log/ +/results/ +/tmp_check/ diff --git a/contrib/pg_term_reason/Makefile b/contrib/pg_term_reason/Makefile new file mode 100644 index 0000000000..fa751801f8 --- /dev/null +++ b/contrib/pg_term_reason/Makefile @@ -0,0 +1,22 @@ +# contrib/pg_term_reason/Makefile + +MODULE_big = pg_term_reason +OBJS = \ + pg_term_reason.o + +EXTENSION = pg_term_reason +DATA = pg_term_reason--1.0.sql +PGFILEDESC = "pg_term_reason - add pg_terminate_backend_reasoned function" + +TAP_TESTS = 1 + +ifdef USE_PGXS +PG_CONFIG = pg_config +PGXS := $(shell $(PG_CONFIG) --pgxs) +include $(PGXS) +else +subdir = contrib/pg_term_reason +top_builddir = ../.. +include $(top_builddir)/src/Makefile.global +include $(top_srcdir)/contrib/contrib-global.mk +endif diff --git a/contrib/pg_term_reason/meson.build b/contrib/pg_term_reason/meson.build new file mode 100644 index 0000000000..ad93b4b5f2 --- /dev/null +++ b/contrib/pg_term_reason/meson.build @@ -0,0 +1,23 @@ +pg_term_sources = files( + 'pg_term_reason.c', +) + +pg_term_reason = shared_module('pg_term_reason', + pg_term_sources, + kwargs: contrib_mod_args, +) +contrib_targets += pg_term_reason + +install_data( + 'pg_term_reason--1.0.sql', + 'pg_term_reason.control', + kwargs: contrib_data_args, +) + +tests += { + 'tap': { + 'tests': [ + 't/001_basic.pl', + ], + }, +} diff --git a/contrib/pg_term_reason/pg_term_reason--1.0.sql b/contrib/pg_term_reason/pg_term_reason--1.0.sql new file mode 100644 index 0000000000..e0f1bd5779 --- /dev/null +++ b/contrib/pg_term_reason/pg_term_reason--1.0.sql @@ -0,0 +1,14 @@ +/* contrib/pg_term_reason/pg_term_reason--1.0.sql */ + +-- complain if script is sourced in psql, rather than via CREATE EXTENSION +\echo Use "CREATE EXTENSION pg_term_reason" to load this file. \quit + +CREATE FUNCTION pg_terminate_backend_msg(pid integer, timeout int8 DEFAULT 0, reason text DEFAULT '') +RETURNS boolean +AS 'MODULE_PATHNAME', 'pg_terminate_backend_msg' +LANGUAGE C STRICT; + +CREATE FUNCTION pg_cancel_backend_msg(pid integer, reason text DEFAULT '') +RETURNS boolean +AS 'MODULE_PATHNAME', 'pg_cancel_backend_msg' +LANGUAGE C STRICT; diff --git a/contrib/pg_term_reason/pg_term_reason.c b/contrib/pg_term_reason/pg_term_reason.c new file mode 100644 index 0000000000..5c6ceb00fd --- /dev/null +++ b/contrib/pg_term_reason/pg_term_reason.c @@ -0,0 +1,89 @@ +/*------------------------------------------------------------------------- + * + * pg_term_reason.c + * Functions to terminate/cancel postgres backends with additional message + * + * IDENTIFICATION + * contrib/pg_term_reason/pg_term_reason.c + * + *------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include "fmgr.h" +#include "varatt.h" +#include "storage/signalfuncs.h" +#include "storage/lock.h" +#include "storage/procarray.h" +#include "storage/proc.h" + +PG_MODULE_MAGIC_EXT( + .name = "pg_term_reason", + .version = PG_VERSION +); + +PG_FUNCTION_INFO_V1(pg_terminate_backend_msg); +PG_FUNCTION_INFO_V1(pg_cancel_backend_msg); + +static void +set_reason(int pid, const char *msg, int msglen) +{ + int len; + PGPROC *proc; + + if (msglen <= 0) { + return; + } + + LWLockAcquire(ProcArrayLock, LW_SHARED); + + proc = BackendPidGetProcWithLock(pid); + + if (proc != NULL) { + len = Min(PROC_TERM_REASON_MAX_LEN - 1, msglen); + strncpy(proc->termReasonStr, msg, len); + } + + LWLockRelease(ProcArrayLock); +} + +Datum +pg_cancel_backend_msg(PG_FUNCTION_ARGS) +{ + int pid; + text *reason; + char *reason_str; + int reason_len; + + pid = PG_GETARG_INT32(0); + reason = PG_GETARG_TEXT_P(1); + + reason_str = VARDATA(reason); + reason_len = VARSIZE(reason) - VARHDRSZ; + + set_reason(pid, reason_str, reason_len); + + return pg_cancel_backend_impl(pid); +} + +Datum +pg_terminate_backend_msg(PG_FUNCTION_ARGS) +{ + int pid; + int timeout; /* milliseconds */ + text *reason; + char *reason_str; + int reason_len; + + pid = PG_GETARG_INT32(0); + timeout = PG_GETARG_INT64(1); + reason = PG_GETARG_TEXT_P(2); + + reason_str = VARDATA(reason); + reason_len = VARSIZE(reason) - VARHDRSZ; + + set_reason(pid, reason_str, reason_len); + + return pg_terminate_backend_impl(pid, timeout); +} diff --git a/contrib/pg_term_reason/pg_term_reason.control b/contrib/pg_term_reason/pg_term_reason.control new file mode 100644 index 0000000000..d3b04459aa --- /dev/null +++ b/contrib/pg_term_reason/pg_term_reason.control @@ -0,0 +1,5 @@ +# pg_term_reason extension +comment = 'extension to terminate backends with additional reason string' +default_version = '1.0' +module_pathname = '$libdir/pg_term_reason' +relocatable = true diff --git a/contrib/pg_term_reason/t/001_basic.pl b/contrib/pg_term_reason/t/001_basic.pl new file mode 100644 index 0000000000..5d9dadc462 --- /dev/null +++ b/contrib/pg_term_reason/t/001_basic.pl @@ -0,0 +1,28 @@ +use strict; +use warnings FATAL => 'all'; +use PostgreSQL::Test::Cluster; +use PostgreSQL::Test::Utils; +use Test::More; + +my $node = PostgreSQL::Test::Cluster->new('primary'); +$node->init(); +$node->start; + +$node->safe_psql('postgres', 'create extension pg_term_reason;'); + +my ($stdout, $stderr); +$node->psql('postgres', + q[select pg_terminate_backend_msg(pg_backend_pid(), 0, 'hello from tap tests!');], + stdout => \$stdout, stderr => \$stderr); +like($stderr, qr/hello from tap tests\!/, "expected message to be passed"); + +$stdout = ''; +$stderr = ''; +$node->psql('postgres', + q[select pg_cancel_backend_msg(pg_backend_pid(), 'hello from tap tests again!');], + stdout => \$stdout, stderr => \$stderr); +like($stderr, qr/hello from tap tests again\!/, "expected message to be passed"); + +$node->stop; + +done_testing(); -- 2.50.1 (Apple Git-155)