SafetyQueue.hpp

 
#ifndef EMX_SAFETYQUEUE_HPP
#define EMX_SAFETYQUEUE_HPP
 
#include <list>
#include <mutex>
#include <condition_variable>
#include <functional>
#include <atomic>
 
namespace Emx {
template<typename T>
class SafetyQueue {
    public:
        using OnCloned = std::function<T(const T& data)>;
 
        using OnFreed = std::function<void(T& data)>;
 
        SafetyQueue(int total = 3, 
            OnCloned clonedCb = nullptr, 
            OnFreed freedCb = nullptr) 
            : m_total(total), m_clonedCb(clonedCb), m_freedCb(freedCb), m_exit(false) {
        }
        
        ~SafetyQueue() {
            Clear();
        }
 
        void Put(const T& data) {
            std::unique_lock<std::mutex> locker(m_mtx);
            T clonedData = (m_clonedCb) ? m_clonedCb(data) : data;
            if (clonedData == nullptr) {
                return;
            }
            m_queue.emplace_back(std::move(clonedData));
 
            while ((int)m_queue.size() > m_total) {
                if (m_freedCb != nullptr) {
                    m_freedCb(m_queue.front());
                }
                m_queue.pop_front();
            }
 
            // 如果存在等待线程，通知一个线程继续执行
            if (!m_waitQueue.empty()) {
                auto wait = m_waitQueue.front();
                m_waitQueue.pop_front();
                (*(wait.first)) = false;
                wait.second->notify_one();
            }
        }
 
        T Get() {
            std::unique_lock<std::mutex> locker(m_mtx);
            while (m_queue.empty()) {
                // 队列为空的时候，线程阻塞等待队列非空状态唤醒。只有当原子位复位才退出等待
                if (m_exit) {
                    return nullptr;
                }
                std::shared_ptr<bool> is_empty = std::make_shared<bool>(true);
                std::shared_ptr<std::condition_variable> cv = std::make_shared<std::condition_variable>();
                m_waitQueue.emplace_back(std::make_pair(is_empty, cv));
                cv->wait(locker, [is_empty, this]() { return !(*is_empty) || m_exit; });
            }
            // 处理stop退出场景，队列被清空的场景
            if (m_exit && m_queue.empty()) {
                return nullptr;
            }
            T data = std::move(m_queue.front());
            m_queue.pop_front();
            return data;
        }
 
        int Size() {
            std::unique_lock<std::mutex> locker(m_mtx);
            return m_queue.size();
        }
 
        void Clear() {
            std::unique_lock<std::mutex> locker(m_mtx);
            if (m_freedCb != nullptr) {
                for(auto &data : m_queue) {
                    if (data != nullptr) {
                        m_freedCb(data);
                    }
                }
            }
            m_queue.clear();
            m_waitQueue.clear();
        }
 
        void Stop() {
            std::unique_lock<std::mutex> locker(m_mtx);
            m_exit = true;
        
            for (auto& wait : m_waitQueue) {
                *(wait.first) = false;
                wait.second->notify_one();
            }
            m_waitQueue.clear();
        }
 
        void SetTotal(int total) {
            std::unique_lock<std::mutex> locker(m_mtx);
            m_total = total;
        }
 
    private:
        int m_total;
        std::mutex m_mtx;
        std::list<T> m_queue;
        std::list<std::pair<std::shared_ptr<bool>, std::shared_ptr<std::condition_variable>>> m_waitQueue;
        OnCloned m_clonedCb;
        OnFreed m_freedCb;
        std::atomic<bool> m_exit;
};
}
 
#endif // EMX_SAFETYQUEUE_HPP