/*
* Fetch a tuple from a tuple queue reader.
*
* The return value is NULL if there are no remaining tuples or if
* nowait = true and no tuple is ready to return. *done, if not NULL,
* is set to true when there are no remaining tuples and otherwise to false.
*
* The returned tuple, if any, is allocated in CurrentMemoryContext.
* Note that this routine must not leak memory! (We used to allow that,
* but not any more.)
*
* Even when shm_mq_receive() returns SHM_MQ_WOULD_BLOCK, this can still
* accumulate bytes from a partially-read message, so it's useful to call
* this with nowait = true even if nothing is returned.
*/
HeapTuple
TupleQueueReaderNext(TupleQueueReader *reader, bool nowait, bool *done)
{
HeapTupleData htup;
shm_mq_result result;
Size nbytes;
void *data;
if (done != NULL)
*done = false;
/* Attempt to read a message. */
result = shm_mq_receive(reader->queue, &nbytes, &data, nowait);
/* If queue is detached, set *done and return NULL. */
if (result == SHM_MQ_DETACHED)
{
if (done != NULL)
*done = true;
return NULL;
}
/* In non-blocking mode, bail out if no message ready yet. */
if (result == SHM_MQ_WOULD_BLOCK)
return NULL;
Assert(result == SHM_MQ_SUCCESS);
/*
* Set up a dummy HeapTupleData pointing to the data from the shm_mq
* (which had better be sufficiently aligned).
*/
ItemPointerSetInvalid(&htup.t_self);
htup.t_tableOid = InvalidOid;
htup.t_len = nbytes;
htup.t_data = data;
return heap_copytuple(&htup);
}
/*
* Loop, receiving and sending messages, until the connection is broken.
*
* This is the "real work" performed by this worker process. Everything that
* happens before this is initialization of one form or another, and everything
* after this point is cleanup.
*/
static void
copy_messages(shm_mq_handle *inqh, shm_mq_handle *outqh)
{
Size len;
void *data;
shm_mq_result res;
for (;;)
{
/* Notice any interrupts that have occurred. */
CHECK_FOR_INTERRUPTS();
/* Receive a message. */
res = shm_mq_receive(inqh, &len, &data, false);
if (res != SHM_MQ_SUCCESS)
break;
/* Send it back out. */
res = shm_mq_send(outqh, len, data, false);
if (res != SHM_MQ_SUCCESS)
break;
}
}
/*
* Simple test of the shared memory message queue infrastructure.
*
* We set up a ring of message queues passing through 1 or more background
* processes and eventually looping back to ourselves. We then send a message
* through the ring a number of times indicated by the loop count. At the end,
* we check whether the final message matches the one we started with.
*/
Datum
test_shm_mq(PG_FUNCTION_ARGS)
{
int64 queue_size = PG_GETARG_INT64(0);
text *message = PG_GETARG_TEXT_PP(1);
char *message_contents = VARDATA_ANY(message);
int message_size = VARSIZE_ANY_EXHDR(message);
int32 loop_count = PG_GETARG_INT32(2);
int32 nworkers = PG_GETARG_INT32(3);
dsm_segment *seg;
shm_mq_handle *outqh;
shm_mq_handle *inqh;
shm_mq_result res;
Size len;
void *data;
/* A negative loopcount is nonsensical. */
if (loop_count < 0)
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("repeat count size must be a non-negative integer")));
/*
* Since this test sends data using the blocking interfaces, it cannot
* send data to itself. Therefore, a minimum of 1 worker is required. Of
* course, a negative worker count is nonsensical.
*/
if (nworkers < 1)
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("number of workers must be a positive integer")));
/* Set up dynamic shared memory segment and background workers. */
test_shm_mq_setup(queue_size, nworkers, &seg, &outqh, &inqh);
/* Send the initial message. */
res = shm_mq_send(outqh, message_size, message_contents, false);
if (res != SHM_MQ_SUCCESS)
ereport(ERROR,
(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
errmsg("could not send message")));
/*
* Receive a message and send it back out again. Do this a number of
* times equal to the loop count.
*/
for (;;)
{
/* Receive a message. */
res = shm_mq_receive(inqh, &len, &data, false);
if (res != SHM_MQ_SUCCESS)
ereport(ERROR,
(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
errmsg("could not receive message")));
/* If this is supposed to be the last iteration, stop here. */
if (--loop_count <= 0)
break;
/* Send it back out. */
res = shm_mq_send(outqh, len, data, false);
if (res != SHM_MQ_SUCCESS)
ereport(ERROR,
(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
errmsg("could not send message")));
}
/*
* Finally, check that we got back the same message from the last
* iteration that we originally sent.
*/
verify_message(message_size, message_contents, len, data);
/* Clean up. */
dsm_detach(seg);
PG_RETURN_VOID();
}
/*
* Pipelined test of the shared memory message queue infrastructure.
*
* As in the basic test, we set up a ring of message queues passing through
* 1 or more background processes and eventually looping back to ourselves.
* Then, we send N copies of the user-specified message through the ring and
* receive them all back. Since this might fill up all message queues in the
* ring and then stall, we must be prepared to begin receiving the messages
* back before we've finished sending them.
*/
Datum
test_shm_mq_pipelined(PG_FUNCTION_ARGS)
{
int64 queue_size = PG_GETARG_INT64(0);
text *message = PG_GETARG_TEXT_PP(1);
char *message_contents = VARDATA_ANY(message);
int message_size = VARSIZE_ANY_EXHDR(message);
int32 loop_count = PG_GETARG_INT32(2);
int32 nworkers = PG_GETARG_INT32(3);
bool verify = PG_GETARG_BOOL(4);
int32 send_count = 0;
int32 receive_count = 0;
dsm_segment *seg;
shm_mq_handle *outqh;
shm_mq_handle *inqh;
shm_mq_result res;
Size len;
void *data;
/* A negative loopcount is nonsensical. */
if (loop_count < 0)
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("repeat count size must be a non-negative integer")));
/*
* Using the nonblocking interfaces, we can even send data to ourselves,
* so the minimum number of workers for this test is zero.
*/
if (nworkers < 0)
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("number of workers must be a non-negative integer")));
/* Set up dynamic shared memory segment and background workers. */
test_shm_mq_setup(queue_size, nworkers, &seg, &outqh, &inqh);
/* Main loop. */
for (;;)
{
bool wait = true;
/*
* If we haven't yet sent the message the requisite number of times,
* try again to send it now. Note that when shm_mq_send() returns
* SHM_MQ_WOULD_BLOCK, the next call to that function must pass the
* same message size and contents; that's not an issue here because
* we're sending the same message every time.
*/
if (send_count < loop_count)
{
res = shm_mq_send(outqh, message_size, message_contents, true);
if (res == SHM_MQ_SUCCESS)
{
++send_count;
wait = false;
}
else if (res == SHM_MQ_DETACHED)
ereport(ERROR,
(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
errmsg("could not send message")));
}
/*
* If we haven't yet received the message the requisite number of
* times, try to receive it again now.
*/
if (receive_count < loop_count)
{
res = shm_mq_receive(inqh, &len, &data, true);
if (res == SHM_MQ_SUCCESS)
{
++receive_count;
/* Verifying every time is slow, so it's optional. */
if (verify)
verify_message(message_size, message_contents, len, data);
wait = false;
}
else if (res == SHM_MQ_DETACHED)
ereport(ERROR,
(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
errmsg("could not receive message")));
}
else
{
/*
* Otherwise, we've received the message enough times. This
* shouldn't happen unless we've also sent it enough times.
*/
if (send_count != receive_count)
ereport(ERROR,
(errcode(ERRCODE_INTERNAL_ERROR),
errmsg("message sent %d times, but received %d times",
send_count, receive_count)));
break;
}
if (wait)
{
/*
* If we made no progress, wait for one of the other processes to
* which we are connected to set our latch, indicating that they
* have read or written data and therefore there may now be work
* for us to do.
*/
WaitLatch(MyLatch, WL_LATCH_SET, 0, PG_WAIT_EXTENSION);
ResetLatch(MyLatch);
CHECK_FOR_INTERRUPTS();
}
}
/* Clean up. */
dsm_detach(seg);
PG_RETURN_VOID();
}
