/*
* Main worker routine. Accepts dsm_handle as an argument
*/
static void
bg_worker_main(Datum main_arg)
{
PartitionArgs *args;
dsm_handle handle = DatumGetInt32(main_arg);
/* Create resource owner */
CurrentResourceOwner = ResourceOwnerCreate(NULL, "CreatePartitionsWorker");
/* Attach to dynamic shared memory */
if (!handle)
{
ereport(WARNING,
(errmsg("pg_pathman worker: invalid dsm_handle")));
}
segment = dsm_attach(handle);
args = dsm_segment_address(segment);
/* Establish connection and start transaction */
BackgroundWorkerInitializeConnectionByOid(args->dbid, InvalidOid);
StartTransactionCommand();
SPI_connect();
PushActiveSnapshot(GetTransactionSnapshot());
/* Create partitions */
args->result = create_partitions(args->relid, PATHMAN_GET_DATUM(args->value, args->by_val), args->value_type, &args->crashed);
/* Cleanup */
SPI_finish();
PopActiveSnapshot();
CommitTransactionCommand();
dsm_detach(segment);
}
/*
* Initialize allocated segment with block structure
*/
void
init_dsm_table(size_t block_size, size_t start, size_t end)
{
int i;
BlockHeaderPtr header;
char *ptr = dsm_segment_address(segment);
/* create blocks */
for (i=start; i<end; i++)
{
header = (BlockHeaderPtr) &ptr[i * block_size];
*header = set_free(header);
*header = set_length(header, 1);
}
return;
}
void
free_dsm_array(DsmArray *arr)
{
int start = arr->offset / dsm_cfg->block_size;
int i = 0;
char *ptr = dsm_segment_address(segment);
BlockHeaderPtr header = (BlockHeaderPtr) &ptr[start * dsm_cfg->block_size];
size_t blocks_count = get_length(header);
/* set blocks free */
for(; i < blocks_count; i++)
{
header = (BlockHeaderPtr) &ptr[(start + i) * dsm_cfg->block_size];
*header = set_free(header);
*header = set_length(header, 1);
}
if (start < dsm_cfg->first_free)
dsm_cfg->first_free = start;
arr->offset = 0;
arr->length = 0;
}
/* attach worker to the shared memory segment, read the job structure */
static void
initialize_worker(uint32 segment)
{
dsm_segment *seg;
ResourceOwner old,
tmp;
/* Connect to dynamic shared memory segment.
*
* In order to attach a dynamic shared memory segment, we need a
* resource owner. We cannot to StartTransactionCommand here, since
* we haven't yet attached to the database: to do this, we need to
* fetch information about connection properties from the shared
* memory segment.
*/
old = CurrentResourceOwner;
CurrentResourceOwner = ResourceOwnerCreate(NULL, "Background Worker");
seg = dsm_attach(segment);
if (seg == NULL)
ereport(ERROR,
(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
errmsg("unable to map dynamic shared memory segment")));
dsm_pin_mapping(seg);
tmp = CurrentResourceOwner;
CurrentResourceOwner = old;
ResourceOwnerDelete(tmp);
job = palloc(sizeof(JobDesc));
/* copy the arguments from shared memory segment */
memcpy(job, dsm_segment_address(seg), sizeof(JobDesc));
/* and detach it right away */
dsm_detach(seg);
Assert(job->magic == JOB_MAGIC);
job_run_function.schema = quote_identifier(job->schemaname);
job_run_function.name = quote_identifier("run_job");
}
void *
dsm_array_get_pointer(const DsmArray *arr)
{
return (char *) dsm_segment_address(segment) + arr->offset + sizeof(BlockHeader);
}
/*
* Allocate array inside dsm_segment
*/
void
alloc_dsm_array(DsmArray *arr, size_t entry_size, size_t length)
{
int i = 0;
int size_requested = entry_size * length;
int min_pos = 0;
int max_pos = 0;
bool found = false;
bool collecting_blocks = false;
size_t offset = -1;
size_t total_length = 0;
BlockHeaderPtr header;
char *ptr = dsm_segment_address(segment);
for (i = dsm_cfg->first_free; i<dsm_cfg->blocks_count; )
{
header = (BlockHeaderPtr) &ptr[i * dsm_cfg->block_size];
if (is_free(header))
{
if (!collecting_blocks)
{
offset = i * dsm_cfg->block_size;
total_length = dsm_cfg->block_size - sizeof(BlockHeader);
min_pos = i;
collecting_blocks = true;
}
else
{
total_length += dsm_cfg->block_size;
}
i++;
}
else
{
collecting_blocks = false;
offset = 0;
total_length = 0;
i += get_length(header);
}
if (total_length >= size_requested)
{
max_pos = i-1;
found = true;
break;
}
}
/*
* If dsm segment size is not enough then resize it (or allocate bigger
* for segment SysV and Windows, not implemented yet)
*/
if (!found)
{
size_t new_blocks_count = dsm_cfg->blocks_count * 2;
dsm_resize(segment, new_blocks_count * dsm_cfg->block_size);
init_dsm_table(dsm_cfg->block_size, dsm_cfg->blocks_count, new_blocks_count);
dsm_cfg->blocks_count = new_blocks_count;
/* try again */
return alloc_dsm_array(arr, entry_size, length);
}
/* look up for first free block */
if (dsm_cfg->first_free == min_pos)
{
for (; i<dsm_cfg->blocks_count; )
{
header = (BlockHeaderPtr) &ptr[i * dsm_cfg->block_size];
if (is_free(header))
{
dsm_cfg->first_free = i;
break;
}
else
{
i += get_length(header);
}
}
}
/* if we found enough of space */
if (total_length >= size_requested)
{
header = (BlockHeaderPtr) &ptr[min_pos * dsm_cfg->block_size];
*header = set_used(header);
*header = set_length(header, max_pos - min_pos + 1);
arr->offset = offset;
arr->length = length;
}
}
/*
* Set up a dynamic shared memory segment.
*
* We set up a small control region that contains only a test_shm_mq_header,
* plus one region per message queue. There are as many message queues as
* the number of workers, plus one.
*/
static void
setup_dynamic_shared_memory(int64 queue_size, int nworkers,
dsm_segment **segp, test_shm_mq_header **hdrp,
shm_mq **outp, shm_mq **inp)
{
shm_toc_estimator e;
int i;
Size segsize;
dsm_segment *seg;
shm_toc *toc;
test_shm_mq_header *hdr;
/* Ensure a valid queue size. */
if (queue_size < 0 || ((uint64) queue_size) < shm_mq_minimum_size)
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("queue size must be at least %zu bytes",
shm_mq_minimum_size)));
if (queue_size != ((Size) queue_size))
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("queue size overflows size_t")));
/*
* Estimate how much shared memory we need.
*
* Because the TOC machinery may choose to insert padding of oddly-sized
* requests, we must estimate each chunk separately.
*
* We need one key to register the location of the header, and we need
* nworkers + 1 keys to track the locations of the message queues.
*/
shm_toc_initialize_estimator(&e);
shm_toc_estimate_chunk(&e, sizeof(test_shm_mq_header));
for (i = 0; i <= nworkers; ++i)
shm_toc_estimate_chunk(&e, (Size) queue_size);
shm_toc_estimate_keys(&e, 2 + nworkers);
segsize = shm_toc_estimate(&e);
/* Create the shared memory segment and establish a table of contents. */
seg = dsm_create(shm_toc_estimate(&e), 0);
toc = shm_toc_create(PG_TEST_SHM_MQ_MAGIC, dsm_segment_address(seg),
segsize);
/* Set up the header region. */
hdr = shm_toc_allocate(toc, sizeof(test_shm_mq_header));
SpinLockInit(&hdr->mutex);
hdr->workers_total = nworkers;
hdr->workers_attached = 0;
hdr->workers_ready = 0;
shm_toc_insert(toc, 0, hdr);
/* Set up one message queue per worker, plus one. */
for (i = 0; i <= nworkers; ++i)
{
shm_mq *mq;
mq = shm_mq_create(shm_toc_allocate(toc, (Size) queue_size),
(Size) queue_size);
shm_toc_insert(toc, i + 1, mq);
if (i == 0)
{
/* We send messages to the first queue. */
shm_mq_set_sender(mq, MyProc);
*outp = mq;
}
if (i == nworkers)
{
/* We receive messages from the last queue. */
shm_mq_set_receiver(mq, MyProc);
*inp = mq;
}
}
/* Return results to caller. */
*segp = seg;
*hdrp = hdr;
}
/*
* Background worker entrypoint.
*
* This is intended to demonstrate how a background worker can be used to
* facilitate a parallel computation. Most of the logic here is fairly
* boilerplate stuff, designed to attach to the shared memory segment,
* notify the user backend that we're alive, and so on. The
* application-specific bits of logic that you'd replace for your own worker
* are attach_to_queues() and copy_messages().
*/
void
test_shm_mq_main(Datum main_arg)
{
dsm_segment *seg;
shm_toc *toc;
shm_mq_handle *inqh;
shm_mq_handle *outqh;
volatile test_shm_mq_header *hdr;
int myworkernumber;
PGPROC *registrant;
/*
* Establish signal handlers.
*
* We want CHECK_FOR_INTERRUPTS() to kill off this worker process just as
* it would a normal user backend. To make that happen, we establish a
* signal handler that is a stripped-down version of die(). We don't have
* any equivalent of the backend's command-read loop, where interrupts can
* be processed immediately, so make sure ImmediateInterruptOK is turned
* off.
*/
pqsignal(SIGTERM, handle_sigterm);
ImmediateInterruptOK = false;
BackgroundWorkerUnblockSignals();
/*
* Connect to the dynamic shared memory segment.
*
* The backend that registered this worker passed us the ID of a shared
* memory segment to which we must attach for further instructions. In
* order to attach to dynamic shared memory, we need a resource owner.
* Once we've mapped the segment in our address space, attach to the table
* of contents so we can locate the various data structures we'll need to
* find within the segment.
*/
CurrentResourceOwner = ResourceOwnerCreate(NULL, "test_shm_mq worker");
seg = dsm_attach(DatumGetInt32(main_arg));
if (seg == NULL)
ereport(ERROR,
(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
errmsg("unable to map dynamic shared memory segment")));
toc = shm_toc_attach(PG_TEST_SHM_MQ_MAGIC, dsm_segment_address(seg));
if (toc == NULL)
ereport(ERROR,
(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
errmsg("bad magic number in dynamic shared memory segment")));
/*
* Acquire a worker number.
*
* By convention, the process registering this background worker should
* have stored the control structure at key 0. We look up that key to
* find it. Our worker number gives our identity: there may be just one
* worker involved in this parallel operation, or there may be many.
*/
hdr = shm_toc_lookup(toc, 0);
SpinLockAcquire(&hdr->mutex);
myworkernumber = ++hdr->workers_attached;
SpinLockRelease(&hdr->mutex);
if (myworkernumber > hdr->workers_total)
ereport(ERROR,
(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
errmsg("too many message queue testing workers already")));
/*
* Attach to the appropriate message queues.
*/
attach_to_queues(seg, toc, myworkernumber, &inqh, &outqh);
/*
* Indicate that we're fully initialized and ready to begin the main part
* of the parallel operation.
*
* Once we signal that we're ready, the user backend is entitled to assume
* that our on_dsm_detach callbacks will fire before we disconnect from
* the shared memory segment and exit. Generally, that means we must have
* attached to all relevant dynamic shared memory data structures by now.
*/
SpinLockAcquire(&hdr->mutex);
++hdr->workers_ready;
SpinLockRelease(&hdr->mutex);
registrant = BackendPidGetProc(MyBgworkerEntry->bgw_notify_pid);
if (registrant == NULL)
{
elog(DEBUG1, "registrant backend has exited prematurely");
proc_exit(1);
}
SetLatch(®istrant->procLatch);
/* Do the work. */
copy_messages(inqh, outqh);
/*
* We're done. Explicitly detach the shared memory segment so that we
* don't get a resource leak warning at commit time. This will fire any
* on_dsm_detach callbacks we've registered, as well. Once that's done,
* we can go ahead and exit.
*/
dsm_detach(seg);
proc_exit(1);
}
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();
}
void
worker_test_main(Datum main_arg)
{
dsm_segment *seg;
volatile test_shm_mq_header *hdr;
PGPROC *registrant;
pqsignal(SIGHUP, handle_sighup);
pqsignal(SIGTERM, handle_sigterm);
BackgroundWorkerUnblockSignals();
printf("worker_test_main: %d\n", DatumGetInt32(main_arg));
CurrentResourceOwner = ResourceOwnerCreate(NULL, "worker test");
seg = dsm_attach(DatumGetInt32(main_arg));
if (seg == NULL)
ereport(ERROR,
(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
errmsg("unable to map dynamic shared memory segment")));
hdr = dsm_segment_address(seg);
/* 開始 */
SpinLockAcquire(&hdr->mutex);
hdr->workers_ready++;
hdr->workers_attached++;
SpinLockRelease(&hdr->mutex);
registrant = BackendPidGetProc(MyBgworkerEntry->bgw_notify_pid);
if (registrant == NULL)
{
elog(DEBUG1, "registrant backend has exited prematurely");
proc_exit(1);
}
SetLatch(®istrant->procLatch);
/* Do the work */
BackgroundWorkerInitializeConnection(hdr->dbname, NULL);
printf("DSM: %p\n", dsm_segment_address);
#if 0
SetCurrentStatementStartTimestamp();
StartTransactionCommand();
SPI_connect();
PushActiveSnapshot(GetTransactionSnapshot());
pgstat_report_activity(STATE_RUNNING, "initializing spi_worker schema");
SPI_finish();
PopActiveSnapshot();
CommitTransactionCommand();
pgstat_report_activity(STATE_IDLE, NULL);
#endif
dsm_detach(seg);
proc_exit(0);
}
/*
* Starts background worker that will create new partitions,
* waits till it finishes the job and returns the result (new partition oid)
*/
Oid
create_partitions_bg_worker(Oid relid, Datum value, Oid value_type, bool *crashed)
{
BackgroundWorker worker;
BackgroundWorkerHandle *worker_handle;
BgwHandleStatus status;
dsm_segment *segment;
dsm_handle segment_handle;
pid_t pid;
PartitionArgs *args;
Oid child_oid;
TypeCacheEntry *tce;
/* Create a dsm segment for the worker to pass arguments */
segment = dsm_create(sizeof(PartitionArgs), 0);
segment_handle = dsm_segment_handle(segment);
tce = lookup_type_cache(value_type, 0);
/* Fill arguments structure */
args = (PartitionArgs *) dsm_segment_address(segment);
args->dbid = MyDatabaseId;
args->relid = relid;
if (tce->typbyval)
args->value = value;
else
memcpy(&args->value, DatumGetPointer(value), sizeof(args->value));
args->by_val = tce->typbyval;
args->value_type = value_type;
args->result = 0;
/* Initialize worker struct */
worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
BGWORKER_BACKEND_DATABASE_CONNECTION;
worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
worker.bgw_restart_time = BGW_NEVER_RESTART;
worker.bgw_main = bg_worker_main;
worker.bgw_main_arg = Int32GetDatum(segment_handle);
worker.bgw_notify_pid = MyProcPid;
/* Start dynamic worker */
if (!RegisterDynamicBackgroundWorker(&worker, &worker_handle))
{
elog(WARNING, "Unable to create background worker for pg_pathman");
}
status = WaitForBackgroundWorkerStartup(worker_handle, &pid);
if (status == BGWH_POSTMASTER_DIED)
{
ereport(WARNING,
(errmsg("Postmaster died during the pg_pathman background worker process"),
errhint("More details may be available in the server log.")));
}
/* Wait till the worker finishes its job */
status = WaitForBackgroundWorkerShutdown(worker_handle);
*crashed = args->crashed;
child_oid = args->result;
/* Free dsm segment */
dsm_detach(segment);
return child_oid;
}