utl_buffer.h

/*
 * Developer : Alexey Zakharov (alexey.zakharov@vectioneer.com)
 * All rights reserved. Copyright (c) 2015 VECTIONEER.
 */
 
#ifndef UTILS_UTL_BUFFER_H
#define UTILS_UTL_BUFFER_H
 
#include "utl_log.h"
#include "utl_tsan.h"
#include <atomic>
#include <cstdint>
#include <cstdio>
#include <shared_mutex>
#include <unistd.h>
 
namespace mcx {
 
namespace utils {
 
template <typename T>
class SpmcBuffer {
  static constexpr uint32_t BUFFER_SIZE = 16;
  static constexpr uint32_t BUFFER_MASK{BUFFER_SIZE - 1};
 
public:
  SpmcBuffer() : buffer_(new Cell[BUFFER_SIZE]()) {}
 
  virtual ~SpmcBuffer() { delete[] buffer_; }
 
  void write(const T& in) {
    uint32_t pos = enqueue_pos_.load(std::memory_order_relaxed);
    Cell* buf;
    do {
      pos += 1;
      buf = &buffer_[pos & BUFFER_MASK];
    } while (!buf->guard.try_lock());
    buf->data = in;
    buf->guard.unlock();
    enqueue_pos_.store(pos, std::memory_order_release);
  }
 
  constexpr uint32_t size() { return BUFFER_SIZE; }
 
  uint32_t read(T* out) const {
    uint32_t pos = enqueue_pos_.load(std::memory_order_acquire);
    Cell* buf = &buffer_[pos & BUFFER_MASK];
    buf->guard.lock_shared();
    *out = buf->data;
    buf->guard.unlock_shared();
    return pos;
  }
 
  uint32_t pos() const { return enqueue_pos_.load(std::memory_order_acquire); }
 
protected:
  struct /*alignas(LEVEL1_DCACHE_LINESIZE)*/ Cell {
#if __cplusplus >= 201703L
    std::shared_mutex guard;
#else
    std::shared_timed_mutex guard;
#endif
    T data;
  };
  Cell* const buffer_;
  std::atomic<uint32_t> enqueue_pos_{};
};
 
template <typename T>
class MpmcQueue {
public:
  explicit MpmcQueue(uint32_t buffer_length) : buffer_(new Cell[buffer_length]{}), buffer_mask_(buffer_length - 1) {
    log_assert((buffer_length >= 2) && ((buffer_length & (buffer_length - 1)) == 0),
               "Queue length must be 2^N") for (uint32_t i = 0; i < buffer_length; ++i) {
      buffer_[i].sequence = i;
    }
  }
 
  virtual ~MpmcQueue() { delete[] buffer_; }
 
  bool write(const T& in, bool no_drop = false) {
    Cell* cell;
    uint32_t pos = enqueue_pos_.load(std::memory_order_relaxed);
    bool is_ok = true;
    for (;;) {
      cell = &buffer_[pos & buffer_mask_];
      uint32_t seq = cell->sequence.load(std::memory_order_acquire);
      intptr_t dif = (intptr_t)seq - (intptr_t)pos;
      if (dif == 0) {
        if (enqueue_pos_.compare_exchange_weak(pos, pos + 1, std::memory_order_relaxed))
          break;
      } else if (dif < 0) {
        if (!no_drop) {
          drop();
          is_ok = false;
        } else {
          return false;
        }
      } else {
        pos = enqueue_pos_.load(std::memory_order_relaxed);
      }
    }
    cell->data = in;
    cell->sequence.store(pos + 1, std::memory_order_release);
    enqueue_pos_size_.store(enqueue_pos_, std::memory_order_relaxed);
    return is_ok;
  }
 
  uint32_t read(T* out) const {
    Cell* cell;
    uint32_t pos = dequeue_pos_.load(std::memory_order_relaxed);
    for (;;) {
      cell = &buffer_[pos & buffer_mask_];
      uint32_t seq = cell->sequence.load(std::memory_order_acquire);
      intptr_t dif = (intptr_t)seq - (intptr_t)(pos + 1);
      if (dif == 0) {
        if (dequeue_pos_.compare_exchange_weak(pos, pos + 1, std::memory_order_relaxed))
          break;
      } else if (dif < 0) {
        return pos;
      } else {
        pos = dequeue_pos_.load(std::memory_order_relaxed);
      }
    }
    *out = cell->data;
    cell->sequence.store(pos + buffer_mask_ + 1, std::memory_order_release);
    dequeue_pos_size_.store(dequeue_pos_, std::memory_order_relaxed);
    return pos + 1;
  }
 
  void drop() {
    Cell* cell;
    uint32_t pos = dequeue_pos_.load(std::memory_order_relaxed);
    for (;;) {
      cell = &buffer_[pos & buffer_mask_];
      uint32_t seq = cell->sequence.load(std::memory_order_acquire);
      intptr_t dif = (intptr_t)seq - (intptr_t)(pos + 1);
      if (dif == 0) {
        if (dequeue_pos_.compare_exchange_weak(pos, pos + 1, std::memory_order_relaxed))
          break;
      } else if (dif < 0) {
        return;
      } else {
        pos = dequeue_pos_.load(std::memory_order_relaxed);
      }
    }
    cell->sequence.store(pos + buffer_mask_ + 1, std::memory_order_release);
    dequeue_pos_size_.store(dequeue_pos_, std::memory_order_relaxed);
  }
 
  uint32_t pos() const { return dequeue_pos_size_.load(std::memory_order_relaxed); }
 
  uint32_t size() const {
    return enqueue_pos_size_.load(std::memory_order_relaxed) - dequeue_pos_size_.load(std::memory_order_relaxed);
  }
 
protected:
  struct Cell {
    std::atomic<uint32_t> sequence;
    T data;
  };
 
  Cell* const buffer_;
  uint32_t const buffer_mask_;
 
  std::atomic<uint32_t> enqueue_pos_{};
  std::atomic<uint32_t> enqueue_pos_size_{};
 
  mutable std::atomic<uint32_t> dequeue_pos_{};
  mutable std::atomic<uint32_t> dequeue_pos_size_{};
};
 
} // namespace utils
 
} // namespace mcx
 
#endif // UTILS_UTL_BUFFER_H