diff --git a/contrib/Makefile b/contrib/Makefile index 25263c0..04ec28a 100644 --- a/contrib/Makefile +++ b/contrib/Makefile @@ -29,6 +29,7 @@ SUBDIRS = \ oid2name \ pageinspect \ passwordcheck \ + pg_background \ pg_buffercache \ pg_freespacemap \ pg_prewarm \ diff --git a/contrib/pg_background/Makefile b/contrib/pg_background/Makefile new file mode 100644 index 0000000..c4e717d --- /dev/null +++ b/contrib/pg_background/Makefile @@ -0,0 +1,18 @@ +# contrib/pg_background/Makefile + +MODULE_big = pg_background +OBJS = pg_background.o + +EXTENSION = pg_background +DATA = pg_background--1.0.sql + +ifdef USE_PGXS +PG_CONFIG = pg_config +PGXS := $(shell $(PG_CONFIG) --pgxs) +include $(PGXS) +else +subdir = contrib/pg_background +top_builddir = ../.. +include $(top_builddir)/src/Makefile.global +include $(top_srcdir)/contrib/contrib-global.mk +endif diff --git a/contrib/pg_background/pg_background--1.0.sql b/contrib/pg_background/pg_background--1.0.sql new file mode 100644 index 0000000..bc8a881 --- /dev/null +++ b/contrib/pg_background/pg_background--1.0.sql @@ -0,0 +1,30 @@ +/* contrib/pg_background/pg_background--1.0.sql */ + +-- complain if script is sourced in psql, rather than via CREATE EXTENSION +\echo Use "CREATE EXTENSION pg_background" to load this file. \quit + +CREATE FUNCTION pg_background_launch() + RETURNS pg_catalog.int4 STRICT + AS 'MODULE_PATHNAME' LANGUAGE C; + +CREATE FUNCTION pg_background_run(pid pg_catalog.int4, sql pg_catalog.text) + RETURNS pg_catalog.void STRICT + AS 'MODULE_PATHNAME' LANGUAGE C; + +CREATE FUNCTION pg_background_result(pid pg_catalog.int4) + RETURNS SETOF pg_catalog.record STRICT + AS 'MODULE_PATHNAME' LANGUAGE C; + +CREATE FUNCTION pg_background_detach(pid pg_catalog.int4) + RETURNS pg_catalog.void STRICT + AS 'MODULE_PATHNAME' LANGUAGE C; + +REVOKE ALL ON FUNCTION pg_background_launch() + FROM public; +REVOKE ALL ON FUNCTION pg_background_run(pid pg_catalog.int4, + sql pg_catalog.text) + FROM public; +REVOKE ALL ON FUNCTION pg_background_result(pg_catalog.int4) + FROM public; +REVOKE ALL ON FUNCTION pg_background_detach(pg_catalog.int4) + FROM public; diff --git a/contrib/pg_background/pg_background.c b/contrib/pg_background/pg_background.c new file mode 100644 index 0000000..8d26b38 --- /dev/null +++ b/contrib/pg_background/pg_background.c @@ -0,0 +1,316 @@ +/*-------------------------------------------------------------------------- + * + * pg_background.c + * Run SQL commands using a background worker. + * + * Copyright (C) 2016, PostgreSQL Global Development Group + * + * IDENTIFICATION + * contrib/pg_background/pg_background.c + * + * ------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include "access/htup_details.h" +#include "catalog/pg_type.h" +#include "funcapi.h" +#include "miscadmin.h" +#include "tcop/bgsession.h" +#include "utils/acl.h" +#include "utils/builtins.h" + +/* Private state maintained by the launching backend for IPC. */ +typedef struct pg_background_worker_info +{ + pid_t pid; + Oid current_user_id; + BackgroundSession *session; + uint32 result_count; /* TODO: uint32? */ +} pg_background_worker_info; + +/* Private state maintained across calls to pg_background_result. */ +typedef struct pg_background_result_state +{ + pg_background_worker_info *info; + BackgroundSessionResult *result; +} pg_background_result_state; + +static HTAB *worker_hash; + +static void remove_worker_info(pid_t pid); +static pg_background_worker_info *find_worker_info(pid_t pid); +static void save_worker_info(pid_t pid, BackgroundSession *session); +static void check_rights(pg_background_worker_info *info); + +PG_MODULE_MAGIC; + +PG_FUNCTION_INFO_V1(pg_background_launch); +PG_FUNCTION_INFO_V1(pg_background_run); +PG_FUNCTION_INFO_V1(pg_background_result); +PG_FUNCTION_INFO_V1(pg_background_detach); + +/* + * Start a dynamic background worker. + */ +Datum +pg_background_launch(PG_FUNCTION_ARGS) +{ + BackgroundSession *session; + pid_t pid; + + session = BackgroundSessionNew(&pid); + + /* Save worker info */ + save_worker_info(pid, session); + + /* Return the worker's PID. */ + PG_RETURN_INT32(pid); +} + +/* + * Run a user-specified SQL command. + */ +Datum +pg_background_run(PG_FUNCTION_ARGS) +{ + int32 pid = PG_GETARG_INT32(0); + char *sql = text_to_cstring(PG_GETARG_TEXT_PP(1)); + pg_background_worker_info *info; + + /* See if we have a connection to the specified PID. */ + if ((info = find_worker_info(pid)) == NULL) + ereport(ERROR, + (errcode(ERRCODE_UNDEFINED_OBJECT), + errmsg("PID %d is not attached to this session", pid))); + check_rights(info); + + /* Execute give SQL query */ + BackgroundSessionExecuteSQL(info->session, sql); + info->result_count++; + + PG_RETURN_VOID(); +} + + +/* + * Retrieve the results of a background query previously launched in this + * session. + */ +Datum +pg_background_result(PG_FUNCTION_ARGS) +{ + int32 pid = PG_GETARG_INT32(0); + FuncCallContext *funcctx; + pg_background_result_state *state; + TupleDesc tupdesc; + BackgroundSessionResult *result; + + /* First-time setup. */ + if (SRF_IS_FIRSTCALL()) + { + MemoryContext oldcontext; + pg_background_worker_info *info; + + funcctx = SRF_FIRSTCALL_INIT(); + oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); + + /* See if we have a connection to the specified PID. */ + if ((info = find_worker_info(pid)) == NULL) + ereport(ERROR, + (errcode(ERRCODE_UNDEFINED_OBJECT), + errmsg("PID %d is not attached to this session", pid))); + check_rights(info); + + /* Can't read results twice. */ + if (info->result_count <= 0) + ereport(ERROR, + (errcode(ERRCODE_UNDEFINED_OBJECT), + errmsg("results for PID %d have already been consumed", pid))); + info->result_count--; + + /* Set up tuple-descriptor based on column definition list. */ + if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE) + ereport(ERROR, + (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), + errmsg("function returning record called in context " + "that cannot accept type record"), + errhint("Try calling the function in the FROM clause " + "using a column definition list."))); + result = BackgroundSessionFetchResult(info->session); + funcctx->tuple_desc = BlessTupleDesc(tupdesc); + + if (result->tupdesc && tupdesc->natts != result->tupdesc->natts) + ereport(ERROR, + (errcode(ERRCODE_DATATYPE_MISMATCH), + errmsg("remote query result rowtype does not match the specified FROM clause rowtype"))); + + /* Cache state that will be needed on every call. */ + state = palloc0(sizeof(pg_background_result_state)); + state->info = info; + state->result = result; + + funcctx->user_fctx = state; + MemoryContextSwitchTo(oldcontext); + } + + funcctx = SRF_PERCALL_SETUP(); + tupdesc = funcctx->tuple_desc; + state = funcctx->user_fctx; + result = state->result; + + if (result->tupdesc) + { + if (result->tuples != NIL) + { + HeapTuple tuple = (HeapTuple) linitial(result->tuples); + Datum *values = (Datum *) palloc(tupdesc->natts * sizeof(Datum)); + bool *isnull = (bool *) palloc(tupdesc->natts * sizeof(bool)); + + /* + * Perform conversion of a tuple according to column + * defination list given in FROM clause + */ + heap_deform_tuple(tuple, result->tupdesc, values, isnull); + tuple = heap_form_tuple(tupdesc, values, isnull); + result->tuples = list_delete_first(result->tuples); + + SRF_RETURN_NEXT(funcctx, HeapTupleGetDatum(tuple)); + } + } + else /* If no data rows, return the command tags instead. */ + { + if (tupdesc->natts != 1 || tupdesc->attrs[0]->atttypid != TEXTOID) + ereport(ERROR, + (errcode(ERRCODE_DATATYPE_MISMATCH), + errmsg("remote query did not return a result set, but result rowtype is not a single text column"))); + + if (result->command != NULL) + { + bool isnull = false; + Datum value = PointerGetDatum(cstring_to_text(result->command)); + HeapTuple tuple = heap_form_tuple(tupdesc, &value, &isnull); + + result->command = NULL; + + SRF_RETURN_NEXT(funcctx, HeapTupleGetDatum(tuple)); + } + } + + SRF_RETURN_DONE(funcctx); +} + +/* + * Detach from the dynamic shared memory segment used for communication with + * a background worker. This prevents the worker from stalling waiting for + * us to read its results. + */ +Datum +pg_background_detach(PG_FUNCTION_ARGS) +{ + int32 pid = PG_GETARG_INT32(0); + pg_background_worker_info *info; + + info = find_worker_info(pid); + if (info == NULL) + ereport(ERROR, + (errcode(ERRCODE_UNDEFINED_OBJECT), + errmsg("PID %d is not attached to this session", pid))); + + check_rights(info); + remove_worker_info(pid); + BackgroundSessionEnd(info->session); + + PG_RETURN_VOID(); +} + +static void +remove_worker_info(pid_t pid) +{ + bool found; + + /* Remove the hashtable entry. */ + hash_search(worker_hash, (void *) &pid, HASH_REMOVE, &found); + if (!found) + elog(ERROR, "pg_background worker_hash table corrupted"); +} + +/* + * Find the background worker information for the worker with a given PID. + */ +static pg_background_worker_info * +find_worker_info(pid_t pid) +{ + pg_background_worker_info *info = NULL; + + if (worker_hash != NULL) + info = hash_search(worker_hash, (void *) &pid, HASH_FIND, NULL); + + return info; +} + +/* + * Save worker info. + */ +static void +save_worker_info(pid_t pid, BackgroundSession *session) +{ + pg_background_worker_info *info; + Oid current_user_id; + int sec_context; + + /* If the hash table hasn't been set up yet, do that now. */ + if (worker_hash == NULL) + { + HASHCTL ctl; + + ctl.keysize = sizeof(pid_t); + ctl.entrysize = sizeof(pg_background_worker_info); + worker_hash = hash_create("pg_background worker_hash", 8, &ctl, + HASH_ELEM); + } + + /* Get current authentication information. */ + GetUserIdAndSecContext(¤t_user_id, &sec_context); + + /* + * In the unlikely event that there's an older worker with this PID, + * just detach it - unless it has a different user ID than the + * currently-active one, in which case someone might be trying to pull + * a fast one. Let's kill the backend to make sure we don't break + * anyone's expectations. + */ + if ((info = find_worker_info(pid)) != NULL) + { + if (current_user_id != info->current_user_id) + ereport(FATAL, + (errcode(ERRCODE_DUPLICATE_OBJECT), + errmsg("background worker with PID \"%d\" already exists", + pid))); + } + + /* Create a new entry for this worker. */ + info = hash_search(worker_hash, (void *) &pid, HASH_ENTER, NULL); + info->session = session; + info->result_count = 0; + info->current_user_id = current_user_id; +} + +/* + * Check whether the current user has rights to manipulate the background + * worker with the given PID. + */ +static void +check_rights(pg_background_worker_info *info) +{ + Oid current_user_id; + int sec_context; + + GetUserIdAndSecContext(¤t_user_id, &sec_context); + if (!has_privs_of_role(current_user_id, info->current_user_id)) + ereport(ERROR, + (errcode(ERRCODE_INSUFFICIENT_PRIVILEGE), + errmsg("permission denied for background worker with PID \"%d\"", + info->pid))); +} diff --git a/contrib/pg_background/pg_background.control b/contrib/pg_background/pg_background.control new file mode 100644 index 0000000..733d0e1 --- /dev/null +++ b/contrib/pg_background/pg_background.control @@ -0,0 +1,4 @@ +comment = 'Run SQL queries in the background' +default_version = '1.0' +module_pathname = '$libdir/pg_background' +relocatable = true diff --git a/src/backend/tcop/bgsession.c b/src/backend/tcop/bgsession.c index 2cc6438..143a244 100644 --- a/src/backend/tcop/bgsession.c +++ b/src/backend/tcop/bgsession.c @@ -115,9 +115,16 @@ static void invalid_protocol_message(char msgtype) pg_attribute_noreturn(); BackgroundSession * BackgroundSessionStart(void) { + pid_t *pid = 0; + + return BackgroundSessionNew(pid); +} + +BackgroundSession * +BackgroundSessionNew(pid_t *pid) +{ ResourceOwner oldowner; BackgroundWorker worker; - pid_t pid; BackgroundSession *session; shm_toc_estimator e; Size segsize; @@ -131,8 +138,11 @@ BackgroundSessionStart(void) BgwHandleStatus bgwstatus; StringInfoData msg; char msgtype; + MemoryContext oldcontext; + oldcontext = MemoryContextSwitchTo(TopMemoryContext); session = palloc(sizeof(*session)); + MemoryContextSwitchTo(oldcontext); session->resowner = ResourceOwnerCreate(NULL, "background session"); @@ -184,8 +194,10 @@ BackgroundSessionStart(void) shm_toc_insert(toc, BGSESSION_KEY_RESPONSE_QUEUE, response_mq); shm_mq_set_receiver(response_mq, MyProc); + oldcontext = MemoryContextSwitchTo(TopMemoryContext); session->command_qh = shm_mq_attach(command_mq, seg, NULL); session->response_qh = shm_mq_attach(response_mq, seg, NULL); + MemoryContextSwitchTo(oldcontext); worker.bgw_flags = BGWORKER_SHMEM_ACCESS | BGWORKER_BACKEND_DATABASE_CONNECTION; @@ -196,16 +208,18 @@ BackgroundSessionStart(void) worker.bgw_main_arg = UInt32GetDatum(dsm_segment_handle(seg)); worker.bgw_notify_pid = MyProcPid; + oldcontext = MemoryContextSwitchTo(TopMemoryContext); if (!RegisterDynamicBackgroundWorker(&worker, &session->worker_handle)) ereport(ERROR, (errcode(ERRCODE_INSUFFICIENT_RESOURCES), errmsg("could not register background process"), errhint("You might need to increase max_worker_processes."))); + MemoryContextSwitchTo(oldcontext); shm_mq_set_handle(session->command_qh, session->worker_handle); shm_mq_set_handle(session->response_qh, session->worker_handle); - bgwstatus = WaitForBackgroundWorkerStartup(session->worker_handle, &pid); + bgwstatus = WaitForBackgroundWorkerStartup(session->worker_handle, pid); if (bgwstatus != BGWH_STARTED) ereport(ERROR, (errcode(ERRCODE_INSUFFICIENT_RESOURCES), @@ -256,24 +270,41 @@ BackgroundSessionEnd(BackgroundSession *session) dsm_detach(session->seg); ResourceOwnerRelease(session->resowner, RESOURCE_RELEASE_BEFORE_LOCKS, false, false); ResourceOwnerDelete(session->resowner); - pfree(session); +//TODO: pfree(session); } - BackgroundSessionResult * BackgroundSessionExecute(BackgroundSession *session, const char *sql) { - StringInfoData msg; - char msgtype; BackgroundSessionResult *result; + BackgroundSessionExecuteSQL(session, sql); + result = BackgroundSessionFetchResult(session); + + return result; +} + +void +BackgroundSessionExecuteSQL(BackgroundSession *session, const char *sql) +{ + StringInfoData msg; + pq_redirect_to_shm_mq(session->seg, session->command_qh); pq_beginmessage(&msg, 'Q'); pq_sendstring(&msg, sql); pq_endmessage(&msg); pq_stop_redirect_to_shm_mq(); +} + +BackgroundSessionResult * +BackgroundSessionFetchResult(BackgroundSession *session) +{ + StringInfoData msg; + char msgtype; + BackgroundSessionResult *result; result = palloc0(sizeof(*result)); + result->command = NULL; do { @@ -322,7 +353,6 @@ BackgroundSessionExecute(BackgroundSession *session, const char *sql) return result; } - BackgroundSessionPreparedStatement * BackgroundSessionPrepare(BackgroundSession *session, const char *sql, int nargs, Oid argtypes[], const char *argnames[]) diff --git a/src/include/tcop/bgsession.h b/src/include/tcop/bgsession.h index 70dad45..f38cb29 100644 --- a/src/include/tcop/bgsession.h +++ b/src/include/tcop/bgsession.h @@ -18,9 +18,11 @@ typedef struct BackgroundSessionResult } BackgroundSessionResult; BackgroundSession *BackgroundSessionStart(void); +BackgroundSession *BackgroundSessionNew(pid_t *pid); void BackgroundSessionEnd(BackgroundSession *session); BackgroundSessionResult *BackgroundSessionExecute(BackgroundSession *session, const char *sql); +void BackgroundSessionExecuteSQL(BackgroundSession *session, const char *sql); +BackgroundSessionResult * BackgroundSessionFetchResult(BackgroundSession *session); BackgroundSessionPreparedStatement *BackgroundSessionPrepare(BackgroundSession *session, const char *sql, int nargs, Oid argtypes[], const char *argnames[]); BackgroundSessionResult *BackgroundSessionExecutePrepared(BackgroundSessionPreparedStatement *stmt, int nargs, Datum values[], bool nulls[]); - #endif /* BGSESSION_H */