From 6a59f94278b4d05f5b7715b11b6dd0bfffc28cce Mon Sep 17 00:00:00 2001
From: 志浩 <neo.dzh@alibaba-inc.com>
Date: 星期四, 15 六月 2023 17:12:15 +0800
Subject: [PATCH] Merge branch 'main' into dev_dzh

---
 funasr/runtime/websocket/funasr-ws-server.cpp |   82 +++++++++++++++++++++++++++-------------
 1 files changed, 55 insertions(+), 27 deletions(-)

diff --git a/funasr/runtime/websocket/websocketmain.cpp b/funasr/runtime/websocket/funasr-ws-server.cpp
similarity index 66%
rename from funasr/runtime/websocket/websocketmain.cpp
rename to funasr/runtime/websocket/funasr-ws-server.cpp
index 4614b51..872f6a1 100644
--- a/funasr/runtime/websocket/websocketmain.cpp
+++ b/funasr/runtime/websocket/funasr-ws-server.cpp
@@ -5,12 +5,12 @@
 /* 2022-2023 by zhaomingwork */
 
 // io server
-// Usage:websocketmain  [--model_thread_num <int>] [--decoder_thread_num <int>]
+// Usage:funasr-ws-server  [--model_thread_num <int>] [--decoder_thread_num <int>]
 //                    [--io_thread_num <int>] [--port <int>] [--listen_ip
 //                    <string>] [--punc-quant <string>] [--punc-dir <string>]
 //                    [--vad-quant <string>] [--vad-dir <string>] [--quantize
 //                    <string>] --model-dir <string> [--] [--version] [-h]
-#include "websocketsrv.h"
+#include "websocket-server.h"
 
 using namespace std;
 void GetValue(TCLAP::ValueArg<std::string>& value_arg, string key,
@@ -25,7 +25,7 @@
     google::InitGoogleLogging(argv[0]);
     FLAGS_logtostderr = true;
 
-    TCLAP::CmdLine cmd("websocketmain", ' ', "1.0");
+    TCLAP::CmdLine cmd("funasr-ws-server", ' ', "1.0");
     TCLAP::ValueArg<std::string> model_dir(
         "", MODEL_DIR,
         "the asr model path, which contains model.onnx, config.yaml, am.mvn",
@@ -64,6 +64,14 @@
     TCLAP::ValueArg<int> model_thread_num("", "model_thread_num",
                                           "model_thread_num", false, 1, "int");
 
+    TCLAP::ValueArg<std::string> certfile("", "certfile", "certfile", false, "",
+                                          "string");
+    TCLAP::ValueArg<std::string> keyfile("", "keyfile", "keyfile", false, "",
+                                         "string");
+
+    cmd.add(certfile);
+    cmd.add(keyfile);
+
     cmd.add(model_dir);
     cmd.add(quantize);
     cmd.add(vad_dir);
@@ -94,46 +102,66 @@
     int s_model_thread_num = model_thread_num.getValue();
 
     asio::io_context io_decoder;  // context for decoding
+    asio::io_context io_server;   // context for server
 
     std::vector<std::thread> decoder_threads;
 
+    std::string s_certfile = certfile.getValue();
+    std::string s_keyfile = keyfile.getValue();
+
+    bool is_ssl = false;
+    if (!s_certfile.empty()) {
+      is_ssl = true;
+    }
+
     auto conn_guard = asio::make_work_guard(
         io_decoder);  // make sure threads can wait in the queue
-
+    auto server_guard = asio::make_work_guard(
+        io_server);  // make sure threads can wait in the queue
     // create threads pool
     for (int32_t i = 0; i < s_decoder_thread_num; ++i) {
       decoder_threads.emplace_back([&io_decoder]() { io_decoder.run(); });
     }
 
-    server server_;       // server for websocket
-    server_.init_asio();  // init asio
-    server_.set_reuse_addr(
-        true);  // reuse address as we create multiple threads
+    server server_;  // server for websocket
+    wss_server wss_server_;
+    if (is_ssl) {
+      wss_server_.init_asio(&io_server);  // init asio
+      wss_server_.set_reuse_addr(
+          true);  // reuse address as we create multiple threads
 
-    // list on port for accept
-    server_.listen(asio::ip::address::from_string(s_listen_ip), s_port);
+      // list on port for accept
+      wss_server_.listen(asio::ip::address::from_string(s_listen_ip), s_port);
+      WebSocketServer websocket_srv(
+          io_decoder, is_ssl, nullptr, &wss_server_, s_certfile,
+          s_keyfile);  // websocket server for asr engine
+      websocket_srv.initAsr(model_path, s_model_thread_num);  // init asr model
 
-    WebSocketServer websocket_srv(io_decoder,
-                                  &server_);  // websocket server for asr engine
-    websocket_srv.initAsr(model_path, s_model_thread_num);  // init asr model
+    } else {
+      server_.init_asio(&io_server);  // init asio
+      server_.set_reuse_addr(
+          true);  // reuse address as we create multiple threads
+
+      // list on port for accept
+      server_.listen(asio::ip::address::from_string(s_listen_ip), s_port);
+      WebSocketServer websocket_srv(
+          io_decoder, is_ssl, &server_, nullptr, s_certfile,
+          s_keyfile);  // websocket server for asr engine
+      websocket_srv.initAsr(model_path, s_model_thread_num);  // init asr model
+    }
+
     std::cout << "asr model init finished. listen on port:" << s_port
               << std::endl;
 
     // Start the ASIO network io_service run loop
-    if (s_io_thread_num == 1) {
-      server_.run();
-    } else {
-      typedef websocketpp::lib::shared_ptr<websocketpp::lib::thread> thread_ptr;
-      std::vector<thread_ptr> ts;
-      // create threads for io network
-      for (size_t i = 0; i < s_io_thread_num; i++) {
-        ts.push_back(websocketpp::lib::make_shared<websocketpp::lib::thread>(
-            &server::run, &server_));
-      }
-      // wait for theads
-      for (size_t i = 0; i < s_io_thread_num; i++) {
-        ts[i]->join();
-      }
+    std::vector<std::thread> ts;
+    // create threads for io network
+    for (size_t i = 0; i < s_io_thread_num; i++) {
+      ts.emplace_back([&io_server]() { io_server.run(); });
+    }
+    // wait for theads
+    for (size_t i = 0; i < s_io_thread_num; i++) {
+      ts[i].join();
     }
 
     // wait for theads

--
Gitblit v1.9.1