worker_spi.c File Reference

#include "postgres.h"
#include "miscadmin.h"
#include "postmaster/bgworker.h"
#include "postmaster/interrupt.h"
#include "storage/latch.h"
#include "access/xact.h"
#include "catalog/pg_database.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"

Include dependency graph for worker_spi.c:

Go to the source code of this file.

Data Structures
struct	worktable

Typedefs
typedef struct worktable	worktable

Functions
	PG_FUNCTION_INFO_V1 (worker_spi_launch)
PGDLLEXPORT pg_noreturn void	worker_spi_main (Datum main_arg)
static void	initialize_worker_spi (worktable *table)
void	_PG_init (void)
Datum	worker_spi_launch (PG_FUNCTION_ARGS)

Variables
	PG_MODULE_MAGIC
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
static uint32	worker_spi_wait_event_main = 0

Typedef Documentation

worktable

typedef struct worktable worktable

Function Documentation

_PG_init()

void _PG_init

(

void

)

Definition at line 302 of file worker_spi.c.

{
    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 database 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);
    }
}

References BackgroundWorker::bgw_flags, BackgroundWorker::bgw_function_name, BackgroundWorker::bgw_library_name, BackgroundWorker::bgw_main_arg, BGW_MAXLEN, BackgroundWorker::bgw_name, BGW_NEVER_RESTART, BackgroundWorker::bgw_notify_pid, BackgroundWorker::bgw_restart_time, BackgroundWorker::bgw_start_time, BackgroundWorker::bgw_type, BGWORKER_BACKEND_DATABASE_CONNECTION, BGWORKER_SHMEM_ACCESS, BgWorkerStart_RecoveryFinished, DefineCustomIntVariable(), DefineCustomStringVariable(), i, Int32GetDatum(), MarkGUCPrefixReserved(), PGC_POSTMASTER, PGC_SIGHUP, process_shared_preload_libraries_in_progress, RegisterBackgroundWorker(), snprintf, sprintf, worker_spi_database, worker_spi_naptime, worker_spi_role, and worker_spi_total_workers.

initialize_worker_spi()

static void initialize_worker_spi ( worktable *  table )

static

Definition at line 69 of file worker_spi.c.

{
    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);
}

References appendStringInfo(), buf, CommitTransactionCommand(), DatumGetInt64(), debug_query_string, elog, FATAL, GetTransactionSnapshot(), initStringInfo(), pgstat_report_activity(), PopActiveSnapshot(), PushActiveSnapshot(), resetStringInfo(), SetCurrentStatementStartTimestamp(), SPI_connect(), SPI_execute(), SPI_finish(), SPI_getbinval(), SPI_OK_SELECT, SPI_OK_UTILITY, SPI_processed, SPI_tuptable, StartTransactionCommand(), STATE_IDLE, STATE_RUNNING, table, SPITupleTable::tupdesc, and SPITupleTable::vals.

Referenced by worker_spi_main().

PG_FUNCTION_INFO_V1()

PG_FUNCTION_INFO_V1

(

worker_spi_launch

)

worker_spi_launch()

Datum worker_spi_launch

(

PG_FUNCTION_ARGS

)

Definition at line 392 of file worker_spi.c.

{
    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);
}

References ARR_ELEMTYPE, ARR_NDIM, array_contains_nulls(), Assert(), BackgroundWorker::bgw_extra, BackgroundWorker::bgw_flags, BackgroundWorker::bgw_function_name, BackgroundWorker::bgw_library_name, BackgroundWorker::bgw_main_arg, BGW_MAXLEN, BackgroundWorker::bgw_name, BGW_NEVER_RESTART, BackgroundWorker::bgw_notify_pid, BackgroundWorker::bgw_restart_time, BackgroundWorker::bgw_start_time, BackgroundWorker::bgw_type, BGWH_POSTMASTER_DIED, BGWH_STARTED, BGWH_STOPPED, BGWORKER_BACKEND_DATABASE_CONNECTION, BGWORKER_BYPASS_ALLOWCONN, BGWORKER_BYPASS_ROLELOGINCHECK, BGWORKER_SHMEM_ACCESS, BgWorkerStart_RecoveryFinished, deconstruct_array_builtin(), ereport, errcode(), errhint(), errmsg(), ERROR, get_database_oid(), get_role_oid(), i, Int32GetDatum(), MyProcPid, OidIsValid, PG_GETARG_ARRAYTYPE_P, PG_GETARG_INT32, PG_GETARG_OID, PG_RETURN_INT32, PG_RETURN_NULL, RegisterDynamicBackgroundWorker(), snprintf, sprintf, TextDatumGetCString, WaitForBackgroundWorkerStartup(), worker_spi_database, and worker_spi_role.

worker_spi_main()

void worker_spi_main

(

Datum

main_arg

)

Definition at line 134 of file worker_spi.c.

{
    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);
 
    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 */
}

References appendStringInfo(), BackgroundWorkerInitializeConnection(), BackgroundWorkerInitializeConnectionByOid(), BackgroundWorkerUnblockSignals(), BackgroundWorker::bgw_extra, BackgroundWorker::bgw_name, buf, CHECK_FOR_INTERRUPTS, CommitTransactionCommand(), ConfigReloadPending, DatumGetInt32(), debug_query_string, die, elog, FATAL, GetTransactionSnapshot(), initialize_worker_spi(), initStringInfo(), LOG, MyBgworkerEntry, MyLatch, name, OidIsValid, palloc(), PGC_SIGHUP, pgstat_report_activity(), pgstat_report_stat(), PopActiveSnapshot(), pqsignal, ProcessConfigFile(), pstrdup(), PushActiveSnapshot(), quote_identifier(), ResetLatch(), SetCurrentStatementStartTimestamp(), SIGHUP, SignalHandlerForConfigReload(), SPI_connect(), SPI_execute(), SPI_finish(), SPI_getbinval(), SPI_OK_UPDATE_RETURNING, SPI_processed, SPI_tuptable, sprintf, StartTransactionCommand(), STATE_IDLE, STATE_RUNNING, table, SPITupleTable::tupdesc, val, SPITupleTable::vals, WaitEventExtensionNew(), WaitLatch(), WL_EXIT_ON_PM_DEATH, WL_LATCH_SET, WL_TIMEOUT, worker_spi_database, worker_spi_naptime, worker_spi_role, and worker_spi_wait_event_main.

Variable Documentation

PG_MODULE_MAGIC

PG_MODULE_MAGIC

Definition at line 43 of file worker_spi.c.

worker_spi_database

char* worker_spi_database = NULL

static

Definition at line 52 of file worker_spi.c.

Referenced by _PG_init(), worker_spi_launch(), and worker_spi_main().

worker_spi_naptime

int worker_spi_naptime = 10

static

Definition at line 50 of file worker_spi.c.

Referenced by _PG_init(), and worker_spi_main().

worker_spi_role

char* worker_spi_role = NULL

static

Definition at line 53 of file worker_spi.c.

Referenced by _PG_init(), worker_spi_launch(), and worker_spi_main().

worker_spi_total_workers

int worker_spi_total_workers = 2

static

Definition at line 51 of file worker_spi.c.

Referenced by _PG_init().

worker_spi_wait_event_main

uint32 worker_spi_wait_event_main = 0

static

Definition at line 56 of file worker_spi.c.

Referenced by worker_spi_main().