[pubsub] Fix lag jumpahead logic

The existing behavior when the subscriber lags (behind the write head by
queue size, since the buffer is circular, new data will have been written
to the location the subscriber is trying to read from) was to move to the
oldest valid location. However, since the jumpahead was happening outside
of any locks, between the time of jumping ahead and acquring the read lock
for that location, the publisher(s) could have written past the new read
location.

The fix was to jump ahead optimistically, rather than to the last valid
location. Currently, this is mid-way between the write head and the last
valid location. Another issue that surfaced was that since the jumpahead
happens outside any lock, in the case of particularly fast publishers, even
here the same issue can be seen, where the write head wraps around and writes
past the new (optimistic) read location.

The second fix was to do a similar jump ahead (and associated check) while
under read lock. This is an expensive operation, since we might potentially
by readig from two buffer elements. However, this is a (very rate) slow path.

To ensure this logic is correct, the benchmark for pubsub has an assert to
check that each subsequent subscribed elements is increasing. The benchmark
has also been rewritten to make it much cleaner.

Author: Squadrick

benchmark/pubsub.cpp

@@ -21,6 +21,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 ==============================================================================*/
 
+#include <cassert>
 #include <chrono>
 #include <iostream>
 
@@ -35,66 +36,69 @@ SOFTWARE.
 #endif
 
 const char topic[] = "raw_benchmark_topic";
-const int SECONDS = 10;
-const int VECTOR_SIZE = 32 * 1024;
 
 shm::stats::Welford per_second_lag;
 
 struct Message {
+  uint64_t count;
   uint64_t timestamp;
   uint8_t *data;
 };
 
+uint64_t current_count = 0;
+
 void callback(shm::memory::Memblock *memblock) {
+  if (memblock->is_empty()) {
+    return;
+  }
   auto *msg = reinterpret_cast<Message *>(memblock->ptr);
   auto lag = shm::current_time() - msg->timestamp;
+  // assert(current_count == msg->count - 1);
+  assert(current_count < msg->count);
+  current_count = msg->count;
   per_second_lag.add(lag);
 }
 
-void subscribe_loop() {
-  std::this_thread::sleep_for(std::chrono::seconds(1));
+void subscribe_loop(int seconds) {
   shm::pubsub::Subscriber sub(topic, callback, nullptr);
-  int seconds = 0;
-
-  auto start = shm::current_time();
-  while (true) {
-    sub.spin_once();
-    auto end = shm::current_time();
-    auto diff = end - start;
-
-    if (diff > TIMESCALE_COUNT) {
-      std::cout << "Lag (" << TIMESCALE_NAME << "): " << per_second_lag
-                << std::endl;
-      if (++seconds == SECONDS) break;
-      start = shm::current_time();
-      per_second_lag.clear();
+
+  for (int sec = 0; sec < seconds; ++sec) {
+    auto start = std::chrono::steady_clock::now();
+    for (auto now = start; now < start + std::chrono::seconds(1);
+         now = std::chrono::steady_clock::now()) {
+      sub.spin_once();
     }
+    std::cout << per_second_lag << std::endl;
+    per_second_lag.clear();
   }
 }
 
-void publish_loop() {
-  shm::pubsub::Publisher pub(topic, nullptr);
+void publish_loop(int seconds, int vector_size) {
+  std::cout << "Number of bytes = " << vector_size << std::endl;
+  std::cout << "Time unit = " << TIMESCALE_NAME << std::endl;
 
-  auto *rawptr = malloc(VECTOR_SIZE);
-  std::memset(rawptr, 255, VECTOR_SIZE);
+  auto *rawptr = malloc(vector_size);
+  std::memset(rawptr, 255, vector_size);
   Message *msg = reinterpret_cast<Message *>(rawptr);
+  msg->count = 0;
 
-  std::cout << "Number of bytes = " << VECTOR_SIZE << std::endl;
+  shm::pubsub::Publisher pub(topic, nullptr);
 
-  auto start = shm::current_time();
-  while (true) {
+  auto start = std::chrono::steady_clock::now();
+  for (auto now = start; now < start + std::chrono::seconds(seconds);
+       now = std::chrono::steady_clock::now()) {
+    msg->count++;
     msg->timestamp = shm::current_time();
-    pub.publish(msg, VECTOR_SIZE);
-    auto end = shm::current_time();
-    auto diff = end - start;
-    if (diff > (SECONDS + 1) * TIMESCALE_COUNT) break;
+    pub.publish(msg, vector_size);
   }
   free(msg);
 }
 
 int main() {
-  std::thread subscribe_thread(subscribe_loop);
-  std::thread publish_thread(publish_loop);
+  const int SECONDS = 10;
+  const int VECTOR_SIZE = 32 * 1024;
+  std::thread subscribe_thread(subscribe_loop, SECONDS);
+  std::thread publish_thread(publish_loop, SECONDS, VECTOR_SIZE);
 
   subscribe_thread.join();
   publish_thread.join();

include/shadesmar/pubsub/subscriber.h

@@ -83,23 +83,29 @@ memory::Memblock Subscriber::get_message() {
     return memory::Memblock();
   }
 
-  if (topic_->counter() - counter_ > topic_->queue_size()) {
+  if (topic_->counter() - counter_ >= topic_->queue_size()) {
     /*
+     * Why is the check >= (not >)? This is because in topic's
+     * `write` we do an `inc` at the end, so the write head
+     * (topic_->counter()) is currently under-write.
+     *
      * If we have fallen behind by the size of the queue
      * in the case of overlap, we go to last existing
      * element in the queue.
      *
-     * Why is the check > (not >=)? This is because in topic's
-     * `write` we do a `fetch_add`, so the queue pointer is already
-     * ahead of where it last wrote.
+     *
+     * Rather than going to the last valid location:
+     *    topic_->counter() - topic_->queue_size()
+     * We move to a more optimistic location (see `jumpahead`), to prevent
+     * hitting a case of always trying to keep up with the publisher.
      */
-    counter_ = topic_->counter() - topic_->queue_size();
+    counter_ = jumpahead(topic_->counter(), topic_->queue_size());
   }
 
   memory::Memblock memblock;
   memblock.free = true;
 
-  if (!topic_->read(&memblock, counter_)) {
+  if (!topic_->read(&memblock, &counter_)) {
     return memory::Memblock();
   }
 

include/shadesmar/pubsub/topic.h

@@ -37,6 +37,13 @@ SOFTWARE.
 
 namespace shm::pubsub {
 
+// The logic for optimistically jumping ahead, if the current read
+// logic has fallen behind (circular wrap around). `counter` is the
+// current write head location.
+inline uint32_t jumpahead(uint32_t counter, uint32_t queue_size) {
+  return counter - queue_size / 2;
+}
+
 template <class LockT>
 struct TopicElemT : public memory::Element {
   LockT mutex;
@@ -81,7 +88,7 @@ class Topic {
       std::cerr << "Increase buffer_size" << std::endl;
       return false;
     }
-    uint32_t q_pos = fetch_add_counter() & (queue_size() - 1);
+    uint32_t q_pos = counter() & (queue_size() - 1);
     TopicElem *elem = &(memory_.shared_queue_->elements[q_pos]);
 
     /*
@@ -120,6 +127,7 @@ class Topic {
     while (!memory_.allocator_->free(old_address)) {
       std::this_thread::sleep_for(std::chrono::microseconds(100));
     }
+    inc_counter();
     return true;
   }
 
@@ -129,41 +137,62 @@ class Topic {
    * defaults to the head of the queue. Reads requires an explicit
    * position in the queue to read(`pos`). We use this position
    * argument since we need to support multiple subscribers each
-   * reading at their own pace. Look at `subscriber.h` to see the
-   * logic to handle circular queue wrap around logic.
+   * reading at their own pace. Between picking the element at pos to
+   * read from and acquiring a read lock, the publisher may write a new
+   * value at pos. In this case, we implement a slow path to jump ahead.
    */
 
-  bool read(memory::Memblock *memblock, uint32_t pos) {
-    /*
-     * Read into a raw array.
-     */
+  bool read(memory::Memblock *memblock, std::atomic<uint32_t> *pos) {
     TopicElem *elem =
-        &(memory_.shared_queue_->elements[pos & (queue_size() - 1)]);
+        &(memory_.shared_queue_->elements[*pos & (queue_size() - 1)]);
 
     /*
-     * Code path:
+     * Code path (without the slow path for lag):
      *  1. Acquire sharable lock
      *  2. Check for emptiness
      *  3. Copy from shared memory to input param `msg`
      *  4. Release sharable lock
      */
     Scope<concurrent::SHARED> _(&elem->mutex);
 
-    if (elem->empty) {
-      return false;
+    // Using a lambda for this reduced throughput.
+    #define MOVE_ELEM(_elem)                                                \
+      if (_elem->empty) {                                                   \
+        return false;                                                       \
+      }                                                                     \
+      auto *dst = memory_.allocator_->handle_to_ptr(_elem->address_handle); \
+      memblock->size = _elem->size;                                         \
+      memblock->ptr = copier_->alloc(memblock->size);                       \
+      copier_->shm_to_user(memblock->ptr, dst, memblock->size);
+
+    if (queue_size() > counter() - *pos) {
+      // Fast path.
+      MOVE_ELEM(elem);
+      return true;
     }
 
-    auto *dst = memory_.allocator_->handle_to_ptr(elem->address_handle);
-    memblock->size = elem->size;
-    memblock->ptr = copier_->alloc(memblock->size);
-
-    copier_->shm_to_user(memblock->ptr, dst, memblock->size);
-
+    // See comment in `pubsub/subscriber.h`, in function `get_function()` for
+    // more info. *pos is outdated, the publisher has already written here
+    // before the reader lock was held. Jump ahead optimisically.
+    //
+    // Q: Why no lock on `next_best_elem`?
+    // A: `elem` is behind `next_best_elem`. With a lock on the former, the
+    //    publisher cannot cross `elem` to get to `next_best_elem`.
+    //
+    // Q: Why is the jump ahead implemented again in `get_message()`?
+    // A: `get_message()` can jump ahead outside of holding a lock. If a
+    //     lag can be detected there, it is more performant. This is a slow
+    //     path under a read lock.
+    *pos = jumpahead(counter(), queue_size());
+    TopicElem *next_best_elem =
+        &(memory_.shared_queue_->elements[*pos & (queue_size() - 1)]);
+    MOVE_ELEM(next_best_elem);
     return true;
+    #undef MOVE_ELEM
   }
 
-  inline __attribute__((always_inline)) uint32_t fetch_add_counter() const {
-    return memory_.shared_queue_->counter.fetch_add(1);
+  inline __attribute__((always_inline)) void inc_counter() const {
+    memory_.shared_queue_->counter++;
   }
 
   inline __attribute__((always_inline)) uint32_t counter() const {

test/pubsub_test.cpp

@@ -289,11 +289,13 @@ TEST_CASE("sub_counter_jump") {
   }
 
   sub.spin_once();
-  REQUIRE(answer == 0);
+  int lookback =
+      shm::pubsub::jumpahead(shm::memory::QUEUE_SIZE, shm::memory::QUEUE_SIZE);
+  REQUIRE(answer == lookback);
 
-  int moveahead = 0.5 * shm::memory::QUEUE_SIZE;
+  int moveahead = shm::memory::QUEUE_SIZE - lookback + 1;
   for (int i = 0; i < moveahead; ++i) {
-    int message = shm::memory::QUEUE_SIZE + i;
+    int message = lookback + i;
     pub.publish(reinterpret_cast<void *>(&message), sizeof(int));
   }