recode添加colse功能和mq功能

Hu Chunming
1 parent beec83ee
Showing 13 changed files with 145 additions and 5 deletions
src/ai_platform/MultiSourceProcess.cpp
src/decoder/dvpp/DvppDecoder.cpp
src/decoder/dvpp/DvppDecoder.h
src/decoder/dvpp/DvppDecoderApi.cpp
src/decoder/dvpp/DvppDecoderApi.h
src/decoder/dvpp/FFRecoder.cpp
src/decoder/dvpp/FFRecoder.h
src/decoder/dvpp/FFRecoderTaskManager.cpp
src/decoder/dvpp/FFRecoderTaskManager.h
src/decoder/interface/AbstractDecoder.h
src/decoder/interface/DecoderManager.cpp
src/decoder/interface/DecoderManager.h
src/decoder/test_recoder.cpp
@@ -303,6 +303,10 @@ bool CMultiSourceProcess::AddTask(task_param _cur_task_param){
       return false;
     }
+#ifdef POST_USE_RABBITMQ
+    pDecManager->set_mq_callback(config.name, std::bind(&mq::Manager::publish, mq_manager_, mq_type, std::placeholders::_1, true));
+#endif
+
     // 人车物跟踪
     if (task_has_vpt_algor(task_id))
       vpt_process.addTaskTracker(task_id, 1, 1, skip_frame_);
@@ -245,6 +245,8 @@ void DvppDecoder::close(){
 	if(m_read_thread != 0){
 		pthread_join(m_read_thread,0);
 	}
+
+    m_recoderManager.close();
 }
 void DvppDecoder::setPostDecArg(const void* postDecArg){
@@ -886,4 +888,8 @@ void DvppDecoder::release_dvpp(){
 void DvppDecoder::doRecode(RecoderInfo& recoderInfo) {
     m_recoderManager.create_recode_task2(recoderInfo);
+}
+
+void DvppDecoder::set_mq_callback(mq_callback_t cb) {
+    m_recoderManager.set_mq_callback(cb);
 }
 \ No newline at end of file
@@ -54,6 +54,8 @@ public:
     void doRecode(RecoderInfo& recoderInfo);
+    void set_mq_callback(mq_callback_t cb);
+
 public:
     void doVdppVdecCallBack(acldvppStreamDesc *input, acldvppPicDesc *output, unsigned long long frame_nb);
     void doProcessReport();
@@ -136,4 +136,10 @@ void DvppDecoderApi::doRecode(RecoderInfo&amp; recoderInfo) {
     if(m_pDecoder != nullptr){
         return m_pDecoder->doRecode(recoderInfo);
     }
+}
+
+void DvppDecoderApi::set_mq_callback(std::function<bool(const char *msg)> mq_publish) {
+    if(m_pDecoder != nullptr){
+        return m_pDecoder->set_mq_callback(mq_publish);
+    }
 }
 \ No newline at end of file
@@ -41,6 +41,8 @@ public:
     void setFinishedDecArg(const void* finishedDecArg);
     void doRecode(RecoderInfo& recoderInfo);
+
+    void set_mq_callback(std::function<bool(const char *msg)> mq_publish);
 private:
     DvppDecoder* m_pDecoder;
 };
 \ No newline at end of file
@@ -123,7 +123,7 @@ bool FFRecoder::init(int w, int h, AVRational time_base, AVCodecContext* avctx, 
 bool FFRecoder::init(AVRational time_base, AVCodecContext* avctx, const char* outfile_name) {
-	codec_ctx_ = new AVCodecContext();
+	codec_ctx_ = (AVCodecContext*)av_malloc(sizeof(AVCodecContext));
 	avcodec_copy_context(codec_ctx_, avctx);
 	codec_ctx_->time_base = time_base;
@@ -244,6 +244,11 @@ void FFRecoder::update_pts(AVPacket* pkt) {
 bool FFRecoder::write_pkt(AVPacket *pkt) {
 	char errbuf[64]{ 0 };
+	// frame_number++;
+	// pkt->pts = av_rescale_q(frame_number, codec_ctx_->time_base, out_stream_->time_base);
+	// pkt->dts = pkt->pts;
+	// pkt->duration = av_rescale_q(1, codec_ctx_->time_base, out_stream_->time_base);
+
 	av_packet_rescale_ts(pkt, codec_ctx_->time_base, out_stream_->time_base);
 	pkt->stream_index = out_stream_->index;
 	update_pts(pkt);
@@ -49,4 +49,6 @@ private:
 	bool bFirstFrame;
 	int64_t last_src_pts;
 	int64_t last_pts;
+
+	int64_t frame_number{0};
 };
 \ No newline at end of file
@@ -56,6 +56,11 @@ bool FFRecoderTaskManager::init(AVRational time_base, AVCodecContext* avctx){
 }
 void FFRecoderTaskManager::cache_pkt(AVPacket* pkt, long long frame_nb){
+    if(m_bExit) {
+        // 任务退出了就不再缓存数据了
+        return;
+    }
+
     std::lock_guard<std::mutex> l_pkt(m_pkt_list_mtx);
     // 考虑到一个AVPacket中的数据并不很大，为减少与解码模块的耦合度，方便管理，这里做一个clone
@@ -131,6 +136,14 @@ list&lt;DataPacket*&gt;::iterator FFRecoderTaskManager::getStartIterator(unsigned long
         return it_first;
     }
+    // auto it_second = m_pkt_list.begin();
+    // for(;it_second != m_pkt_list.end(); it_second++) {
+    //     DataPacket* dataPkt = *it_second;
+    //     if (dataPkt->frame_nb >= start_frame_nb){
+    //         return it_second;
+    //     }
+    // }
+
     auto it_second = m_pkt_list.begin();
     for(;it_second != m_pkt_list.end(); it_second++) {
         DataPacket* dataPkt = *it_second;
@@ -179,6 +192,10 @@ void FFRecoderTaskManager::create_recode_task(AVRational time_base, AVCodecConte
 }
 void FFRecoderTaskManager::create_recode_task2(RecoderInfo& recoderInfo) {
+    if(m_bExit) {
+        // 任务退出了就不再接收录制任务
+        return;
+    }
     save_intask_recoderinfo(recoderInfo);
 }
@@ -219,8 +236,12 @@ void FFRecoderTaskManager::recode_thread(RecodeParam recodeParam){
 }
 void FFRecoderTaskManager::recode_thread2() {
-
+    LOG_INFO("recode_thread2 start...");
     while(true) {
+        if(m_bExit) {
+            break;
+        }
+
         m_recoderinfo_list_mtx.lock();
         if(m_recoderinfo_list.size() <= 0){
             m_recoderinfo_list_mtx.unlock();
@@ -258,6 +279,7 @@ void FFRecoderTaskManager::recode_thread2() {
         bool bInit = ffrecoder.init(m_time_base, m_avctx, file_name.c_str());
         if (!bInit) {
             LOG_ERROR("ffrecoder init error : {} {} {}", recoderinfo.task_id, recoderinfo.object_id, recoderinfo.frame_nb);
+            ffrecoder.uninit();
             continue;
         }
         LOG_INFO("record start, pkt_list size: {} id: {}", m_pkt_list.size(), id);
@@ -277,7 +299,34 @@ void FFRecoderTaskManager::recode_thread2() {
         // ffrecoder.flush();
         ffrecoder.uninit();
+        // 发送mq消息
+        if(mq_publish_func) {
+            mq_publish_func(recoderinfo.mq_info.c_str());
+        }
+
         LOG_INFO("record end, total save: {} start_frame_nb: {} end_frame_nb: {}  file_path: {}", count, start_frame_nb, end_frame_nb, file_name);
     }
+    LOG_INFO("recode_thread2 end.");
+}
+
+void FFRecoderTaskManager::close() {
+    m_bExit = true;
+
+    if (m_recoder_thread) {
+        m_recoder_thread->join();
+        m_recoder_thread = nullptr;
+    }
+
+    // 清空数据
+    while(!m_pkt_list.empty()) {
+        DataPacket* dataPkt = m_pkt_list.front();
+        delete dataPkt;
+        dataPkt = nullptr;
+        m_pkt_list.pop_front();
+    }
+}
+
+void FFRecoderTaskManager::set_mq_callback(mq_callback_t cb) {
+    mq_publish_func = cb;
 }
 \ No newline at end of file
@@ -18,6 +18,8 @@ struct RecodeParam {
     AVCodecContext* avctx;
 };
+typedef std::function<bool(const char *msg)> mq_callback_t;
+
 class FFRecoderTaskManager {
 public:
     FFRecoderTaskManager();
@@ -29,6 +31,10 @@ public:
     bool init(AVRational time_base, AVCodecContext* avctx);
     void create_recode_task2(RecoderInfo& recoderInfo);
+    void close();
+
+    void set_mq_callback(mq_callback_t cb);
+
 public:
     void recode_thread(RecodeParam param);
     list<DataPacket*>::iterator getStartIterator(unsigned long long frame_nb);
@@ -59,5 +65,7 @@ private:
     AVRational m_time_base;
     AVCodecContext* m_avctx;
-    thread* m_recoder_thread {nullptr};
+    thread* m_recoder_thread{nullptr};
+
+    mq_callback_t mq_publish_func;
 };
 \ No newline at end of file
@@ -5,6 +5,8 @@
 #include "../../ai_platform/common_header.h"
+#include <functional>  
+
 using namespace std;
 class AbstractDecoder{
@@ -42,6 +44,8 @@ public:
     virtual void doRecode(RecoderInfo& recoderInfo) = 0;
+    virtual void set_mq_callback(std::function<bool(const char *msg)> mq_publish) = 0;
+
 public:
     bool isSnapTime();
@@ -536,4 +536,22 @@ void DecoderManager::doRecode(RecoderInfo&amp; recoderInfo) {
     LOG_ERROR("没有找到name为{}的解码器",name);
     return;
+}
+
+void DecoderManager::set_mq_callback(const string name, std::function<bool(const char *msg)> mq_publish) {
+    if (name.empty()){
+        LOG_ERROR("name 为空!");
+        return;
+    }
+
+    std::lock_guard<std::mutex> l(m_mutex);
+
+    auto dec = decoderMap.find(name);
+    if (dec != decoderMap.end()){
+        dec->second->set_mq_callback(mq_publish);
+        return;
+    }
+    
+    LOG_ERROR("没有找到name为{}的解码器",name);
+    return;
 }
 \ No newline at end of file
@@ -275,6 +275,8 @@ public:
 	void doRecode(RecoderInfo& recoderInfo);
+	void set_mq_callback(const string name, std::function<bool(const char *msg)> mq_publish);
+
 private:
     DecoderManager(){}
@@ -15,13 +15,18 @@ deque&lt;DeviceMemory*&gt; m_RgbDataList;
 mutex m_DataListMtx; 
 thread* m_pAlgorthimThread{nullptr};
+thread* m_recodeThread{nullptr};
 bool m_bfinish{false};
 int m_devId = 0;
 const char* task_id = "test0";
 int skip_frame_ = 5;
 int m_batch_size = 20;
+deque<RecoderInfo> m_recoderinfo_queue;
+mutex m_recoderinfo_queue_mtx;
+
 void algorthim_process_thread();
+void recode_thread();
 void algorthim_face_detect(vector<DeviceMemory*> vec_gpuMem);
 void post_decod_cbk(const void * userPtr, DeviceMemory* devFrame){
@@ -90,6 +95,12 @@ int main(){
     }
     , nullptr);
+    m_recodeThread = new thread([](void* arg) {
+      recode_thread();
+      return (void*)0;
+    }
+    , nullptr);
+
     pDecManager->startDecodeByName(config.name);
     while (getchar() != 'q');
@@ -162,10 +173,31 @@ void algorthim_face_detect(vector&lt;DeviceMemory*&gt; vec_gpuMem) {
     recoderInfo.object_id = std::to_string(obj_id);
     recoderInfo.recoderDir = "./res/recode";
     recoderInfo.frame_nb = mem->getFrameNb();
-    DecoderManager* pDecManager = DecoderManager::getInstance();
-    pDecManager->doRecode(recoderInfo);
+
+    m_recoderinfo_queue_mtx.lock();
+    m_recoderinfo_queue.push_back(recoderInfo);
+    m_recoderinfo_queue_mtx.unlock();
     obj_id++;
   }
+}
+
+void recode_thread() {
+  while(true) {
+
+    m_recoderinfo_queue_mtx.lock();
+    if(m_recoderinfo_queue.size() <= 0) {
+      m_recoderinfo_queue_mtx.unlock();
+      std::this_thread::sleep_for(std::chrono::milliseconds(5));
+      continue;
+    }
+
+    RecoderInfo info = m_recoderinfo_queue.front();
+    m_recoderinfo_queue.pop_front();
+    m_recoderinfo_queue_mtx.unlock();
+
+    DecoderManager* pDecManager = DecoderManager::getInstance();
+    pDecManager->doRecode(info);
+  }
 }
 \ No newline at end of file