145 lines
3.4 KiB
C++
145 lines
3.4 KiB
C++
/**
|
|
* Copyright (c) 2016-present, Facebook, Inc.
|
|
* All rights reserved.
|
|
*
|
|
* This source code is licensed under the BSD-style license found in the
|
|
* LICENSE file in the root directory of this source tree. An additional grant
|
|
* of patent rights can be found in the PATENTS file in the same directory.
|
|
*/
|
|
#pragma once
|
|
|
|
#include "utils/Buffer.h"
|
|
|
|
#include <atomic>
|
|
#include <cassert>
|
|
#include <condition_variable>
|
|
#include <cstddef>
|
|
#include <functional>
|
|
#include <mutex>
|
|
#include <queue>
|
|
|
|
namespace pzstd {
|
|
|
|
/// Unbounded thread-safe work queue.
|
|
template <typename T>
|
|
class WorkQueue {
|
|
// Protects all member variable access
|
|
std::mutex mutex_;
|
|
std::condition_variable cv_;
|
|
|
|
std::queue<T> queue_;
|
|
bool done_;
|
|
|
|
public:
|
|
/// Constructs an empty work queue.
|
|
WorkQueue() : done_(false) {}
|
|
|
|
/**
|
|
* Push an item onto the work queue. Notify a single thread that work is
|
|
* available. If `finish()` has been called, do nothing and return false.
|
|
*
|
|
* @param item Item to push onto the queue.
|
|
* @returns True upon success, false if `finish()` has been called. An
|
|
* item was pushed iff `push()` returns true.
|
|
*/
|
|
bool push(T item) {
|
|
{
|
|
std::lock_guard<std::mutex> lock(mutex_);
|
|
if (done_) {
|
|
return false;
|
|
}
|
|
queue_.push(std::move(item));
|
|
}
|
|
cv_.notify_one();
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Attempts to pop an item off the work queue. It will block until data is
|
|
* available or `finish()` has been called.
|
|
*
|
|
* @param[out] item If `pop` returns `true`, it contains the popped item.
|
|
* If `pop` returns `false`, it is unmodified.
|
|
* @returns True upon success. False if the queue is empty and
|
|
* `finish()` has been called.
|
|
*/
|
|
bool pop(T& item) {
|
|
std::unique_lock<std::mutex> lock(mutex_);
|
|
while (queue_.empty() && !done_) {
|
|
cv_.wait(lock);
|
|
}
|
|
if (queue_.empty()) {
|
|
assert(done_);
|
|
return false;
|
|
}
|
|
item = std::move(queue_.front());
|
|
queue_.pop();
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Promise that `push()` won't be called again, so once the queue is empty
|
|
* there will never any more work.
|
|
*/
|
|
void finish() {
|
|
{
|
|
std::lock_guard<std::mutex> lock(mutex_);
|
|
assert(!done_);
|
|
done_ = true;
|
|
}
|
|
cv_.notify_all();
|
|
}
|
|
|
|
/// Blocks until `finish()` has been called (but the queue may not be empty).
|
|
void waitUntilFinished() {
|
|
std::unique_lock<std::mutex> lock(mutex_);
|
|
while (!done_) {
|
|
cv_.wait(lock);
|
|
// If we were woken by a push, we need to wake a thread waiting on pop().
|
|
if (!done_) {
|
|
lock.unlock();
|
|
cv_.notify_one();
|
|
lock.lock();
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
/// Work queue for `Buffer`s that knows the total number of bytes in the queue.
|
|
class BufferWorkQueue {
|
|
WorkQueue<Buffer> queue_;
|
|
std::atomic<std::size_t> size_;
|
|
|
|
public:
|
|
BufferWorkQueue() : size_(0) {}
|
|
|
|
void push(Buffer buffer) {
|
|
size_.fetch_add(buffer.size());
|
|
queue_.push(std::move(buffer));
|
|
}
|
|
|
|
bool pop(Buffer& buffer) {
|
|
bool result = queue_.pop(buffer);
|
|
if (result) {
|
|
size_.fetch_sub(buffer.size());
|
|
}
|
|
return result;
|
|
}
|
|
|
|
void finish() {
|
|
queue_.finish();
|
|
}
|
|
|
|
/**
|
|
* Blocks until `finish()` has been called.
|
|
*
|
|
* @returns The total number of bytes of all the `Buffer`s currently in the
|
|
* queue.
|
|
*/
|
|
std::size_t size() {
|
|
queue_.waitUntilFinished();
|
|
return size_.load();
|
|
}
|
|
};
|
|
}
|