nim/server/src/app/application.cpp

#include "app/application.h"

#include <atomic>
#include <chrono>
#include <csignal>
#include <cstring>
#include <memory>
#include <string>
#include <thread>
#include <unistd.h>

#include "common/logger.h"
#include "third/cos.h"
#include "util/file.h"
#include "im/handlers/account_handler.h"
#include "im/handlers/favorite_handler.h"
#include "im/handlers/friend_handler.h"
#include "im/handlers/group_handler.h"
#include "im/handlers/message_handler.h"
#include "im/handlers/sticker_handler.h"
#include "im/handlers/system_handler.h"
#include "db/sticker_store.h"

namespace im {
namespace {

std::atomic<bool> g_running{true};

void on_signal(int)
{
    g_running = false;
}

std::string find_init_sql(const std::string& cfg_path)
{
    const std::string parent = ylib::file::parent_dir(cfg_path);
    const std::string candidates[] = {
        parent + "/../sql/init.sql",
        parent + "/sql/init.sql",
        "sql/init.sql",
        "../sql/init.sql",
    };
    for (const auto& path : candidates) {
        if (ylib::file::exist(path)) {
            return path;
        }
    }
    return parent + "/../sql/init.sql";
}

}  // namespace

int Application::run(int argc, char** argv)
{
    if (!init(argc, argv)) {
        return 1;
    }
    wait_shutdown();
    shutdown();
    return 0;
}

bool Application::init(int argc, char** argv)
{
    bool init_db = false;
    std::string cfg_path;
    for (int i = 1; i < argc; ++i) {
        const std::string arg = argv[i] != nullptr ? argv[i] : "";
        if (arg.empty()) {
            continue;
        }
        if (arg == "--init-db") {
            init_db = true;
        } else if (arg[0] == '-') {
            Logger::error("unknown option: " + arg);
            Logger::error("usage: server [--init-db] [config.ini]");
            return false;
        } else {
            cfg_path = arg;
        }
    }
    if (cfg_path.empty()) {
        cfg_path = AppConfig::resolve_path(0, nullptr);
    }
    if (!config_.load(cfg_path)) {
        return false;
    }

    if (init_db) {
        if (!config_.mysql().enable) {
            Logger::error("--init-db requires [mysql] enable=1");
            return false;
        }
        const std::string sql_path = find_init_sql(cfg_path);
        if (!ylib::file::exist(sql_path)) {
            Logger::error("init.sql not found, tried near " + cfg_path);
            return false;
        }
        Logger::info("init database from " + sql_path);
        if (!mysql_.init_schema(config_.mysql(), sql_path)) {
            return false;
        }
    }

    if (config_.cos().enable) {
        if (third::cos::enabled(config_.cos())) {
            Logger::info("cos storage bucket=" + config_.cos().bucket +
                         " region=" + config_.cos().region +
                         " prefix=" + config_.cos().prefix);
        } else {
            Logger::error("cos enable=1 but secret_id/secret_key/bucket/region incomplete");
            return false;
        }
    } else {
        Logger::error("COS storage is required");
        return false;
    }
    if (!config_.cdn_domain().empty()) {
        Logger::info("cdn domain=" + config_.cdn_domain());
    }

    if (config_.mysql().enable) {
        if (!mysql_.start(config_.mysql())) {
            if (config_.mysql().required) {
                Logger::error("mysql is required, abort");
                return false;
            }
            Logger::warn("mysql unavailable, continue without db");
        }
    }

    dispatcher_.add(Type::System, [] { return std::make_unique<SystemHandler>(); });
    dispatcher_.add(Type::Account, [] { return std::make_unique<AccountHandler>(); });
    dispatcher_.add(Type::Friend, [] { return std::make_unique<FriendHandler>(); });
    dispatcher_.add(Type::Message, [] { return std::make_unique<MessageHandler>(); });
    dispatcher_.add(Type::Sticker, [] { return std::make_unique<StickerHandler>(); });
    dispatcher_.add(Type::Favorite, [] { return std::make_unique<FavoriteHandler>(); });
    dispatcher_.add(Type::Group,    [] { return std::make_unique<GroupHandler>(); });

    if (mysql_.started()) {
        StickerStore stickers(mysql_);
        stickers.seed();
    }

    std::string node_id = config_.server().node_id;
    if (node_id.empty()) {
        char host[256] = {0};
        if (gethostname(host, sizeof(host) - 1) != 0 || host[0] == '\0') {
            std::strcpy(host, "im");
        }
        node_id = std::string(host) + ":" + std::to_string(config_.server().port);
    }
    if (config_.redis().enable) {
        if (!redis_.start(config_.redis(), node_id)) {
            if (config_.redis().required) {
                Logger::error("redis is required, abort");
                return false;
            }
            Logger::warn("redis unavailable, continue in single-node mode");
        }
    } else {
        Logger::info("redis disabled, single-node node_id=" + node_id);
    }

    server_ = std::make_unique<ImServer>(config_, sessions_, dispatcher_, mysql_, &redis_);
    if (!server_->start()) {
        return false;
    }

    Logger::info("im server ready, packet=magic:type:cmd:seq:len + body");
    return true;
}

void Application::wait_shutdown()
{
    std::signal(SIGINT, on_signal);
    std::signal(SIGTERM, on_signal);
    Logger::info("press Ctrl+C to stop");
    while (g_running) {
        std::this_thread::sleep_for(std::chrono::milliseconds(200));
    }
}

void Application::shutdown()
{
    Logger::info("shutting down");
    if (server_) {
        server_->stop();
        server_.reset();
    }
    redis_.stop();
    mysql_.stop();
}

}  // namespace im