concurrent_vector.hpp

#pragma once
 
#include <deque>
#include <vector>
#include <mutex>
#include <atomic>
#include <optional>
#include <condition_variable>
 
namespace glim {
 
struct DataStorePolicy {
public:
  template <typename T, typename Alloc>
  void regulate(std::deque<T, Alloc>& queue) const {
    if (queue.size() < max_size) {
      return;
    }
 
    const size_t num_erase = queue.size() - max_size;
    if (pop_front) {
      queue.erase(queue.begin(), queue.begin() + num_erase);
    } else {
      queue.erase(queue.end() - num_erase, queue.end());
    }
  }
 
  static DataStorePolicy UNLIMITED() { return DataStorePolicy(); }
  static DataStorePolicy UPTO(const size_t max_size, const bool pop_front = true) { return DataStorePolicy{max_size, pop_front}; }
 
public:
  const size_t max_size = std::numeric_limits<size_t>::max();
  const bool pop_front = true;
};
 
template <typename T, typename Alloc = std::allocator<T>>
class ConcurrentVector {
public:
  ConcurrentVector(const DataStorePolicy& policy = DataStorePolicy::UNLIMITED()) : policy(policy) { end_of_data = false; }
 
  void submit_end_of_data() {
    end_of_data = true;
    cond.notify_all();
  }
 
  bool empty() const {
    std::lock_guard<std::mutex> lock(mutex);
    return values.empty();
  }
 
  int size() const {
    std::lock_guard<std::mutex> lock(mutex);
    return values.size();
  }
 
  void reserve(int n) {
    std::lock_guard<std::mutex> lock(mutex);
    values.reserve(n);
  }
 
  void push_back(const T& value) {
    std::lock_guard<std::mutex> lock(mutex);
    values.push_back(value);
    policy.regulate(values);
    cond.notify_one();
  }
 
  void clear() {
    std::lock_guard<std::mutex> lock(mutex);
    values.clear();
  }
 
  T front() const {
    std::lock_guard<std::mutex> lock(mutex);
    return values.front();
  }
 
  T back() const {
    std::lock_guard<std::mutex> lock(mutex);
    return values.back();
  }
 
  template <typename Container>
  void insert(const Container& new_values) {
    if (new_values.empty()) {
      return;
    }
 
    std::lock_guard<std::mutex> lock(mutex);
    values.insert(values.end(), new_values.begin(), new_values.end());
    policy.regulate(values);
    cond.notify_all();
  }
 
  std::optional<T> pop() {
    std::lock_guard<std::mutex> lock(mutex);
    if (values.empty()) {
      return std::nullopt;
    }
 
    const T data = values.front();
    values.pop_front();
    return data;
  }
 
  std::optional<T> pop_wait() {
    std::unique_lock<std::mutex> lock(mutex);
 
    std::optional<T> data;
    cond.wait(lock, [this, &data] {
      if (values.empty()) {
        return static_cast<bool>(end_of_data);
      }
 
      data = values.front();
      values.pop_front();
      return true;
    });
 
    return data;
  }
 
  std::vector<T, Alloc> get_all_and_clear_wait() {
    std::unique_lock<std::mutex> lock(mutex);
 
    std::vector<T, Alloc> buffer;
    cond.wait(lock, [this, &buffer] {
      if (values.empty()) {
        return static_cast<bool>(end_of_data);
      }
 
      buffer.assign(values.begin(), values.end());
      values.clear();
      return true;
    });
 
    return buffer;
  }
 
  std::vector<T, Alloc> get_all_and_clear() {
    std::vector<T, Alloc> buffer;
    std::lock_guard<std::mutex> lock(mutex);
    buffer.assign(values.begin(), values.end());
    values.clear();
 
    return buffer;
  }
 
  std::vector<T, Alloc> get_and_clear(int num_max) {
    std::vector<T, Alloc> buffer;
    std::lock_guard<std::mutex> lock(mutex);
    if (values.size() <= num_max) {
      buffer.assign(values.begin(), values.end());
      values.clear();
    } else {
      buffer.assign(values.begin(), values.begin() + num_max);
      values.erase(values.begin(), values.begin() + num_max);
    }
 
    return buffer;
  }
 
private:
  const DataStorePolicy policy;
 
  std::atomic_bool end_of_data;
  std::condition_variable cond;
 
  mutable std::mutex mutex;
  std::deque<T, Alloc> values;
};
 
}  // namespace  glim