LCOV - code coverage report
Current view: top level - src/test/modules/test_shm_mq - worker.c (source / functions) Hit Total Coverage
Test: PostgreSQL 13devel Lines: 38 52 73.1 %
Date: 2019-11-15 22:06:47 Functions: 3 4 75.0 %
Legend: Lines: hit not hit

          Line data    Source code
       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-2019, 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             : 
      28             : #include "test_shm_mq.h"
      29             : 
      30             : static void handle_sigterm(SIGNAL_ARGS);
      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
      47          14 : 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 establish a
      62             :      * signal handler that is a stripped-down version of die().
      63             :      */
      64          14 :     pqsignal(SIGTERM, handle_sigterm);
      65          14 :     BackgroundWorkerUnblockSignals();
      66             : 
      67             :     /*
      68             :      * Connect to the dynamic shared memory segment.
      69             :      *
      70             :      * The backend that registered this worker passed us the ID of a shared
      71             :      * memory segment to which we must attach for further instructions.  Once
      72             :      * we've mapped the segment in our address space, attach to the table of
      73             :      * contents so we can locate the various data structures we'll need to
      74             :      * find within the segment.
      75             :      *
      76             :      * Note: at this point, we have not created any ResourceOwner in this
      77             :      * process.  This will result in our DSM mapping surviving until process
      78             :      * exit, which is fine.  If there were a ResourceOwner, it would acquire
      79             :      * ownership of the mapping, but we have no need for that.
      80             :      */
      81          14 :     seg = dsm_attach(DatumGetInt32(main_arg));
      82          14 :     if (seg == NULL)
      83           0 :         ereport(ERROR,
      84             :                 (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
      85             :                  errmsg("unable to map dynamic shared memory segment")));
      86          14 :     toc = shm_toc_attach(PG_TEST_SHM_MQ_MAGIC, dsm_segment_address(seg));
      87          14 :     if (toc == NULL)
      88           0 :         ereport(ERROR,
      89             :                 (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
      90             :                  errmsg("bad magic number in dynamic shared memory segment")));
      91             : 
      92             :     /*
      93             :      * Acquire a worker number.
      94             :      *
      95             :      * By convention, the process registering this background worker should
      96             :      * have stored the control structure at key 0.  We look up that key to
      97             :      * find it.  Our worker number gives our identity: there may be just one
      98             :      * worker involved in this parallel operation, or there may be many.
      99             :      */
     100          14 :     hdr = shm_toc_lookup(toc, 0, false);
     101          14 :     SpinLockAcquire(&hdr->mutex);
     102          14 :     myworkernumber = ++hdr->workers_attached;
     103          14 :     SpinLockRelease(&hdr->mutex);
     104          14 :     if (myworkernumber > hdr->workers_total)
     105           0 :         ereport(ERROR,
     106             :                 (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
     107             :                  errmsg("too many message queue testing workers already")));
     108             : 
     109             :     /*
     110             :      * Attach to the appropriate message queues.
     111             :      */
     112          14 :     attach_to_queues(seg, toc, myworkernumber, &inqh, &outqh);
     113             : 
     114             :     /*
     115             :      * Indicate that we're fully initialized and ready to begin the main part
     116             :      * of the parallel operation.
     117             :      *
     118             :      * Once we signal that we're ready, the user backend is entitled to assume
     119             :      * that our on_dsm_detach callbacks will fire before we disconnect from
     120             :      * the shared memory segment and exit.  Generally, that means we must have
     121             :      * attached to all relevant dynamic shared memory data structures by now.
     122             :      */
     123          14 :     SpinLockAcquire(&hdr->mutex);
     124          14 :     ++hdr->workers_ready;
     125          14 :     SpinLockRelease(&hdr->mutex);
     126          14 :     registrant = BackendPidGetProc(MyBgworkerEntry->bgw_notify_pid);
     127          14 :     if (registrant == NULL)
     128             :     {
     129           0 :         elog(DEBUG1, "registrant backend has exited prematurely");
     130           0 :         proc_exit(1);
     131             :     }
     132          14 :     SetLatch(&registrant->procLatch);
     133             : 
     134             :     /* Do the work. */
     135          14 :     copy_messages(inqh, outqh);
     136             : 
     137             :     /*
     138             :      * We're done.  For cleanliness, explicitly detach from the shared memory
     139             :      * segment (that would happen anyway during process exit, though).
     140             :      */
     141          14 :     dsm_detach(seg);
     142          14 :     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
     155          14 : 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          14 :     inq = shm_toc_lookup(toc, myworkernumber, false);
     162          14 :     shm_mq_set_receiver(inq, MyProc);
     163          14 :     *inqhp = shm_mq_attach(inq, seg, NULL);
     164          14 :     outq = shm_toc_lookup(toc, myworkernumber + 1, false);
     165          14 :     shm_mq_set_sender(outq, MyProc);
     166          14 :     *outqhp = shm_mq_attach(outq, seg, NULL);
     167          14 : }
     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
     177       49216 : 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       98418 :         CHECK_FOR_INTERRUPTS();
     187             : 
     188             :         /* Receive a message. */
     189       49216 :         res = shm_mq_receive(inqh, &len, &data, false);
     190       49216 :         if (res != SHM_MQ_SUCCESS)
     191          14 :             break;
     192             : 
     193             :         /* Send it back out. */
     194       49202 :         res = shm_mq_send(outqh, len, data, false);
     195       49202 :         if (res != SHM_MQ_SUCCESS)
     196           0 :             break;
     197             :     }
     198          14 : }
     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
     206           0 : handle_sigterm(SIGNAL_ARGS)
     207             : {
     208           0 :     int         save_errno = errno;
     209             : 
     210           0 :     SetLatch(MyLatch);
     211             : 
     212           0 :     if (!proc_exit_inprogress)
     213             :     {
     214           0 :         InterruptPending = true;
     215           0 :         ProcDiePending = true;
     216             :     }
     217             : 
     218           0 :     errno = save_errno;
     219           0 : }

Generated by: LCOV version 1.13