PostgreSQL Source Code  git master
worker_spi.c
Go to the documentation of this file.
1 /* -------------------------------------------------------------------------
2  *
3  * worker_spi.c
4  * Sample background worker code that demonstrates various coding
5  * patterns: establishing a database connection; starting and committing
6  * transactions; using GUC variables, and heeding SIGHUP to reread
7  * the configuration file; reporting to pg_stat_activity; using the
8  * process latch to sleep and exit in case of postmaster death.
9  *
10  * This code connects to a database, creates a schema and table, and summarizes
11  * the numbers contained therein. To see it working, insert an initial value
12  * with "total" type and some initial value; then insert some other rows with
13  * "delta" type. Delta rows will be deleted by this worker and their values
14  * aggregated into the total.
15  *
16  * Copyright (c) 2013-2024, PostgreSQL Global Development Group
17  *
18  * IDENTIFICATION
19  * src/test/modules/worker_spi/worker_spi.c
20  *
21  * -------------------------------------------------------------------------
22  */
23 #include "postgres.h"
24 
25 /* These are always necessary for a bgworker */
26 #include "miscadmin.h"
27 #include "postmaster/bgworker.h"
28 #include "postmaster/interrupt.h"
29 #include "storage/ipc.h"
30 #include "storage/latch.h"
31 #include "storage/lwlock.h"
32 #include "storage/proc.h"
33 #include "storage/shmem.h"
34 
35 /* these headers are used by this particular worker's code */
36 #include "access/xact.h"
37 #include "commands/dbcommands.h"
38 #include "executor/spi.h"
39 #include "fmgr.h"
40 #include "lib/stringinfo.h"
41 #include "pgstat.h"
42 #include "utils/acl.h"
43 #include "utils/builtins.h"
44 #include "utils/snapmgr.h"
45 #include "tcop/utility.h"
46 
48 
50 
52 
53 /* GUC variables */
54 static int worker_spi_naptime = 10;
55 static int worker_spi_total_workers = 2;
56 static char *worker_spi_database = NULL;
57 static char *worker_spi_role = NULL;
58 
59 /* value cached, fetched from shared memory */
61 
62 typedef struct worktable
63 {
64  const char *schema;
65  const char *name;
67 
68 /*
69  * Initialize workspace for a worker process: create the schema if it doesn't
70  * already exist.
71  */
72 static void
74 {
75  int ret;
76  int ntup;
77  bool isnull;
79 
82  SPI_connect();
84  pgstat_report_activity(STATE_RUNNING, "initializing worker_spi schema");
85 
86  /* XXX could we use CREATE SCHEMA IF NOT EXISTS? */
88  appendStringInfo(&buf, "select count(*) from pg_namespace where nspname = '%s'",
89  table->schema);
90 
91  debug_query_string = buf.data;
92  ret = SPI_execute(buf.data, true, 0);
93  if (ret != SPI_OK_SELECT)
94  elog(FATAL, "SPI_execute failed: error code %d", ret);
95 
96  if (SPI_processed != 1)
97  elog(FATAL, "not a singleton result");
98 
101  1, &isnull));
102  if (isnull)
103  elog(FATAL, "null result");
104 
105  if (ntup == 0)
106  {
107  debug_query_string = NULL;
110  "CREATE SCHEMA \"%s\" "
111  "CREATE TABLE \"%s\" ("
112  " type text CHECK (type IN ('total', 'delta')), "
113  " value integer)"
114  "CREATE UNIQUE INDEX \"%s_unique_total\" ON \"%s\" (type) "
115  "WHERE type = 'total'",
116  table->schema, table->name, table->name, table->name);
117 
118  /* set statement start time */
120 
121  debug_query_string = buf.data;
122  ret = SPI_execute(buf.data, false, 0);
123 
124  if (ret != SPI_OK_UTILITY)
125  elog(FATAL, "failed to create my schema");
126 
127  debug_query_string = NULL; /* rest is not statement-specific */
128  }
129 
130  SPI_finish();
133  debug_query_string = NULL;
135 }
136 
137 void
139 {
140  int index = DatumGetInt32(main_arg);
141  worktable *table;
143  char name[20];
144  Oid dboid;
145  Oid roleoid;
146  char *p;
147  bits32 flags = 0;
148 
149  table = palloc(sizeof(worktable));
150  sprintf(name, "schema%d", index);
151  table->schema = pstrdup(name);
152  table->name = pstrdup("counted");
153 
154  /* fetch database and role OIDs, these are set for a dynamic worker */
156  memcpy(&dboid, p, sizeof(Oid));
157  p += sizeof(Oid);
158  memcpy(&roleoid, p, sizeof(Oid));
159  p += sizeof(Oid);
160  memcpy(&flags, p, sizeof(bits32));
161 
162  /* Establish signal handlers before unblocking signals. */
164  pqsignal(SIGTERM, die);
165 
166  /* We're now ready to receive signals */
168 
169  /* Connect to our database */
170  if (OidIsValid(dboid))
171  BackgroundWorkerInitializeConnectionByOid(dboid, roleoid, flags);
172  else
174  worker_spi_role, flags);
175 
176  /*
177  * Disable parallel query for workers started with BYPASS_ALLOWCONN or
178  * BGWORKER_BYPASS_ALLOWCONN so as these don't attempt connections using a
179  * database or a role that may not allow that.
180  */
182  SetConfigOption("max_parallel_workers_per_gather", "0",
184 
185  elog(LOG, "%s initialized with %s.%s",
186  MyBgworkerEntry->bgw_name, table->schema, table->name);
187  initialize_worker_spi(table);
188 
189  /*
190  * Quote identifiers passed to us. Note that this must be done after
191  * initialize_worker_spi, because that routine assumes the names are not
192  * quoted.
193  *
194  * Note some memory might be leaked here.
195  */
196  table->schema = quote_identifier(table->schema);
197  table->name = quote_identifier(table->name);
198 
201  "WITH deleted AS (DELETE "
202  "FROM %s.%s "
203  "WHERE type = 'delta' RETURNING value), "
204  "total AS (SELECT coalesce(sum(value), 0) as sum "
205  "FROM deleted) "
206  "UPDATE %s.%s "
207  "SET value = %s.value + total.sum "
208  "FROM total WHERE type = 'total' "
209  "RETURNING %s.value",
210  table->schema, table->name,
211  table->schema, table->name,
212  table->name,
213  table->name);
214 
215  /*
216  * Main loop: do this until SIGTERM is received and processed by
217  * ProcessInterrupts.
218  */
219  for (;;)
220  {
221  int ret;
222 
223  /* First time, allocate or get the custom wait event */
226 
227  /*
228  * Background workers mustn't call usleep() or any direct equivalent:
229  * instead, they may wait on their process latch, which sleeps as
230  * necessary, but is awakened if postmaster dies. That way the
231  * background process goes away immediately in an emergency.
232  */
233  (void) WaitLatch(MyLatch,
235  worker_spi_naptime * 1000L,
238 
240 
241  /*
242  * In case of a SIGHUP, just reload the configuration.
243  */
245  {
246  ConfigReloadPending = false;
248  }
249 
250  /*
251  * Start a transaction on which we can run queries. Note that each
252  * StartTransactionCommand() call should be preceded by a
253  * SetCurrentStatementStartTimestamp() call, which sets both the time
254  * for the statement we're about the run, and also the transaction
255  * start time. Also, each other query sent to SPI should probably be
256  * preceded by SetCurrentStatementStartTimestamp(), so that statement
257  * start time is always up to date.
258  *
259  * The SPI_connect() call lets us run queries through the SPI manager,
260  * and the PushActiveSnapshot() call creates an "active" snapshot
261  * which is necessary for queries to have MVCC data to work on.
262  *
263  * The pgstat_report_activity() call makes our activity visible
264  * through the pgstat views.
265  */
268  SPI_connect();
270  debug_query_string = buf.data;
272 
273  /* We can now execute queries via SPI */
274  ret = SPI_execute(buf.data, false, 0);
275 
276  if (ret != SPI_OK_UPDATE_RETURNING)
277  elog(FATAL, "cannot select from table %s.%s: error code %d",
278  table->schema, table->name, ret);
279 
280  if (SPI_processed > 0)
281  {
282  bool isnull;
283  int32 val;
284 
287  1, &isnull));
288  if (!isnull)
289  elog(LOG, "%s: count in %s.%s is now %d",
291  table->schema, table->name, val);
292  }
293 
294  /*
295  * And finish our transaction.
296  */
297  SPI_finish();
300  debug_query_string = NULL;
301  pgstat_report_stat(true);
303  }
304 
305  /* Not reachable */
306 }
307 
308 /*
309  * Entrypoint of this module.
310  *
311  * We register more than one worker process here, to demonstrate how that can
312  * be done.
313  */
314 void
315 _PG_init(void)
316 {
317  BackgroundWorker worker;
318 
319  /* get the configuration */
320 
321  /*
322  * These GUCs are defined even if this library is not loaded with
323  * shared_preload_libraries, for worker_spi_launch().
324  */
325  DefineCustomIntVariable("worker_spi.naptime",
326  "Duration between each check (in seconds).",
327  NULL,
329  10,
330  1,
331  INT_MAX,
332  PGC_SIGHUP,
333  0,
334  NULL,
335  NULL,
336  NULL);
337 
338  DefineCustomStringVariable("worker_spi.database",
339  "Database to connect to.",
340  NULL,
342  "postgres",
343  PGC_SIGHUP,
344  0,
345  NULL, NULL, NULL);
346 
347  DefineCustomStringVariable("worker_spi.role",
348  "Role to connect with.",
349  NULL,
351  NULL,
352  PGC_SIGHUP,
353  0,
354  NULL, NULL, NULL);
355 
357  return;
358 
359  DefineCustomIntVariable("worker_spi.total_workers",
360  "Number of workers.",
361  NULL,
363  2,
364  1,
365  100,
367  0,
368  NULL,
369  NULL,
370  NULL);
371 
372  MarkGUCPrefixReserved("worker_spi");
373 
374  /* set up common data for all our workers */
375  memset(&worker, 0, sizeof(worker));
380  sprintf(worker.bgw_library_name, "worker_spi");
381  sprintf(worker.bgw_function_name, "worker_spi_main");
382  worker.bgw_notify_pid = 0;
383 
384  /*
385  * Now fill in worker-specific data, and do the actual registrations.
386  *
387  * bgw_extra can optionally include a dabatase OID, a role OID and a set
388  * of flags. This is left empty here to fallback to the related GUCs at
389  * startup (0 for the bgworker flags).
390  */
391  for (int i = 1; i <= worker_spi_total_workers; i++)
392  {
393  snprintf(worker.bgw_name, BGW_MAXLEN, "worker_spi worker %d", i);
394  snprintf(worker.bgw_type, BGW_MAXLEN, "worker_spi");
395  worker.bgw_main_arg = Int32GetDatum(i);
396 
397  RegisterBackgroundWorker(&worker);
398  }
399 }
400 
401 /*
402  * Dynamically launch an SPI worker.
403  */
404 Datum
406 {
407  int32 i = PG_GETARG_INT32(0);
408  Oid dboid = PG_GETARG_OID(1);
409  Oid roleoid = PG_GETARG_OID(2);
410  BackgroundWorker worker;
411  BackgroundWorkerHandle *handle;
412  BgwHandleStatus status;
413  pid_t pid;
414  char *p;
415  bits32 flags = 0;
417  Size ndim;
418  int nelems;
419  Datum *datum_flags;
420 
421  memset(&worker, 0, sizeof(worker));
426  sprintf(worker.bgw_library_name, "worker_spi");
427  sprintf(worker.bgw_function_name, "worker_spi_main");
428  snprintf(worker.bgw_name, BGW_MAXLEN, "worker_spi dynamic worker %d", i);
429  snprintf(worker.bgw_type, BGW_MAXLEN, "worker_spi dynamic");
430  worker.bgw_main_arg = Int32GetDatum(i);
431  /* set bgw_notify_pid so that we can use WaitForBackgroundWorkerStartup */
432  worker.bgw_notify_pid = MyProcPid;
433 
434  /* extract flags, if any */
435  ndim = ARR_NDIM(arr);
436  if (ndim > 1)
437  ereport(ERROR,
438  (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
439  errmsg("flags array must be one-dimensional")));
440 
441  if (array_contains_nulls(arr))
442  ereport(ERROR,
443  (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
444  errmsg("flags array must not contain nulls")));
445 
446  Assert(ARR_ELEMTYPE(arr) == TEXTOID);
447  deconstruct_array_builtin(arr, TEXTOID, &datum_flags, NULL, &nelems);
448 
449  for (i = 0; i < nelems; i++)
450  {
451  char *optname = TextDatumGetCString(datum_flags[i]);
452 
453  if (strcmp(optname, "ALLOWCONN") == 0)
454  flags |= BGWORKER_BYPASS_ALLOWCONN;
455  else if (strcmp(optname, "ROLELOGINCHECK") == 0)
457  else
458  ereport(ERROR,
459  (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
460  errmsg("incorrect flag value found in array")));
461  }
462 
463  /*
464  * Register database and role to use for the worker started in bgw_extra.
465  * If none have been provided, this will fall back to the GUCs at startup.
466  */
467  if (!OidIsValid(dboid))
468  dboid = get_database_oid(worker_spi_database, false);
469 
470  /*
471  * worker_spi_role is NULL by default, so this gives to worker_spi_main()
472  * an invalid OID in this case.
473  */
474  if (!OidIsValid(roleoid) && worker_spi_role)
475  roleoid = get_role_oid(worker_spi_role, false);
476 
477  p = worker.bgw_extra;
478  memcpy(p, &dboid, sizeof(Oid));
479  p += sizeof(Oid);
480  memcpy(p, &roleoid, sizeof(Oid));
481  p += sizeof(Oid);
482  memcpy(p, &flags, sizeof(bits32));
483 
484  if (!RegisterDynamicBackgroundWorker(&worker, &handle))
485  PG_RETURN_NULL();
486 
487  status = WaitForBackgroundWorkerStartup(handle, &pid);
488 
489  if (status == BGWH_STOPPED)
490  ereport(ERROR,
491  (errcode(ERRCODE_INSUFFICIENT_RESOURCES),
492  errmsg("could not start background process"),
493  errhint("More details may be available in the server log.")));
494  if (status == BGWH_POSTMASTER_DIED)
495  ereport(ERROR,
496  (errcode(ERRCODE_INSUFFICIENT_RESOURCES),
497  errmsg("cannot start background processes without postmaster"),
498  errhint("Kill all remaining database processes and restart the database.")));
499  Assert(status == BGWH_STARTED);
500 
501  PG_RETURN_INT32(pid);
502 }
Oid get_role_oid(const char *rolname, bool missing_ok)
Definition: acl.c:5346
#define ARR_NDIM(a)
Definition: array.h:290
#define PG_GETARG_ARRAYTYPE_P(n)
Definition: array.h:263
#define ARR_ELEMTYPE(a)
Definition: array.h:292
bool array_contains_nulls(ArrayType *array)
Definition: arrayfuncs.c:3749
void deconstruct_array_builtin(ArrayType *array, Oid elmtype, Datum **elemsp, bool **nullsp, int *nelemsp)
Definition: arrayfuncs.c:3679
void pgstat_report_activity(BackendState state, const char *cmd_str)
@ STATE_IDLE
@ STATE_RUNNING
void RegisterBackgroundWorker(BackgroundWorker *worker)
Definition: bgworker.c:852
BgwHandleStatus WaitForBackgroundWorkerStartup(BackgroundWorkerHandle *handle, pid_t *pidp)
Definition: bgworker.c:1127
bool RegisterDynamicBackgroundWorker(BackgroundWorker *worker, BackgroundWorkerHandle **handle)
Definition: bgworker.c:960
#define BGW_NEVER_RESTART
Definition: bgworker.h:85
#define BGWORKER_BYPASS_ROLELOGINCHECK
Definition: bgworker.h:157
BgwHandleStatus
Definition: bgworker.h:104
@ BGWH_POSTMASTER_DIED
Definition: bgworker.h:108
@ BGWH_STARTED
Definition: bgworker.h:105
@ BGWH_STOPPED
Definition: bgworker.h:107
@ BgWorkerStart_RecoveryFinished
Definition: bgworker.h:81
#define BGWORKER_BACKEND_DATABASE_CONNECTION
Definition: bgworker.h:60
#define BGWORKER_BYPASS_ALLOWCONN
Definition: bgworker.h:156
#define BGWORKER_SHMEM_ACCESS
Definition: bgworker.h:53
#define BGW_MAXLEN
Definition: bgworker.h:86
#define TextDatumGetCString(d)
Definition: builtins.h:98
unsigned int uint32
Definition: c.h:495
signed int int32
Definition: c.h:483
#define pg_attribute_noreturn()
Definition: c.h:206
#define PGDLLEXPORT
Definition: c.h:1320
uint32 bits32
Definition: c.h:504
#define OidIsValid(objectId)
Definition: c.h:764
size_t Size
Definition: c.h:594
Oid get_database_oid(const char *dbname, bool missing_ok)
Definition: dbcommands.c:3042
int errhint(const char *fmt,...)
Definition: elog.c:1322
int errcode(int sqlerrcode)
Definition: elog.c:860
int errmsg(const char *fmt,...)
Definition: elog.c:1075
#define LOG
Definition: elog.h:31
#define FATAL
Definition: elog.h:41
#define ERROR
Definition: elog.h:39
#define ereport(elevel,...)
Definition: elog.h:149
#define PG_GETARG_OID(n)
Definition: fmgr.h:275
#define PG_RETURN_NULL()
Definition: fmgr.h:345
#define PG_RETURN_INT32(x)
Definition: fmgr.h:354
#define PG_GETARG_INT32(n)
Definition: fmgr.h:269
#define PG_FUNCTION_ARGS
Definition: fmgr.h:193
int MyProcPid
Definition: globals.c:45
struct Latch * MyLatch
Definition: globals.c:59
void SetConfigOption(const char *name, const char *value, GucContext context, GucSource source)
Definition: guc.c:4256
void DefineCustomStringVariable(const char *name, const char *short_desc, const char *long_desc, char **valueAddr, const char *bootValue, GucContext context, int flags, GucStringCheckHook check_hook, GucStringAssignHook assign_hook, GucShowHook show_hook)
Definition: guc.c:5137
void MarkGUCPrefixReserved(const char *className)
Definition: guc.c:5198
void DefineCustomIntVariable(const char *name, const char *short_desc, const char *long_desc, int *valueAddr, int bootValue, int minValue, int maxValue, GucContext context, int flags, GucIntCheckHook check_hook, GucIntAssignHook assign_hook, GucShowHook show_hook)
Definition: guc.c:5077
@ PGC_S_OVERRIDE
Definition: guc.h:119
@ PGC_USERSET
Definition: guc.h:75
@ PGC_POSTMASTER
Definition: guc.h:70
@ PGC_SIGHUP
Definition: guc.h:71
void ProcessConfigFile(GucContext context)
long val
Definition: informix.c:664
volatile sig_atomic_t ConfigReloadPending
Definition: interrupt.c:27
void SignalHandlerForConfigReload(SIGNAL_ARGS)
Definition: interrupt.c:61
int i
Definition: isn.c:73
void ResetLatch(Latch *latch)
Definition: latch.c:725
int WaitLatch(Latch *latch, int wakeEvents, long timeout, uint32 wait_event_info)
Definition: latch.c:518
#define WL_TIMEOUT
Definition: latch.h:130
#define WL_EXIT_ON_PM_DEATH
Definition: latch.h:132
#define WL_LATCH_SET
Definition: latch.h:127
Assert(fmt[strlen(fmt) - 1] !='\n')
char * pstrdup(const char *in)
Definition: mcxt.c:1619
void * palloc(Size size)
Definition: mcxt.c:1201
#define CHECK_FOR_INTERRUPTS()
Definition: miscadmin.h:122
bool process_shared_preload_libraries_in_progress
Definition: miscinit.c:1779
#define die(msg)
static char * buf
Definition: pg_test_fsync.c:73
long pgstat_report_stat(bool force)
Definition: pgstat.c:582
#define sprintf
Definition: port.h:240
pqsigfunc pqsignal(int signo, pqsigfunc func)
#define snprintf
Definition: port.h:238
const char * debug_query_string
Definition: postgres.c:86
static int64 DatumGetInt64(Datum X)
Definition: postgres.h:385
uintptr_t Datum
Definition: postgres.h:64
static Datum Int32GetDatum(int32 X)
Definition: postgres.h:212
static int32 DatumGetInt32(Datum X)
Definition: postgres.h:202
unsigned int Oid
Definition: postgres_ext.h:31
void BackgroundWorkerInitializeConnection(const char *dbname, const char *username, uint32 flags)
Definition: postmaster.c:5622
void BackgroundWorkerUnblockSignals(void)
Definition: postmaster.c:5696
void BackgroundWorkerInitializeConnectionByOid(Oid dboid, Oid useroid, uint32 flags)
Definition: postmaster.c:5656
BackgroundWorker * MyBgworkerEntry
Definition: postmaster.c:194
const char * quote_identifier(const char *ident)
Definition: ruleutils.c:11975
Snapshot GetTransactionSnapshot(void)
Definition: snapmgr.c:223
void PushActiveSnapshot(Snapshot snapshot)
Definition: snapmgr.c:655
void PopActiveSnapshot(void)
Definition: snapmgr.c:750
uint64 SPI_processed
Definition: spi.c:45
SPITupleTable * SPI_tuptable
Definition: spi.c:46
int SPI_connect(void)
Definition: spi.c:95
int SPI_finish(void)
Definition: spi.c:183
int SPI_execute(const char *src, bool read_only, long tcount)
Definition: spi.c:594
Datum SPI_getbinval(HeapTuple tuple, TupleDesc tupdesc, int fnumber, bool *isnull)
Definition: spi.c:1250
#define SPI_OK_UTILITY
Definition: spi.h:85
#define SPI_OK_UPDATE_RETURNING
Definition: spi.h:94
#define SPI_OK_SELECT
Definition: spi.h:86
void resetStringInfo(StringInfo str)
Definition: stringinfo.c:78
void appendStringInfo(StringInfo str, const char *fmt,...)
Definition: stringinfo.c:97
void initStringInfo(StringInfo str)
Definition: stringinfo.c:59
char bgw_function_name[BGW_MAXLEN]
Definition: bgworker.h:97
Datum bgw_main_arg
Definition: bgworker.h:98
char bgw_name[BGW_MAXLEN]
Definition: bgworker.h:91
int bgw_restart_time
Definition: bgworker.h:95
char bgw_type[BGW_MAXLEN]
Definition: bgworker.h:92
BgWorkerStartTime bgw_start_time
Definition: bgworker.h:94
char bgw_extra[BGW_EXTRALEN]
Definition: bgworker.h:99
pid_t bgw_notify_pid
Definition: bgworker.h:100
char bgw_library_name[MAXPGPATH]
Definition: bgworker.h:96
TupleDesc tupdesc
Definition: spi.h:25
HeapTuple * vals
Definition: spi.h:26
Definition: type.h:95
const char * name
Definition: worker_spi.c:65
const char * schema
Definition: worker_spi.c:64
uint32 WaitEventExtensionNew(const char *wait_event_name)
Definition: wait_event.c:164
const char * name
#define SIGHUP
Definition: win32_port.h:168
static int worker_spi_naptime
Definition: worker_spi.c:54
static void initialize_worker_spi(worktable *table)
Definition: worker_spi.c:73
void _PG_init(void)
Definition: worker_spi.c:315
PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn()
Definition: worker_spi.c:138
PG_MODULE_MAGIC
Definition: worker_spi.c:47
Datum worker_spi_launch(PG_FUNCTION_ARGS)
Definition: worker_spi.c:405
static uint32 worker_spi_wait_event_main
Definition: worker_spi.c:60
static char * worker_spi_database
Definition: worker_spi.c:56
static int worker_spi_total_workers
Definition: worker_spi.c:55
struct worktable worktable
static char * worker_spi_role
Definition: worker_spi.c:57
PG_FUNCTION_INFO_V1(worker_spi_launch)
void StartTransactionCommand(void)
Definition: xact.c:2953
void SetCurrentStatementStartTimestamp(void)
Definition: xact.c:899
void CommitTransactionCommand(void)
Definition: xact.c:3050