nim/server/src/net/im_server.cpp

#include "net/im_server.h"

#include <exception>
#include <memory>
#include <string>
#include <vector>

#include "common/i18n.h"
#include "common/logger.h"
#include "common/protocol_def.h"
#include "db/friend_store.h"
#include "db/redis_service.h"
#include "net/packet_log.h"
#include "util/json.h"
#include "HPSocket/HPSocket.h"

namespace im {

ImServer::ImServer(const AppConfig& cfg, SessionMgr& sessions, Dispatcher& dispatcher,
                   MysqlService& mysql, RedisService* redis)
    : cfg_(cfg),
      sessions_(sessions),
      dispatcher_(dispatcher),
      mysql_(mysql),
      redis_(redis),
      codec_(cfg.protocol().magic, cfg.protocol().max_packet)
{
}

ImServer::~ImServer()
{
    stop();
}

bool ImServer::start()
{
    stopping_.store(false, std::memory_order_release);

    if (redis_ != nullptr) {
        redis_->set_message_handler([this](const std::string& payload) {
            on_cluster_payload(payload);
        });
    }

    if (!start_thread_pool()) {
        return false;
    }

    server_.on_accept([this](ylib::network::tcp::server*, uint64 connid) {
        on_accept(connid);
    });
    server_.on_recv(
        [this](ylib::network::tcp::server*, uint64 connid, const char* data, uint32 len) {
            on_recv(connid, data, len);
        });
    server_.on_close([this](ylib::network::tcp::server*, uint64 connid) {
        on_close(connid);
    });

    configure_hp();

    ylib::AddressPort ap;
    ap.address = cfg_.server().host;
    ap.port = cfg_.server().port;

    if (!server_.start(ap, ylib::PUSH_DEFAULT)) {
        Logger::error("tcp listen failed: " + server_.last_error());
        stop_thread_pool();
        return false;
    }

    ITcpServer* hp = server_.getHP();
    Logger::info("tcp listen " + cfg_.server().host + ":" +
                 std::to_string(cfg_.server().port) +
                 " io_workers=" +
                 std::to_string(hp != nullptr ? hp->GetWorkerThreadCount() : 0));
    return true;
}

void ImServer::stop()
{
    const bool first = !stopping_.exchange(true, std::memory_order_acq_rel);
    if (server_.started()) {
        server_.close();
    }
    stop_thread_pool();
    sessions_.clear();
    if (first) {
        Logger::info("tcp server stopped");
    }
}

void ImServer::configure_hp()
{
    ITcpServer* hp = server_.getHP();
    if (hp == nullptr) {
        return;
    }
    hp->SetMaxConnectionCount(cfg_.server().max_connections);
    if (cfg_.server().worker_threads != 0) {
        hp->SetWorkerThreadCount(static_cast<DWORD>(cfg_.server().worker_threads));
    }
}

bool ImServer::start_thread_pool()
{
    thread_pool_ = HP_Create_ThreadPool();
    if (thread_pool_ == nullptr) {
        Logger::error("hp thread pool create failed");
        return false;
    }
    if (!thread_pool_->Start(static_cast<DWORD>(cfg_.server().thread_pool),
                             cfg_.server().thread_pool_queue, TRP_CALL_FAIL)) {
        Logger::error("hp thread pool start failed");
        HP_Destroy_ThreadPool(thread_pool_);
        thread_pool_ = nullptr;
        return false;
    }
    Logger::info("hp thread pool workers=" +
                 std::to_string(thread_pool_->GetThreadCount()) +
                 " queue_limit=" + std::to_string(cfg_.server().thread_pool_queue));
    return true;
}

void ImServer::stop_thread_pool()
{
    if (thread_pool_ == nullptr) {
        return;
    }
    thread_pool_->Stop(INFINITE);
    HP_Destroy_ThreadPool(thread_pool_);
    thread_pool_ = nullptr;
}

bool ImServer::send(ConnId connid, const Packet& packet)
{
    log_send(connid, packet);
    Packet out = packet;
    out.magic = cfg_.protocol().magic;
    return server_.send(connid, codec_.encode(out));
}

bool ImServer::send(const SessionPtr& session, const Packet& packet)
{
    if (!conn_matches(session)) {
        return false;
    }
    return send(session->connid(), packet);
}

void ImServer::disconnect(ConnId connid)
{
    server_.disConnect(connid);
}

bool ImServer::conn_matches(const SessionPtr& session) const
{
    if (!session || session->closed() || stopping_.load(std::memory_order_acquire)) {
        return false;
    }
    if (!sessions_.is_live(session)) {
        return false;
    }
    ITcpServer* hp = const_cast<ylib::network::tcp::server&>(server_).getHP();
    if (hp == nullptr) {
        return false;
    }
    PVOID extra = nullptr;
    if (!hp->GetConnectionExtra(session->connid(), &extra) || extra != session.get()) {
        return false;
    }
    return true;
}

void ImServer::notify_friend_presence(UserId uid, int online)
{
    if (uid <= 0 || uid == kFileHelperUid) {
        return;
    }
    ylib::json body;
    body["uid"] = static_cast<int64>(uid);
    body["online"] = online ? 1 : 0;
    const Packet pkt = Packet::make(Type::Friend, Cmd::PresenceNotify, 0, body.to_string());
    try {
        FriendStore store(mysql_);
        for (const auto& f : store.list_friends(uid)) {
            if (f.uid <= 0 || f.uid == kFileHelperUid) {
                continue;
            }
            push_to_uid(f.uid, pkt);
        }
    } catch (const std::exception& e) {
        Logger::warn(std::string("presence notify uid=") + std::to_string(uid) + ": " + e.what());
    }
}

void ImServer::on_accept(ConnId connid)
{
    if (stopping_.load(std::memory_order_acquire)) {
        disconnect(connid);
        return;
    }
    auto session = std::make_shared<Session>(connid, remote_of(connid),
                                             cfg_.protocol().magic,
                                             cfg_.protocol().max_packet);
    sessions_.add(session);
    ITcpServer* hp = server_.getHP();
    if (hp != nullptr) {
        hp->SetConnectionExtra(connid, session.get());
    }
    Logger::info("accept conn=" + std::to_string(connid) + " remote=" + session->remote() +
                 " online=" + std::to_string(sessions_.size()));
}

void ImServer::on_recv(ConnId connid, const char* data, uint32_t len)
{
    if (stopping_.load(std::memory_order_acquire)) {
        return;
    }
    auto session = sessions_.get(connid);
    if (!session || session->closed()) {
        return;
    }

    std::vector<Packet> packets;
    const DecodeStatus st = session->feed(data, len, packets);
    if (st == DecodeStatus::BadMagic || st == DecodeStatus::TooLarge) {
        Logger::warn("bad packet conn=" + std::to_string(connid) +
                     (st == DecodeStatus::BadMagic ? " magic" : " too large"));
        disconnect(connid);
        return;
    }

    for (auto& pkt : packets) {
        log_recv(connid, pkt);
        bool start_task = false;
        if (!session->enqueue(std::move(pkt), &start_task)) {
            return;
        }
        if (start_task && !submit_session(session)) {
            session->clear_running();
            Logger::error("thread pool submit failed conn=" + std::to_string(connid));
            disconnect(connid);
            return;
        }
    }
}

void ImServer::on_close(ConnId connid)
{
    auto session = sessions_.get(connid);
    if (session) {
        session->mark_closed();
    }
    ITcpServer* hp = server_.getHP();
    if (hp != nullptr) {
        hp->SetConnectionExtra(connid, nullptr);
    }
    const UserId uid = session ? session->uid() : 0;
    const std::string remote = session ? session->remote() : "";
    sessions_.remove(connid);
    if (uid > 0 && !sessions_.find_by_uid(uid)) {
        unbind_cluster(uid);
        if (!is_online(uid)) {
            notify_friend_presence(uid, 0);
        }
    }
    Logger::info("close conn=" + std::to_string(connid) + " remote=" + remote +
                 " online=" + std::to_string(sessions_.size()));
}

bool ImServer::submit_session(const SessionPtr& session)
{
    if (thread_pool_ == nullptr || stopping_.load(std::memory_order_acquire)) {
        return false;
    }
    auto* task = new PoolTask{this, session};
    if (!thread_pool_->Submit(reinterpret_cast<Fn_TaskProc>(&ImServer::on_pool_task), task)) {
        delete task;
        return false;
    }
    return true;
}

void ImServer::on_pool_task(void* arg)
{
    std::unique_ptr<PoolTask> task(static_cast<PoolTask*>(arg));
    if (!task || task->server == nullptr) {
        return;
    }
    try {
        task->server->drain_session(task->session);
    } catch (const std::exception& e) {
        Logger::error(std::string("thread pool task exception: ") + e.what());
    } catch (...) {
        Logger::error("thread pool task unknown exception");
    }
}

void ImServer::drain_session(const SessionPtr& session)
{
    if (!session) {
        return;
    }
    while (true) {
        Packet pkt;
        if (!session->take_next(pkt)) {
            return;
        }
        if (!conn_matches(session)) {
            session->clear_running();
            return;
        }
        try {
            HandlerContext ctx;
            ctx.session = session;
            ctx.server = this;
            ctx.mysql = &mysql_;
            ctx.config = &cfg_;
            ctx.packet = std::move(pkt);
            dispatcher_.dispatch(std::move(ctx));
        } catch (const std::exception& e) {
            Logger::error(std::string("dispatch exception conn=") +
                          std::to_string(session->connid()) + ": " + e.what());
        } catch (...) {
            Logger::error("dispatch unknown exception conn=" +
                          std::to_string(session->connid()));
        }
    }
}

std::string ImServer::remote_of(ConnId connid) const
{
    ITcpServer* hp = const_cast<ylib::network::tcp::server&>(server_).getHP();
    if (hp == nullptr) {
        return "";
    }
    char addr[256] = {0};
    int addr_len = static_cast<int>(sizeof(addr));
    USHORT port = 0;
    if (!hp->GetRemoteAddress(connid, addr, addr_len, port)) {
        return "";
    }
    return std::string(addr) + ":" + std::to_string(port);
}

bool ImServer::is_online(UserId uid)
{
    if (uid <= 0) {
        return false;
    }
    if (sessions_.find_by_uid(uid)) {
        return true;
    }
    if (redis_ == nullptr || !redis_->enabled()) {
        return false;
    }
    return !redis_->online_node(uid).empty();
}

bool ImServer::push_to_uid(UserId uid, const Packet& packet)
{
    if (uid <= 0) {
        return false;
    }
    if (auto local = sessions_.find_by_uid(uid)) {
        return send(local, packet);
    }
    if (redis_ == nullptr || !redis_->enabled()) {
        return false;
    }
    const std::string node = redis_->online_node(uid);
    if (node.empty() || node == redis_->node_id()) {
        return false;
    }
    ylib::json msg;
    msg["op"] = "send";
    msg["uid"] = static_cast<int64>(uid);
    msg["type"] = static_cast<int64>(packet.type);
    msg["cmd"] = static_cast<int64>(packet.cmd);
    msg["body"] = packet.body.to_string();
    return redis_->publish_to_node(node, msg.to_string());
}

void ImServer::bind_cluster(UserId uid)
{
    if (redis_ == nullptr || !redis_->enabled() || uid <= 0) {
        return;
    }
    const std::string old = redis_->bind_online(uid);
    if (old.empty() || old == redis_->node_id()) {
        return;
    }
    ylib::json msg;
    msg["op"] = "kick";
    msg["uid"] = static_cast<int64>(uid);
    redis_->publish_to_node(old, msg.to_string());
}

void ImServer::unbind_cluster(UserId uid)
{
    if (redis_ == nullptr || !redis_->enabled() || uid <= 0) {
        return;
    }
    redis_->unbind_online(uid);
}

void ImServer::kick_local(UserId uid)
{
    auto session = sessions_.find_by_uid(uid);
    if (!session) {
        return;
    }
    ylib::json body;
    body["code"] = Err::Ok;
    body["msg"] = i18n::text(i18n::Msg::Kicked, session->lang());
    send(session, Packet::make(Type::System, Cmd::Kicked, 0, body.to_string()));
    disconnect(session->connid());
}

void ImServer::on_cluster_payload(const std::string& payload)
{
    if (payload.empty() || payload == "{}") {
        return;
    }
    ylib::json msg;
    if (!msg.parse(payload) || !msg.exist("op")) {
        return;
    }
    const std::string op = msg["op"].to<std::string>();
    const UserId uid = msg.exist("uid") ? static_cast<UserId>(msg["uid"].to<int64>()) : 0;
    if (uid <= 0) {
        return;
    }
    if (op == "kick") {
        kick_local(uid);
        return;
    }
    if (op != "send") {
        return;
    }
    auto session = sessions_.find_by_uid(uid);
    if (!session) {
        return;
    }
    Packet pkt;
    pkt.type = static_cast<uint16_t>(msg.exist("type") ? msg["type"].to<int64>() : 0);
    pkt.cmd = static_cast<uint16_t>(msg.exist("cmd") ? msg["cmd"].to<int64>() : 0);
    if (msg.exist("body")) {
        pkt.body = msg["body"].to<std::string>();
    }
    send(session, pkt);
}

}  // namespace im