503 lines
14 KiB
C
503 lines
14 KiB
C
/* -------------------------------------------------------------------------
|
|
*
|
|
* worker_spi.c
|
|
* Sample background worker code that demonstrates various coding
|
|
* patterns: establishing a database connection; starting and committing
|
|
* transactions; using GUC variables, and heeding SIGHUP to reread
|
|
* the configuration file; reporting to pg_stat_activity; using the
|
|
* process latch to sleep and exit in case of postmaster death.
|
|
*
|
|
* This code connects to a database, creates a schema and table, and summarizes
|
|
* the numbers contained therein. To see it working, insert an initial value
|
|
* with "total" type and some initial value; then insert some other rows with
|
|
* "delta" type. Delta rows will be deleted by this worker and their values
|
|
* aggregated into the total.
|
|
*
|
|
* Copyright (c) 2013-2024, PostgreSQL Global Development Group
|
|
*
|
|
* IDENTIFICATION
|
|
* src/test/modules/worker_spi/worker_spi.c
|
|
*
|
|
* -------------------------------------------------------------------------
|
|
*/
|
|
#include "postgres.h"
|
|
|
|
/* These are always necessary for a bgworker */
|
|
#include "miscadmin.h"
|
|
#include "postmaster/bgworker.h"
|
|
#include "postmaster/interrupt.h"
|
|
#include "storage/ipc.h"
|
|
#include "storage/latch.h"
|
|
#include "storage/lwlock.h"
|
|
#include "storage/proc.h"
|
|
#include "storage/shmem.h"
|
|
|
|
/* these headers are used by this particular worker's code */
|
|
#include "access/xact.h"
|
|
#include "commands/dbcommands.h"
|
|
#include "executor/spi.h"
|
|
#include "fmgr.h"
|
|
#include "lib/stringinfo.h"
|
|
#include "pgstat.h"
|
|
#include "tcop/utility.h"
|
|
#include "utils/acl.h"
|
|
#include "utils/builtins.h"
|
|
#include "utils/snapmgr.h"
|
|
|
|
PG_MODULE_MAGIC;
|
|
|
|
PG_FUNCTION_INFO_V1(worker_spi_launch);
|
|
|
|
PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn();
|
|
|
|
/* GUC variables */
|
|
static int worker_spi_naptime = 10;
|
|
static int worker_spi_total_workers = 2;
|
|
static char *worker_spi_database = NULL;
|
|
static char *worker_spi_role = NULL;
|
|
|
|
/* value cached, fetched from shared memory */
|
|
static uint32 worker_spi_wait_event_main = 0;
|
|
|
|
typedef struct worktable
|
|
{
|
|
const char *schema;
|
|
const char *name;
|
|
} worktable;
|
|
|
|
/*
|
|
* Initialize workspace for a worker process: create the schema if it doesn't
|
|
* already exist.
|
|
*/
|
|
static void
|
|
initialize_worker_spi(worktable *table)
|
|
{
|
|
int ret;
|
|
int ntup;
|
|
bool isnull;
|
|
StringInfoData buf;
|
|
|
|
SetCurrentStatementStartTimestamp();
|
|
StartTransactionCommand();
|
|
SPI_connect();
|
|
PushActiveSnapshot(GetTransactionSnapshot());
|
|
pgstat_report_activity(STATE_RUNNING, "initializing worker_spi schema");
|
|
|
|
/* XXX could we use CREATE SCHEMA IF NOT EXISTS? */
|
|
initStringInfo(&buf);
|
|
appendStringInfo(&buf, "select count(*) from pg_namespace where nspname = '%s'",
|
|
table->schema);
|
|
|
|
debug_query_string = buf.data;
|
|
ret = SPI_execute(buf.data, true, 0);
|
|
if (ret != SPI_OK_SELECT)
|
|
elog(FATAL, "SPI_execute failed: error code %d", ret);
|
|
|
|
if (SPI_processed != 1)
|
|
elog(FATAL, "not a singleton result");
|
|
|
|
ntup = DatumGetInt64(SPI_getbinval(SPI_tuptable->vals[0],
|
|
SPI_tuptable->tupdesc,
|
|
1, &isnull));
|
|
if (isnull)
|
|
elog(FATAL, "null result");
|
|
|
|
if (ntup == 0)
|
|
{
|
|
debug_query_string = NULL;
|
|
resetStringInfo(&buf);
|
|
appendStringInfo(&buf,
|
|
"CREATE SCHEMA \"%s\" "
|
|
"CREATE TABLE \"%s\" ("
|
|
" type text CHECK (type IN ('total', 'delta')), "
|
|
" value integer)"
|
|
"CREATE UNIQUE INDEX \"%s_unique_total\" ON \"%s\" (type) "
|
|
"WHERE type = 'total'",
|
|
table->schema, table->name, table->name, table->name);
|
|
|
|
/* set statement start time */
|
|
SetCurrentStatementStartTimestamp();
|
|
|
|
debug_query_string = buf.data;
|
|
ret = SPI_execute(buf.data, false, 0);
|
|
|
|
if (ret != SPI_OK_UTILITY)
|
|
elog(FATAL, "failed to create my schema");
|
|
|
|
debug_query_string = NULL; /* rest is not statement-specific */
|
|
}
|
|
|
|
SPI_finish();
|
|
PopActiveSnapshot();
|
|
CommitTransactionCommand();
|
|
debug_query_string = NULL;
|
|
pgstat_report_activity(STATE_IDLE, NULL);
|
|
}
|
|
|
|
void
|
|
worker_spi_main(Datum main_arg)
|
|
{
|
|
int index = DatumGetInt32(main_arg);
|
|
worktable *table;
|
|
StringInfoData buf;
|
|
char name[20];
|
|
Oid dboid;
|
|
Oid roleoid;
|
|
char *p;
|
|
bits32 flags = 0;
|
|
|
|
table = palloc(sizeof(worktable));
|
|
sprintf(name, "schema%d", index);
|
|
table->schema = pstrdup(name);
|
|
table->name = pstrdup("counted");
|
|
|
|
/* fetch database and role OIDs, these are set for a dynamic worker */
|
|
p = MyBgworkerEntry->bgw_extra;
|
|
memcpy(&dboid, p, sizeof(Oid));
|
|
p += sizeof(Oid);
|
|
memcpy(&roleoid, p, sizeof(Oid));
|
|
p += sizeof(Oid);
|
|
memcpy(&flags, p, sizeof(bits32));
|
|
|
|
/* Establish signal handlers before unblocking signals. */
|
|
pqsignal(SIGHUP, SignalHandlerForConfigReload);
|
|
pqsignal(SIGTERM, die);
|
|
|
|
/* We're now ready to receive signals */
|
|
BackgroundWorkerUnblockSignals();
|
|
|
|
/* Connect to our database */
|
|
if (OidIsValid(dboid))
|
|
BackgroundWorkerInitializeConnectionByOid(dboid, roleoid, flags);
|
|
else
|
|
BackgroundWorkerInitializeConnection(worker_spi_database,
|
|
worker_spi_role, flags);
|
|
|
|
/*
|
|
* Disable parallel query for workers started with BYPASS_ALLOWCONN or
|
|
* BGWORKER_BYPASS_ALLOWCONN so as these don't attempt connections using a
|
|
* database or a role that may not allow that.
|
|
*/
|
|
if ((flags & (BGWORKER_BYPASS_ALLOWCONN | BGWORKER_BYPASS_ROLELOGINCHECK)))
|
|
SetConfigOption("max_parallel_workers_per_gather", "0",
|
|
PGC_USERSET, PGC_S_OVERRIDE);
|
|
|
|
elog(LOG, "%s initialized with %s.%s",
|
|
MyBgworkerEntry->bgw_name, table->schema, table->name);
|
|
initialize_worker_spi(table);
|
|
|
|
/*
|
|
* Quote identifiers passed to us. Note that this must be done after
|
|
* initialize_worker_spi, because that routine assumes the names are not
|
|
* quoted.
|
|
*
|
|
* Note some memory might be leaked here.
|
|
*/
|
|
table->schema = quote_identifier(table->schema);
|
|
table->name = quote_identifier(table->name);
|
|
|
|
initStringInfo(&buf);
|
|
appendStringInfo(&buf,
|
|
"WITH deleted AS (DELETE "
|
|
"FROM %s.%s "
|
|
"WHERE type = 'delta' RETURNING value), "
|
|
"total AS (SELECT coalesce(sum(value), 0) as sum "
|
|
"FROM deleted) "
|
|
"UPDATE %s.%s "
|
|
"SET value = %s.value + total.sum "
|
|
"FROM total WHERE type = 'total' "
|
|
"RETURNING %s.value",
|
|
table->schema, table->name,
|
|
table->schema, table->name,
|
|
table->name,
|
|
table->name);
|
|
|
|
/*
|
|
* Main loop: do this until SIGTERM is received and processed by
|
|
* ProcessInterrupts.
|
|
*/
|
|
for (;;)
|
|
{
|
|
int ret;
|
|
|
|
/* First time, allocate or get the custom wait event */
|
|
if (worker_spi_wait_event_main == 0)
|
|
worker_spi_wait_event_main = WaitEventExtensionNew("WorkerSpiMain");
|
|
|
|
/*
|
|
* Background workers mustn't call usleep() or any direct equivalent:
|
|
* instead, they may wait on their process latch, which sleeps as
|
|
* necessary, but is awakened if postmaster dies. That way the
|
|
* background process goes away immediately in an emergency.
|
|
*/
|
|
(void) WaitLatch(MyLatch,
|
|
WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
|
|
worker_spi_naptime * 1000L,
|
|
worker_spi_wait_event_main);
|
|
ResetLatch(MyLatch);
|
|
|
|
CHECK_FOR_INTERRUPTS();
|
|
|
|
/*
|
|
* In case of a SIGHUP, just reload the configuration.
|
|
*/
|
|
if (ConfigReloadPending)
|
|
{
|
|
ConfigReloadPending = false;
|
|
ProcessConfigFile(PGC_SIGHUP);
|
|
}
|
|
|
|
/*
|
|
* Start a transaction on which we can run queries. Note that each
|
|
* StartTransactionCommand() call should be preceded by a
|
|
* SetCurrentStatementStartTimestamp() call, which sets both the time
|
|
* for the statement we're about the run, and also the transaction
|
|
* start time. Also, each other query sent to SPI should probably be
|
|
* preceded by SetCurrentStatementStartTimestamp(), so that statement
|
|
* start time is always up to date.
|
|
*
|
|
* The SPI_connect() call lets us run queries through the SPI manager,
|
|
* and the PushActiveSnapshot() call creates an "active" snapshot
|
|
* which is necessary for queries to have MVCC data to work on.
|
|
*
|
|
* The pgstat_report_activity() call makes our activity visible
|
|
* through the pgstat views.
|
|
*/
|
|
SetCurrentStatementStartTimestamp();
|
|
StartTransactionCommand();
|
|
SPI_connect();
|
|
PushActiveSnapshot(GetTransactionSnapshot());
|
|
debug_query_string = buf.data;
|
|
pgstat_report_activity(STATE_RUNNING, buf.data);
|
|
|
|
/* We can now execute queries via SPI */
|
|
ret = SPI_execute(buf.data, false, 0);
|
|
|
|
if (ret != SPI_OK_UPDATE_RETURNING)
|
|
elog(FATAL, "cannot select from table %s.%s: error code %d",
|
|
table->schema, table->name, ret);
|
|
|
|
if (SPI_processed > 0)
|
|
{
|
|
bool isnull;
|
|
int32 val;
|
|
|
|
val = DatumGetInt32(SPI_getbinval(SPI_tuptable->vals[0],
|
|
SPI_tuptable->tupdesc,
|
|
1, &isnull));
|
|
if (!isnull)
|
|
elog(LOG, "%s: count in %s.%s is now %d",
|
|
MyBgworkerEntry->bgw_name,
|
|
table->schema, table->name, val);
|
|
}
|
|
|
|
/*
|
|
* And finish our transaction.
|
|
*/
|
|
SPI_finish();
|
|
PopActiveSnapshot();
|
|
CommitTransactionCommand();
|
|
debug_query_string = NULL;
|
|
pgstat_report_stat(true);
|
|
pgstat_report_activity(STATE_IDLE, NULL);
|
|
}
|
|
|
|
/* Not reachable */
|
|
}
|
|
|
|
/*
|
|
* Entrypoint of this module.
|
|
*
|
|
* We register more than one worker process here, to demonstrate how that can
|
|
* be done.
|
|
*/
|
|
void
|
|
_PG_init(void)
|
|
{
|
|
BackgroundWorker worker;
|
|
|
|
/* get the configuration */
|
|
|
|
/*
|
|
* These GUCs are defined even if this library is not loaded with
|
|
* shared_preload_libraries, for worker_spi_launch().
|
|
*/
|
|
DefineCustomIntVariable("worker_spi.naptime",
|
|
"Duration between each check (in seconds).",
|
|
NULL,
|
|
&worker_spi_naptime,
|
|
10,
|
|
1,
|
|
INT_MAX,
|
|
PGC_SIGHUP,
|
|
0,
|
|
NULL,
|
|
NULL,
|
|
NULL);
|
|
|
|
DefineCustomStringVariable("worker_spi.database",
|
|
"Database to connect to.",
|
|
NULL,
|
|
&worker_spi_database,
|
|
"postgres",
|
|
PGC_SIGHUP,
|
|
0,
|
|
NULL, NULL, NULL);
|
|
|
|
DefineCustomStringVariable("worker_spi.role",
|
|
"Role to connect with.",
|
|
NULL,
|
|
&worker_spi_role,
|
|
NULL,
|
|
PGC_SIGHUP,
|
|
0,
|
|
NULL, NULL, NULL);
|
|
|
|
if (!process_shared_preload_libraries_in_progress)
|
|
return;
|
|
|
|
DefineCustomIntVariable("worker_spi.total_workers",
|
|
"Number of workers.",
|
|
NULL,
|
|
&worker_spi_total_workers,
|
|
2,
|
|
1,
|
|
100,
|
|
PGC_POSTMASTER,
|
|
0,
|
|
NULL,
|
|
NULL,
|
|
NULL);
|
|
|
|
MarkGUCPrefixReserved("worker_spi");
|
|
|
|
/* set up common data for all our workers */
|
|
memset(&worker, 0, sizeof(worker));
|
|
worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
|
|
BGWORKER_BACKEND_DATABASE_CONNECTION;
|
|
worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
|
|
worker.bgw_restart_time = BGW_NEVER_RESTART;
|
|
sprintf(worker.bgw_library_name, "worker_spi");
|
|
sprintf(worker.bgw_function_name, "worker_spi_main");
|
|
worker.bgw_notify_pid = 0;
|
|
|
|
/*
|
|
* Now fill in worker-specific data, and do the actual registrations.
|
|
*
|
|
* bgw_extra can optionally include a dabatase OID, a role OID and a set
|
|
* of flags. This is left empty here to fallback to the related GUCs at
|
|
* startup (0 for the bgworker flags).
|
|
*/
|
|
for (int i = 1; i <= worker_spi_total_workers; i++)
|
|
{
|
|
snprintf(worker.bgw_name, BGW_MAXLEN, "worker_spi worker %d", i);
|
|
snprintf(worker.bgw_type, BGW_MAXLEN, "worker_spi");
|
|
worker.bgw_main_arg = Int32GetDatum(i);
|
|
|
|
RegisterBackgroundWorker(&worker);
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Dynamically launch an SPI worker.
|
|
*/
|
|
Datum
|
|
worker_spi_launch(PG_FUNCTION_ARGS)
|
|
{
|
|
int32 i = PG_GETARG_INT32(0);
|
|
Oid dboid = PG_GETARG_OID(1);
|
|
Oid roleoid = PG_GETARG_OID(2);
|
|
BackgroundWorker worker;
|
|
BackgroundWorkerHandle *handle;
|
|
BgwHandleStatus status;
|
|
pid_t pid;
|
|
char *p;
|
|
bits32 flags = 0;
|
|
ArrayType *arr = PG_GETARG_ARRAYTYPE_P(3);
|
|
Size ndim;
|
|
int nelems;
|
|
Datum *datum_flags;
|
|
|
|
memset(&worker, 0, sizeof(worker));
|
|
worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
|
|
BGWORKER_BACKEND_DATABASE_CONNECTION;
|
|
worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
|
|
worker.bgw_restart_time = BGW_NEVER_RESTART;
|
|
sprintf(worker.bgw_library_name, "worker_spi");
|
|
sprintf(worker.bgw_function_name, "worker_spi_main");
|
|
snprintf(worker.bgw_name, BGW_MAXLEN, "worker_spi dynamic worker %d", i);
|
|
snprintf(worker.bgw_type, BGW_MAXLEN, "worker_spi dynamic");
|
|
worker.bgw_main_arg = Int32GetDatum(i);
|
|
/* set bgw_notify_pid so that we can use WaitForBackgroundWorkerStartup */
|
|
worker.bgw_notify_pid = MyProcPid;
|
|
|
|
/* extract flags, if any */
|
|
ndim = ARR_NDIM(arr);
|
|
if (ndim > 1)
|
|
ereport(ERROR,
|
|
(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
|
|
errmsg("flags array must be one-dimensional")));
|
|
|
|
if (array_contains_nulls(arr))
|
|
ereport(ERROR,
|
|
(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
|
|
errmsg("flags array must not contain nulls")));
|
|
|
|
Assert(ARR_ELEMTYPE(arr) == TEXTOID);
|
|
deconstruct_array_builtin(arr, TEXTOID, &datum_flags, NULL, &nelems);
|
|
|
|
for (i = 0; i < nelems; i++)
|
|
{
|
|
char *optname = TextDatumGetCString(datum_flags[i]);
|
|
|
|
if (strcmp(optname, "ALLOWCONN") == 0)
|
|
flags |= BGWORKER_BYPASS_ALLOWCONN;
|
|
else if (strcmp(optname, "ROLELOGINCHECK") == 0)
|
|
flags |= BGWORKER_BYPASS_ROLELOGINCHECK;
|
|
else
|
|
ereport(ERROR,
|
|
(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
|
|
errmsg("incorrect flag value found in array")));
|
|
}
|
|
|
|
/*
|
|
* Register database and role to use for the worker started in bgw_extra.
|
|
* If none have been provided, this will fall back to the GUCs at startup.
|
|
*/
|
|
if (!OidIsValid(dboid))
|
|
dboid = get_database_oid(worker_spi_database, false);
|
|
|
|
/*
|
|
* worker_spi_role is NULL by default, so this gives to worker_spi_main()
|
|
* an invalid OID in this case.
|
|
*/
|
|
if (!OidIsValid(roleoid) && worker_spi_role)
|
|
roleoid = get_role_oid(worker_spi_role, false);
|
|
|
|
p = worker.bgw_extra;
|
|
memcpy(p, &dboid, sizeof(Oid));
|
|
p += sizeof(Oid);
|
|
memcpy(p, &roleoid, sizeof(Oid));
|
|
p += sizeof(Oid);
|
|
memcpy(p, &flags, sizeof(bits32));
|
|
|
|
if (!RegisterDynamicBackgroundWorker(&worker, &handle))
|
|
PG_RETURN_NULL();
|
|
|
|
status = WaitForBackgroundWorkerStartup(handle, &pid);
|
|
|
|
if (status == BGWH_STOPPED)
|
|
ereport(ERROR,
|
|
(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
|
|
errmsg("could not start background process"),
|
|
errhint("More details may be available in the server log.")));
|
|
if (status == BGWH_POSTMASTER_DIED)
|
|
ereport(ERROR,
|
|
(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
|
|
errmsg("cannot start background processes without postmaster"),
|
|
errhint("Kill all remaining database processes and restart the database.")));
|
|
Assert(status == BGWH_STARTED);
|
|
|
|
PG_RETURN_INT32(pid);
|
|
}
|