#include "postgres.h"
#include "miscadmin.h"
#include "pgstat.h"
#include "postmaster/bgworker.h"
#include "storage/shm_toc.h"
#include "test_shm_mq.h"
#include "utils/memutils.h"

Include dependency graph for setup.c:

Data Structures
struct	worker_state

Functions
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)

static worker_state *	setup_background_workers (int nworkers, dsm_segment *seg)

static void	cleanup_background_workers (dsm_segment *seg, Datum arg)

static void	wait_for_workers_to_become_ready (worker_state wstate, volatile test_shm_mq_header hdr)

static bool	check_worker_status (worker_state *wstate)

void	test_shm_mq_setup (int64 queue_size, int32 nworkers, dsm_segment segp, shm_mq_handle output, shm_mq_handle **input)

Variables
static uint32	we_bgworker_startup = 0

Function Documentation

◆ check_worker_status()

static bool check_worker_status ( worker_state * wstate )

static

Definition at line 305 of file setup.c.

 {
     int         n;
  
     /* If any workers (or the postmaster) have died, we have failed. */
     for (n = 0; n < wstate->nworkers; ++n)
     {
         BgwHandleStatus status;
         pid_t       pid;
  
         status = GetBackgroundWorkerPid(wstate->handle[n], &pid);
         if (status == BGWH_STOPPED || status == BGWH_POSTMASTER_DIED)
             return false;
     }
  
     /* Otherwise, things still look OK. */
     return true;
 }

References BGWH_POSTMASTER_DIED, BGWH_STOPPED, GetBackgroundWorkerPid(), worker_state::handle, and worker_state::nworkers.

Referenced by wait_for_workers_to_become_ready().

◆ cleanup_background_workers()

static void cleanup_background_workers	(	dsm_segment *	seg,
		Datum	arg
	)

static

Definition at line 245 of file setup.c.

 {
     worker_state *wstate = (worker_state *) DatumGetPointer(arg);
  
     while (wstate->nworkers > 0)
     {
         --wstate->nworkers;
         TerminateBackgroundWorker(wstate->handle[wstate->nworkers]);
     }
 }

References arg, DatumGetPointer(), worker_state::handle, worker_state::nworkers, and TerminateBackgroundWorker().

Referenced by setup_background_workers(), and test_shm_mq_setup().

◆ setup_background_workers()

static worker_state * setup_background_workers	(	int	nworkers,
		dsm_segment *	seg
	)

static

Definition at line 174 of file setup.c.

 {
     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. */
     memset(&worker, 0, sizeof(worker));
     worker.bgw_flags = BGWORKER_SHMEM_ACCESS;
     worker.bgw_start_time = BgWorkerStart_ConsistentState;
     worker.bgw_restart_time = BGW_NEVER_RESTART;
     sprintf(worker.bgw_library_name, "test_shm_mq");
     sprintf(worker.bgw_function_name, "test_shm_mq_main");
     snprintf(worker.bgw_type, 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;
 }

Referenced by test_shm_mq_setup().

◆ setup_dynamic_shared_memory()

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
	)

static

Definition at line 91 of file setup.c.

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

References dsm_create(), dsm_segment_address(), ereport, errcode(), errmsg(), ERROR, i, test_shm_mq_header::mutex, MyProc, PG_TEST_SHM_MQ_MAGIC, shm_mq_create(), shm_mq_minimum_size, shm_mq_set_receiver(), shm_mq_set_sender(), shm_toc_allocate(), shm_toc_create(), shm_toc_estimate(), shm_toc_estimate_chunk, shm_toc_estimate_keys, shm_toc_initialize_estimator, shm_toc_insert(), SpinLockInit, test_shm_mq_header::workers_attached, test_shm_mq_header::workers_ready, and test_shm_mq_header::workers_total.

Referenced by test_shm_mq_setup().

◆ test_shm_mq_setup()

void test_shm_mq_setup	(	int64	queue_size,
		int32	nworkers,
		dsm_segment **	segp,
		shm_mq_handle **	output,
		shm_mq_handle **	input
	)

Definition at line 50 of file setup.c.

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

References cancel_on_dsm_detach(), cleanup_background_workers(), worker_state::handle, input, output, pfree(), PointerGetDatum(), setup_background_workers(), setup_dynamic_shared_memory(), shm_mq_attach(), and wait_for_workers_to_become_ready().

Referenced by test_shm_mq(), and test_shm_mq_pipelined().

◆ wait_for_workers_to_become_ready()

static void wait_for_workers_to_become_ready	(	worker_state *	wstate,
		volatile test_shm_mq_header *	hdr
	)

static

Definition at line 257 of file setup.c.

 {
     bool        result = false;
  
     for (;;)
     {
         int         workers_ready;
  
         /* If all the workers are ready, we have succeeded. */
         SpinLockAcquire(&hdr->mutex);
         workers_ready = hdr->workers_ready;
         SpinLockRelease(&hdr->mutex);
         if (workers_ready >= wstate->nworkers)
         {
             result = true;
             break;
         }
  
         /* If any workers (or the postmaster) have died, we have failed. */
         if (!check_worker_status(wstate))
         {
             result = false;
             break;
         }
  
         /* first time, allocate or get the custom wait event */
         if (we_bgworker_startup == 0)
             we_bgworker_startup = WaitEventExtensionNew("TestShmMqBgWorkerStartup");
  
         /* Wait to be signaled. */
         (void) WaitLatch(MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0,
                          we_bgworker_startup);
  
         /* Reset the latch so we don't spin. */
         ResetLatch(MyLatch);
  
         /* An interrupt may have occurred while we were waiting. */
         CHECK_FOR_INTERRUPTS();
     }
  
     if (!result)
         ereport(ERROR,
                 (errcode(ERRCODE_INSUFFICIENT_RESOURCES),
                  errmsg("one or more background workers failed to start")));
 }