ClickHouse/utils/zookeeper-test/main.cpp

#include <IO/ReadBufferFromString.h>
#include <IO/ReadHelpers.h>
#include <Poco/ConsoleChannel.h>
#include <Common/ZooKeeper/KeeperException.h>
#include <Common/ZooKeeper/ZooKeeper.h>
#include <base/LineReader.h>
#include <Common/logger_useful.h>
#include <fmt/format.h>
#include <random>
#include <iterator>
#include <algorithm>
#include <chrono>
#include <Common/Stopwatch.h>

#include <iostream>
#include <sstream>
#include <exception>
#include <future>

using namespace std;

/// TODO: Remove ME

void checkEq(zkutil::ZooKeeper & zk, const std::string & path, const std::string & expected)
{
    auto result = zk.get(path);
    if (result != expected)
        throw std::runtime_error(fmt::format("Data on path '{}' = '{}' doesn't match expected '{}'",
                                             path, result, expected));
}

void checkExists(zkutil::ZooKeeper & zk, const std::string & path)
{
    if (!zk.exists(path))
        throw std::runtime_error(fmt::format("Path '{}' doesn't exists", path));
}

void testCreateGetExistsNode(zkutil::ZooKeeper & zk)
{
    zk.create("/data", "test_string", zkutil::CreateMode::Persistent);
    zk.create("/data/seq-", "another_string", zkutil::CreateMode::PersistentSequential);
    checkEq(zk, "/data", "test_string");
    checkExists(zk, "/data/seq-0000000000");
    checkEq(zk, "/data/seq-0000000000", "another_string");
}

void testCreateSetNode(zkutil::ZooKeeper & zk)
{
    zk.create("/data/set", "sssss", zkutil::CreateMode::Persistent);
    checkEq(zk, "/data/set", "sssss");
    zk.set("/data/set", "qqqqq");
    checkEq(zk, "/data/set", "qqqqq");
}

void testCreateList(zkutil::ZooKeeper & zk)
{
    zk.create("/data/lst", "", zkutil::CreateMode::Persistent);
    zk.create("/data/lst/d1", "", zkutil::CreateMode::Persistent);
    zk.create("/data/lst/d2", "", zkutil::CreateMode::Persistent);
    zk.create("/data/lst/d3", "", zkutil::CreateMode::Persistent);
    auto children = zk.getChildren("/data/lst");
    if (children.size() != 3)
        throw std::runtime_error("Children of /data/lst doesn't equal to three");
    for (size_t i = 0; i < children.size(); ++i)
    {
        std::cerr << "children:" << children[i] << std::endl;
        std::cerr << "children size:" << children[i].size() << std::endl;
        if (children[i] != "d" + std::to_string(i + 1))
            throw std::runtime_error(fmt::format("Incorrect children #{} got {}, expected {}", i, children[i], "d" + std::to_string(i + 1)));
    }
}

void testCreateSetVersionRequest(zkutil::ZooKeeper & zk)
{
    zk.create("/data/check_data", "d", zkutil::CreateMode::Persistent);
    Coordination::Stat stat{};
    try
    {
        zk.set("/data/check_data", "e", stat.version + 2);
        std::terminate();
    }
    catch (...)
    {
        std::cerr << "Got exception on incorrect version (it's ok)\n";
    }

    checkEq(zk, "/data/check_data", "d");
    zk.set("/data/check_data", "e", stat.version);

    checkEq(zk, "/data/check_data", "e");
}

void testCreateSetWatchEvent(zkutil::ZooKeeper & zk)
{

    std::shared_ptr<Poco::Event> event = std::make_shared<Poco::Event>();
    zk.create("/data/nodeforwatch", "", zkutil::CreateMode::Persistent);
    Coordination::Stat stat;
    zk.get("/data/nodeforwatch", &stat, event);

    if (event->tryWait(300))
        throw std::runtime_error(fmt::format("Event for path {} was set without any actions", "/data/nodeforwatch"));

    zk.set("/data/nodeforwatch", "x");
    if (!event->tryWait(300))
        throw std::runtime_error(fmt::format("Event for path {} was not set after set", "/data/nodeforwatch"));
    else
        std::cerr << "Event was set well\n";
}

void testCreateListWatchEvent(zkutil::ZooKeeper & zk)
{
    std::shared_ptr<Poco::Event> event = std::make_shared<Poco::Event>();
    std::string path = "/data/pathforwatch";
    zk.create(path, "", zkutil::CreateMode::Persistent);
    zk.create(path + "/n1", "", zkutil::CreateMode::Persistent);
    zk.create(path + "/n2", "", zkutil::CreateMode::Persistent);
    zk.getChildren(path, nullptr, event);

    if (event->tryWait(300))
        throw std::runtime_error(fmt::format("ListEvent for path {} was set without any actions", path));

    zk.create(path + "/n3", "", zkutil::CreateMode::Persistent);
    if (!event->tryWait(300))
        throw std::runtime_error(fmt::format("ListEvent for path {} was not set after create", path));
    else
        std::cerr << "ListEvent was set well\n";
}

void testMultiRequest(zkutil::ZooKeeper & zk)
{
    std::cerr << "Testing multi request\n";
    Coordination::Requests requests;
    requests.push_back(zkutil::makeCreateRequest("/data/multirequest", "aaa", zkutil::CreateMode::Persistent));
    requests.push_back(zkutil::makeSetRequest("/data/multirequest", "bbb", -1));
    zk.multi(requests);
    std::cerr << "Multi executed\n";

    try
    {
        requests.clear();
        std::cerr << "Testing bad multi\n";
        requests.push_back(zkutil::makeCreateRequest("/data/multirequest", "qweqwe", zkutil::CreateMode::Persistent));
        requests.push_back(zkutil::makeSetRequest("/data/multirequest", "bbb", -1));
        requests.push_back(zkutil::makeSetRequest("/data/multirequest", "ccc", -1));
        zk.multi(requests);
        std::cerr << "Bad multi executed\n";
        std::terminate();
    }
    catch (...)
    {
        std::cerr << "Got exception on multy request (it's ok)\n";
    }

    checkEq(zk, "/data/multirequest", "bbb");
    std::cerr << "Multi request finished\n";
}

std::mutex elements_mutex;
std::vector<int> current_elements;
std::atomic<int> watches_triggered = 0;

void triggerWatch(const Coordination::WatchResponse &)
{
    watches_triggered++;
}

template<typename Iter, typename RandomGenerator>
Iter select_randomly(Iter start, Iter end, RandomGenerator& g)
{
    std::uniform_int_distribution<> dis(0, std::distance(start, end) - 1);
    std::advance(start, dis(g));
    return start;
}

template<typename Iter>
Iter select_randomly(Iter start, Iter end)
{
    static std::random_device rd;
    static std::mt19937 gen(rd());
    return select_randomly(start, end, gen);
}

std::atomic<int> element_counter = 0;
std::atomic<int> failed_setup_counter = 0;

void createPathAndSetWatch(zkutil::ZooKeeper & zk, const String & path_prefix, size_t total)
{
    for (size_t i = 0; i < total; ++i)
    {
        int element = element_counter++;
        zk.createIfNotExists(path_prefix + "/" + std::to_string(element), "");

        std::string result;
        if (!zk.tryGetWatch(path_prefix + "/" + std::to_string(element), result, nullptr, triggerWatch))
            failed_setup_counter++;

        {
            std::lock_guard lock(elements_mutex);
            current_elements.push_back(element);
        }

        std::this_thread::sleep_for(std::chrono::milliseconds(200));

        {
            std::lock_guard lock(elements_mutex);
            if (current_elements.empty())
                continue;
            element = *select_randomly(current_elements.begin(), current_elements.end());
            current_elements.erase(std::remove(current_elements.begin(), current_elements.end(), element), current_elements.end());
        }
        zk.tryRemove(path_prefix + "/" + std::to_string(element));
    }

}

std::string random_string(size_t length)
{
    auto randchar = []() -> char
    {
        const char charset[] =
        "0123456789"
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "abcdefghijklmnopqrstuvwxyz";
        const size_t max_index = (sizeof(charset) - 1);
        return charset[rand() % max_index]; /// NOLINT
    };
    std::string str(length, 0);
    std::generate_n(str.begin(), length, randchar);
    return str;
}

std::string currentDateTime()
{
    time_t now = time(nullptr);
    tm tstruct;
    char buf[80];
    tstruct = *localtime(&now);
    // Visit http://en.cppreference.com/w/cpp/chrono/c/strftime
    // for more information about date/time format
    strftime(buf, sizeof(buf), "%Y-%m-%d %H:%M:%S", &tstruct);

    return buf;
}


void createOnPrefix(const std::string & zkhost, const String & path_prefix, size_t datasize, size_t total)
{
    zkutil::ZooKeeper zk(zkhost);
    std::vector<std::future<Coordination::CreateResponse>> holder_futures;
    using namespace std::chrono;
    try
    {
        for (size_t i = 0; i < total; ++i)
        {
            std::cerr << currentDateTime() << "] Request:" << i << std::endl;
            std::string path = path_prefix + "/element" + std::to_string(i);
            holder_futures.push_back(zk.asyncCreate(path, random_string(datasize), zkutil::CreateMode::Persistent));
        }

        for (auto & future : holder_futures)
            future.get();
    }
    catch (...)
    {
        ::exit(-1);
    }
}


void createConcurrent(zkutil::ZooKeeper & testzk, const std::string & zkhost, size_t threads, size_t requests, size_t blobsize)
{
    std::vector<std::future<void>> asyncs;
    for (size_t i = 0; i < threads; ++i)
    {
        std::string path_prefix = "/data/create_test" + std::to_string(i);
        testzk.createIfNotExists(path_prefix, "");
        auto callback = [&zkhost, path_prefix, requests, blobsize] ()
        {
            createOnPrefix(zkhost, path_prefix, blobsize, requests);
        };
        asyncs.push_back(std::async(std::launch::async, callback));
    }

    size_t i = 0;
    for (auto & async : asyncs)
    {
        async.wait();
        i++;
    }
}

void tryConcurrentWatches(zkutil::ZooKeeper & zk)
{
    std::string path_prefix = "/concurrent_watches";
    std::vector<std::future<void>> asyncs;
    zk.createIfNotExists(path_prefix, "");
    for (size_t i = 0; i < 100; ++i)
    {
        auto callback = [&zk, path_prefix] ()
        {
            createPathAndSetWatch(zk, path_prefix, 100);
        };
        asyncs.push_back(std::async(std::launch::async, callback));
    }

    for (auto & async : asyncs)
    {
        async.wait();
    }

    size_t counter = 0;
    while (watches_triggered != 100 * 100)
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
        if (counter++ > 20)
            break;
    }

    std::cerr << "Failed setup counter:" << failed_setup_counter << std::endl;
    std::cerr << "Current elements size:" << current_elements.size() << std::endl;
    std::cerr << "WatchesTriggered:" << watches_triggered << std::endl;
}


int main(int argc, char *argv[])
{
    if (argc != 2)
    {
        std::cerr << "usage: " << argv[0] << " hosts" << std::endl;
        return 2;
    }
    Poco::AutoPtr<Poco::ConsoleChannel> channel = new Poco::ConsoleChannel(std::cerr);
    Poco::Logger::root().setChannel(channel);
    Poco::Logger::root().setLevel("trace");

    zkutil::ZooKeeper zk(argv[1]);

    try
    {
        std::cerr << "Removing\n";
        zk.tryRemoveRecursive("/data");
        std::cerr << "Creating\n";
        zk.createIfNotExists("/data", "");
        std::cerr << "Created\n";

        Stopwatch watch;
        createConcurrent(zk, argv[1], 1, 1005000, 10);
        std::cerr << "Finished in: " << watch.elapsedMilliseconds() << "ms" << std::endl;

        //testCreateGetExistsNode(zk);
        //testCreateSetNode(zk);
        //testCreateList(zk);
        //testCreateSetVersionRequest(zk);
        //testMultiRequest(zk);
        //testCreateSetWatchEvent(zk);
        //testCreateListWatchEvent(zk);
        //tryConcurrentWatches(zk);
    }
    catch (...)
    {
        zk.tryRemoveRecursive("/data");
        throw;
    }
    return 0;
}