ClickHouse/programs/keeper-bench/KeeperBench.cpp

#include <csignal>

#include <iostream>
#include <boost/program_options.hpp>
#include <GeneratedRunner.h>
#include <LogRunner.h>
#include <StorageRunner.h>
#include <Common/ErrnoException.h>
#include <Common/Exception.h>
#include <Common/TerminalSize.h>
#include <Common/ThreadPool.h>
#include <IO/SharedThreadPools.h>
#include <Common/scope_guard_safe.h>
#include <Core/Types.h>
#include <boost/program_options/variables_map.hpp>

namespace DB::ErrorCodes
{
    extern const int CANNOT_BLOCK_SIGNAL;
}

namespace
{

template <typename T>
std::optional<T> valueToOptional(const boost::program_options::variable_value & value)
{
    if (value.empty())
        return std::nullopt;

    return value.as<T>();
}

}

int mainEntryClickHouseKeeperBench(int argc, char ** argv);
int mainEntryClickHouseKeeperBench(int argc, char ** argv)
{

    bool print_stacktrace = true;

    //Poco::AutoPtr<Poco::ConsoleChannel> channel(new Poco::ConsoleChannel(std::cerr));
    //Poco::Logger::root().setChannel(channel);
    //Poco::Logger::root().setLevel("trace");

    /// Join global-pool threads before the statics they may have accessed are destroyed.
    /// That way, accesses happen-before destruction.
    SCOPE_EXIT_SAFE({
        DB::StaticThreadPool::shutdownAll();
        GlobalThreadPool::shutdown();
    });

    try
    {
        /// Block SIGINT in the main thread before any other threads are created,
        /// so that every thread inherits the blocked mask. Otherwise a SIGINT
        /// delivered to a thread that doesn't block it (e.g. a global thread pool
        /// thread) terminates the process immediately. The runners poll the
        /// pending signal with `InterruptListener` and shut down gracefully.
        sigset_t sig_set;
        if (sigemptyset(&sig_set)
            || sigaddset(&sig_set, SIGINT)
            || pthread_sigmask(SIG_BLOCK, &sig_set, nullptr))
        {
            throw DB::ErrnoException(DB::ErrorCodes::CANNOT_BLOCK_SIGNAL, "Cannot block signal");
        }

        using boost::program_options::value;

        boost::program_options::options_description desc = createOptionsDescription("Allowed options", getTerminalWidth());
        desc.add_options()
            ("help",                                                                         "produce help message")
            ("config",            value<std::string>()->default_value(""),                      "yaml/xml file containing configuration")
            ("input-request-log", value<std::string>()->default_value(""),                      "log of requests that will be replayed (enables replay mode, overriding the config's `generator` section)")
            ("setup-nodes-snapshot-path", value<std::string>()->default_value(""),                      "directory containing snapshots with starting state")
            ("concurrency,c",     value<unsigned>(),                                            "number of parallel queries")
            ("report-delay,d",    value<double>(),                                              "delay between intermediate reports in seconds (set 0 to disable reports)")
            ("iterations,i",      value<size_t>(),                                              "amount of queries to be executed")
            ("time-limit,t",      value<double>(),                                              "stop launch of queries after specified time limit")
            ("hosts,h",           value<Strings>()->multitoken()->default_value(Strings{}, ""), "")
            ("continue_on_errors", "continue testing even if a query fails")
            ("storage",           "benchmark KeeperStorage in-process, without the rest of the keeper server")
        ;

        boost::program_options::variables_map options;
        boost::program_options::store(boost::program_options::parse_command_line(argc, argv, desc), options);
        boost::program_options::notify(options);

        if (options.contains("help"))
        {
            std::cout << "Usage: " << argv[0] << " [options] < queries.txt\n";
            std::cout << desc << "\n";
            return 1;
        }

        if (options.contains("storage"))
        {
            StorageRunner storage_runner(
                options["config"].as<std::string>(),
                valueToOptional<unsigned>(options["concurrency"]).transform([](unsigned v) { return static_cast<size_t>(v); }),
                valueToOptional<double>(options["time-limit"]),
                valueToOptional<double>(options["report-delay"]),
                valueToOptional<size_t>(options["iterations"]),
                options.contains("continue_on_errors") ? std::optional<bool>(true) : std::nullopt);

            try
            {
                storage_runner.runBenchmark();
            }
            catch (...)
            {
                std::cout << "Got exception while trying to run storage benchmark: " << DB::getCurrentExceptionMessage(true) << std::endl;
            }

            return 0;
        }

        if (const auto input_request_log = options["input-request-log"].as<std::string>(); !input_request_log.empty())
        {
            LogRunner runner(valueToOptional<unsigned>(options["concurrency"]),
                             options["config"].as<std::string>(),
                             input_request_log,
                             options["setup-nodes-snapshot-path"].as<std::string>(),
                             options["hosts"].as<Strings>(),
                             valueToOptional<double>(options["report-delay"]));

            try
            {
                runner.runBenchmark();
            }
            catch (...)
            {
                std::cout << "Got exception while trying to run benchmark: " << DB::getCurrentExceptionMessage(true) << std::endl;
            }

            return 0;
        }

        GeneratedRunner runner(valueToOptional<unsigned>(options["concurrency"]),
                               options["config"].as<std::string>(),
                               options["hosts"].as<Strings>(),
                               valueToOptional<double>(options["time-limit"]),
                               valueToOptional<double>(options["report-delay"]),
                               options.contains("continue_on_errors") ? std::optional<bool>(true) : std::nullopt,
                               valueToOptional<size_t>(options["iterations"]));

        try
        {
            runner.runBenchmark();
        }
        catch (...)
        {
            std::cout << "Got exception while trying to run benchmark: " << DB::getCurrentExceptionMessage(true) << std::endl;
        }

        return 0;
    }
    catch (...)
    {
        std::cerr << DB::getCurrentExceptionMessage(print_stacktrace, true) << std::endl;
        return DB::getCurrentExceptionCode();
    }
}