ClickHouse/utils/zookeeper-test/main.cpp
2020-12-17 19:21:46 +03:00

274 lines
8.5 KiB
C++

#include <IO/ReadBufferFromString.h>
#include <IO/ReadHelpers.h>
#include <Poco/ConsoleChannel.h>
#include <Common/ZooKeeper/KeeperException.h>
#include <Common/ZooKeeper/ZooKeeper.h>
#include <common/LineReader.h>
#include <common/logger_useful.h>
#include <fmt/format.h>
#include <random>
#include <iterator>
#include <algorithm>
#include <chrono>
#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)
{
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;
std::string result = zk.get("/data/check_data", &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)
{
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);
try
{
requests.clear();
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::terminate();
}
catch (...)
{
std::cerr << "Got exception on multy request (it's ok)\n";
}
checkEq(zk, "/data/multirequest", "bbb");
}
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));
}
}
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
{
zk.tryRemoveRecursive("/data");
testCreateGetExistsNode(zk);
testCreateSetNode(zk);
testCreateList(zk);
testCreateSetVersionRequest(zk);
testMultiRequest(zk);
testCreateSetWatchEvent(zk);
testCreateListWatchEvent(zk);
tryConcurrentWatches(zk);
}
catch (...)
{
zk.tryRemoveRecursive("/data");
throw;
}
return 0;
}