ClickHouse/utils/keeper-bench/Generator.h
2023-04-14 18:49:28 +00:00

200 lines
6.4 KiB
C++

#pragma once
#include <Common/ZooKeeper/KeeperException.h>
#include <Common/ZooKeeper/ZooKeeper.h>
#include <Common/ZooKeeper/ZooKeeperImpl.h>
#include <Common/ZooKeeper/ZooKeeperCommon.h>
#include <functional>
#include <optional>
#include <pcg-random/pcg_random.hpp>
#include <Poco/Util/AbstractConfiguration.h>
#include <Common/randomSeed.h>
struct NumberGetter
{
static NumberGetter fromConfig(const std::string & key, const Poco::Util::AbstractConfiguration & config, std::optional<uint64_t> default_value = std::nullopt);
uint64_t getNumber() const;
std::string description() const;
private:
struct NumberRange
{
uint64_t min_value;
uint64_t max_value;
};
std::variant<uint64_t, NumberRange> value;
};
struct StringGetter
{
explicit StringGetter(NumberGetter number_getter)
: value(std::move(number_getter))
{}
StringGetter() = default;
static StringGetter fromConfig(const std::string & key, const Poco::Util::AbstractConfiguration & config);
void setString(std::string name);
std::string getString() const;
std::string description() const;
bool isRandom() const;
private:
std::variant<std::string, NumberGetter> value;
};
struct PathGetter
{
static PathGetter fromConfig(const std::string & key, const Poco::Util::AbstractConfiguration & config);
std::string getPath() const;
std::string description() const;
void initialize(Coordination::ZooKeeper & zookeeper);
private:
std::vector<std::string> parent_paths;
bool initialized = false;
std::vector<std::string> paths;
mutable std::uniform_int_distribution<size_t> path_picker;
};
struct RequestGenerator
{
virtual ~RequestGenerator() = default;
void getFromConfig(const std::string & key, const Poco::Util::AbstractConfiguration & config);
Coordination::ZooKeeperRequestPtr generate(const Coordination::ACLs & acls);
std::string description();
void startup(Coordination::ZooKeeper & zookeeper);
size_t getWeight() const;
private:
virtual void getFromConfigImpl(const std::string & key, const Poco::Util::AbstractConfiguration & config) = 0;
virtual std::string descriptionImpl() = 0;
virtual Coordination::ZooKeeperRequestPtr generateImpl(const Coordination::ACLs & acls) = 0;
virtual void startupImpl(Coordination::ZooKeeper &) {}
size_t weight = 1;
};
using RequestGeneratorPtr = std::shared_ptr<RequestGenerator>;
struct CreateRequestGenerator final : public RequestGenerator
{
CreateRequestGenerator();
private:
void getFromConfigImpl(const std::string & key, const Poco::Util::AbstractConfiguration & config) override;
std::string descriptionImpl() override;
Coordination::ZooKeeperRequestPtr generateImpl(const Coordination::ACLs & acls) override;
void startupImpl(Coordination::ZooKeeper & zookeeper) override;
PathGetter parent_path;
StringGetter name;
std::optional<StringGetter> data;
std::optional<double> remove_factor;
pcg64 rng;
std::uniform_real_distribution<double> remove_picker;
std::unordered_set<std::string> paths_created;
};
struct SetRequestGenerator final : public RequestGenerator
{
private:
void getFromConfigImpl(const std::string & key, const Poco::Util::AbstractConfiguration & config) override;
std::string descriptionImpl() override;
Coordination::ZooKeeperRequestPtr generateImpl(const Coordination::ACLs & acls) override;
void startupImpl(Coordination::ZooKeeper & zookeeper) override;
PathGetter path;
StringGetter data;
};
struct GetRequestGenerator final : public RequestGenerator
{
private:
void getFromConfigImpl(const std::string & key, const Poco::Util::AbstractConfiguration & config) override;
std::string descriptionImpl() override;
Coordination::ZooKeeperRequestPtr generateImpl(const Coordination::ACLs & acls) override;
void startupImpl(Coordination::ZooKeeper & zookeeper) override;
PathGetter path;
};
struct ListRequestGenerator final : public RequestGenerator
{
private:
void getFromConfigImpl(const std::string & key, const Poco::Util::AbstractConfiguration & config) override;
std::string descriptionImpl() override;
Coordination::ZooKeeperRequestPtr generateImpl(const Coordination::ACLs & acls) override;
void startupImpl(Coordination::ZooKeeper & zookeeper) override;
PathGetter path;
};
struct RequestGetter
{
explicit RequestGetter(std::vector<RequestGeneratorPtr> request_generators_);
RequestGetter() = default;
static RequestGetter fromConfig(const std::string & key, const Poco::Util::AbstractConfiguration & config, bool for_multi = false);
RequestGeneratorPtr getRequestGenerator() const;
std::string description() const;
void startup(Coordination::ZooKeeper & zookeeper);
const std::vector<RequestGeneratorPtr> & requestGenerators() const;
private:
std::vector<RequestGeneratorPtr> request_generators;
std::vector<size_t> weights;
mutable std::uniform_int_distribution<size_t> request_generator_picker;
};
struct MultiRequestGenerator final : public RequestGenerator
{
private:
void getFromConfigImpl(const std::string & key, const Poco::Util::AbstractConfiguration & config) override;
std::string descriptionImpl() override;
Coordination::ZooKeeperRequestPtr generateImpl(const Coordination::ACLs & acls) override;
void startupImpl(Coordination::ZooKeeper & zookeeper) override;
std::optional<NumberGetter> size;
RequestGetter request_getter;
};
class Generator
{
public:
explicit Generator(const Poco::Util::AbstractConfiguration & config);
void startup(Coordination::ZooKeeper & zookeeper);
Coordination::ZooKeeperRequestPtr generate();
void cleanup(Coordination::ZooKeeper & zookeeper);
private:
struct Node
{
StringGetter name;
std::optional<StringGetter> data;
std::vector<std::shared_ptr<Node>> children;
size_t repeat_count = 0;
std::shared_ptr<Node> clone() const;
void createNode(Coordination::ZooKeeper & zookeeper, const std::string & parent_path, const Coordination::ACLs & acls) const;
void dumpTree(int level = 0) const;
};
static std::shared_ptr<Node> parseNode(const std::string & key, const Poco::Util::AbstractConfiguration & config);
std::uniform_int_distribution<size_t> request_picker;
std::vector<std::shared_ptr<Node>> root_nodes;
RequestGetter request_getter;
Coordination::ACLs default_acls;
};
std::optional<Generator> getGenerator(const std::string & name);