From 61f28d3ab2af7a2b3b7e0d746fde2113b14fc036 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Botond=20D=C3=A9nes?= Date: Wed, 25 Jan 2023 03:44:08 -0500 Subject: [PATCH] test/lib: add key_utils.hh Contains methods to generate partition and clustering keys. In the case of the former, one can specify the shard to generate keys for. We currently have some methods to generate these but they are not generic. Therefore the tests are littered by open-coded variants. The methods introduced here are completely generic: they can generate keys for any schema. --- test/lib/key_utils.cc | 109 ++++++++++++++++++++++++++++++++++++++ test/lib/key_utils.hh | 48 +++++++++++++++++ test/lib/simple_schema.hh | 4 +- 3 files changed, 158 insertions(+), 3 deletions(-) create mode 100644 test/lib/key_utils.cc create mode 100644 test/lib/key_utils.hh diff --git a/test/lib/key_utils.cc b/test/lib/key_utils.cc new file mode 100644 index 0000000000..1b5bb22330 --- /dev/null +++ b/test/lib/key_utils.cc @@ -0,0 +1,109 @@ +/* + * Copyright (C) 2023-present ScyllaDB + */ + +/* + * SPDX-License-Identifier: AGPL-3.0-or-later + */ + +#include "test/lib/key_utils.hh" +#include "test/lib/random_schema.hh" +#include "test/lib/random_utils.hh" +#include "test/lib/test_utils.hh" + +namespace tests { + +namespace { + +template +std::vector generate_keys( + size_t n, + schema_ptr s, + Comparator cmp, + const std::vector& types, + std::function(const RawKey&)> decorate_fun, + bool allow_prefixes, + std::optional size) { + auto keys = std::set(cmp); + const auto effective_size = size.value_or(tests::key_size{1, 128}); + + std::mt19937 engine(tests::random::get_int()); + std::uniform_int_distribution component_count_dist(1, types.size()); + tests::value_generator value_gen; + + std::vector components; + components.reserve(types.size()); + + while (keys.size() != n) { + components.clear(); + auto component_count = allow_prefixes ? component_count_dist(engine) : types.size(); + for (size_t i = 0; i < component_count; ++i) { + components.emplace_back(value_gen.generate_atomic_value(engine, *types.at(i), effective_size.min, effective_size.max)); + } + auto raw_key = RawKey::from_deeply_exploded(*s, components); + // discard empty keys on the off chance that we generate one + if (raw_key.is_empty() || (types.size() == 1 && raw_key.begin(*s)->empty())) { + continue; + } + if constexpr (std::is_same_v) { + keys.emplace(std::move(raw_key)); + } else if (auto decorated_key_opt = decorate_fun(raw_key); decorated_key_opt) { + keys.emplace(std::move(*decorated_key_opt)); + } + } + + return std::vector(keys.begin(), keys.end()); +} + +} + +std::vector generate_partition_keys(size_t n, schema_ptr s, std::optional shard, std::optional size) { + return generate_keys( + n, + s, + dht::decorated_key::less_comparator(s), + s->partition_key_type()->types(), + [s, shard, tokens = std::set()] (const partition_key& pkey) mutable -> std::optional { + auto dkey = dht::decorate_key(*s, pkey); + if (shard && *shard != dht::shard_of(*s, dkey.token())) { + return {}; + } + if (!tokens.insert(dkey.token()).second) { + return {}; + } + return dkey; + }, + false, + size); +} + +std::vector generate_partition_keys(size_t n, schema_ptr s, local_shard_only lso, std::optional size) { + return generate_partition_keys(n, std::move(s), lso == local_shard_only::yes ? std::optional(this_shard_id()) : std::nullopt, size); +} + +dht::decorated_key generate_partition_key(schema_ptr s, std::optional shard, std::optional size) { + auto&& keys = generate_partition_keys(1, std::move(s), shard, size); + return std::move(keys.front()); +} + +dht::decorated_key generate_partition_key(schema_ptr s, local_shard_only lso, std::optional size) { + return generate_partition_key(std::move(s), lso == local_shard_only::yes ? std::optional(this_shard_id()) : std::nullopt, size); +} + +std::vector generate_clustering_keys(size_t n, schema_ptr s, bool allow_prefixes, std::optional size) { + return generate_keys( + n, + s, + clustering_key::less_compare(*s), + s->clustering_key_type()->types(), + {}, + allow_prefixes, + size); +} + +clustering_key generate_clustering_key(schema_ptr s, bool allow_prefix, std::optional size) { + auto&& keys = generate_clustering_keys(1, std::move(s), allow_prefix, size); + return std::move(keys.front()); +} + +} // namespace tests diff --git a/test/lib/key_utils.hh b/test/lib/key_utils.hh new file mode 100644 index 0000000000..b2ec6979f4 --- /dev/null +++ b/test/lib/key_utils.hh @@ -0,0 +1,48 @@ +/* + * Copyright (C) 2023-present ScyllaDB + */ + +/* + * SPDX-License-Identifier: AGPL-3.0-or-later + */ + +#pragma once + +#include "dht/i_partitioner.hh" + +struct local_shard_only_tag { }; +using local_shard_only = bool_class; + +namespace tests { + +struct key_size { + size_t min; + size_t max; +}; + +// Generate n partition keys for the given schema. +// +// Returned keys are unique (their token too), ordered and never empty. +// Parameters: +// * n - number of keys +// * s - schema of the keys, used also to obtain the sharder +// * shard - only generate keys for this shard (if engaged) +// * size - the min and max size of the key in bytes, if disengaged default +// limits (1-128) are used. If you want exactly sized keys, use +// ascii or bytes types only as the key types. +std::vector generate_partition_keys(size_t n, schema_ptr s, std::optional shard = this_shard_id(), std::optional size = {}); +std::vector generate_partition_keys(size_t n, schema_ptr s, local_shard_only lso, std::optional size = {}); +// Overload for a single key +dht::decorated_key generate_partition_key(schema_ptr s, std::optional shard = this_shard_id(), std::optional size = {}); +dht::decorated_key generate_partition_key(schema_ptr s, local_shard_only lso, std::optional size = {}); + +// Generate n clustering keys +// +// Returned keys are unique, ordered and never empty. +// Parameters are the same as that of generate_partition_keys(). +// If allow_prefixes is true, prefix keys may be generated too. +std::vector generate_clustering_keys(size_t n, schema_ptr s, bool allow_prefixes = false, std::optional size = {}); +// Overload for a single key +clustering_key generate_clustering_key(schema_ptr s, bool allow_prefix = false, std::optional size = {}); + +} // namespace tests diff --git a/test/lib/simple_schema.hh b/test/lib/simple_schema.hh index a3a95b1010..2d3de6a596 100644 --- a/test/lib/simple_schema.hh +++ b/test/lib/simple_schema.hh @@ -20,11 +20,9 @@ #include "schema_builder.hh" #include "reader_permit.hh" #include "types/map.hh" +#include "test/lib/key_utils.hh" #include "atomic_cell_or_collection.hh" -struct local_shard_only_tag { }; -using local_shard_only = bool_class; - // // Make set of keys sorted by token for current or remote shard. //