rtsp_tensorrt/pipeline/input/frame_queue.cpp

#include "frame_queue.hpp"

namespace pipeline {

FrameQueue::FrameQueue(size_t max_size)
    : max_size_(max_size) {
}

FrameQueue::~FrameQueue() {
    close();
}

void FrameQueue::close() {
    std::unique_lock<std::mutex> lock(mutex_);
    is_closed_ = true;
    // 通知所有等待的线程
    not_empty_.notify_all();
    not_full_.notify_all();
}

bool FrameQueue::push(const cv::Mat& frame, std::chrono::milliseconds timeout) {
    std::unique_lock<std::mutex> lock(mutex_);
    
    // 如果队列已关闭，直接返回
    if (is_closed_) {
        return false;
    }
    
    // 等待队列未满或超时
    bool not_full = not_full_.wait_for(lock, timeout, [this] {
        return queue_.size() < max_size_ || is_closed_;
    });
    
    if (!not_full || is_closed_) {
        return false;
    }
    
    // 深拷贝帧数据并推入队列
    queue_.push(frame.clone());
    
    // 通知可能在等待的消费者
    not_empty_.notify_one();
    
    return true;
}

bool FrameQueue::pop(cv::Mat& frame, std::chrono::milliseconds timeout) {
    std::unique_lock<std::mutex> lock(mutex_);
    
    // 等待队列非空或超时
    bool not_empty = not_empty_.wait_for(lock, timeout, [this] {
        return !queue_.empty() || is_closed_;
    });
    
    if (!not_empty || (queue_.empty() && is_closed_)) {
        return false;
    }
    
    // 获取队首帧
    frame = queue_.front().clone();  // 确保深拷贝
    queue_.pop();
    
    // 通知可能在等待的生产者
    not_full_.notify_all();  // 通知所有等待的生产者
    
    return true;
}

void FrameQueue::clear() {
    std::unique_lock<std::mutex> lock(mutex_);
    std::queue<cv::Mat> empty;
    std::swap(queue_, empty);
    not_full_.notify_all();  // 通知所有等待的生产者
}

bool FrameQueue::empty() const {
    std::unique_lock<std::mutex> lock(mutex_);
    return queue_.empty();
}

bool FrameQueue::full() const {
    std::unique_lock<std::mutex> lock(mutex_);
    return queue_.size() >= max_size_;
}

size_t FrameQueue::size() const {
    std::unique_lock<std::mutex> lock(mutex_);
    return queue_.size();
}

} // namespace pipeline