/*
* Initialize a space for temporary files that can be opened for read-only
* access by other backends. Other backends must attach to it before
* accessing it. Associate this SharedFileSet with 'seg'. Any contained
* files will be deleted when the last backend detaches.
*
* Files will be distributed over the tablespaces configured in
* temp_tablespaces.
*
* Under the covers the set is one or more directories which will eventually
* be deleted when there are no backends attached.
*/
void
SharedFileSetInit(SharedFileSet *fileset, dsm_segment *seg)
{
static uint32 counter = 0;
SpinLockInit(&fileset->mutex);
fileset->refcnt = 1;
fileset->creator_pid = MyProcPid;
fileset->number = counter;
counter = (counter + 1) % INT_MAX;
/* Capture the tablespace OIDs so that all backends agree on them. */
PrepareTempTablespaces();
fileset->ntablespaces =
GetTempTablespaces(&fileset->tablespaces[0],
lengthof(fileset->tablespaces));
if (fileset->ntablespaces == 0)
{
fileset->tablespaces[0] = DEFAULTTABLESPACE_OID;
fileset->ntablespaces = 1;
}
/* Register our cleanup callback. */
on_dsm_detach(seg, SharedFileSetOnDetach, PointerGetDatum(fileset));
}
/*
* Arrange to redirect frontend/backend protocol messages to a shared-memory
* message queue.
*/
void
pq_redirect_to_shm_mq(dsm_segment *seg, shm_mq_handle *mqh)
{
PqCommMethods = &PqCommMqMethods;
pq_mq_handle = mqh;
whereToSendOutput = DestRemote;
FrontendProtocol = PG_PROTOCOL_LATEST;
on_dsm_detach(seg, pq_cleanup_redirect_to_shm_mq, (Datum) 0);
}
/*
* Attach to a shared message queue so we can send or receive messages.
*
* The memory context in effect at the time this function is called should
* be one which will last for at least as long as the message queue itself.
* We'll allocate the handle in that context, and future allocations that
* are needed to buffer incoming data will happen in that context as well.
*
* If seg != NULL, the queue will be automatically detached when that dynamic
* shared memory segment is detached.
*
* If handle != NULL, the queue can be read or written even before the
* other process has attached. We'll wait for it to do so if needed. The
* handle must be for a background worker initialized with bgw_notify_pid
* equal to our PID.
*
* shm_mq_detach() should be called when done. This will free the
* shm_mq_handle and mark the queue itself as detached, so that our
* counterpart won't get stuck waiting for us to fill or drain the queue
* after we've already lost interest.
*/
shm_mq_handle *
shm_mq_attach(shm_mq *mq, dsm_segment *seg, BackgroundWorkerHandle *handle)
{
shm_mq_handle *mqh = palloc(sizeof(shm_mq_handle));
Assert(mq->mq_receiver == MyProc || mq->mq_sender == MyProc);
mqh->mqh_queue = mq;
mqh->mqh_segment = seg;
mqh->mqh_buffer = NULL;
mqh->mqh_handle = handle;
mqh->mqh_buflen = 0;
mqh->mqh_consume_pending = 0;
mqh->mqh_context = CurrentMemoryContext;
mqh->mqh_partial_bytes = 0;
mqh->mqh_length_word_complete = false;
mqh->mqh_counterparty_attached = false;
if (seg != NULL)
on_dsm_detach(seg, shm_mq_detach_callback, PointerGetDatum(mq));
return mqh;
}
/*
* Attach to a set of directories that was created with SharedFileSetInit.
*/
void
SharedFileSetAttach(SharedFileSet *fileset, dsm_segment *seg)
{
bool success;
SpinLockAcquire(&fileset->mutex);
if (fileset->refcnt == 0)
success = false;
else
{
++fileset->refcnt;
success = true;
}
SpinLockRelease(&fileset->mutex);
if (!success)
ereport(ERROR,
(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
errmsg("could not attach to a SharedFileSet that is already destroyed")));
/* Register our cleanup callback. */
on_dsm_detach(seg, SharedFileSetOnDetach, PointerGetDatum(fileset));
}
/*
* Register background workers.
*/
static worker_state *
setup_background_workers(int nworkers, dsm_segment *seg)
{
MemoryContext oldcontext;
BackgroundWorker worker;
worker_state *wstate;
int i;
/*
* We need the worker_state object and the background worker handles to
* which it points to be allocated in CurTransactionContext rather than
* ExprContext; otherwise, they'll be destroyed before the on_dsm_detach
* hooks run.
*/
oldcontext = MemoryContextSwitchTo(CurTransactionContext);
/* Create worker state object. */
wstate = MemoryContextAlloc(TopTransactionContext,
offsetof(worker_state, handle) +
sizeof(BackgroundWorkerHandle *) * nworkers);
wstate->nworkers = 0;
/*
* Arrange to kill all the workers if we abort before all workers are
* finished hooking themselves up to the dynamic shared memory segment.
*
* If we die after all the workers have finished hooking themselves up to
* the dynamic shared memory segment, we'll mark the two queues to which
* we're directly connected as detached, and the worker(s) connected to
* those queues will exit, marking any other queues to which they are
* connected as detached. This will cause any as-yet-unaware workers
* connected to those queues to exit in their turn, and so on, until
* everybody exits.
*
* But suppose the workers which are supposed to connect to the queues to
* which we're directly attached exit due to some error before they
* actually attach the queues. The remaining workers will have no way of
* knowing this. From their perspective, they're still waiting for those
* workers to start, when in fact they've already died.
*/
on_dsm_detach(seg, cleanup_background_workers,
PointerGetDatum(wstate));
/* Configure a worker. */
worker.bgw_flags = BGWORKER_SHMEM_ACCESS;
worker.bgw_start_time = BgWorkerStart_ConsistentState;
worker.bgw_restart_time = BGW_NEVER_RESTART;
worker.bgw_main = NULL; /* new worker might not have library loaded */
sprintf(worker.bgw_library_name, "test_shm_mq");
sprintf(worker.bgw_function_name, "test_shm_mq_main");
snprintf(worker.bgw_name, BGW_MAXLEN, "test_shm_mq");
worker.bgw_main_arg = UInt32GetDatum(dsm_segment_handle(seg));
/* set bgw_notify_pid, so we can detect if the worker stops */
worker.bgw_notify_pid = MyProcPid;
/* Register the workers. */
for (i = 0; i < nworkers; ++i)
{
if (!RegisterDynamicBackgroundWorker(&worker, &wstate->handle[i]))
ereport(ERROR,
(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
errmsg("could not register background process"),
errhint("You may need to increase max_worker_processes.")));
++wstate->nworkers;
}
/* All done. */
MemoryContextSwitchTo(oldcontext);
return wstate;
}
Datum
worker_test(PG_FUNCTION_ARGS)
{
int i, nworkers;
dsm_segment *seg;
test_shm_mq_header *hdr;
MemoryContext oldcontext;
worker_state *wstate;
nworkers = PG_GETARG_INT32(0);
#if PG_VERSION_NUM >= 90500
seg = dsm_create(sizeof(test_shm_mq_header), 0);
#else
seg = dsm_create(sizeof(test_shm_mq_header)
#endif
hdr = dsm_segment_address(seg);
printf("begin worker_test: %d, %p\n", dsm_segment_handle(seg), hdr);
MemSet(hdr, 0, sizeof(test_shm_mq_header));
SpinLockInit(&hdr->mutex);
strncpy(hdr->dbname, get_database_name(MyDatabaseId), sizeof(hdr->dbname));
oldcontext = MemoryContextSwitchTo(CurTransactionContext);
wstate = MemoryContextAlloc(TopTransactionContext,
offsetof(worker_state, handle) +
sizeof(BackgroundWorkerHandle *) * nworkers);
MemSet(wstate, 0, offsetof(worker_state, handle) + sizeof(BackgroundWorkerHandle *) * nworkers);
on_dsm_detach(seg, cleanup_background_workers,
PointerGetDatum(wstate));
for (i = 0; i < nworkers; i++)
{
BackgroundWorker worker;
MemSet(&worker, 0, sizeof(worker));
worker.bgw_flags = BGWORKER_SHMEM_ACCESS | BGWORKER_BACKEND_DATABASE_CONNECTION;
worker.bgw_start_time = BgWorkerStart_ConsistentState;
worker.bgw_restart_time = BGW_NEVER_RESTART;
worker.bgw_main = NULL; /* new worker might not have library loaded */
sprintf(worker.bgw_library_name, "worker_test");
sprintf(worker.bgw_function_name, "worker_test_main");
snprintf(worker.bgw_name, BGW_MAXLEN, "worker_test %d", i);
worker.bgw_main_arg = UInt32GetDatum(dsm_segment_handle(seg));
worker.bgw_notify_pid = MyProcPid;
if (!RegisterDynamicBackgroundWorker(&worker, &wstate->handle[i]))
ereport(ERROR,
(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
errmsg("could not register background process"),
errhint("You may need to increase max_worker_processes.")));
++wstate->nworkers;
}
for (i = 0; i < nworkers; i++)
{
BgwHandleStatus status;
pid_t pid;
status = WaitForBackgroundWorkerStartup(wstate->handle[i], &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);
}
wait_for_workers_to_become_ready(wstate, hdr);
cancel_on_dsm_detach(seg, cleanup_background_workers,
PointerGetDatum(wstate));
dsm_detach(seg);
pfree(wstate);
MemoryContextSwitchTo(oldcontext);
PG_RETURN_VOID();
}
