Files

590 lines
21 KiB
C++

#include <NodesSetup.h>
#include <fmt/ranges.h>
#include <deque>
#include <filesystem>
#include <iostream>
#include <ranges>
#include <Common/Exception.h>
#include <Common/ZooKeeper/KeeperException.h>
namespace DB::ErrorCodes
{
extern const int BAD_ARGUMENTS;
}
namespace
{
constexpr size_t WRITE_BATCH_SIZE = 10000;
constexpr size_t READ_BATCH_SIZE = 10000;
/// How often RequestBatcher prints a progress line, in batches.
constexpr size_t REPORT_EVERY_BATCHES = 10;
Coordination::MultiResponse executeMulti(Coordination::ZooKeeper & zookeeper, const Coordination::Requests & batch)
{
auto promise = std::make_shared<std::promise<Coordination::MultiResponse>>();
auto future = promise->get_future();
zookeeper.multi(batch, [promise](const Coordination::MultiResponse & response)
{
promise->set_value(response);
});
return future.get();
}
void flushMulti(Coordination::ZooKeeper & zookeeper, Coordination::Requests & batch)
{
if (batch.empty())
return;
auto response = executeMulti(zookeeper, batch);
if (response.error != Coordination::Error::ZOK)
throw zkutil::KeeperException(response.error, "Multi request failed");
batch.clear();
}
/// Accumulates requests and executes them in Multi requests of WRITE_BATCH_SIZE,
/// printing a progress line every REPORT_EVERY_BATCHES executed batches.
class RequestBatcher
{
public:
/// `action` is a past-tense verb for progress messages: "<action> <count>[/<total>] nodes".
/// Pass `total` of 0 if the total is not known in advance.
RequestBatcher(Coordination::ZooKeeper & zookeeper_, std::string action_, size_t total_)
: zookeeper(zookeeper_), action(std::move(action_)), total(total_)
{
}
void add(Coordination::RequestPtr request)
{
batch.push_back(std::move(request));
if (batch.size() >= WRITE_BATCH_SIZE)
flush();
}
/// Execute the last incomplete batch and print the final count.
void finish()
{
flush();
if (executed != last_reported)
report();
}
private:
void flush()
{
if (batch.empty())
return;
size_t batch_size = batch.size();
flushMulti(zookeeper, batch);
executed += batch_size;
++batches;
if (batches % REPORT_EVERY_BATCHES == 0)
report();
}
void report()
{
if (total != 0)
std::cerr << action << " " << executed << "/" << total << " nodes" << std::endl;
else
std::cerr << action << " " << executed << " nodes" << std::endl;
last_reported = executed;
}
Coordination::ZooKeeper & zookeeper;
std::string action;
size_t total = 0;
Coordination::Requests batch;
size_t executed = 0;
size_t batches = 0;
size_t last_reported = 0;
};
void removeRecursiveManual(Coordination::ZooKeeper & zookeeper, const std::string & path, RequestBatcher & remove_batcher)
{
namespace fs = std::filesystem;
auto promise = std::make_shared<std::promise<Coordination::Error>>();
auto future = promise->get_future();
std::vector<std::string> children;
auto list_callback = [promise, &children](const Coordination::ListResponse & response)
{
children = response.names;
promise->set_value(response.error);
};
zookeeper.list(path, Coordination::ListRequestType::ALL, list_callback, {}, false, false);
auto error = future.get();
if (error == Coordination::Error::ZNONODE)
return;
if (error != Coordination::Error::ZOK)
throw zkutil::KeeperException(error, "Failed to list children of {}", path);
for (const auto & child : children)
removeRecursiveManual(zookeeper, fs::path(path) / child, remove_batcher);
remove_batcher.add(zkutil::makeRemoveRequest(path, -1));
}
/// Like removeRecursiveManual, but suitable for huge subtrees (~100M nodes):
/// * List requests are batched into MultiRead requests instead of one round trip per node,
/// * nodes that their parent's listing showed to be childless are removed without listing.
/// Requires the MULTI_READ, FILTERED_LIST, and LIST_WITH_STAT_AND_DATA feature flags.
void removeRecursiveManualFast(Coordination::ZooKeeper & zookeeper, const std::string & root_path, RequestBatcher & remove_batcher)
{
namespace fs = std::filesystem;
/// Nodes that had children when their parent was listed (plus the root).
/// The walk below appends parents before their children (BFS pre-order), so
/// removing in reverse order removes children before parents. Nodes without
/// children are removed right away instead.
std::vector<std::string> listed;
/// Nodes to list: path and the number of children reported by the parent's listing.
std::deque<std::pair<std::string, size_t>> queue;
queue.emplace_back(root_path, 1);
while (!queue.empty())
{
/// Group List requests into one MultiRead, capping the expected total
/// number of listed children to bound the response size.
std::vector<std::string> batch_paths;
size_t expected_children = 0;
while (!queue.empty() && (batch_paths.empty() || expected_children + queue.front().second <= READ_BATCH_SIZE))
{
expected_children += queue.front().second;
batch_paths.push_back(std::move(queue.front().first));
queue.pop_front();
}
Coordination::Requests list_batch;
list_batch.reserve(batch_paths.size());
for (const auto & path : batch_paths)
list_batch.push_back(zkutil::makeListRequest(path, Coordination::ListRequestType::ALL, /*with_stat=*/ true, /*with_data=*/ false));
auto multi_response = executeMulti(zookeeper, list_batch);
/// An error of a subrequest (e.g. ZNONODE, which we tolerate) is propagated to the
/// error of the whole MultiRead. So check the subresponse errors first, and treat the
/// multi-level error as fatal only if no subresponse explains it (e.g. a connection
/// loss, where the subresponses are left blank and must not be trusted).
chassert(multi_response.responses.size() == batch_paths.size());
bool sub_error_found = false;
for (size_t i = 0; i < batch_paths.size(); ++i)
{
auto sub_error = multi_response.responses.at(i)->error;
if (sub_error == Coordination::Error::ZNONODE)
sub_error_found = true;
else if (sub_error != Coordination::Error::ZOK)
throw zkutil::KeeperException(sub_error, "Failed to list children of {}", batch_paths[i]);
}
if (multi_response.error != Coordination::Error::ZOK && !sub_error_found)
throw zkutil::KeeperException(multi_response.error, "MultiRead of {} List requests failed", list_batch.size());
for (size_t i = 0; i < batch_paths.size(); ++i)
{
const auto & response = *multi_response.responses.at(i);
/// The node was removed concurrently (e.g. an expired ephemeral); nothing to do.
if (response.error == Coordination::Error::ZNONODE)
continue;
const auto & list_response = dynamic_cast<const Coordination::ListResponse &>(response);
chassert(list_response.stats.size() == list_response.names.size());
for (size_t j = 0; j < list_response.names.size(); ++j)
{
std::string child_path = fs::path(batch_paths[i]) / list_response.names[j];
if (list_response.stats[j].numChildren == 0)
remove_batcher.add(zkutil::makeRemoveRequest(child_path, -1));
else
queue.emplace_back(std::move(child_path), list_response.stats[j].numChildren);
}
listed.push_back(std::move(batch_paths[i]));
}
}
for (auto it = listed.rbegin(); it != listed.rend(); ++it)
remove_batcher.add(zkutil::makeRemoveRequest(*it, -1));
}
void removeRecursive(Coordination::ZooKeeper & zookeeper, const std::string & path, bool allow_native)
{
if (allow_native && zookeeper.isFeatureEnabled(DB::KeeperFeatureFlag::REMOVE_RECURSIVE))
{
auto promise = std::make_shared<std::promise<Coordination::Error>>();
auto future = promise->get_future();
/// A native RemoveRecursive request is indivisible: the whole subtree is removed
/// in one raft entry, blocking all other writes for its duration. Limit its size
/// and fall back to manual removal (batched into many multis) for bigger subtrees.
zookeeper.removeRecursive(path, /*remove_nodes_limit=*/ 1000000,
[promise](const Coordination::RemoveRecursiveResponse & response)
{
promise->set_value(response.error);
});
auto error = future.get();
if (error == Coordination::Error::ZNONODE || error == Coordination::Error::ZOK)
return;
/// ZNOTEMPTY means the subtree exceeded remove_nodes_limit (nothing was removed).
if (error != Coordination::Error::ZNOTEMPTY)
throw zkutil::KeeperException(error, "Failed to recursively remove {}", path);
std::cerr << "Subtree at " << path << " is too large for one native RemoveRecursive request, removing manually" << std::endl;
}
RequestBatcher remove_batcher(zookeeper, "Removed", /*total=*/ 0);
if (zookeeper.isFeatureEnabled(DB::KeeperFeatureFlag::MULTI_READ)
&& zookeeper.isFeatureEnabled(DB::KeeperFeatureFlag::FILTERED_LIST)
&& zookeeper.isFeatureEnabled(DB::KeeperFeatureFlag::LIST_WITH_STAT_AND_DATA))
removeRecursiveManualFast(zookeeper, path, remove_batcher);
else
removeRecursiveManual(zookeeper, path, remove_batcher);
remove_batcher.finish();
}
}
void NodesSetup::initializeFromConfig(const Poco::Util::AbstractConfiguration & config)
{
default_acls = getDefaultACLs();
use_remove_recursive = config.getBool("use_remove_recursive", true);
std::cerr << "---- Parsing setup ---- " << std::endl;
static const std::string setup_key = "setup";
Poco::Util::AbstractConfiguration::Keys keys;
config.keys(setup_key, keys);
for (const auto & key : keys)
{
if (key.starts_with("node"))
{
const auto node = root_nodes.emplace_back(parseNode(setup_key + "." + key, config));
std::cerr << "Tree to create:" << std::endl;
node->dumpTree();
std::cerr << std::endl;
}
}
std::cerr << "---- Done parsing data setup ----\n" << std::endl;
}
std::shared_ptr<NodesSetup::Node> NodesSetup::parseNode(const std::string & key, const Poco::Util::AbstractConfiguration & config)
{
auto node = std::make_shared<NodesSetup::Node>();
node->name = StringGetter::fromConfig(key + ".name", config);
if (config.has(key + ".data"))
node->data = StringGetter::fromConfig(key + ".data", config);
if (config.has(key + ".tag"))
{
node->tag = config.getString(key + ".tag");
node->tag_set = getOrCreateTagSet(*node->tag);
node->tag_set->is_setup_tag = true;
}
if (config.has(key + ".repeat"))
{
if (!node->name.isRandom())
throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Repeating node creation for key {}, but name is not randomly generated", key);
node->repeat_count = config.getUInt64(key + ".repeat");
if (node->repeat_count == 0)
throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "repeat must be >= 1 for key {}", key);
}
Poco::Util::AbstractConfiguration::Keys node_keys;
config.keys(key, node_keys);
for (const auto & node_key : node_keys)
{
if (!node_key.starts_with("node"))
continue;
node->children.push_back(parseNode(key + "." + node_key, config));
}
return node;
}
size_t NodesSetup::Node::countNodes() const
{
size_t count = 1;
for (const auto & child : children)
count += child->repeat_count * child->countNodes();
return count;
}
void NodesSetup::Node::dumpTree(int level) const
{
std::string data_string
= data.has_value() ? fmt::format("{}", data->description()) : "no data";
std::string repeat_count_string = repeat_count != 1 ? fmt::format(", repeated {} times", repeat_count) : "";
std::string tag_string = tag.has_value() ? fmt::format(", tag: \"{}\"", *tag) : "";
std::cerr << fmt::format("{}name: {}, data: {}{}{}", std::string(level, '\t'), name.description(), data_string, repeat_count_string, tag_string) << std::endl;
for (const auto & child : children)
child->dumpTree(level + 1);
}
void NodesSetup::Node::createAt(const std::string & path, const CreateRequestSink & sink, const Coordination::ACLs & acls, pcg64 & rng_) const
{
auto request = std::make_shared<Coordination::ZooKeeperCreateRequest>();
request->path = path;
request->data = data ? data->getString(rng_) : "";
request->acls = acls;
sink(std::move(request));
if (tag_set && tag_set->used_as_input)
tag_set->populate(path);
for (const auto & child : children)
for (size_t i = 0; i < child->repeat_count; ++i)
child->createAt(std::filesystem::path(path) / child->name.getString(rng_), sink, acls, rng_);
}
PathSetPtr NodesSetup::getOrCreateTagSet(const std::string & tag)
{
auto & set = tag_sets[tag];
if (!set)
{
set = std::make_shared<PathSet>();
set->name = fmt::format("tag \"{}\"", tag);
}
return set;
}
PathSetPtr NodesSetup::getOrCreateChildrenOfSet(const std::string & parent_path)
{
auto & set = children_of_sets[parent_path];
if (!set)
{
set = std::make_shared<PathSet>();
set->name = fmt::format("children of {}", parent_path);
set->is_children_of = true;
set->children_of_parent = parent_path;
}
return set;
}
PathSetPtr NodesSetup::createLiteralSet(std::vector<std::string> paths)
{
auto set = std::make_shared<PathSet>();
set->name = fmt::format("paths [{}]", fmt::join(paths, ", "));
set->is_literal = true;
for (auto & path : paths)
set->populate(std::move(path));
literal_sets.push_back(set);
return set;
}
PathSetPtr NodesSetup::createAnonymousSet(std::string display_name)
{
auto set = std::make_shared<PathSet>();
set->name = std::move(display_name);
anonymous_sets.push_back(set);
return set;
}
void NodesSetup::registerTagChildrenOfConflict(std::string parent_path, std::string tag_name)
{
tag_children_of_conflicts.emplace_back(std::move(parent_path), std::move(tag_name));
}
void NodesSetup::finalizePathSets(size_t num_threads)
{
auto finalize = [&](const PathSetPtr & set)
{
/// A set some generator writes to changes while the benchmark runs.
set->is_dynamic = set->used_as_output;
set->finalize(num_threads);
};
for (const auto & [_, set] : tag_sets)
finalize(set);
for (const auto & [_, set] : children_of_sets)
finalize(set);
for (const auto & set : literal_sets)
finalize(set);
for (const auto & set : anonymous_sets)
finalize(set);
}
void NodesSetup::resolveChildrenOf(const ListChildrenFn & list_children)
{
for (const auto & [parent_path, set] : children_of_sets)
{
if (!set->used_as_input)
continue;
for (const auto & child : list_children(parent_path))
set->populate(std::filesystem::path(parent_path) / child);
}
}
void NodesSetup::validatePathSets()
{
for (const auto & [parent_path, tag_name] : tag_children_of_conflicts)
{
auto it = children_of_sets.find(parent_path);
if (it != children_of_sets.end() && it->second->used_as_input)
throw DB::Exception(
DB::ErrorCodes::BAD_ARGUMENTS,
"A create generator with parent {} outputs to tag \"{}\", but another generator reads `children_of` {}: "
"the created nodes would be tracked in one set and read from another. Use the tag consistently "
"(reference `tagged: {}` instead of `children_of`)",
parent_path, tag_name, parent_path, tag_name);
}
bool printed_header = false;
auto validate = [&](const PathSetPtr & set)
{
bool defined = set->is_setup_tag || set->is_children_of || set->is_literal || set->used_as_output;
if (set->used_as_input && !defined)
throw DB::Exception(
DB::ErrorCodes::BAD_ARGUMENTS,
"{} is used by a request generator, but no setup node has this tag "
"and no create generator outputs to it. Available tags: {}",
set->name,
tag_sets.empty() ? "(none)" : fmt::to_string(fmt::join(tag_sets | std::views::keys, ", ")));
if (set->used_as_input && !set->is_dynamic && set->totalSize() == 0)
throw DB::Exception(
DB::ErrorCodes::BAD_ARGUMENTS,
"{} is empty: check that the `children_of` target has children or that setup nodes carry the tag",
set->name);
if (set->keep_count)
{
size_t num_shards = set->shards.size();
size_t total_target = *set->keep_count == 0 ? set->totalSize() : *set->keep_count;
set->target_count_per_shard = std::max<size_t>(1, total_target / num_shards);
}
if (set->used_as_input)
{
if (!printed_header)
{
std::cerr << "Path sets:" << std::endl;
printed_header = true;
}
std::string target_string
= set->target_count_per_shard != 0 ? fmt::format(", target {} per thread", set->target_count_per_shard) : "";
std::cerr << fmt::format(" {}: {} paths{}{}", set->name, set->totalSize(), set->is_dynamic ? " (dynamic)" : "", target_string)
<< std::endl;
}
};
for (const auto & [_, set] : tag_sets)
validate(set);
for (const auto & [_, set] : children_of_sets)
validate(set);
for (const auto & set : literal_sets)
validate(set);
for (const auto & set : anonymous_sets)
validate(set);
}
std::string NodesSetup::describeDynamicPathSets() const
{
std::string result;
auto append = [&](const PathSetPtr & set)
{
if (!set->is_dynamic || !set->used_as_input)
return;
if (!result.empty())
result += ", ";
result += fmt::format("{}: {}", set->name, set->totalSize());
if (set->target_count_per_shard != 0)
result += fmt::format(" (target {})", set->target_count_per_shard * set->shards.size());
};
for (const auto & [_, set] : tag_sets)
append(set);
for (const auto & [_, set] : children_of_sets)
append(set);
for (const auto & set : literal_sets)
append(set);
for (const auto & set : anonymous_sets)
append(set);
return result;
}
const std::vector<std::string> & NodesSetup::prepareRootPaths()
{
if (created_root_paths.empty())
{
for (const auto & node : root_nodes)
for (size_t i = 0; i < node->repeat_count; ++i)
created_root_paths.push_back(std::filesystem::path("/") / node->name.getString(rng));
}
return created_root_paths;
}
void NodesSetup::createNodes(const CreateRequestSink & sink)
{
const auto & root_paths = prepareRootPaths();
size_t path_idx = 0;
for (const auto & node : root_nodes)
for (size_t i = 0; i < node->repeat_count; ++i)
node->createAt(root_paths.at(path_idx++), sink, default_acls, rng);
}
size_t NodesSetup::countTotalNodes() const
{
size_t total = 0;
for (const auto & node : root_nodes)
total += node->repeat_count * node->countNodes();
return total;
}
void NodesSetup::startup(Coordination::ZooKeeper & zookeeper)
{
if (root_nodes.empty())
return;
std::cerr << "---- Creating test data ----" << std::endl;
/// Remove leftovers from previous runs.
for (const auto & root_path : prepareRootPaths())
{
std::cerr << "Cleaning up " << root_path << std::endl;
removeRecursive(zookeeper, root_path, use_remove_recursive);
}
size_t total_nodes = countTotalNodes();
std::cerr << "Creating " << total_nodes << " nodes" << std::endl;
RequestBatcher create_batcher(zookeeper, "Created", total_nodes);
createNodes([&](std::shared_ptr<Coordination::ZooKeeperCreateRequest> request)
{
create_batcher.add(std::move(request));
});
create_batcher.finish();
std::cerr << "---- Created test data ----\n" << std::endl;
}
void NodesSetup::cleanup(Coordination::ZooKeeper & zookeeper)
{
if (created_root_paths.empty())
return;
std::cerr << "---- Cleaning up test data ----" << std::endl;
for (const auto & root_path : created_root_paths)
{
std::cerr << "Cleaning up " << root_path << std::endl;
removeRecursive(zookeeper, root_path, use_remove_recursive);
}
}