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-2021, 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 "tcop/tcopprot.h"
28 
29 #include "test_shm_mq.h"
30 
31 static void attach_to_queues(dsm_segment *seg, shm_toc *toc,
32 							 int myworkernumber, shm_mq_handle **inqhp,
33 							 shm_mq_handle **outqhp);
34 static void copy_messages(shm_mq_handle *inqh, shm_mq_handle *outqh);
35 
36 /*
37  * Background worker entrypoint.
38  *
39  * This is intended to demonstrate how a background worker can be used to
40  * facilitate a parallel computation.  Most of the logic here is fairly
41  * boilerplate stuff, designed to attach to the shared memory segment,
42  * notify the user backend that we're alive, and so on.  The
43  * application-specific bits of logic that you'd replace for your own worker
44  * are attach_to_queues() and copy_messages().
45  */
46 void
test_shm_mq_main(Datum main_arg)47 test_shm_mq_main(Datum main_arg)
48 {
49 	dsm_segment *seg;
50 	shm_toc    *toc;
51 	shm_mq_handle *inqh;
52 	shm_mq_handle *outqh;
53 	volatile test_shm_mq_header *hdr;
54 	int			myworkernumber;
55 	PGPROC	   *registrant;
56 
57 	/*
58 	 * Establish signal handlers.
59 	 *
60 	 * We want CHECK_FOR_INTERRUPTS() to kill off this worker process just as
61 	 * it would a normal user backend.  To make that happen, we use die().
62 	 */
63 	pqsignal(SIGTERM, die);
64 	BackgroundWorkerUnblockSignals();
65 
66 	/*
67 	 * Connect to the dynamic shared memory segment.
68 	 *
69 	 * The backend that registered this worker passed us the ID of a shared
70 	 * memory segment to which we must attach for further instructions.  Once
71 	 * we've mapped the segment in our address space, attach to the table of
72 	 * contents so we can locate the various data structures we'll need to
73 	 * find within the segment.
74 	 *
75 	 * Note: at this point, we have not created any ResourceOwner in this
76 	 * process.  This will result in our DSM mapping surviving until process
77 	 * exit, which is fine.  If there were a ResourceOwner, it would acquire
78 	 * ownership of the mapping, but we have no need for that.
79 	 */
80 	seg = dsm_attach(DatumGetInt32(main_arg));
81 	if (seg == NULL)
82 		ereport(ERROR,
83 				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
84 				 errmsg("unable to map dynamic shared memory segment")));
85 	toc = shm_toc_attach(PG_TEST_SHM_MQ_MAGIC, dsm_segment_address(seg));
86 	if (toc == NULL)
87 		ereport(ERROR,
88 				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
89 				 errmsg("bad magic number in dynamic shared memory segment")));
90 
91 	/*
92 	 * Acquire a worker number.
93 	 *
94 	 * By convention, the process registering this background worker should
95 	 * have stored the control structure at key 0.  We look up that key to
96 	 * find it.  Our worker number gives our identity: there may be just one
97 	 * worker involved in this parallel operation, or there may be many.
98 	 */
99 	hdr = shm_toc_lookup(toc, 0, false);
100 	SpinLockAcquire(&hdr->mutex);
101 	myworkernumber = ++hdr->workers_attached;
102 	SpinLockRelease(&hdr->mutex);
103 	if (myworkernumber > hdr->workers_total)
104 		ereport(ERROR,
105 				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
106 				 errmsg("too many message queue testing workers already")));
107 
108 	/*
109 	 * Attach to the appropriate message queues.
110 	 */
111 	attach_to_queues(seg, toc, myworkernumber, &inqh, &outqh);
112 
113 	/*
114 	 * Indicate that we're fully initialized and ready to begin the main part
115 	 * of the parallel operation.
116 	 *
117 	 * Once we signal that we're ready, the user backend is entitled to assume
118 	 * that our on_dsm_detach callbacks will fire before we disconnect from
119 	 * the shared memory segment and exit.  Generally, that means we must have
120 	 * attached to all relevant dynamic shared memory data structures by now.
121 	 */
122 	SpinLockAcquire(&hdr->mutex);
123 	++hdr->workers_ready;
124 	SpinLockRelease(&hdr->mutex);
125 	registrant = BackendPidGetProc(MyBgworkerEntry->bgw_notify_pid);
126 	if (registrant == NULL)
127 	{
128 		elog(DEBUG1, "registrant backend has exited prematurely");
129 		proc_exit(1);
130 	}
131 	SetLatch(&registrant->procLatch);
132 
133 	/* Do the work. */
134 	copy_messages(inqh, outqh);
135 
136 	/*
137 	 * We're done.  For cleanliness, explicitly detach from the shared memory
138 	 * segment (that would happen anyway during process exit, though).
139 	 */
140 	dsm_detach(seg);
141 	proc_exit(1);
142 }
143 
144 /*
145  * Attach to shared memory message queues.
146  *
147  * We use our worker number to determine to which queue we should attach.
148  * The queues are registered at keys 1..<number-of-workers>.  The user backend
149  * writes to queue #1 and reads from queue #<number-of-workers>; each worker
150  * reads from the queue whose number is equal to its worker number and writes
151  * to the next higher-numbered queue.
152  */
153 static void
attach_to_queues(dsm_segment * seg,shm_toc * toc,int myworkernumber,shm_mq_handle ** inqhp,shm_mq_handle ** outqhp)154 attach_to_queues(dsm_segment *seg, shm_toc *toc, int myworkernumber,
155 				 shm_mq_handle **inqhp, shm_mq_handle **outqhp)
156 {
157 	shm_mq	   *inq;
158 	shm_mq	   *outq;
159 
160 	inq = shm_toc_lookup(toc, myworkernumber, false);
161 	shm_mq_set_receiver(inq, MyProc);
162 	*inqhp = shm_mq_attach(inq, seg, NULL);
163 	outq = shm_toc_lookup(toc, myworkernumber + 1, false);
164 	shm_mq_set_sender(outq, MyProc);
165 	*outqhp = shm_mq_attach(outq, seg, NULL);
166 }
167 
168 /*
169  * Loop, receiving and sending messages, until the connection is broken.
170  *
171  * This is the "real work" performed by this worker process.  Everything that
172  * happens before this is initialization of one form or another, and everything
173  * after this point is cleanup.
174  */
175 static void
copy_messages(shm_mq_handle * inqh,shm_mq_handle * outqh)176 copy_messages(shm_mq_handle *inqh, shm_mq_handle *outqh)
177 {
178 	Size		len;
179 	void	   *data;
180 	shm_mq_result res;
181 
182 	for (;;)
183 	{
184 		/* Notice any interrupts that have occurred. */
185 		CHECK_FOR_INTERRUPTS();
186 
187 		/* Receive a message. */
188 		res = shm_mq_receive(inqh, &len, &data, false);
189 		if (res != SHM_MQ_SUCCESS)
190 			break;
191 
192 		/* Send it back out. */
193 		res = shm_mq_send(outqh, len, data, false);
194 		if (res != SHM_MQ_SUCCESS)
195 			break;
196 	}
197 }
198