diff --git a/mooncake-store/AGENTS.md b/mooncake-store/AGENTS.md new file mode 100644 index 0000000000..7e2480637e --- /dev/null +++ b/mooncake-store/AGENTS.md @@ -0,0 +1,13 @@ +# Mooncake Store Instructions + +## Random Numbers + +- Production code in `include/` and `src/` must use `include/random.h`; do not + create random engines, seeds, or distributions at call sites. +- Extend `random.h` and its tests when an operation is missing. Keep bounded + sampling unbiased and reject invalid bounds. +- Tests and benchmarks may use explicit seeds for reproducibility, but should + use the shared sampling helpers. +- Do not migrate `src/cachelib_memory_allocator/` unless explicitly requested. +- The shared engine is not cryptographically secure; do not use it for secrets + or authentication tokens. diff --git a/mooncake-store/include/allocation_strategy.h b/mooncake-store/include/allocation_strategy.h index 558a4c3272..2417369079 100644 --- a/mooncake-store/include/allocation_strategy.h +++ b/mooncake-store/include/allocation_strategy.h @@ -2,7 +2,6 @@ #include #include -#include #include #include #include @@ -13,6 +12,7 @@ #include "allocator.h" // Contains BufferAllocator declaration #include "replica.h" #include "types.h" +#include "random.h" namespace mooncake { @@ -242,9 +242,6 @@ class RandomAllocationStrategy : public AllocationStrategy { return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } - // Random number generator. - static thread_local std::mt19937 generator(std::random_device{}()); - std::vector replicas; replicas.reserve(replica_num); @@ -254,8 +251,8 @@ class RandomAllocationStrategy : public AllocationStrategy { return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } - auto buffer = allocateSingle(allocator_manager, names[0], - slice_length, generator); + auto buffer = + allocateSingle(allocator_manager, names[0], slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -275,7 +272,7 @@ class RandomAllocationStrategy : public AllocationStrategy { } auto buffer = allocateSingle(allocator_manager, preferred_segment, - slice_length, generator); + slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -290,8 +287,7 @@ class RandomAllocationStrategy : public AllocationStrategy { // If replica_num is not satisfied, allocate the remaining replicas // randomly. - std::uniform_int_distribution distribution(0, names.size() - 1); - size_t start_idx = distribution(generator); + size_t start_idx = randomIndex(names.size()); const size_t max_retry = std::min(kMaxRetryLimit, names.size()); size_t try_count = 0; @@ -307,8 +303,8 @@ class RandomAllocationStrategy : public AllocationStrategy { continue; } - auto buffer = allocateSingle(allocator_manager, names[index], - slice_length, generator); + auto buffer = + allocateSingle(allocator_manager, names[index], slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -328,9 +324,6 @@ class RandomAllocationStrategy : public AllocationStrategy { tl::expected AllocateFrom( const AllocatorManager& allocator_manager, const size_t slice_length, const std::string& segment_name) { - // Random number generator. - static thread_local std::mt19937 generator(std::random_device{}()); - // Validate input parameters if (slice_length == 0) { return tl::make_unexpected(ErrorCode::INVALID_PARAMS); @@ -341,8 +334,8 @@ class RandomAllocationStrategy : public AllocationStrategy { return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND); } - auto buffer = allocateSingle(allocator_manager, segment_name, - slice_length, generator); + auto buffer = + allocateSingle(allocator_manager, segment_name, slice_length); if (buffer == nullptr) { return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } @@ -352,7 +345,7 @@ class RandomAllocationStrategy : public AllocationStrategy { std::unique_ptr allocateSingle( const AllocatorManager& allocator_manager, const std::string& name, - const size_t slice_length, std::mt19937& generator) { + const size_t slice_length) { const auto allocators = allocator_manager.getAllocators(name); if (allocators == nullptr || allocators->size() == 0) { return nullptr; @@ -366,9 +359,8 @@ class RandomAllocationStrategy : public AllocationStrategy { // Randomly select a start point to distribute // allocations across all segments - std::uniform_int_distribution dist(0, num_segs - 1); - size_t seg_offset = - dist(generator); // select a start segment to place replica + // Select a start segment to place the replica. + size_t seg_offset = randomIndex(num_segs); for (size_t i = 0; i < num_segs; i++) { // only allocate one replica auto& allocator = (*allocators)[(i + seg_offset) % num_segs]; if (auto buffer = allocator->allocate(slice_length)) { @@ -420,8 +412,6 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } - static thread_local std::mt19937 generator(std::random_device{}()); - std::vector replicas; replicas.reserve(replica_num); std::set used_segments; @@ -434,7 +424,7 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { } auto buffer = allocateSingle(allocator_manager, preferred_segment, - slice_length, generator); + slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -452,8 +442,7 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { size_t sample_count = std::min(kCandidateMultiplier * remaining, names.size()); - std::uniform_int_distribution start_dist(0, names.size() - 1); - size_t start_idx = start_dist(generator); + size_t start_idx = randomIndex(names.size()); struct Candidate { size_t name_idx; @@ -489,8 +478,7 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { continue; } - auto buffer = allocateSingle(allocator_manager, name, slice_length, - generator); + auto buffer = allocateSingle(allocator_manager, name, slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -503,8 +491,7 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { } // --- Fallback: Random allocation for any remaining replicas --- - std::uniform_int_distribution distribution(0, names.size() - 1); - size_t fallback_idx = distribution(generator); + size_t fallback_idx = randomIndex(names.size()); const size_t max_retry = std::min(kMaxRetryLimit, names.size()); size_t try_count = 0; @@ -519,8 +506,8 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { continue; } - auto buffer = allocateSingle(allocator_manager, names[index], - slice_length, generator); + auto buffer = + allocateSingle(allocator_manager, names[index], slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -578,8 +565,6 @@ class SsdFreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } - static thread_local std::mt19937 generator(std::random_device{}()); - std::vector replicas; replicas.reserve(replica_num); std::set used_segments; @@ -592,7 +577,7 @@ class SsdFreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { } auto buffer = allocateSingle(allocator_manager, preferred_segment, - slice_length, generator); + slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -609,8 +594,7 @@ class SsdFreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { size_t sample_count = std::min(kCandidateMultiplier * remaining, names.size()); - std::uniform_int_distribution start_dist(0, names.size() - 1); - size_t start_idx = start_dist(generator); + size_t start_idx = randomIndex(names.size()); struct Candidate { size_t name_idx; @@ -643,8 +627,7 @@ class SsdFreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { } const auto& name = names[candidate.name_idx]; - auto buffer = allocateSingle(allocator_manager, name, slice_length, - generator); + auto buffer = allocateSingle(allocator_manager, name, slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -657,8 +640,7 @@ class SsdFreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { } // Fallback: Random allocation for remaining replicas - std::uniform_int_distribution distribution(0, names.size() - 1); - size_t fallback_idx = distribution(generator); + size_t fallback_idx = randomIndex(names.size()); const size_t max_retry = std::min(kMaxRetryLimit, names.size()); size_t try_count = 0; @@ -674,8 +656,7 @@ class SsdFreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { continue; } - auto buffer = allocateSingle(allocator_manager, name, slice_length, - generator); + auto buffer = allocateSingle(allocator_manager, name, slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); diff --git a/mooncake-store/include/random.h b/mooncake-store/include/random.h new file mode 100644 index 0000000000..8a16c590db --- /dev/null +++ b/mooncake-store/include/random.h @@ -0,0 +1,87 @@ +#pragma once + +#include +#include +#include +#include + +namespace mooncake { + +using RandomEngine = std::mt19937_64; + +namespace detail { + +template +concept UniformDistributionInteger = + std::same_as || std::same_as || + std::same_as || std::same_as || + std::same_as || + std::same_as || + std::same_as || + std::same_as; + +template +Result sampleUniform(Result lower_bound, Result upper_bound, + Generator& generator) { + std::uniform_int_distribution distribution( + static_cast(lower_bound), + static_cast(upper_bound)); + return static_cast(distribution(generator)); +} + +} // namespace detail + +// Returns the pseudo-random engine shared by random helpers on this thread. +// The engine is seeded once per thread and is not safe to use from another +// thread. +inline RandomEngine& threadLocalRandomEngine() { + thread_local RandomEngine engine = [] { + std::random_device device; + std::seed_seq seed{device(), device(), device(), device(), + device(), device(), device(), device()}; + return RandomEngine(seed); + }(); + return engine; +} + +template +size_t randomIndex(size_t upper_bound, Generator& generator) { + if (upper_bound == 0) { + throw std::invalid_argument("randomIndex upper bound must be positive"); + } + std::uniform_int_distribution distribution(0, upper_bound - 1); + return distribution(generator); +} + +// Returns an unbiased random index in [0, upper_bound). +inline size_t randomIndex(size_t upper_bound) { + return randomIndex(upper_bound, threadLocalRandomEngine()); +} + +template +Integer randomUniform(Integer lower_bound, Integer upper_bound, + Generator& generator) { + if (lower_bound > upper_bound) { + throw std::invalid_argument( + "randomUniform lower bound must not exceed upper bound"); + } + if constexpr (detail::UniformDistributionInteger) { + return detail::sampleUniform(lower_bound, upper_bound, + generator); + } else if constexpr (std::signed_integral) { + return detail::sampleUniform( + lower_bound, upper_bound, generator); + } else { + return detail::sampleUniform( + lower_bound, upper_bound, generator); + } +} + +// Returns an unbiased random integer in [lower_bound, upper_bound]. +template +Integer randomUniform(Integer lower_bound, Integer upper_bound) { + return randomUniform(lower_bound, upper_bound, threadLocalRandomEngine()); +} + +} // namespace mooncake diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 13beddbe37..343e7f19ea 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -9,7 +9,6 @@ #include #include #include -#include #include #include #include @@ -44,6 +43,7 @@ #include "ha/snapshot/snapshot_logger.h" #include "utils/zstd_util.h" #include "utils/file_util.h" +#include "random.h" #include "utils.h" #include "kv_event/kv_event_config.h" #include "master_snapshot_manager.h" @@ -80,12 +80,6 @@ tl::expected ParseSnapshotCatalogKind( std::string(store_type)); } -size_t RandomIndex(size_t upper_bound) { - static thread_local std::mt19937 generator(std::random_device{}()); - std::uniform_int_distribution dist(0, upper_bound - 1); - return dist(generator); -} - uint64_t SaturatingAdd(uint64_t lhs, uint64_t rhs) { if (lhs > std::numeric_limits::max() - rhs) { return std::numeric_limits::max(); @@ -2119,7 +2113,7 @@ std::vector> MasterService::BatchExistKey( } } - const size_t start_shard = RandomIndex(kNumShards); + const size_t start_shard = randomIndex(kNumShards); for (size_t scanned = 0; scanned < kNumShards; ++scanned) { const size_t shard_idx = (start_shard + kNumShards - scanned) % kNumShards; @@ -2647,7 +2641,7 @@ MasterService::BatchGetReplicaList(const std::vector& keys, } } - const size_t start_shard = RandomIndex(kNumShards); + const size_t start_shard = randomIndex(kNumShards); for (size_t scanned = 0; scanned < kNumShards; ++scanned) { const size_t shard_idx = (start_shard + kNumShards - scanned) % kNumShards; @@ -2768,7 +2762,7 @@ MasterService::BatchGetReplicaListForAdmin(const std::vector& keys, } } - const size_t start_shard = RandomIndex(kNumShards); + const size_t start_shard = randomIndex(kNumShards); for (size_t scanned = 0; scanned < kNumShards; ++scanned) { const size_t shard_idx = (start_shard + kNumShards - scanned) % kNumShards; @@ -6283,7 +6277,7 @@ MasterService::EvictTenantMemoryForQuota(const std::string& tenant_id, }; auto pass = [&](bool allow_soft_pinned) { - const size_t start_shard = RandomIndex(kNumShards); + const size_t start_shard = randomIndex(kNumShards); for (size_t scanned = 0; scanned < kNumShards && total.freed_bytes < target_bytes; ++scanned) { @@ -6553,7 +6547,7 @@ void MasterService::BatchEvict(double evict_ratio_target, // Randomly select a starting shard to avoid imbalance eviction between // shards. - size_t start_idx = RandomIndex(kNumShards); + size_t start_idx = randomIndex(kNumShards); std::shared_lock shared_lock(snapshot_mutex_); // ===== Phase 1: Parallel candidate collection ===== @@ -6940,7 +6934,7 @@ void MasterService::NoFBatchEvict(double evict_ratio_target, long object_count = 0; uint64_t total_freed_size = 0; - size_t start_idx = RandomIndex(metadata_shards_.size()); + size_t start_idx = randomIndex(metadata_shards_.size()); for (size_t i = 0; i < metadata_shards_.size(); i++) { MetadataShardAccessorRW shard( this, (start_idx + i) % metadata_shards_.size()); @@ -7988,9 +7982,8 @@ tl::expected MasterService::CreateCopyTask( } // Randomly pick a segment from the source replicas - static thread_local std::mt19937 gen(std::random_device{}()); - std::uniform_int_distribution dis(0, segment_names.size() - 1); - std::string selected_source_segment = segment_names[dis(gen)]; + std::string selected_source_segment = + segment_names[randomIndex(segment_names.size())]; UUID select_client; ErrorCode error = segment_accessor.GetClientIdBySegmentName( selected_source_segment, select_client); diff --git a/mooncake-store/src/utils.cpp b/mooncake-store/src/utils.cpp index 3f923b1b80..a07ca23209 100644 --- a/mooncake-store/src/utils.cpp +++ b/mooncake-store/src/utils.cpp @@ -1,4 +1,5 @@ #include "utils.h" +#include "random.h" #include "mmap_arena.h" #include "config.h" #include "common.h" @@ -25,7 +26,6 @@ #include #include #include -#include #include #include #include @@ -79,12 +79,8 @@ bool isPortAvailable(int port) { // AutoPortBinder implementation AutoPortBinder::AutoPortBinder(int min_port, int max_port) : socket_fd_(-1), port_(-1) { - static std::random_device rand_gen; - std::mt19937 gen(rand_gen()); - std::uniform_int_distribution<> rand_dist(min_port, max_port); - for (int attempt = 0; attempt < 20; ++attempt) { - int port = rand_dist(gen); + int port = randomUniform(min_port, max_port); socket_fd_ = socket(AF_INET, SOCK_STREAM, 0); if (socket_fd_ < 0) continue; diff --git a/mooncake-store/tests/utils_test.cpp b/mooncake-store/tests/utils_test.cpp index b3102fd9fa..b9b2af8623 100644 --- a/mooncake-store/tests/utils_test.cpp +++ b/mooncake-store/tests/utils_test.cpp @@ -1,13 +1,84 @@ #include "utils.h" +#include "random.h" +#include +#include #include #include +#include #include +#include #include #include using namespace mooncake; +TEST(RandomTest, ReusesEngineWithinThread) { + EXPECT_EQ(&threadLocalRandomEngine(), &threadLocalRandomEngine()); +} + +TEST(RandomTest, UsesDifferentEngineAcrossThreads) { + auto* main_engine = &threadLocalRandomEngine(); + std::promise worker_address; + std::promise release_worker; + auto release_future = release_worker.get_future(); + + std::thread worker([&] { + worker_address.set_value(&threadLocalRandomEngine()); + release_future.wait(); + }); + + EXPECT_NE(main_engine, worker_address.get_future().get()); + release_worker.set_value(); + worker.join(); +} + +TEST(RandomTest, RandomIndexStaysWithinBounds) { + std::mt19937_64 engine(42); + for (size_t i = 0; i < 10000; ++i) { + EXPECT_LT(randomIndex(17, engine), 17); + } +} + +TEST(RandomTest, ExplicitEngineIsDeterministic) { + std::mt19937_64 first(42); + std::mt19937_64 second(42); + for (size_t i = 0; i < 100; ++i) { + EXPECT_EQ(randomIndex(1000, first), randomIndex(1000, second)); + } +} + +TEST(RandomTest, RandomUniformIncludesRequestedBounds) { + std::mt19937_64 engine(42); + for (size_t i = 0; i < 10000; ++i) { + int value = randomUniform(-5, 9, engine); + EXPECT_GE(value, -5); + EXPECT_LE(value, 9); + } + EXPECT_EQ(randomUniform(7, 7, engine), 7); +} + +TEST(RandomTest, RandomUniformSupportsNarrowIntegers) { + std::mt19937_64 engine(42); + for (size_t i = 0; i < 1000; ++i) { + auto signed_value = randomUniform(-5, 9, engine); + EXPECT_GE(signed_value, -5); + EXPECT_LE(signed_value, 9); + + auto unsigned_value = randomUniform(2, 7, engine); + EXPECT_GE(unsigned_value, 2); + EXPECT_LE(unsigned_value, 7); + } + EXPECT_FALSE(randomUniform(false, false, engine)); + EXPECT_TRUE(randomUniform(true, true, engine)); +} + +TEST(RandomTest, RejectsInvalidBounds) { + std::mt19937_64 engine(42); + EXPECT_THROW(randomIndex(0, engine), std::invalid_argument); + EXPECT_THROW(randomUniform(2, 1, engine), std::invalid_argument); +} + TEST(UtilsTest, ByteSizeToString) { EXPECT_EQ(byte_size_to_string(999), "999 B"); EXPECT_EQ(byte_size_to_string(2048), "2.00 KB");