Add thread safe queue
Reviewers: dgleich, florijan Reviewed By: dgleich Differential Revision: https://phabricator.memgraph.io/D979
This commit is contained in:
parent
85ef12def5
commit
5c0f378394
65
src/data_structures/queue.hpp
Normal file
65
src/data_structures/queue.hpp
Normal file
@ -0,0 +1,65 @@
|
|||||||
|
#pragma once
|
||||||
|
|
||||||
|
#include <condition_variable>
|
||||||
|
#include <cstdint>
|
||||||
|
#include <experimental/optional>
|
||||||
|
#include <iostream>
|
||||||
|
#include <mutex>
|
||||||
|
#include <queue>
|
||||||
|
|
||||||
|
// Thread safe queue. Probably doesn't perform very well, but it works.
|
||||||
|
template <typename T>
|
||||||
|
class Queue {
|
||||||
|
public:
|
||||||
|
Queue() = default;
|
||||||
|
Queue(const Queue &) = delete;
|
||||||
|
Queue &operator=(const Queue &) = delete;
|
||||||
|
Queue(Queue &&) = delete;
|
||||||
|
Queue &operator=(Queue &&) = delete;
|
||||||
|
|
||||||
|
void Push(T x) {
|
||||||
|
std::unique_lock<std::mutex> guard(mutex_);
|
||||||
|
queue_.emplace(std::move(x));
|
||||||
|
guard.unlock();
|
||||||
|
cvar_.notify_one();
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename... Args>
|
||||||
|
void Emplace(Args &&... args) {
|
||||||
|
std::unique_lock<std::mutex> guard(mutex_);
|
||||||
|
queue_.emplace(std::forward<Args>(args)...);
|
||||||
|
guard.unlock();
|
||||||
|
cvar_.notify_one();
|
||||||
|
}
|
||||||
|
|
||||||
|
int64_t size() const {
|
||||||
|
std::unique_lock<std::mutex> guard(mutex_);
|
||||||
|
return queue_.size();
|
||||||
|
}
|
||||||
|
|
||||||
|
bool empty() const {
|
||||||
|
std::unique_lock<std::mutex> guard(mutex_);
|
||||||
|
return queue_.empty();
|
||||||
|
}
|
||||||
|
|
||||||
|
T AwaitPop() {
|
||||||
|
std::unique_lock<std::mutex> guard(mutex_);
|
||||||
|
cvar_.wait(guard, [this]() { return !queue_.empty(); });
|
||||||
|
auto x = std::move(queue_.front());
|
||||||
|
queue_.pop();
|
||||||
|
return x;
|
||||||
|
}
|
||||||
|
|
||||||
|
std::experimental::optional<T> MaybePop() {
|
||||||
|
std::unique_lock<std::mutex> guard(mutex_);
|
||||||
|
if (queue_.empty()) return std::experimental::nullopt;
|
||||||
|
auto x = std::move(queue_.front());
|
||||||
|
queue_.pop();
|
||||||
|
return x;
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
std::queue<T> queue_;
|
||||||
|
std::condition_variable cvar_;
|
||||||
|
mutable std::mutex mutex_;
|
||||||
|
};
|
139
tests/unit/queue.cpp
Normal file
139
tests/unit/queue.cpp
Normal file
@ -0,0 +1,139 @@
|
|||||||
|
#include <atomic>
|
||||||
|
#include <chrono>
|
||||||
|
#include <experimental/optional>
|
||||||
|
#include <string>
|
||||||
|
#include <thread>
|
||||||
|
#include <utility>
|
||||||
|
|
||||||
|
#include "gmock/gmock.h"
|
||||||
|
#include "gtest/gtest.h"
|
||||||
|
|
||||||
|
#include "data_structures/queue.hpp"
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
using namespace std::literals::chrono_literals;
|
||||||
|
|
||||||
|
TEST(Queue, PushMaybePop) {
|
||||||
|
Queue<int> q;
|
||||||
|
q.Push(1);
|
||||||
|
EXPECT_EQ(*q.MaybePop(), 1);
|
||||||
|
EXPECT_EQ(q.MaybePop(), std::experimental::nullopt);
|
||||||
|
|
||||||
|
q.Push(2);
|
||||||
|
q.Push(3);
|
||||||
|
EXPECT_EQ(*q.MaybePop(), 2);
|
||||||
|
|
||||||
|
q.Push(4);
|
||||||
|
q.Push(5);
|
||||||
|
EXPECT_EQ(*q.MaybePop(), 3);
|
||||||
|
EXPECT_EQ(*q.MaybePop(), 4);
|
||||||
|
EXPECT_EQ(*q.MaybePop(), 5);
|
||||||
|
EXPECT_EQ(q.MaybePop(), std::experimental::nullopt);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(Queue, Emplace) {
|
||||||
|
Queue<std::pair<std::string, int>> q;
|
||||||
|
q.Emplace("abc", 123);
|
||||||
|
EXPECT_THAT(*q.MaybePop(), testing::Pair("abc", 123));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(Queue, Size) {
|
||||||
|
Queue<int> q;
|
||||||
|
EXPECT_EQ(q.size(), 0);
|
||||||
|
|
||||||
|
q.Push(1);
|
||||||
|
EXPECT_EQ(q.size(), 1);
|
||||||
|
|
||||||
|
q.Push(1);
|
||||||
|
EXPECT_EQ(q.size(), 2);
|
||||||
|
|
||||||
|
q.MaybePop();
|
||||||
|
EXPECT_EQ(q.size(), 1);
|
||||||
|
|
||||||
|
q.MaybePop();
|
||||||
|
EXPECT_EQ(q.size(), 0);
|
||||||
|
q.MaybePop();
|
||||||
|
EXPECT_EQ(q.size(), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(Queue, Empty) {
|
||||||
|
Queue<int> q;
|
||||||
|
EXPECT_TRUE(q.empty());
|
||||||
|
|
||||||
|
q.Push(1);
|
||||||
|
EXPECT_FALSE(q.empty());
|
||||||
|
|
||||||
|
q.MaybePop();
|
||||||
|
EXPECT_TRUE(q.empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(Queue, AwaitPop) {
|
||||||
|
Queue<int> q;
|
||||||
|
std::thread t([&] {
|
||||||
|
q.Push(1);
|
||||||
|
q.Push(2);
|
||||||
|
std::this_thread::sleep_for(200ms);
|
||||||
|
q.Push(3);
|
||||||
|
q.Push(4);
|
||||||
|
});
|
||||||
|
|
||||||
|
EXPECT_EQ(q.AwaitPop(), 1);
|
||||||
|
EXPECT_EQ(q.AwaitPop(), 2);
|
||||||
|
EXPECT_EQ(q.AwaitPop(), 3);
|
||||||
|
EXPECT_EQ(q.AwaitPop(), 4);
|
||||||
|
t.join();
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(Queue, Concurrent) {
|
||||||
|
Queue<int> q;
|
||||||
|
|
||||||
|
const int kNumProducers = 10;
|
||||||
|
const int kNumConsumers = 10;
|
||||||
|
const int kNumElementsPerProducer = 300000;
|
||||||
|
|
||||||
|
std::vector<std::thread> producers;
|
||||||
|
std::atomic<int> next{0};
|
||||||
|
|
||||||
|
for (int i = 0; i < kNumProducers; ++i) {
|
||||||
|
producers.emplace_back([&] {
|
||||||
|
for (int i = 0; i < kNumElementsPerProducer; ++i) {
|
||||||
|
q.Push(next++);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
std::vector<std::thread> consumers;
|
||||||
|
std::vector<int> retrieved[kNumConsumers];
|
||||||
|
std::atomic<int> num_retrieved{0};
|
||||||
|
for (int i = 0; i < kNumConsumers; ++i) {
|
||||||
|
consumers.emplace_back(
|
||||||
|
[&](int thread_id) {
|
||||||
|
while (true) {
|
||||||
|
int count = num_retrieved++;
|
||||||
|
if (count >= kNumProducers * kNumElementsPerProducer) break;
|
||||||
|
retrieved[thread_id].push_back(q.AwaitPop());
|
||||||
|
}
|
||||||
|
},
|
||||||
|
i);
|
||||||
|
}
|
||||||
|
|
||||||
|
for (auto &t : consumers) {
|
||||||
|
t.join();
|
||||||
|
}
|
||||||
|
for (auto &t : producers) {
|
||||||
|
t.join();
|
||||||
|
}
|
||||||
|
|
||||||
|
EXPECT_EQ(q.MaybePop(), std::experimental::nullopt);
|
||||||
|
|
||||||
|
std::set<int> all_elements;
|
||||||
|
for (auto &r : retrieved) {
|
||||||
|
all_elements.insert(r.begin(), r.end());
|
||||||
|
}
|
||||||
|
EXPECT_EQ(all_elements.size(), kNumProducers * kNumElementsPerProducer);
|
||||||
|
EXPECT_EQ(*all_elements.begin(), 0);
|
||||||
|
EXPECT_EQ(*all_elements.rbegin(),
|
||||||
|
kNumProducers * kNumElementsPerProducer - 1);
|
||||||
|
}
|
||||||
|
}
|
Loading…
Reference in New Issue
Block a user