/* * Set up a dynamic shared memory segment and zero or more background workers * for a test run. */ void test_shm_mq_setup(int64 queue_size, int32 nworkers, dsm_segment **segp, shm_mq_handle **output, shm_mq_handle **input) { dsm_segment *seg; test_shm_mq_header *hdr; shm_mq *outq = NULL; /* placate compiler */ shm_mq *inq = NULL; /* placate compiler */ worker_state *wstate; /* Set up a dynamic shared memory segment. */ setup_dynamic_shared_memory(queue_size, nworkers, &seg, &hdr, &outq, &inq); *segp = seg; /* Register background workers. */ wstate = setup_background_workers(nworkers, seg); /* Attach the queues. */ *output = shm_mq_attach(outq, seg, wstate->handle[0]); *input = shm_mq_attach(inq, seg, wstate->handle[nworkers - 1]); /* Wait for workers to become ready. */ wait_for_workers_to_become_ready(wstate, hdr); /* * Once we reach this point, all workers are ready. We no longer need to * kill them if we die; they'll die on their own as the message queues * shut down. */ cancel_on_dsm_detach(seg, cleanup_background_workers, PointerGetDatum(wstate)); pfree(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(); }