wehub-resource-sync
77 lines
1.5 KiB
C++
77 lines
1.5 KiB
C++
//
|
|
// BlockingQueue.hpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/11/19.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#ifndef BlockingQueue_hpp
|
|
#define BlockingQueue_hpp
|
|
#include <MNN/MNNDefine.h>
|
|
#include <condition_variable>
|
|
#include <mutex>
|
|
#include <queue>
|
|
|
|
namespace MNN {
|
|
namespace Train {
|
|
|
|
template <typename T>
|
|
class BlockingQueue {
|
|
public:
|
|
BlockingQueue() = default;
|
|
BlockingQueue(size_t maxSize) : mMaxSize(maxSize) {
|
|
}
|
|
|
|
bool isFull() {
|
|
return mQueue.size() == mMaxSize;
|
|
}
|
|
|
|
bool isEmpty() {
|
|
return mQueue.empty();
|
|
}
|
|
|
|
void push(T value) {
|
|
{
|
|
std::unique_lock<std::mutex> lock(mMutex);
|
|
mCondVar.wait(lock, [&] { return !isFull(); });
|
|
MNN_ASSERT(!isFull());
|
|
mQueue.push(std::move(value));
|
|
lock.unlock();
|
|
}
|
|
mCondVar.notify_one();
|
|
}
|
|
|
|
T pop() {
|
|
std::unique_lock<std::mutex> lock(mMutex);
|
|
mCondVar.wait(lock, [&] { return !isEmpty(); });
|
|
MNN_ASSERT(!isEmpty());
|
|
T value = mQueue.front();
|
|
mQueue.pop();
|
|
mCondVar.notify_one();
|
|
lock.unlock();
|
|
|
|
return std::move(value);
|
|
}
|
|
|
|
size_t clear() {
|
|
std::lock_guard<std::mutex> lock(mMutex);
|
|
const auto size = mQueue.size();
|
|
while (!isEmpty()) {
|
|
mQueue.pop();
|
|
}
|
|
return size;
|
|
}
|
|
|
|
private:
|
|
size_t mMaxSize;
|
|
std::queue<T> mQueue;
|
|
std::mutex mMutex;
|
|
std::condition_variable_any mCondVar;
|
|
};
|
|
|
|
} // namespace Train
|
|
} // namespace MNN
|
|
|
|
#endif // BlockingQueue_hpp
|