Потокобезопасная очередь

#pragma once
 
// Очередь реализуется согласано книге Anthony Williams -
// "C++ Concurrency in Action: Practical Multithreading".
// Смотреть главу 6, раздел 6.2
 
#include <condition_variable>
#include <memory>
#include <mutex>
 
template <typename T>
class ThreadSafeQueue {
 private:
  using LockGuard = std::lock_guard<std::mutex>;
  using UniqueLock = std::unique_lock<std::mutex>;
 
  // Очередь реализуется в виде односвязного списка
  struct Node {
    // Данные в узле списка
    std::shared_ptr<T> data = nullptr;
 
    // Указатель на следующей элемент в списке
    std::unique_ptr<Node> next = nullptr;
  };
 
 private:
  // Потокобезопасный доступ к концу очереди
  const Node* get_tail() const {
    LockGuard tail_lock(m_tail_mutex);
    return m_tail;
  }
 
  // Извлечение из начала очереди
  std::unique_ptr<Node> pop_head() {
    std::unique_ptr<Node> old_head = std::move(m_head);
    m_head = std::move(old_head->next);
    return std::move(old_head);
  }
 
  // Ожидание появления данных в очереди
  UniqueLock wait_for_data() {
    UniqueLock head_lock(m_head_mutex);
    m_data_cond.wait(head_lock, [this] { return m_head.get() != get_tail(); });
    return std::move(head_lock);
  }
 
 public:
  // Создаем фиктивный узел, чтобы next никогда не указывал на nullptr.
  // Это необходимо для потокобезопасной работы и удобства.
  ThreadSafeQueue() : m_head(std::make_unique<Node>()), m_tail(m_head.get()) {}
  ThreadSafeQueue(const ThreadSafeQueue&) = delete;
  ThreadSafeQueue& operator=(const ThreadSafeQueue&) = delete;
 
  // Проверка наличия элемента и если он есть, то извлечение
  std::shared_ptr<T> try_pop() {
    LockGuard head_lock(m_head_mutex);
    if (m_head.get() == get_tail()) {
      return nullptr;
    }
 
    std::unique_ptr<Node> old_head = pop_head();
    return old_head->data;
  }
 
  // Проверка наличия элемента и если он есть, то извлечение
  bool try_pop(T& value) {
    LockGuard head_lock(m_head_mutex);
    if (m_head.get() == get_tail()) {
      return false;
    }
 
    value = std::move(*m_head->data);
    pop_head();
    return true;
  }
 
  // Извлечение элемента из очереди с ожиданием
  std::shared_ptr<T> wait_and_pop() {
    UniqueLock head_lock(wait_for_data());
    std::unique_ptr<Node> old_head = pop_head();
    return old_head->data;
  }
 
  // Извлечение элемента из очереди с ожиданием
  void wait_and_pop(T& value) {
    UniqueLock head_lock(wait_for_data());
    value = std::move(*m_head->data);
    pop_head();
  }
 
  // Вставка нового элемента в очередь. Надо понимать тот факт,
  // что фиктивная нода всегда находится в конце связного списка.
  // Если на фиктивную ноду указывают начало и конец очереди,
  // значит очередь пуста.
  void push(T new_value) {
    std::shared_ptr<T> new_data = std::make_shared<T>(std::move(new_value));
    std::unique_ptr<Node> new_dummy_node = std::make_unique<Node>();
    {
      LockGuard tail_lock(m_tail_mutex);
      Node* old_tail = m_tail;
      m_tail = new_dummy_node.get();
      old_tail->data = std::move(new_data);
      old_tail->next = std::move(new_dummy_node);
    }
    m_data_cond.notify_one();
  }
 
  bool empty() const {
    LockGuard head_lock(m_head_mutex);
    return m_head.get() == get_tail();
  }
 
 private:
  // Начало очереди
  mutable std::mutex m_head_mutex;
  std::unique_ptr<Node> m_head;
 
  // Конец очереди
  mutable std::mutex m_tail_mutex;
  Node* m_tail;
 
  std::condition_variable m_data_cond;
};

#include <algorithm>
#include <future>
#include <iostream>
#include <utility>
 
#include <gtest/gtest.h>
#include "thread_safe_queue.h"
 
class ThreadSafeQueueTest : public ::testing::Test {
 public:
  using Pair = std::pair<std::size_t, std::size_t>;
 
  void push_to_queue(ThreadSafeQueue<Pair>& queue, std::size_t number,
                     std::size_t count) {
    for (std::size_t i = 0; i < count; ++i) {
      queue.push(std::make_pair(number, i));
    }
  }
 
  std::vector<Pair> wait_and_pop_from_queue(ThreadSafeQueue<Pair>& queue,
                                            std::size_t count) {
    std::vector<Pair> pairs(count);
    for (std::size_t i = 0; i < count; ++i) {
      queue.wait_and_pop(pairs[i]);
    }
    return pairs;
  }
 
  std::vector<Pair> try_pop_from_queue(ThreadSafeQueue<Pair>& queue,
                                       std::size_t count) {
    std::vector<Pair> pairs;
    pairs.reserve(count);
    for (std::size_t i = 0; i < count; ++i) {
      auto result = queue.try_pop();
      if (result) {
        pairs.push_back(*result);
      }
    }
    return pairs;
  }
};
 
TEST_F(ThreadSafeQueueTest, one_push_pop) {
  ThreadSafeQueue<std::size_t> queue;
  ASSERT_TRUE(queue.empty());
 
  queue.push(1);
  ASSERT_FALSE(queue.empty());
 
  std::size_t pop_result = *queue.try_pop();
  ASSERT_TRUE(queue.empty());
  ASSERT_EQ(queue.try_pop(), nullptr);
  ASSERT_EQ(pop_result, 1);
 
  queue.push(2);
  ASSERT_FALSE(queue.empty());
 
  ASSERT_TRUE(queue.try_pop(pop_result));
  ASSERT_TRUE(queue.empty());
  ASSERT_EQ(pop_result, 2);
}
 
TEST_F(ThreadSafeQueueTest, many_push_pop) {
  ThreadSafeQueue<std::size_t> queue;
  ASSERT_TRUE(queue.empty());
 
  queue.push(1);
  queue.push(2);
  ASSERT_FALSE(queue.empty());
 
  std::size_t pop_result = *queue.try_pop();
  ASSERT_FALSE(queue.empty());
  ASSERT_EQ(pop_result, 1);
 
  ASSERT_TRUE(queue.try_pop(pop_result));
  ASSERT_TRUE(queue.empty());
  ASSERT_EQ(pop_result, 2);
}
 
TEST_F(ThreadSafeQueueTest, wait_and_pop_from_queue) {
  ThreadSafeQueue<Pair> queue;
 
  std::future<void> push1 = std::async(&ThreadSafeQueueTest::push_to_queue,
                                       this, std::ref(queue), 0, 5);
 
  std::future<void> push2 = std::async(&ThreadSafeQueueTest::push_to_queue,
                                       this, std::ref(queue), 1, 5);
 
  std::future<std::vector<Pair>> pop1 = std::async(
      &ThreadSafeQueueTest::wait_and_pop_from_queue, this, std::ref(queue), 10);
 
  push1.get();
  push2.get();
  auto result = pop1.get();
  std::stable_partition(result.begin(), result.end(),
                        [](const Pair& pair) { return pair.first == 0; });
 
  std::vector<Pair> expected = {{0, 0},
                                {0, 1},
                                {0, 2},
                                {0, 3},
                                {0, 4},
                                {1, 0},
                                {1, 1},
                                {1, 2},
                                {1, 3},
                                {1, 4}};
 
  ASSERT_TRUE(queue.empty());
  ASSERT_EQ(queue.try_pop(), nullptr);
  ASSERT_EQ(result, expected);
}
 
TEST_F(ThreadSafeQueueTest, try_pop_from_queue) {
  ThreadSafeQueue<Pair> queue;
 
  std::future<void> push1 = std::async(&ThreadSafeQueueTest::push_to_queue,
                                       this, std::ref(queue), 0, 5);
 
  std::future<void> push2 = std::async(&ThreadSafeQueueTest::push_to_queue,
                                       this, std::ref(queue), 1, 5);
 
  std::future<std::vector<Pair>> pop1 = std::async(
      &ThreadSafeQueueTest::try_pop_from_queue, this, std::ref(queue), 10);
 
  push1.get();
  push2.get();
  auto result = pop1.get();
 
  while (!queue.empty()) {
    result.push_back(*queue.try_pop());
  }
 
  std::stable_partition(result.begin(), result.end(),
                        [](const Pair& pair) { return pair.first == 0; });
 
  std::vector<Pair> expected = {{0, 0},
                                {0, 1},
                                {0, 2},
                                {0, 3},
                                {0, 4},
                                {1, 0},
                                {1, 1},
                                {1, 2},
                                {1, 3},
                                {1, 4}};
 
  ASSERT_TRUE(queue.empty());
  ASSERT_EQ(queue.try_pop(), nullptr);
  ASSERT_EQ(result, expected);
}
 
int main(int argc, char** argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}