From 2a66366be4c2715870e4859fd5a5db6e8a9dc00a Mon Sep 17 00:00:00 2001
From: chenmengzheAAA <123789350+chenmengzheAAA@users.noreply.github.com>
Date: 星期四, 14 九月 2023 19:00:17 +0800
Subject: [PATCH] Merge pull request #956 from alibaba-damo-academy/chenmengzheAAA-patch-4

---
 funasr/runtime/websocket/websocket-server-2pass.cpp |  190 +++++++++++++++++++++++++++++++----------------
 1 files changed, 124 insertions(+), 66 deletions(-)

diff --git a/funasr/runtime/websocket/websocket-server-2pass.cpp b/funasr/runtime/websocket/websocket-server-2pass.cpp
index 50f0edc..028879a 100644
--- a/funasr/runtime/websocket/websocket-server-2pass.cpp
+++ b/funasr/runtime/websocket/websocket-server-2pass.cpp
@@ -15,7 +15,7 @@
 #include <thread>
 #include <utility>
 #include <vector>
-
+#include <chrono>
 context_ptr WebSocketServer::on_tls_init(tls_mode mode,
                                          websocketpp::connection_hdl hdl,
                                          std::string& s_certfile,
@@ -83,29 +83,30 @@
   scoped_lock guard(thread_lock);
   if(!tpass_online_handle){
 	  LOG(INFO) << "tpass_online_handle  is free, return";
+	  msg["access_num"]=(int)msg["access_num"]-1;
 	  return;
   }
-  FUNASR_RESULT Result = nullptr;
-  int asr_mode_ = 2;
-  if (msg.contains("mode")) {
-    std::string modeltype = msg["mode"];
-    if (modeltype == "offline") {
-      asr_mode_ = 0;
-    } else if (modeltype == "online") {
-      asr_mode_ = 1;
-    } else if (modeltype == "2pass") {
+  try {
+    FUNASR_RESULT Result = nullptr;
+    int asr_mode_ = 2;
+    if (msg.contains("mode")) {
+      std::string modeltype = msg["mode"];
+      if (modeltype == "offline") {
+        asr_mode_ = 0;
+      } else if (modeltype == "online") {
+        asr_mode_ = 1;
+      } else if (modeltype == "2pass") {
+        asr_mode_ = 2;
+      }
+    } else {
+      // default value
+      msg["mode"] = "2pass";
       asr_mode_ = 2;
     }
-  } else {
-    // default value
-    msg["mode"] = "2pass";
-    asr_mode_ = 2;
-  }
 
-  try {
     // loop to send chunk_size 800*2 data to asr engine.   TODO: chunk_size need
     // get from client
-    while (buffer.size() >= 800 * 2) {
+    while (buffer.size() >= 800 * 2 && !msg["is_eof"]) {
       std::vector<char> subvector = {buffer.begin(), buffer.begin() + 800 * 2};
       buffer.erase(buffer.begin(), buffer.begin() + 800 * 2);
 
@@ -117,10 +118,13 @@
                                        msg["wav_format"], (ASR_TYPE)asr_mode_);
 
         } else {
+          msg["access_num"]=(int)msg["access_num"]-1;
           return;
         }
       } catch (std::exception const& e) {
         LOG(ERROR) << e.what();
+        msg["access_num"]=(int)msg["access_num"]-1;
+        return;
       }
       if (Result) {
         websocketpp::lib::error_code ec;
@@ -139,7 +143,7 @@
         FunASRFreeResult(Result);
       }
     }
-    if (is_final) {
+    if (is_final && !msg["is_eof"]) {
       try {
         if (tpass_online_handle) {
           Result = FunTpassInferBuffer(tpass_handle, tpass_online_handle,
@@ -147,10 +151,13 @@
                                        is_final, msg["audio_fs"],
                                        msg["wav_format"], (ASR_TYPE)asr_mode_);
         } else {
+          msg["access_num"]=(int)msg["access_num"]-1;	 
           return;
         }
       } catch (std::exception const& e) {
         LOG(ERROR) << e.what();
+        msg["access_num"]=(int)msg["access_num"]-1;
+        return;
       }
       if(punc_cache.size()>0){
         for (auto& vec : punc_cache) {
@@ -170,35 +177,53 @@
                         websocketpp::frame::opcode::text, ec);
         }
         FunASRFreeResult(Result);
+      }else{
+        if(msg["wav_format"] != "pcm" && msg["wav_format"] != "PCM"){
+          websocketpp::lib::error_code ec;
+          nlohmann::json jsonresult;
+          jsonresult["text"] = "ERROR. Real-time transcription service ONLY SUPPORT wav_format pcm.";
+          jsonresult["wav_name"] = wav_name;
+          jsonresult["is_final"] = true;
+          if (is_ssl) {
+            wss_server_->send(hdl, jsonresult.dump(),
+                              websocketpp::frame::opcode::text, ec);
+          } else {
+            server_->send(hdl, jsonresult.dump(),
+                          websocketpp::frame::opcode::text, ec);
+          }
+        }
       }
     }
 
   } catch (std::exception const& e) {
     std::cerr << "Error: " << e.what() << std::endl;
   }
+  msg["access_num"]=(int)msg["access_num"]-1;
+ 
 }
 
 void WebSocketServer::on_open(websocketpp::connection_hdl hdl) {
   scoped_lock guard(m_lock);     // for threads safty
-  check_and_clean_connection();  // remove closed connection
+  try{
+    std::shared_ptr<FUNASR_MESSAGE> data_msg =
+        std::make_shared<FUNASR_MESSAGE>();  // put a new data vector for new
+                                            // connection
+    data_msg->samples = std::make_shared<std::vector<char>>();
+    data_msg->thread_lock = std::make_shared<websocketpp::lib::mutex>();  
 
-  std::shared_ptr<FUNASR_MESSAGE> data_msg =
-      std::make_shared<FUNASR_MESSAGE>();  // put a new data vector for new
-                                           // connection
-  data_msg->samples = std::make_shared<std::vector<char>>();
-  data_msg->thread_lock = std::make_shared<websocketpp::lib::mutex>();  
+    data_msg->msg = nlohmann::json::parse("{}");
+    data_msg->msg["wav_format"] = "pcm";
+    data_msg->msg["audio_fs"] = 16000;
+    data_msg->msg["access_num"] = 0; // the number of access for this object, when it is 0, we can free it saftly
+    data_msg->msg["is_eof"]=false; // if this connection is closed
+    data_msg->punc_cache =
+        std::make_shared<std::vector<std::vector<std::string>>>(2);
+  	data_msg->strand_ =	std::make_shared<asio::io_context::strand>(io_decoder_);
 
-  data_msg->msg = nlohmann::json::parse("{}");
-  data_msg->msg["wav_format"] = "pcm";
-  data_msg->msg["audio_fs"] = 16000;
-  data_msg->punc_cache =
-      std::make_shared<std::vector<std::vector<std::string>>>(2);
-  // std::vector<int> chunk_size = {5, 10, 5};  //TODO, need get from client
-  // FUNASR_HANDLE tpass_online_handle =
-  //     FunTpassOnlineInit(tpass_handle, chunk_size);
-  // data_msg->tpass_online_handle = tpass_online_handle;
-  data_map.emplace(hdl, data_msg);
-  LOG(INFO) << "on_open, active connections: " << data_map.size();
+    data_map.emplace(hdl, data_msg);
+  }catch (std::exception const& e) {
+    std::cerr << "Error: " << e.what() << std::endl;
+  }
 }
 
 void remove_hdl(
@@ -216,47 +241,75 @@
   // scoped_lock guard_decoder(*(data_msg->thread_lock));  //wait for do_decoder
   // finished and avoid access freed tpass_online_handle
   unique_lock guard_decoder(*(data_msg->thread_lock));
-  if (data_msg->tpass_online_handle) {
-    LOG(INFO) << "----------------FunTpassOnlineUninit----------------------";
+  if (data_msg->msg["access_num"]==0 && data_msg->msg["is_eof"]==true) {
     FunTpassOnlineUninit(data_msg->tpass_online_handle);
     data_msg->tpass_online_handle = nullptr;
+	  data_map.erase(hdl);
   }
-  
  
   guard_decoder.unlock();
-  data_map.erase(hdl);  // remove data vector when  connection is closed
 }
 
 void WebSocketServer::on_close(websocketpp::connection_hdl hdl) {
   scoped_lock guard(m_lock);
-  remove_hdl(hdl, data_map);
-  LOG(INFO) << "on_close, active connections: " << data_map.size();
+  std::shared_ptr<FUNASR_MESSAGE> data_msg = nullptr;
+  auto it_data = data_map.find(hdl);
+  if (it_data != data_map.end()) {
+    data_msg = it_data->second;
+  } else {
+    return;
+  }
+  unique_lock guard_decoder(*(data_msg->thread_lock));
+  data_msg->msg["is_eof"]=true;
+  guard_decoder.unlock();
 }
-
+ 
 // remove closed connection
 void WebSocketServer::check_and_clean_connection() {
-  std::vector<websocketpp::connection_hdl> to_remove;  // remove list
-  auto iter = data_map.begin();
-  while (iter != data_map.end()) {  // loop to find closed connection
-    websocketpp::connection_hdl hdl = iter->first;
-
-    if (is_ssl) {
-      wss_server::connection_ptr con = wss_server_->get_con_from_hdl(hdl);
-      if (con->get_state() != 1) {  // session::state::open ==1
-        to_remove.push_back(hdl);
+  while(true){
+    std::this_thread::sleep_for(std::chrono::milliseconds(5000));
+    std::vector<websocketpp::connection_hdl> to_remove;  // remove list
+    auto iter = data_map.begin();
+    while (iter != data_map.end()) {  // loop to find closed connection
+      websocketpp::connection_hdl hdl = iter->first;
+      try{
+        if (is_ssl) {
+          wss_server::connection_ptr con = wss_server_->get_con_from_hdl(hdl);
+          if (con->get_state() != 1) {  // session::state::open ==1
+            to_remove.push_back(hdl);
+          }
+        } else {
+          server::connection_ptr con = server_->get_con_from_hdl(hdl);
+          if (con->get_state() != 1) {  // session::state::open ==1
+            to_remove.push_back(hdl);
+          }
+        }
       }
-    } else {
-      server::connection_ptr con = server_->get_con_from_hdl(hdl);
-      if (con->get_state() != 1) {  // session::state::open ==1
+      catch (std::exception const &e)
+      {
+        // if connection is close, we set is_eof = true
+        std::shared_ptr<FUNASR_MESSAGE> data_msg = nullptr;
+        auto it_data = data_map.find(hdl);
+        if (it_data != data_map.end()) {
+          data_msg = it_data->second;
+        } else {
+            continue;
+        }
+        unique_lock guard_decoder(*(data_msg->thread_lock));
+        data_msg->msg["is_eof"]=true;
+        guard_decoder.unlock();
         to_remove.push_back(hdl);
+        LOG(INFO)<<"connection is closed: "<<e.what();
+        
+      }
+      iter++;
+    }
+    for (auto hdl : to_remove) {
+      {
+        unique_lock lock(m_lock);
+        remove_hdl(hdl, data_map);
       }
     }
-
-    iter++;
-  }
-  for (auto hdl : to_remove) {
-    remove_hdl(hdl, data_map);
-    LOG(INFO) << "remove one connection ";
   }
 }
 void WebSocketServer::on_message(websocketpp::connection_hdl hdl,
@@ -287,7 +340,7 @@
   }
 
   const std::string& payload = msg->get_payload();  // get msg type
-
+  unique_lock guard_decoder(*(thread_lock_p)); // mutex for one connection
   switch (msg->get_opcode()) {
     case websocketpp::frame::opcode::text: {
       nlohmann::json jsonresult = nlohmann::json::parse(payload);
@@ -308,8 +361,6 @@
         if (msg_data->tpass_online_handle == NULL) {
           std::vector<int> chunk_size_vec =
               jsonresult["chunk_size"].get<std::vector<int>>();
-          LOG(INFO)
-              << "----------------FunTpassOnlineInit----------------------";
           FUNASR_HANDLE tpass_online_handle =
               FunTpassOnlineInit(tpass_handle, chunk_size_vec);
           msg_data->tpass_online_handle = tpass_online_handle;
@@ -323,14 +374,15 @@
 
         // if it is in final message, post the sample_data to decode
         try{
-          asio::post(
-              io_decoder_,
+		  
+          msg_data->strand_->post(
               std::bind(&WebSocketServer::do_decoder, this,
                         std::move(*(sample_data_p.get())), std::move(hdl),
                         std::ref(msg_data->msg), std::ref(*(punc_cache_p.get())),
                         std::ref(*thread_lock_p), std::move(true),
                         msg_data->msg["wav_name"],
                         std::ref(msg_data->tpass_online_handle)));
+		      msg_data->msg["access_num"]=(int)(msg_data->msg["access_num"])+1;
         }
         catch (std::exception const &e)
         {
@@ -362,7 +414,7 @@
 
           try{
             // post to decode
-            asio::post(io_decoder_,
+            msg_data->strand_->post( 
                       std::bind(&WebSocketServer::do_decoder, this,
                                 std::move(subvector), std::move(hdl),
                                 std::ref(msg_data->msg),
@@ -370,6 +422,7 @@
                                 std::ref(*thread_lock_p), std::move(false),
                                 msg_data->msg["wav_name"],
                                 std::ref(msg_data->tpass_online_handle)));
+            msg_data->msg["access_num"]=(int)(msg_data->msg["access_num"])+1;
           }
           catch (std::exception const &e)
           {
@@ -385,6 +438,7 @@
     default:
       break;
   }
+  guard_decoder.unlock();
 }
 
 // init asr model
@@ -396,6 +450,10 @@
       LOG(ERROR) << "FunTpassInit init failed";
       exit(-1);
     }
+    LOG(INFO) << "initAsr run check_and_clean_connection";
+    std::thread clean_thread(&WebSocketServer::check_and_clean_connection,this);  
+    clean_thread.detach();
+    LOG(INFO) << "initAsr run check_and_clean_connection finished";
 
   } catch (const std::exception& e) {
     LOG(INFO) << e.what();

--
Gitblit v1.9.1