Enable termination in the wait free buffer

This change adds a field to the shared wait free buffer analogous to the
ring_cmd field in the ring buffer IPC mechanism. It is set to
PGM_TERMINATE when the consumer has no more data to send.

3 files changed, 32 insertions, 15 deletions

diff --git a/include/wait_free.h b/include/wait_free.h
index 38e2442..1a366b2 100644
--- a/include/wait_free.h
+++ b/include/wait_free.h
@@ -7,6 +7,10 @@
 
 // Implements a wait-free buffer where each buffer is an 8 byte uint64_t.
 struct wait_free {
+  // This fulfills the same purpose as edge->ring_cmd in pgm.cpp. It must only
+  // ever be 0 or PGM_TERMINATE. However, it must reside in the same colored
+  // shared memory as the ring buffer.
+  pgm_command_t command;
   std::atomic_uchar latest;
   std::atomic_uchar reading;
   int wdcnt;
@@ -20,6 +24,7 @@ static inline int InitWaitFree(struct wait_free *b) {
   b->latest = 0;
   b->wdcnt = -1;
   b->reading = 0;
+  b->command = 0;
   return 0;
 }
 
diff --git a/src/pgm.cpp b/src/pgm.cpp
index 3a2c463..93c8533 100644
--- a/src/pgm.cpp
+++ b/src/pgm.cpp
@@ -3561,12 +3561,20 @@ static int pgm_send_ring_data(struct pgm_edge* e, pgm_command_t tag)
         return 0;
 }
 
+static int pgm_send_waitfree_data(struct pgm_edge *e, pgm_command_t tag) {
+        if (tag & PGM_TERMINATE) {
+                e->wait_free_buffer->command = tag;
+        } else {
+                e->ops->write(e, pgm_get_user_ptr(e->buf_out), e->attr.nr_produce);
+        }
+        return 0;
+}
+
 static int pgm_send_data(struct pgm_edge* e, pgm_command_t tag = PGM_NORMAL)
 {
-        if(!(e->attr.type & __PGM_EDGE_RING))
-                return pgm_send_std_data(e, tag);
-        else
-                return pgm_send_ring_data(e, tag);
+        if (e->attr.type & __PGM_EDGE_RING) return pgm_send_ring_data(e, tag);
+        if (e->attr.type & __PGM_EDGE_WAITFREE) return pgm_send_waitfree_data(e, tag);
+        return pgm_send_std_data(e, tag);
 }
 
 static eWaitStatus pgm_wait_for_data(pgm_fd_mask_t* to_wait,
@@ -3707,6 +3715,20 @@ wait_for_data: // jump here if we would block on read
                         continue;
                 }
 
+                if (e->attr.type & __PGM_EDGE_WAITFREE) {
+                        // Shortcut for the wait-free buffer IPC, similar to the previous block
+                        // for ring buffers. (The only difference is that the wait free buffer
+                        // has no notion of being full or empty)
+                        if (e->wait_free_buffer->command & PGM_TERMINATE) {
+                                n->nr_terminate_msgs++;
+                        }
+                        // This is INTENTIONALLY re-read if a terminate is received; to make sure
+                        // that the final data is read. Since it will only re-read the final data
+                        // this can't result in anything being read that wasn't written.
+                        e->ops->read(e, dest_ptrs[i], e->attr.nr_consume);
+                        continue;
+                }
+
 read_more: // jump to here if we need to read more bytes into our buffer
 
                 remaining = e->attr.nr_consume - (dest_ptrs[i] - (char*)pgm_get_user_ptr(e->buf_in));
@@ -3830,7 +3852,7 @@ read_more: // jump to here if we need to read more bytes into our buffer
 out:
         // signal terminte if everyone has checked in
         if(n->nr_terminate_msgs &&
-           (n->nr_terminate_msgs == (n->nr_in_data - n->nr_in_data_backedges)))
+           (n->nr_terminate_msgs >= (n->nr_in_data - n->nr_in_data_backedges)))
                 wait_status = WaitExhaustedAndTerminate;
 
         return wait_status;
diff --git a/tools/wait_free_test.cpp b/tools/wait_free_test.cpp
index f4e3376..f62f952 100644
--- a/tools/wait_free_test.cpp
+++ b/tools/wait_free_test.cpp
@@ -107,16 +107,6 @@ int main(void) {
   waitfree_attr.nr_consume = sizeof(uint32_t);
   waitfree_attr.nr_threshold = sizeof(uint32_t);
 
-  /*
-  edge_attr_t ring_attr;
-  memset(&ring_attr, 0, sizeof(ring_attr));
-  ring_attr.type = pgm_ring_edge;
-  ring_attr.nr_produce = sizeof(uint32_t);
-  ring_attr.nr_consume = sizeof(uint32_t);
-  ring_attr.nr_threshold = sizeof(uint32_t);
-  ring_attr.nmemb = 32;
-  */
-
   CheckError(pgm_init_process_local());
   CheckError(pgm_init_graph(&g, "demo"));