rtsp_tensorrt/pipeline/input/input_manager.cpp

#include "input_manager.hpp"
#include <chrono>
#include <iostream>

namespace pipeline {

InputManager::InputManager(size_t queue_capacity)
    : running_(true) {
}

InputManager::~InputManager() {
    running_ = false;
    clear();
}

bool InputManager::addSource(const std::string& name, const RtspReader::Config& config, const std::string& url) {
    std::lock_guard<std::mutex> lock(sources_mutex_);
    
    // 检查源是否已存在
    if (source_types_.find(name) != source_types_.end()) {
        return false;
    }
    
    // 创建新的RTSP读取器
    auto reader = std::make_unique<RtspReader>(config);
    if (!reader->connect(url)) {
        return false;
    }
    
    // 创建源状态
    SourceStatus status;
    status.is_connected = true;
    source_status_[name] = status;
    
    // 保存读取器和类型
    rtsp_sources_[name] = std::move(reader);
    source_types_[name] = SourceType::RTSP;
    
    // 创建并启动源线程
    active_threads_++;
    source_threads_[name] = std::make_unique<std::thread>(&InputManager::sourceThread, this, name);
    
    return true;
}

bool InputManager::addVideoSource(const std::string& name, const VideoReader::Config& config, const std::string& path) {
    std::lock_guard<std::mutex> lock(sources_mutex_);
    
    // 检查源是否已存在
    if (source_types_.find(name) != source_types_.end()) {
        return false;
    }
    
    // 创建新的视频读取器
    auto reader = std::make_unique<VideoReader>(config);
    if (!reader->open(path)) {
        return false;
    }
    
    // 创建源状态
    SourceStatus status;
    status.is_connected = true;
    source_status_[name] = status;
    
    // 创建帧队列
    source_queues_[name] = std::make_unique<FrameQueue>(config.buffer_size);
    
    // 保存读取器和类型
    video_sources_[name] = std::move(reader);
    source_types_[name] = SourceType::VIDEO;
    
    // 创建并启动源线程
    active_threads_++;
    source_threads_[name] = std::make_unique<std::thread>(&InputManager::videoSourceThread, this, name);
    
    return true;
}

bool InputManager::removeSource(const std::string& name) {
    std::unique_lock<std::mutex> lock(sources_mutex_);
    
    // 检查源是否存在
    auto type_it = source_types_.find(name);
    if (type_it == source_types_.end()) {
        return false;
    }
    
    // 获取源类型
    SourceType type = type_it->second;
    
    // 停止并等待线程结束
    if (source_threads_.find(name) != source_threads_.end()) {
        auto thread = std::move(source_threads_[name]);
        // 在解锁前先移除源，这样线程就不会访问已删除的资源
        source_threads_.erase(name);
        source_status_.erase(name);
        source_types_.erase(name);
        if (type == SourceType::RTSP) {
            rtsp_sources_.erase(name);
        } else {
            video_sources_.erase(name);
        }
        lock.unlock();
        if (thread && thread->joinable()) {
            thread->join();
        }
    } else {
        // 如果没有线程，直接移除源
        source_threads_.erase(name);
        source_status_.erase(name);
        source_types_.erase(name);
        if (type == SourceType::RTSP) {
            rtsp_sources_.erase(name);
        } else {
            video_sources_.erase(name);
        }
    }
    
    active_threads_--;
    return true;
}

bool InputManager::getNextBatch(std::vector<cv::Mat>& frames, int timeout_ms) {
    std::cout << "Trying to get next batch with timeout " << timeout_ms << "ms" << std::endl;
    
    // 获取所有源的名称
    std::vector<std::string> source_names;
    {
        std::lock_guard<std::mutex> lock(sources_mutex_);
        for (const auto& [name, _] : source_types_) {
            source_names.push_back(name);
        }
    }
    
    if (source_names.empty()) {
        return false;
    }
    
    // 从每个源获取一帧
    bool got_any_frame = false;
    for (const auto& name : source_names) {
        cv::Mat frame;
        auto queue_it = source_queues_.find(name);
        if (queue_it != source_queues_.end() && queue_it->second) {
            if (queue_it->second->pop(frame, std::chrono::milliseconds(timeout_ms))) {
                frames.push_back(frame);
                got_any_frame = true;
                std::cout << "Got frame from source: " << name << std::endl;
            }
        }
    }
    
    std::cout << "Got " << frames.size() << " frames in total" << std::endl;
    return got_any_frame;
}

bool InputManager::getSourceStatus(const std::string& name, SourceStatus& status) {
    std::lock_guard<std::mutex> lock(sources_mutex_);
    auto it = source_status_.find(name);
    if (it == source_status_.end()) {
        return false;
    }
    status = it->second;
    return true;
}

std::vector<std::string> InputManager::getSourceNames() const {
    std::lock_guard<std::mutex> lock(sources_mutex_);
    std::vector<std::string> names;
    names.reserve(source_types_.size());
    for (const auto& [name, _] : source_types_) {
        names.push_back(name);
    }
    return names;
}

size_t InputManager::getSourceCount() const {
    std::lock_guard<std::mutex> lock(sources_mutex_);
    return source_types_.size();
}

void InputManager::clear() {
    std::cout << "Starting to clear input manager..." << std::endl;
    
    // 首先停止所有线程
    running_ = false;
    
    // 关闭所有帧队列并等待线程结束
    {
        std::unique_lock<std::mutex> lock(sources_mutex_);
        
        // 关闭所有队列
        for (auto& [name, queue] : source_queues_) {
            if (queue) {
                queue->close();
            }
        }
        
        // 等待所有线程结束
        for (auto& [name, thread] : source_threads_) {
            std::cout << "Waiting for thread to finish: " << name << std::endl;
            if (thread && thread->joinable()) {
                // 在等待线程时释放锁，避免死锁
                lock.unlock();
                thread->join();
                lock.lock();
            }
        }
        
        // 清理所有资源
        source_threads_.clear();
        rtsp_sources_.clear();
        video_sources_.clear();
        source_types_.clear();
        source_status_.clear();
        source_queues_.clear();
    }
    
    // 重置状态
    active_threads_ = 0;
    running_ = true;
    
    std::cout << "Input manager cleared successfully" << std::endl;
}

void InputManager::sourceThread(const std::string& name) {
    cv::Mat frame;
    SourceStatus status;
    status.is_connected = true;
    status.error_count = 0;
    status.last_error = "";
    
    auto start_time = std::chrono::steady_clock::now();
    int frame_count = 0;
    const int PUSH_TIMEOUT_MS = 1000;
    const int MAX_RETRIES = 10;
    
    std::cout << "Starting RTSP source thread for: " << name << std::endl;
    
    while (running_) {
        // 获取RTSP读取器和帧队列
        std::unique_lock<std::mutex> lock(sources_mutex_);
        auto reader_it = rtsp_sources_.find(name);
        auto queue_it = source_queues_.find(name);
        if (reader_it == rtsp_sources_.end() || queue_it == source_queues_.end()) {
            std::cout << "Source " << name << " not found, exiting thread" << std::endl;
            break;
        }
        auto reader = reader_it->second.get();
        auto& queue = queue_it->second;
        lock.unlock();
        
        // 读取帧
        if (!reader->read(frame)) {
            if (!running_) {
                std::cout << "Source " << name << " stopped" << std::endl;
                break;
            }
            status.error_count++;
            status.last_error = "Failed to read frame";
            updateSourceStatus(name, status);
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            continue;
        }
        
        // 更新状态
        frame_count++;
        auto current_time = std::chrono::steady_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::seconds>(
            current_time - start_time).count();
        if (duration > 0) {
            status.current_fps = static_cast<float>(frame_count) / duration;
        }
        status.frame_count = frame_count;
        updateSourceStatus(name, status);
        
        // 尝试推送帧到队列
        int retries = 0;
        bool pushed = false;
        while (running_ && retries < MAX_RETRIES) {
            if (queue->push(frame, std::chrono::milliseconds(100))) {
                pushed = true;
                break;
            }
            retries++;
            if (retries == 1) {
                std::cout << "Source " << name << ": Queue full, retrying..." << std::endl;
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(10 * retries));
        }
        
        if (!pushed && running_) {
            status.error_count++;
            status.last_error = "Failed to push frame to queue after " + std::to_string(MAX_RETRIES) + " retries";
            updateSourceStatus(name, status);
            std::cout << "Source " << name << ": " << status.last_error << std::endl;
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }
    }
    
    std::cout << "RTSP source thread exiting for: " << name << std::endl;
    active_threads_--;
}

void InputManager::videoSourceThread(const std::string& name) {
    cv::Mat frame;
    SourceStatus status;
    status.is_connected = true;
    status.error_count = 0;
    status.last_error = "";
    
    auto start_time = std::chrono::steady_clock::now();
    auto last_frame_time = start_time;
    int frame_count = 0;
    const int MAX_RETRIES = 10;
    const float TARGET_FPS = 15.0f;
    const float FRAME_TIME_US = 1000000.0f / TARGET_FPS;
    
    std::cout << "Starting video source thread for: " << name << std::endl;
    
    VideoReader* reader = nullptr;
    FrameQueue* queue = nullptr;
    
    // 获取视频读取器和帧队列
    {
        std::unique_lock<std::mutex> lock(sources_mutex_);
        auto reader_it = video_sources_.find(name);
        auto queue_it = source_queues_.find(name);
        if (reader_it == video_sources_.end() || queue_it == source_queues_.end()) {
            std::cout << "Source " << name << " not found, exiting thread" << std::endl;
            return;
        }
        reader = reader_it->second.get();
        queue = queue_it->second.get();
    }
    
    while (running_) {
        // 读取帧
        if (!reader->read(frame)) {
            if (!running_) {
                std::cout << "Source " << name << " stopped" << std::endl;
                break;
            }
            status.error_count++;
            status.last_error = "Failed to read frame";
            updateSourceStatus(name, status);
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            continue;
        }
        
        // 更新状态
        frame_count++;
        auto current_time = std::chrono::steady_clock::now();
        
        // 控制帧处理速率
        auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(
            current_time - last_frame_time).count();
        if (elapsed < FRAME_TIME_US) {
            auto sleep_time = std::chrono::microseconds(
                static_cast<long long>(FRAME_TIME_US - elapsed));
            std::this_thread::sleep_for(sleep_time);
            current_time = std::chrono::steady_clock::now();
        }
        last_frame_time = current_time;
        
        // 更新FPS
        auto duration = std::chrono::duration_cast<std::chrono::seconds>(
            current_time - start_time).count();
        if (duration > 0) {
            status.current_fps = static_cast<float>(frame_count) / duration;
        }
        status.frame_count = frame_count;
        updateSourceStatus(name, status);
        
        // 检查是否需要退出
        if (!running_) {
            break;
        }
        
        // 尝试推送帧到队列
        int retries = 0;
        bool pushed = false;
        while (running_ && retries < MAX_RETRIES) {
            if (queue->push(frame, std::chrono::milliseconds(100))) {
                pushed = true;
                break;
            }
            retries++;
            if (retries == 1) {
                std::cout << "Source " << name << ": Queue full, retrying..." << std::endl;
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(10 * retries));
        }
        
        if (!pushed && running_) {
            status.error_count++;
            status.last_error = "Failed to push frame to queue after " + std::to_string(MAX_RETRIES) + " retries";
            updateSourceStatus(name, status);
            std::cout << "Source " << name << ": " << status.last_error << std::endl;
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }
    }
    
    std::cout << "Video source thread exiting for: " << name << std::endl;
    active_threads_--;
}

void InputManager::updateSourceStatus(const std::string& name, const SourceStatus& status) {
    std::lock_guard<std::mutex> lock(sources_mutex_);
    source_status_[name] = status;
}

} // namespace pipeline