1 /*--------------------------------------------------------------------------
2 *
3 * worker.c
4 * Code for sample worker making use of shared memory message queues.
5 * Our test worker simply reads messages from one message queue and
6 * writes them back out to another message queue. In a real
7 * application, you'd presumably want the worker to do some more
8 * complex calculation rather than simply returning the input,
9 * but it should be possible to use much of the control logic just
10 * as presented here.
11 *
12 * Copyright (c) 2013-2016, PostgreSQL Global Development Group
13 *
14 * IDENTIFICATION
15 * src/test/modules/test_shm_mq/worker.c
16 *
17 * -------------------------------------------------------------------------
18 */
19
20 #include "postgres.h"
21
22 #include "miscadmin.h"
23 #include "storage/ipc.h"
24 #include "storage/procarray.h"
25 #include "storage/shm_mq.h"
26 #include "storage/shm_toc.h"
27 #include "utils/resowner.h"
28
29 #include "test_shm_mq.h"
30
31 static void handle_sigterm(SIGNAL_ARGS);
32 static void attach_to_queues(dsm_segment *seg, shm_toc *toc,
33 int myworkernumber, shm_mq_handle **inqhp,
34 shm_mq_handle **outqhp);
35 static void copy_messages(shm_mq_handle *inqh, shm_mq_handle *outqh);
36
37 /*
38 * Background worker entrypoint.
39 *
40 * This is intended to demonstrate how a background worker can be used to
41 * facilitate a parallel computation. Most of the logic here is fairly
42 * boilerplate stuff, designed to attach to the shared memory segment,
43 * notify the user backend that we're alive, and so on. The
44 * application-specific bits of logic that you'd replace for your own worker
45 * are attach_to_queues() and copy_messages().
46 */
47 void
test_shm_mq_main(Datum main_arg)48 test_shm_mq_main(Datum main_arg)
49 {
50 dsm_segment *seg;
51 shm_toc *toc;
52 shm_mq_handle *inqh;
53 shm_mq_handle *outqh;
54 volatile test_shm_mq_header *hdr;
55 int myworkernumber;
56 PGPROC *registrant;
57
58 /*
59 * Establish signal handlers.
60 *
61 * We want CHECK_FOR_INTERRUPTS() to kill off this worker process just as
62 * it would a normal user backend. To make that happen, we establish a
63 * signal handler that is a stripped-down version of die().
64 */
65 pqsignal(SIGTERM, handle_sigterm);
66 BackgroundWorkerUnblockSignals();
67
68 /*
69 * Connect to the dynamic shared memory segment.
70 *
71 * The backend that registered this worker passed us the ID of a shared
72 * memory segment to which we must attach for further instructions. In
73 * order to attach to dynamic shared memory, we need a resource owner.
74 * Once we've mapped the segment in our address space, attach to the table
75 * of contents so we can locate the various data structures we'll need to
76 * find within the segment.
77 */
78 CurrentResourceOwner = ResourceOwnerCreate(NULL, "test_shm_mq worker");
79 seg = dsm_attach(DatumGetInt32(main_arg));
80 if (seg == NULL)
81 ereport(ERROR,
82 (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
83 errmsg("unable to map dynamic shared memory segment")));
84 toc = shm_toc_attach(PG_TEST_SHM_MQ_MAGIC, dsm_segment_address(seg));
85 if (toc == NULL)
86 ereport(ERROR,
87 (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
88 errmsg("bad magic number in dynamic shared memory segment")));
89
90 /*
91 * Acquire a worker number.
92 *
93 * By convention, the process registering this background worker should
94 * have stored the control structure at key 0. We look up that key to
95 * find it. Our worker number gives our identity: there may be just one
96 * worker involved in this parallel operation, or there may be many.
97 */
98 hdr = shm_toc_lookup(toc, 0);
99 SpinLockAcquire(&hdr->mutex);
100 myworkernumber = ++hdr->workers_attached;
101 SpinLockRelease(&hdr->mutex);
102 if (myworkernumber > hdr->workers_total)
103 ereport(ERROR,
104 (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
105 errmsg("too many message queue testing workers already")));
106
107 /*
108 * Attach to the appropriate message queues.
109 */
110 attach_to_queues(seg, toc, myworkernumber, &inqh, &outqh);
111
112 /*
113 * Indicate that we're fully initialized and ready to begin the main part
114 * of the parallel operation.
115 *
116 * Once we signal that we're ready, the user backend is entitled to assume
117 * that our on_dsm_detach callbacks will fire before we disconnect from
118 * the shared memory segment and exit. Generally, that means we must have
119 * attached to all relevant dynamic shared memory data structures by now.
120 */
121 SpinLockAcquire(&hdr->mutex);
122 ++hdr->workers_ready;
123 SpinLockRelease(&hdr->mutex);
124 registrant = BackendPidGetProc(MyBgworkerEntry->bgw_notify_pid);
125 if (registrant == NULL)
126 {
127 elog(DEBUG1, "registrant backend has exited prematurely");
128 proc_exit(1);
129 }
130 SetLatch(®istrant->procLatch);
131
132 /* Do the work. */
133 copy_messages(inqh, outqh);
134
135 /*
136 * We're done. Explicitly detach the shared memory segment so that we
137 * don't get a resource leak warning at commit time. This will fire any
138 * on_dsm_detach callbacks we've registered, as well. Once that's done,
139 * we can go ahead and exit.
140 */
141 dsm_detach(seg);
142 proc_exit(1);
143 }
144
145 /*
146 * Attach to shared memory message queues.
147 *
148 * We use our worker number to determine to which queue we should attach.
149 * The queues are registered at keys 1..<number-of-workers>. The user backend
150 * writes to queue #1 and reads from queue #<number-of-workers>; each worker
151 * reads from the queue whose number is equal to its worker number and writes
152 * to the next higher-numbered queue.
153 */
154 static void
attach_to_queues(dsm_segment * seg,shm_toc * toc,int myworkernumber,shm_mq_handle ** inqhp,shm_mq_handle ** outqhp)155 attach_to_queues(dsm_segment *seg, shm_toc *toc, int myworkernumber,
156 shm_mq_handle **inqhp, shm_mq_handle **outqhp)
157 {
158 shm_mq *inq;
159 shm_mq *outq;
160
161 inq = shm_toc_lookup(toc, myworkernumber);
162 shm_mq_set_receiver(inq, MyProc);
163 *inqhp = shm_mq_attach(inq, seg, NULL);
164 outq = shm_toc_lookup(toc, myworkernumber + 1);
165 shm_mq_set_sender(outq, MyProc);
166 *outqhp = shm_mq_attach(outq, seg, NULL);
167 }
168
169 /*
170 * Loop, receiving and sending messages, until the connection is broken.
171 *
172 * This is the "real work" performed by this worker process. Everything that
173 * happens before this is initialization of one form or another, and everything
174 * after this point is cleanup.
175 */
176 static void
copy_messages(shm_mq_handle * inqh,shm_mq_handle * outqh)177 copy_messages(shm_mq_handle *inqh, shm_mq_handle *outqh)
178 {
179 Size len;
180 void *data;
181 shm_mq_result res;
182
183 for (;;)
184 {
185 /* Notice any interrupts that have occurred. */
186 CHECK_FOR_INTERRUPTS();
187
188 /* Receive a message. */
189 res = shm_mq_receive(inqh, &len, &data, false);
190 if (res != SHM_MQ_SUCCESS)
191 break;
192
193 /* Send it back out. */
194 res = shm_mq_send(outqh, len, data, false);
195 if (res != SHM_MQ_SUCCESS)
196 break;
197 }
198 }
199
200 /*
201 * When we receive a SIGTERM, we set InterruptPending and ProcDiePending just
202 * like a normal backend. The next CHECK_FOR_INTERRUPTS() will do the right
203 * thing.
204 */
205 static void
handle_sigterm(SIGNAL_ARGS)206 handle_sigterm(SIGNAL_ARGS)
207 {
208 int save_errno = errno;
209
210 SetLatch(MyLatch);
211
212 if (!proc_exit_inprogress)
213 {
214 InterruptPending = true;
215 ProcDiePending = true;
216 }
217
218 errno = save_errno;
219 }
220