PostgreSQL Source Code  git master
setup.c
Go to the documentation of this file.
1 /*--------------------------------------------------------------------------
2  *
3  * setup.c
4  * Code to set up a dynamic shared memory segments and a specified
5  * number of background workers for shared memory message queue
6  * testing.
7  *
8  * Copyright (c) 2013-2024, PostgreSQL Global Development Group
9  *
10  * IDENTIFICATION
11  * src/test/modules/test_shm_mq/setup.c
12  *
13  * -------------------------------------------------------------------------
14  */
15 
16 #include "postgres.h"
17 
18 #include "miscadmin.h"
19 #include "pgstat.h"
20 #include "postmaster/bgworker.h"
21 #include "storage/procsignal.h"
22 #include "storage/shm_toc.h"
23 #include "test_shm_mq.h"
24 #include "utils/memutils.h"
25 
26 typedef struct
27 {
28  int nworkers;
30 } worker_state;
31 
32 static void setup_dynamic_shared_memory(int64 queue_size, int nworkers,
33  dsm_segment **segp,
34  test_shm_mq_header **hdrp,
35  shm_mq **outp, shm_mq **inp);
36 static worker_state *setup_background_workers(int nworkers,
37  dsm_segment *seg);
40  volatile test_shm_mq_header *hdr);
41 static bool check_worker_status(worker_state *wstate);
42 
43 /* value cached, fetched from shared memory */
45 
46 /*
47  * Set up a dynamic shared memory segment and zero or more background workers
48  * for a test run.
49  */
50 void
51 test_shm_mq_setup(int64 queue_size, int32 nworkers, dsm_segment **segp,
53 {
54  dsm_segment *seg;
55  test_shm_mq_header *hdr;
56  shm_mq *outq = NULL; /* placate compiler */
57  shm_mq *inq = NULL; /* placate compiler */
58  worker_state *wstate;
59 
60  /* Set up a dynamic shared memory segment. */
61  setup_dynamic_shared_memory(queue_size, nworkers, &seg, &hdr, &outq, &inq);
62  *segp = seg;
63 
64  /* Register background workers. */
65  wstate = setup_background_workers(nworkers, seg);
66 
67  /* Attach the queues. */
68  *output = shm_mq_attach(outq, seg, wstate->handle[0]);
69  *input = shm_mq_attach(inq, seg, wstate->handle[nworkers - 1]);
70 
71  /* Wait for workers to become ready. */
73 
74  /*
75  * Once we reach this point, all workers are ready. We no longer need to
76  * kill them if we die; they'll die on their own as the message queues
77  * shut down.
78  */
80  PointerGetDatum(wstate));
81  pfree(wstate);
82 }
83 
84 /*
85  * Set up a dynamic shared memory segment.
86  *
87  * We set up a small control region that contains only a test_shm_mq_header,
88  * plus one region per message queue. There are as many message queues as
89  * the number of workers, plus one.
90  */
91 static void
92 setup_dynamic_shared_memory(int64 queue_size, int nworkers,
93  dsm_segment **segp, test_shm_mq_header **hdrp,
94  shm_mq **outp, shm_mq **inp)
95 {
97  int i;
98  Size segsize;
99  dsm_segment *seg;
100  shm_toc *toc;
101  test_shm_mq_header *hdr;
102 
103  /* Ensure a valid queue size. */
104  if (queue_size < 0 || ((uint64) queue_size) < shm_mq_minimum_size)
105  ereport(ERROR,
106  (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
107  errmsg("queue size must be at least %zu bytes",
109  if (queue_size != ((Size) queue_size))
110  ereport(ERROR,
111  (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
112  errmsg("queue size overflows size_t")));
113 
114  /*
115  * Estimate how much shared memory we need.
116  *
117  * Because the TOC machinery may choose to insert padding of oddly-sized
118  * requests, we must estimate each chunk separately.
119  *
120  * We need one key to register the location of the header, and we need
121  * nworkers + 1 keys to track the locations of the message queues.
122  */
125  for (i = 0; i <= nworkers; ++i)
126  shm_toc_estimate_chunk(&e, (Size) queue_size);
127  shm_toc_estimate_keys(&e, 2 + nworkers);
128  segsize = shm_toc_estimate(&e);
129 
130  /* Create the shared memory segment and establish a table of contents. */
131  seg = dsm_create(shm_toc_estimate(&e), 0);
133  segsize);
134 
135  /* Set up the header region. */
136  hdr = shm_toc_allocate(toc, sizeof(test_shm_mq_header));
137  SpinLockInit(&hdr->mutex);
138  hdr->workers_total = nworkers;
139  hdr->workers_attached = 0;
140  hdr->workers_ready = 0;
141  shm_toc_insert(toc, 0, hdr);
142 
143  /* Set up one message queue per worker, plus one. */
144  for (i = 0; i <= nworkers; ++i)
145  {
146  shm_mq *mq;
147 
148  mq = shm_mq_create(shm_toc_allocate(toc, (Size) queue_size),
149  (Size) queue_size);
150  shm_toc_insert(toc, i + 1, mq);
151 
152  if (i == 0)
153  {
154  /* We send messages to the first queue. */
156  *outp = mq;
157  }
158  if (i == nworkers)
159  {
160  /* We receive messages from the last queue. */
162  *inp = mq;
163  }
164  }
165 
166  /* Return results to caller. */
167  *segp = seg;
168  *hdrp = hdr;
169 }
170 
171 /*
172  * Register background workers.
173  */
174 static worker_state *
176 {
177  MemoryContext oldcontext;
178  BackgroundWorker worker;
179  worker_state *wstate;
180  int i;
181 
182  /*
183  * We need the worker_state object and the background worker handles to
184  * which it points to be allocated in CurTransactionContext rather than
185  * ExprContext; otherwise, they'll be destroyed before the on_dsm_detach
186  * hooks run.
187  */
189 
190  /* Create worker state object. */
192  offsetof(worker_state, handle) +
193  sizeof(BackgroundWorkerHandle *) * nworkers);
194  wstate->nworkers = 0;
195 
196  /*
197  * Arrange to kill all the workers if we abort before all workers are
198  * finished hooking themselves up to the dynamic shared memory segment.
199  *
200  * If we die after all the workers have finished hooking themselves up to
201  * the dynamic shared memory segment, we'll mark the two queues to which
202  * we're directly connected as detached, and the worker(s) connected to
203  * those queues will exit, marking any other queues to which they are
204  * connected as detached. This will cause any as-yet-unaware workers
205  * connected to those queues to exit in their turn, and so on, until
206  * everybody exits.
207  *
208  * But suppose the workers which are supposed to connect to the queues to
209  * which we're directly attached exit due to some error before they
210  * actually attach the queues. The remaining workers will have no way of
211  * knowing this. From their perspective, they're still waiting for those
212  * workers to start, when in fact they've already died.
213  */
215  PointerGetDatum(wstate));
216 
217  /* Configure a worker. */
218  memset(&worker, 0, sizeof(worker));
222  sprintf(worker.bgw_library_name, "test_shm_mq");
223  sprintf(worker.bgw_function_name, "test_shm_mq_main");
224  snprintf(worker.bgw_type, BGW_MAXLEN, "test_shm_mq");
226  /* set bgw_notify_pid, so we can detect if the worker stops */
227  worker.bgw_notify_pid = MyProcPid;
228 
229  /* Register the workers. */
230  for (i = 0; i < nworkers; ++i)
231  {
232  if (!RegisterDynamicBackgroundWorker(&worker, &wstate->handle[i]))
233  ereport(ERROR,
234  (errcode(ERRCODE_INSUFFICIENT_RESOURCES),
235  errmsg("could not register background process"),
236  errhint("You may need to increase max_worker_processes.")));
237  ++wstate->nworkers;
238  }
239 
240  /* All done. */
241  MemoryContextSwitchTo(oldcontext);
242  return wstate;
243 }
244 
245 static void
247 {
249 
250  while (wstate->nworkers > 0)
251  {
252  --wstate->nworkers;
253  TerminateBackgroundWorker(wstate->handle[wstate->nworkers]);
254  }
255 }
256 
257 static void
259  volatile test_shm_mq_header *hdr)
260 {
261  bool result = false;
262 
263  for (;;)
264  {
265  int workers_ready;
266 
267  /* If all the workers are ready, we have succeeded. */
268  SpinLockAcquire(&hdr->mutex);
269  workers_ready = hdr->workers_ready;
270  SpinLockRelease(&hdr->mutex);
271  if (workers_ready >= wstate->nworkers)
272  {
273  result = true;
274  break;
275  }
276 
277  /* If any workers (or the postmaster) have died, we have failed. */
278  if (!check_worker_status(wstate))
279  {
280  result = false;
281  break;
282  }
283 
284  /* first time, allocate or get the custom wait event */
285  if (we_bgworker_startup == 0)
286  we_bgworker_startup = WaitEventExtensionNew("TestShmMqBgWorkerStartup");
287 
288  /* Wait to be signaled. */
291 
292  /* Reset the latch so we don't spin. */
294 
295  /* An interrupt may have occurred while we were waiting. */
297  }
298 
299  if (!result)
300  ereport(ERROR,
301  (errcode(ERRCODE_INSUFFICIENT_RESOURCES),
302  errmsg("one or more background workers failed to start")));
303 }
304 
305 static bool
307 {
308  int n;
309 
310  /* If any workers (or the postmaster) have died, we have failed. */
311  for (n = 0; n < wstate->nworkers; ++n)
312  {
313  BgwHandleStatus status;
314  pid_t pid;
315 
316  status = GetBackgroundWorkerPid(wstate->handle[n], &pid);
317  if (status == BGWH_STOPPED || status == BGWH_POSTMASTER_DIED)
318  return false;
319  }
320 
321  /* Otherwise, things still look OK. */
322  return true;
323 }
void TerminateBackgroundWorker(BackgroundWorkerHandle *handle)
Definition: bgworker.c:1221
BgwHandleStatus GetBackgroundWorkerPid(BackgroundWorkerHandle *handle, pid_t *pidp)
Definition: bgworker.c:1082
bool RegisterDynamicBackgroundWorker(BackgroundWorker *worker, BackgroundWorkerHandle **handle)
Definition: bgworker.c:970
#define BGW_NEVER_RESTART
Definition: bgworker.h:85
BgwHandleStatus
Definition: bgworker.h:104
@ BGWH_POSTMASTER_DIED
Definition: bgworker.h:108
@ BGWH_STOPPED
Definition: bgworker.h:107
@ BgWorkerStart_ConsistentState
Definition: bgworker.h:80
#define BGWORKER_SHMEM_ACCESS
Definition: bgworker.h:53
#define BGW_MAXLEN
Definition: bgworker.h:86
unsigned int uint32
Definition: c.h:493
signed int int32
Definition: c.h:481
#define FLEXIBLE_ARRAY_MEMBER
Definition: c.h:385
size_t Size
Definition: c.h:592
dsm_handle dsm_segment_handle(dsm_segment *seg)
Definition: dsm.c:1123
void * dsm_segment_address(dsm_segment *seg)
Definition: dsm.c:1095
void on_dsm_detach(dsm_segment *seg, on_dsm_detach_callback function, Datum arg)
Definition: dsm.c:1132
dsm_segment * dsm_create(Size size, int flags)
Definition: dsm.c:516
void cancel_on_dsm_detach(dsm_segment *seg, on_dsm_detach_callback function, Datum arg)
Definition: dsm.c:1147
int errhint(const char *fmt,...)
Definition: elog.c:1319
int errcode(int sqlerrcode)
Definition: elog.c:859
int errmsg(const char *fmt,...)
Definition: elog.c:1072
#define ERROR
Definition: elog.h:39
#define ereport(elevel,...)
Definition: elog.h:149
int MyProcPid
Definition: globals.c:45
struct Latch * MyLatch
Definition: globals.c:60
FILE * input
FILE * output
int i
Definition: isn.c:73
void ResetLatch(Latch *latch)
Definition: latch.c:724
int WaitLatch(Latch *latch, int wakeEvents, long timeout, uint32 wait_event_info)
Definition: latch.c:517
#define WL_EXIT_ON_PM_DEATH
Definition: latch.h:132
#define WL_LATCH_SET
Definition: latch.h:127
MemoryContext TopTransactionContext
Definition: mcxt.c:142
void pfree(void *pointer)
Definition: mcxt.c:1508
MemoryContext CurTransactionContext
Definition: mcxt.c:143
void * MemoryContextAlloc(MemoryContext context, Size size)
Definition: mcxt.c:1168
#define CHECK_FOR_INTERRUPTS()
Definition: miscadmin.h:122
static MemoryContext MemoryContextSwitchTo(MemoryContext context)
Definition: palloc.h:124
void * arg
#define sprintf
Definition: port.h:240
#define snprintf
Definition: port.h:238
static Datum PointerGetDatum(const void *X)
Definition: postgres.h:322
uintptr_t Datum
Definition: postgres.h:64
static Pointer DatumGetPointer(Datum X)
Definition: postgres.h:312
static Datum UInt32GetDatum(uint32 X)
Definition: postgres.h:232
e
Definition: preproc-init.c:82
static worker_state * setup_background_workers(int nworkers, dsm_segment *seg)
Definition: setup.c:175
static void cleanup_background_workers(dsm_segment *seg, Datum arg)
Definition: setup.c:246
static void wait_for_workers_to_become_ready(worker_state *wstate, volatile test_shm_mq_header *hdr)
Definition: setup.c:258
void test_shm_mq_setup(int64 queue_size, int32 nworkers, dsm_segment **segp, shm_mq_handle **output, shm_mq_handle **input)
Definition: setup.c:51
static uint32 we_bgworker_startup
Definition: setup.c:44
static bool check_worker_status(worker_state *wstate)
Definition: setup.c:306
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)
Definition: setup.c:92
shm_mq_handle * shm_mq_attach(shm_mq *mq, dsm_segment *seg, BackgroundWorkerHandle *handle)
Definition: shm_mq.c:290
void shm_mq_set_sender(shm_mq *mq, PGPROC *proc)
Definition: shm_mq.c:224
shm_mq * shm_mq_create(void *address, Size size)
Definition: shm_mq.c:177
void shm_mq_set_receiver(shm_mq *mq, PGPROC *proc)
Definition: shm_mq.c:206
const Size shm_mq_minimum_size
Definition: shm_mq.c:168
shm_toc * shm_toc_create(uint64 magic, void *address, Size nbytes)
Definition: shm_toc.c:40
Size shm_toc_estimate(shm_toc_estimator *e)
Definition: shm_toc.c:263
void shm_toc_insert(shm_toc *toc, uint64 key, void *address)
Definition: shm_toc.c:171
void * shm_toc_allocate(shm_toc *toc, Size nbytes)
Definition: shm_toc.c:88
#define shm_toc_estimate_chunk(e, sz)
Definition: shm_toc.h:51
#define shm_toc_initialize_estimator(e)
Definition: shm_toc.h:49
#define shm_toc_estimate_keys(e, cnt)
Definition: shm_toc.h:53
#define SpinLockInit(lock)
Definition: spin.h:60
#define SpinLockRelease(lock)
Definition: spin.h:64
#define SpinLockAcquire(lock)
Definition: spin.h:62
PGPROC * MyProc
Definition: proc.c:66
char bgw_function_name[BGW_MAXLEN]
Definition: bgworker.h:97
Datum bgw_main_arg
Definition: bgworker.h:98
int bgw_restart_time
Definition: bgworker.h:95
char bgw_type[BGW_MAXLEN]
Definition: bgworker.h:92
BgWorkerStartTime bgw_start_time
Definition: bgworker.h:94
pid_t bgw_notify_pid
Definition: bgworker.h:100
char bgw_library_name[MAXPGPATH]
Definition: bgworker.h:96
Definition: shm_mq.c:72
int nworkers
Definition: setup.c:28
BackgroundWorkerHandle * handle[FLEXIBLE_ARRAY_MEMBER]
Definition: setup.c:29
#define PG_TEST_SHM_MQ_MAGIC
Definition: test_shm_mq.h:22
uint32 WaitEventExtensionNew(const char *wait_event_name)
Definition: wait_event.c:162