志浩
2023-06-15 6a59f94278b4d05f5b7715b11b6dd0bfffc28cce
funasr/runtime/websocket/funasr-ws-server.cpp
File was renamed from funasr/runtime/websocket/websocketmain.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