ECAP_Server/module/Local_Server/server/WebSocket_Manager.cpp

#include "WebSocket_Manager.h"
#include "../global_include.h"
#include "Global.h"
#include <drogon/HttpAppFramework.h>
#include <drogon/WebSocketController.h>
#include <unordered_set>
#include <string_view>
class DSP_Ws_Mgr : public drogon::WebSocketController<DSP_Ws_Mgr> {
public:
  static DSP_Ws_Mgr *instance() { return inst_; }
  WS_PATH_LIST_BEGIN
  WS_ADD_PATH_VIA_REGEX(R"(^/ws/.*$)", drogon::Get);
  WS_PATH_LIST_END
  void handleNewMessage(const drogon::WebSocketConnectionPtr &, std::string &&,
                        const drogon::WebSocketMessageType &) override;
  void
  handleConnectionClosed(const drogon::WebSocketConnectionPtr &conn) override;
  void handleNewConnection(const drogon::HttpRequestPtr &req,
                           const drogon::WebSocketConnectionPtr &conn) override;
  DSP_Ws_Mgr();

protected:
  static DSP_Ws_Mgr *inst_;
};
DSP_Ws_Mgr *DSP_Ws_Mgr::inst_ = nullptr;
struct WebSocket_Connection {
  std::string type;
};

WebSocket_Manager::WebSocket_Manager(std::string_view path)
    : path(path) {}
void WebSocket_Manager::remove(const drogon::WebSocketConnectionPtr &conn) {
  std::lock_guard<std::mutex> g(mtx);
  connections.erase(conn);
}
void WebSocket_Manager::add(const drogon::WebSocketConnectionPtr &conn) {
  std::lock_guard<std::mutex> g(mtx);
  connections.insert(conn);
}
void WebSocket_Manager::visit(
    const std::function<void(const drogon::WebSocketConnectionPtr &conn)>
        &func) {
  std::lock_guard<std::mutex> g(mtx);
  for (const auto &coon : connections) {
    func(coon);
  }
}
void WebSocket_Manager::push_data(const char *buf, uint64_t len) {
  std::lock_guard<std::mutex> g(mtx);
  for (const auto &conn : connections) {
    conn->send(buf, len, drogon::WebSocketMessageType::Binary);
  }
}
void WebSocket_Manager::set_message_handler(Message_Handler handler) {
  message_handler = std::move(handler);
}
void WebSocket_Manager::set_close_handler(Close_Handler handler) {
  close_handler = std::move(handler);
}
void WebSocket_Manager::handle_message(const drogon::WebSocketConnectionPtr &conn,
                                       std::string &&message,
                                       const drogon::WebSocketMessageType &type) {
  if (message_handler) {
    message_handler(conn, std::move(message), type);
  }
}
void WebSocket_Manager::handle_close(const drogon::WebSocketConnectionPtr &conn) {
  if (close_handler) {
    close_handler(conn);
  }
}
size_t WebSocket_Manager::get_websocket_connect_num() {
  std::lock_guard<std::mutex> g(mtx);
  return connections.size();
}

std::shared_ptr<WebSocket_Manager> Ws_Mgr::register_ws(std::string_view path) {
  auto path_string = std::string(path);
  auto it = map.find(path_string);
  if (it != map.end()) {
    Psc::fail_fast_core_dump("register_ws:" + VAR_STR_1(path_string) + "重复创建！");
    return it->second;
  }
  std::shared_ptr<WebSocket_Manager> cur =
      std::make_shared<WebSocket_Manager>(path_string);
  map[path_string] = cur;
  return cur;
}

std::shared_ptr<WebSocket_Manager> Ws_Mgr::get_ws(std::string_view path) {
  auto it = map.find(std::string(path));
  if (it != map.end()) {
    return it->second;
  }
  return nullptr;
}

void DSP_Ws_Mgr::handleNewMessage(const drogon::WebSocketConnectionPtr &conn,
                                  std::string &&message,
                                  const drogon::WebSocketMessageType &type) {
  auto &s = conn->getContextRef<WebSocket_Connection>();
  auto ws = Ws_Mgr::instance()->get_ws(s.type);
  if (ws) {
    ws->handle_message(conn, std::move(message), type);
  }
}
void DSP_Ws_Mgr::handleConnectionClosed(
    const drogon::WebSocketConnectionPtr &conn) {
  auto &s = conn->getContextRef<WebSocket_Connection>();
  auto rp = s.type;
  auto ws = Ws_Mgr::instance()->get_ws(rp);
  if (!ws) {
    return;
  }
  ws->handle_close(conn);
  // LOG_DEBUG
  std::ostringstream oss;
  oss << "已有连接数:" << ws->get_websocket_connect_num() << "  "
      << "关闭websocket连接:" << rp << "\n"
      << "local_addr:" << inet_address_to_string(conn->localAddr()) << "\n"
      << "peer_addr:" << inet_address_to_string(conn->peerAddr()) << "\n";
  dsp_logger->c_debug("websocket", {}, oss.str());
  ws->remove(conn);
}
void DSP_Ws_Mgr::handleNewConnection(
    const drogon::HttpRequestPtr &req,
    const drogon::WebSocketConnectionPtr &conn) {
  auto rpt = req->path();
  auto rp = rpt.substr(std::string("/ws").size());
  auto ws = Ws_Mgr::instance()->get_ws(rp);
  if (!ws) {
    conn->shutdown();
    return;
  }
  std::ostringstream oss;
  oss << "已有连接数:" << ws->get_websocket_connect_num() << "  "
      << "新建websocket连接:" << rp << "  "
      << "local_addr:" << inet_address_to_string(conn->localAddr()) << "  "
      << "peer_addr:" << inet_address_to_string(conn->peerAddr()) << "\n";

  dsp_logger->c_debug("websocket", {}, oss.str());

  WebSocket_Connection s;
  s.type = rp;
  conn->setContext(std::make_shared<WebSocket_Connection>(std::move(s)));
  ws->add(conn);
}
DSP_Ws_Mgr::DSP_Ws_Mgr() { inst_ = this; }