fix(set): fix random in SRANDMEMBER and SPOP commands (#3022)

Signed-off-by: Stepan Bagritsevich <sbagritsevich@quantumbrains.com>
This commit is contained in:
Stepan Bagritsevich 2024-05-13 11:08:01 +04:00 committed by GitHub
parent 4cd142d42c
commit d3a585113f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 297 additions and 164 deletions

View file

@ -4,6 +4,7 @@
#pragma once
#include <absl/random/random.h>
#include <absl/strings/ascii.h>
#include <absl/strings/str_cat.h>
#include <absl/types/span.h>
@ -309,4 +310,47 @@ struct MemoryBytesFlag {
bool AbslParseFlag(std::string_view in, dfly::MemoryBytesFlag* flag, std::string* err);
std::string AbslUnparseFlag(const dfly::MemoryBytesFlag& flag);
using RandomPick = std::uint32_t;
class PicksGenerator {
public:
virtual RandomPick Generate() = 0;
virtual ~PicksGenerator() = default;
};
class NonUniquePicksGenerator : public PicksGenerator {
public:
/* The generated value will be within the closed-open interval [0, max_range) */
NonUniquePicksGenerator(RandomPick max_range);
RandomPick Generate() override;
private:
const RandomPick max_range_;
absl::BitGen bitgen_{};
};
/*
* Generates unique index in O(1).
*
* picks_count specifies the number of random indexes to be generated.
* In other words, this is the number of times the Generate() function is called.
*
* The class uses Robert Floyd's sampling algorithm
* https://dl.acm.org/doi/pdf/10.1145/30401.315746
* */
class UniquePicksGenerator : public PicksGenerator {
public:
/* The generated value will be within the closed-open interval [0, max_range) */
UniquePicksGenerator(std::uint32_t picks_count, RandomPick max_range);
RandomPick Generate() override;
private:
RandomPick current_random_limit_;
std::uint32_t remaining_picks_count_;
std::unordered_set<RandomPick> picked_indexes_;
absl::BitGen bitgen_{};
};
} // namespace dfly