64 lines
1.6 KiB
C++
64 lines
1.6 KiB
C++
/*
|
|
* A simple thread-safe multi-producer, multi-consumer queue.
|
|
*
|
|
* This is a wrapper around std::deque that provides non-blocking
|
|
* queue operations like readIfNotEmpty and writeIfNotFull using
|
|
* std mutexes and the underlying queue can only be accessed
|
|
* with synchronized sections.
|
|
*
|
|
* For now the goal is to provide a simple implementation that
|
|
* works in all cases and produces no surprises to users.
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include <deque>
|
|
#include <mutex>
|
|
#include <type_traits>
|
|
|
|
namespace torch::nativert::detail {
|
|
|
|
// TODO(zhxchen17) Add wrapper for concurrentqueue.
|
|
template <typename T>
|
|
class MPMCQueue {
|
|
static_assert(!std::is_reference_v<T>);
|
|
|
|
public:
|
|
explicit MPMCQueue(size_t capacity) : capacity_(capacity) {}
|
|
|
|
/**
|
|
* Read from the queue if it is not empty.
|
|
* @param out The value to read into.
|
|
* @return true if the read succeeded, false if the queue is empty.
|
|
*/
|
|
bool readIfNotEmpty(T& out) {
|
|
std::lock_guard<std::mutex> lock(mutex_);
|
|
if (storage_.empty()) {
|
|
return false;
|
|
}
|
|
out = std::move(storage_.front());
|
|
storage_.pop_front();
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Write to the queue if it is not full.
|
|
* @param in The value to write. For now we only support moveable types.
|
|
* @return true if the write succeeded, false if the queue is full.
|
|
*/
|
|
bool writeIfNotFull(T in) {
|
|
std::lock_guard<std::mutex> lock(mutex_);
|
|
if (storage_.size() == capacity_) {
|
|
return false;
|
|
}
|
|
storage_.push_back(std::move(in));
|
|
return true;
|
|
}
|
|
|
|
private:
|
|
std::mutex mutex_;
|
|
std::deque<T> storage_;
|
|
size_t capacity_;
|
|
};
|
|
} // namespace torch::nativert::detail
|