Branch data Line data Source code
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-2026, 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/latch.h"
30 : :
31 : : /* these headers are used by this particular worker's code */
32 : : #include "access/xact.h"
33 : : #include "catalog/pg_database.h"
34 : : #include "executor/spi.h"
35 : : #include "fmgr.h"
36 : : #include "lib/stringinfo.h"
37 : : #include "pgstat.h"
38 : : #include "tcop/utility.h"
39 : : #include "utils/acl.h"
40 : : #include "utils/builtins.h"
41 : : #include "utils/snapmgr.h"
42 : : #include "utils/wait_event.h"
43 : :
44 : 16 : PG_MODULE_MAGIC;
45 : :
46 : 15 : PG_FUNCTION_INFO_V1(worker_spi_launch);
47 : :
48 : : pg_noreturn PGDLLEXPORT void worker_spi_main(Datum main_arg);
49 : :
50 : : /* GUC variables */
51 : : static int worker_spi_naptime = 10;
52 : : static int worker_spi_total_workers = 2;
53 : : static char *worker_spi_database = NULL;
54 : : static char *worker_spi_role = NULL;
55 : :
56 : : /* value cached, fetched from shared memory */
57 : : static uint32 worker_spi_wait_event_main = 0;
58 : :
59 : : typedef struct worktable
60 : : {
61 : : const char *schema;
62 : : const char *name;
63 : : } worktable;
64 : :
65 : : /*
66 : : * Initialize workspace for a worker process: create the schema if it doesn't
67 : : * already exist.
68 : : */
69 : : static void
70 : 6 : initialize_worker_spi(worktable *table)
71 : : {
72 : : int ret;
73 : : int ntup;
74 : : bool isnull;
75 : : StringInfoData buf;
76 : :
77 : 6 : SetCurrentStatementStartTimestamp();
78 : 6 : StartTransactionCommand();
79 : 6 : SPI_connect();
80 : 6 : PushActiveSnapshot(GetTransactionSnapshot());
81 : 6 : pgstat_report_activity(STATE_RUNNING, "initializing worker_spi schema");
82 : :
83 : : /* XXX could we use CREATE SCHEMA IF NOT EXISTS? */
84 : 6 : initStringInfo(&buf);
85 : 6 : appendStringInfo(&buf, "select count(*) from pg_namespace where nspname = '%s'",
86 : : table->schema);
87 : :
88 : 6 : debug_query_string = buf.data;
89 : 6 : ret = SPI_execute(buf.data, true, 0);
90 [ - + ]: 6 : if (ret != SPI_OK_SELECT)
91 [ # # ]: 0 : elog(FATAL, "SPI_execute failed: error code %d", ret);
92 : :
93 [ - + ]: 6 : if (SPI_processed != 1)
94 [ # # ]: 0 : elog(FATAL, "not a singleton result");
95 : :
96 : 6 : ntup = DatumGetInt64(SPI_getbinval(SPI_tuptable->vals[0],
97 : 6 : SPI_tuptable->tupdesc,
98 : : 1, &isnull));
99 [ - + ]: 6 : if (isnull)
100 [ # # ]: 0 : elog(FATAL, "null result");
101 : :
102 [ + - ]: 6 : if (ntup == 0)
103 : : {
104 : 6 : debug_query_string = NULL;
105 : 6 : resetStringInfo(&buf);
106 : 6 : appendStringInfo(&buf,
107 : : "CREATE SCHEMA \"%s\" "
108 : : "CREATE TABLE \"%s\" ("
109 : : " type text CHECK (type IN ('total', 'delta')), "
110 : : " value integer)"
111 : : "CREATE UNIQUE INDEX \"%s_unique_total\" ON \"%s\" (type) "
112 : : "WHERE type = 'total'",
113 : : table->schema, table->name, table->name, table->name);
114 : :
115 : : /* set statement start time */
116 : 6 : SetCurrentStatementStartTimestamp();
117 : :
118 : 6 : debug_query_string = buf.data;
119 : 6 : ret = SPI_execute(buf.data, false, 0);
120 : :
121 [ - + ]: 6 : if (ret != SPI_OK_UTILITY)
122 [ # # ]: 0 : elog(FATAL, "failed to create my schema");
123 : :
124 : 6 : debug_query_string = NULL; /* rest is not statement-specific */
125 : : }
126 : :
127 : 6 : SPI_finish();
128 : 6 : PopActiveSnapshot();
129 : 6 : CommitTransactionCommand();
130 : 6 : debug_query_string = NULL;
131 : 6 : pgstat_report_activity(STATE_IDLE, NULL);
132 : 6 : }
133 : :
134 : : void
135 : 8 : worker_spi_main(Datum main_arg)
136 : : {
137 : 8 : int index = DatumGetInt32(main_arg);
138 : : worktable *table;
139 : : StringInfoData buf;
140 : : char name[20];
141 : : Oid dboid;
142 : : Oid roleoid;
143 : : char *p;
144 : 8 : uint32 flags = 0;
145 : :
146 : 8 : table = palloc_object(worktable);
147 : 8 : sprintf(name, "schema%d", index);
148 : 8 : table->schema = pstrdup(name);
149 : 8 : table->name = pstrdup("counted");
150 : :
151 : : /* fetch database and role OIDs, these are set for a dynamic worker */
152 : 8 : p = MyBgworkerEntry->bgw_extra;
153 : 8 : memcpy(&dboid, p, sizeof(Oid));
154 : 8 : p += sizeof(Oid);
155 : 8 : memcpy(&roleoid, p, sizeof(Oid));
156 : 8 : p += sizeof(Oid);
157 : 8 : memcpy(&flags, p, sizeof(uint32));
158 : :
159 : : /* Establish signal handlers before unblocking signals. */
160 : 8 : pqsignal(SIGHUP, SignalHandlerForConfigReload);
161 : 8 : pqsignal(SIGTERM, die);
162 : :
163 : : /* We're now ready to receive signals */
164 : 8 : BackgroundWorkerUnblockSignals();
165 : :
166 : : /* Connect to our database */
167 [ + - ]: 8 : if (OidIsValid(dboid))
168 : 8 : BackgroundWorkerInitializeConnectionByOid(dboid, roleoid, flags);
169 : : else
170 : 0 : BackgroundWorkerInitializeConnection(worker_spi_database,
171 : : worker_spi_role, flags);
172 : :
173 [ + - ]: 6 : elog(LOG, "%s initialized with %s.%s",
174 : : MyBgworkerEntry->bgw_name, table->schema, table->name);
175 : 6 : initialize_worker_spi(table);
176 : :
177 : : /*
178 : : * Quote identifiers passed to us. Note that this must be done after
179 : : * initialize_worker_spi, because that routine assumes the names are not
180 : : * quoted.
181 : : *
182 : : * Note some memory might be leaked here.
183 : : */
184 : 6 : table->schema = quote_identifier(table->schema);
185 : 6 : table->name = quote_identifier(table->name);
186 : :
187 : 6 : initStringInfo(&buf);
188 : 6 : appendStringInfo(&buf,
189 : : "WITH deleted AS (DELETE "
190 : : "FROM %s.%s "
191 : : "WHERE type = 'delta' RETURNING value), "
192 : : "total AS (SELECT coalesce(sum(value), 0) as sum "
193 : : "FROM deleted) "
194 : : "UPDATE %s.%s "
195 : : "SET value = %s.value + total.sum "
196 : : "FROM total WHERE type = 'total' "
197 : : "RETURNING %s.value",
198 : : table->schema, table->name,
199 : : table->schema, table->name,
200 : : table->name,
201 : : table->name);
202 : :
203 : : /*
204 : : * Main loop: do this until SIGTERM is received and processed by
205 : : * ProcessInterrupts.
206 : : */
207 : : for (;;)
208 : 6 : {
209 : : int ret;
210 : :
211 : : /* First time, allocate or get the custom wait event */
212 [ + + ]: 12 : if (worker_spi_wait_event_main == 0)
213 : 6 : worker_spi_wait_event_main = WaitEventExtensionNew("WorkerSpiMain");
214 : :
215 : : /*
216 : : * Background workers mustn't call usleep() or any direct equivalent:
217 : : * instead, they may wait on their process latch, which sleeps as
218 : : * necessary, but is awakened if postmaster dies. That way the
219 : : * background process goes away immediately in an emergency.
220 : : */
221 : 12 : (void) WaitLatch(MyLatch,
222 : : WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
223 : : worker_spi_naptime * 1000L,
224 : : worker_spi_wait_event_main);
225 : 12 : ResetLatch(MyLatch);
226 : :
227 [ + + ]: 12 : CHECK_FOR_INTERRUPTS();
228 : :
229 : : /*
230 : : * In case of a SIGHUP, just reload the configuration.
231 : : */
232 [ + + ]: 6 : if (ConfigReloadPending)
233 : : {
234 : 1 : ConfigReloadPending = false;
235 : 1 : ProcessConfigFile(PGC_SIGHUP);
236 : : }
237 : :
238 : : /*
239 : : * Start a transaction on which we can run queries. Note that each
240 : : * StartTransactionCommand() call should be preceded by a
241 : : * SetCurrentStatementStartTimestamp() call, which sets both the time
242 : : * for the statement we're about the run, and also the transaction
243 : : * start time. Also, each other query sent to SPI should probably be
244 : : * preceded by SetCurrentStatementStartTimestamp(), so that statement
245 : : * start time is always up to date.
246 : : *
247 : : * The SPI_connect() call lets us run queries through the SPI manager,
248 : : * and the PushActiveSnapshot() call creates an "active" snapshot
249 : : * which is necessary for queries to have MVCC data to work on.
250 : : *
251 : : * The pgstat_report_activity() call makes our activity visible
252 : : * through the pgstat views.
253 : : */
254 : 6 : SetCurrentStatementStartTimestamp();
255 : 6 : StartTransactionCommand();
256 : 6 : SPI_connect();
257 : 6 : PushActiveSnapshot(GetTransactionSnapshot());
258 : 6 : debug_query_string = buf.data;
259 : 6 : pgstat_report_activity(STATE_RUNNING, buf.data);
260 : :
261 : : /* We can now execute queries via SPI */
262 : 6 : ret = SPI_execute(buf.data, false, 0);
263 : :
264 [ - + ]: 6 : if (ret != SPI_OK_UPDATE_RETURNING)
265 [ # # ]: 0 : elog(FATAL, "cannot select from table %s.%s: error code %d",
266 : : table->schema, table->name, ret);
267 : :
268 [ + + ]: 6 : if (SPI_processed > 0)
269 : : {
270 : : bool isnull;
271 : : int32 val;
272 : :
273 : 1 : val = DatumGetInt32(SPI_getbinval(SPI_tuptable->vals[0],
274 : 1 : SPI_tuptable->tupdesc,
275 : : 1, &isnull));
276 [ + - ]: 1 : if (!isnull)
277 [ + - ]: 1 : elog(LOG, "%s: count in %s.%s is now %d",
278 : : MyBgworkerEntry->bgw_name,
279 : : table->schema, table->name, val);
280 : : }
281 : :
282 : : /*
283 : : * And finish our transaction.
284 : : */
285 : 6 : SPI_finish();
286 : 6 : PopActiveSnapshot();
287 : 6 : CommitTransactionCommand();
288 : 6 : debug_query_string = NULL;
289 : 6 : pgstat_report_stat(true);
290 : 6 : pgstat_report_activity(STATE_IDLE, NULL);
291 : : }
292 : :
293 : : /* Not reachable */
294 : : }
295 : :
296 : : /*
297 : : * Entrypoint of this module.
298 : : *
299 : : * We register more than one worker process here, to demonstrate how that can
300 : : * be done.
301 : : */
302 : : void
303 : 16 : _PG_init(void)
304 : : {
305 : : BackgroundWorker worker;
306 : :
307 : : /* get the configuration */
308 : :
309 : : /*
310 : : * These GUCs are defined even if this library is not loaded with
311 : : * shared_preload_libraries, for worker_spi_launch().
312 : : */
313 : 16 : DefineCustomIntVariable("worker_spi.naptime",
314 : : "Duration between each check (in seconds).",
315 : : NULL,
316 : : &worker_spi_naptime,
317 : : 10,
318 : : 1,
319 : : INT_MAX,
320 : : PGC_SIGHUP,
321 : : 0,
322 : : NULL,
323 : : NULL,
324 : : NULL);
325 : :
326 : 16 : DefineCustomStringVariable("worker_spi.database",
327 : : "Database to connect to.",
328 : : NULL,
329 : : &worker_spi_database,
330 : : "postgres",
331 : : PGC_SIGHUP,
332 : : 0,
333 : : NULL, NULL, NULL);
334 : :
335 : 16 : DefineCustomStringVariable("worker_spi.role",
336 : : "Role to connect with.",
337 : : NULL,
338 : : &worker_spi_role,
339 : : NULL,
340 : : PGC_SIGHUP,
341 : : 0,
342 : : NULL, NULL, NULL);
343 : :
344 [ + + ]: 16 : if (!process_shared_preload_libraries_in_progress)
345 : 15 : return;
346 : :
347 : 1 : DefineCustomIntVariable("worker_spi.total_workers",
348 : : "Number of workers.",
349 : : NULL,
350 : : &worker_spi_total_workers,
351 : : 2,
352 : : 1,
353 : : 100,
354 : : PGC_POSTMASTER,
355 : : 0,
356 : : NULL,
357 : : NULL,
358 : : NULL);
359 : :
360 : 1 : MarkGUCPrefixReserved("worker_spi");
361 : :
362 : : /* set up common data for all our workers */
363 : 1 : memset(&worker, 0, sizeof(worker));
364 : 1 : worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
365 : : BGWORKER_BACKEND_DATABASE_CONNECTION;
366 : 1 : worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
367 : 1 : worker.bgw_restart_time = BGW_NEVER_RESTART;
368 : 1 : sprintf(worker.bgw_library_name, "worker_spi");
369 : 1 : sprintf(worker.bgw_function_name, "worker_spi_main");
370 : 1 : worker.bgw_notify_pid = 0;
371 : :
372 : : /*
373 : : * Now fill in worker-specific data, and do the actual registrations.
374 : : *
375 : : * bgw_extra can optionally include a database OID, a role OID and a set
376 : : * of flags. This is left empty here to fallback to the related GUCs at
377 : : * startup (0 for the bgworker flags).
378 : : */
379 [ + + ]: 4 : for (int i = 1; i <= worker_spi_total_workers; i++)
380 : : {
381 : 3 : snprintf(worker.bgw_name, BGW_MAXLEN, "worker_spi worker %d", i);
382 : 3 : snprintf(worker.bgw_type, BGW_MAXLEN, "worker_spi");
383 : 3 : worker.bgw_main_arg = Int32GetDatum(i);
384 : :
385 : 3 : RegisterBackgroundWorker(&worker);
386 : : }
387 : : }
388 : :
389 : : /*
390 : : * Dynamically launch an SPI worker.
391 : : */
392 : : Datum
393 : 12 : worker_spi_launch(PG_FUNCTION_ARGS)
394 : : {
395 : 12 : int32 i = PG_GETARG_INT32(0);
396 : 12 : Oid dboid = PG_GETARG_OID(1);
397 : 12 : Oid roleoid = PG_GETARG_OID(2);
398 : : BackgroundWorker worker;
399 : : BackgroundWorkerHandle *handle;
400 : : BgwHandleStatus status;
401 : : pid_t pid;
402 : : char *p;
403 : 12 : uint32 flags = 0;
404 : 12 : ArrayType *arr = PG_GETARG_ARRAYTYPE_P(3);
405 : : Size ndim;
406 : : int nelems;
407 : : Datum *datum_flags;
408 : 12 : bool interruptible = PG_GETARG_BOOL(4);
409 : :
410 : 12 : memset(&worker, 0, sizeof(worker));
411 : 12 : worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
412 : : BGWORKER_BACKEND_DATABASE_CONNECTION;
413 : :
414 [ + + ]: 12 : if (interruptible)
415 : 4 : worker.bgw_flags |= BGWORKER_INTERRUPTIBLE;
416 : :
417 : 12 : worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
418 : 12 : worker.bgw_restart_time = BGW_NEVER_RESTART;
419 : 12 : sprintf(worker.bgw_library_name, "worker_spi");
420 : 12 : sprintf(worker.bgw_function_name, "worker_spi_main");
421 : 12 : snprintf(worker.bgw_name, BGW_MAXLEN, "worker_spi dynamic worker %d", i);
422 : 12 : snprintf(worker.bgw_type, BGW_MAXLEN, "worker_spi dynamic");
423 : 12 : worker.bgw_main_arg = Int32GetDatum(i);
424 : : /* set bgw_notify_pid so that we can use WaitForBackgroundWorkerStartup */
425 : 12 : worker.bgw_notify_pid = MyProcPid;
426 : :
427 : : /* extract flags, if any */
428 : 12 : ndim = ARR_NDIM(arr);
429 [ - + ]: 12 : if (ndim > 1)
430 [ # # ]: 0 : ereport(ERROR,
431 : : (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
432 : : errmsg("flags array must be one-dimensional")));
433 : :
434 [ - + ]: 12 : if (array_contains_nulls(arr))
435 [ # # ]: 0 : ereport(ERROR,
436 : : (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
437 : : errmsg("flags array must not contain nulls")));
438 : :
439 : : Assert(ARR_ELEMTYPE(arr) == TEXTOID);
440 : 12 : deconstruct_array_builtin(arr, TEXTOID, &datum_flags, NULL, &nelems);
441 : :
442 [ + + ]: 14 : for (i = 0; i < nelems; i++)
443 : : {
444 : 2 : char *optname = TextDatumGetCString(datum_flags[i]);
445 : :
446 [ + + ]: 2 : if (strcmp(optname, "ALLOWCONN") == 0)
447 : 1 : flags |= BGWORKER_BYPASS_ALLOWCONN;
448 [ + - ]: 1 : else if (strcmp(optname, "ROLELOGINCHECK") == 0)
449 : 1 : flags |= BGWORKER_BYPASS_ROLELOGINCHECK;
450 : : else
451 [ # # ]: 0 : ereport(ERROR,
452 : : (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
453 : : errmsg("incorrect flag value found in array")));
454 : : }
455 : :
456 : : /*
457 : : * Register database and role to use for the worker started in bgw_extra.
458 : : * If none have been provided, this will fall back to the GUCs at startup.
459 : : */
460 [ + + ]: 12 : if (!OidIsValid(dboid))
461 : 1 : dboid = get_database_oid(worker_spi_database, false);
462 : :
463 : : /*
464 : : * worker_spi_role is NULL by default, so this gives to worker_spi_main()
465 : : * an invalid OID in this case.
466 : : */
467 [ + + - + ]: 12 : if (!OidIsValid(roleoid) && worker_spi_role)
468 : 0 : roleoid = get_role_oid(worker_spi_role, false);
469 : :
470 : 12 : p = worker.bgw_extra;
471 : 12 : memcpy(p, &dboid, sizeof(Oid));
472 : 12 : p += sizeof(Oid);
473 : 12 : memcpy(p, &roleoid, sizeof(Oid));
474 : 12 : p += sizeof(Oid);
475 : 12 : memcpy(p, &flags, sizeof(uint32));
476 : :
477 [ - + ]: 12 : if (!RegisterDynamicBackgroundWorker(&worker, &handle))
478 : 0 : PG_RETURN_NULL();
479 : :
480 : 12 : status = WaitForBackgroundWorkerStartup(handle, &pid);
481 : :
482 [ - + ]: 12 : if (status == BGWH_STOPPED)
483 [ # # ]: 0 : ereport(ERROR,
484 : : (errcode(ERRCODE_INSUFFICIENT_RESOURCES),
485 : : errmsg("could not start background process"),
486 : : errhint("More details may be available in the server log.")));
487 [ - + ]: 12 : if (status == BGWH_POSTMASTER_DIED)
488 [ # # ]: 0 : ereport(ERROR,
489 : : (errcode(ERRCODE_INSUFFICIENT_RESOURCES),
490 : : errmsg("cannot start background processes without postmaster"),
491 : : errhint("Kill all remaining database processes and restart the database.")));
492 : : Assert(status == BGWH_STARTED);
493 : :
494 : 12 : PG_RETURN_INT32(pid);
495 : : }
|
