From b3672ca79102cc7ec0edeef4e05650851909df9b Mon Sep 17 00:00:00 2001 From: lixu Date: Thu, 25 Jun 2026 19:50:40 +0800 Subject: [PATCH] Fix offload bucket tail flush --- mooncake-store/include/storage_backend.h | 11 +- mooncake-store/src/storage_backend.cpp | 168 ++++++++++++--------- mooncake-store/tests/file_storage_test.cpp | 76 ++++++++-- 3 files changed, 172 insertions(+), 83 deletions(-) diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h index dfeabfb54a..16346b223e 100644 --- a/mooncake-store/include/storage_backend.h +++ b/mooncake-store/include/storage_backend.h @@ -200,6 +200,10 @@ struct BucketBackendConfig { // eviction_policy. Set via // MOONCAKE_OFFLOAD_DISABLE_SSD_EVICTION. + int64_t tail_flush_heartbeat_threshold = + 3; // Flush a partial tail bucket after this many consecutive + // grouping rounds without new keys. + bool Validate() const; static BucketBackendConfig FromEnvironment(); @@ -785,7 +789,9 @@ class BucketStorageBackend : public StorageBackendInterface { * @param offloading_objects Input map of object keys and their sizes * (bytes). * @param buckets_keys Output: bucketized keys; each inner vector is a - * bucket. + * bucket. Full buckets are returned immediately. The final partial bucket + * is retained briefly for aggregation and flushed after consecutive idle + * heartbeats so pending offloads cannot wait forever for more keys. * @return tl::expected indicating operation status. */ tl::expected AllocateOffloadingBuckets( @@ -980,6 +986,7 @@ class BucketStorageBackend : public StorageBackendInterface { mutable Mutex offloading_mutex_; std::unordered_map GUARDED_BY(offloading_mutex_) ungrouped_offloading_objects_; + int64_t tail_idle_heartbeats_ GUARDED_BY(offloading_mutex_) = 0; // File handle cache for UringFile to avoid repeated open/close overhead mutable Mutex file_cache_mutex_; @@ -1214,4 +1221,4 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { tl::expected, ErrorCode> CreateStorageBackend(const FileStorageConfig& config); -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 7a6792351e..aff6dfaa65 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -42,6 +42,11 @@ bool BucketBackendConfig::Validate() const { LOG(ERROR) << "BucketBackendConfig: bucket_size_limit must > 0"; return false; } + if (tail_flush_heartbeat_threshold <= 0) { + LOG(ERROR) + << "BucketBackendConfig: tail_flush_heartbeat_threshold must > 0"; + return false; + } return true; } @@ -66,6 +71,10 @@ BucketBackendConfig BucketBackendConfig::FromEnvironment() { config.bucket_size_limit = GetEnvOr( "MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", config.bucket_size_limit); + config.tail_flush_heartbeat_threshold = GetEnvOr( + "MOONCAKE_OFFLOAD_BUCKET_TAIL_FLUSH_HEARTBEATS", + config.tail_flush_heartbeat_threshold); + config.max_total_size = GetEnvOr("MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE", GetEnvOr("MOONCAKE_BUCKET_MAX_TOTAL_SIZE", @@ -1836,6 +1845,7 @@ tl::expected BucketStorageBackend::AllocateOffloadingBuckets( void BucketStorageBackend::ClearUngroupedOffloadingObjects() { MutexLocker locker(&offloading_mutex_); ungrouped_offloading_objects_.clear(); + tail_idle_heartbeats_ = 0; } size_t BucketStorageBackend::UngroupedOffloadingObjectsSize() const { @@ -1847,92 +1857,108 @@ tl::expected BucketStorageBackend::GroupOffloadingKeysByBucket( const std::unordered_map& offloading_objects, std::vector>& buckets_keys) { MutexLocker offloading_locker(&offloading_mutex_); - auto& ungrouped_offloading_objects = ungrouped_offloading_objects_; - auto it = offloading_objects.cbegin(); - int64_t residue_count = static_cast( - offloading_objects.size() + ungrouped_offloading_objects.size()); - int64_t total_count = residue_count; + int64_t new_key_count = 0; + for (const auto& [key, size] : offloading_objects) { + auto [_, inserted] = ungrouped_offloading_objects_.emplace(key, size); + if (inserted) { + ++new_key_count; + } + } + if (ungrouped_offloading_objects_.empty()) { + tail_idle_heartbeats_ = 0; + return {}; + } + if (new_key_count > 0) { + tail_idle_heartbeats_ = 0; + } else { + ++tail_idle_heartbeats_; + } + + int64_t grouped_count = 0; + const int64_t total_count = + static_cast(ungrouped_offloading_objects_.size()); + std::vector grouped_keys; auto is_exist_func = [this](const std::string& key) -> tl::expected { return IsExist(key); }; - while (it != offloading_objects.cend()) { - std::vector bucket_keys; - std::unordered_map bucket_objects; - int64_t bucket_data_size = 0; - - if (!ungrouped_offloading_objects.empty()) { - for (const auto& ungrouped_it : ungrouped_offloading_objects) { - bucket_data_size += ungrouped_it.second; - bucket_keys.push_back(ungrouped_it.first); - bucket_objects.emplace(ungrouped_it.first, ungrouped_it.second); - } - VLOG(1) << "Ungrouped offloading objects have been processed and " - "cleared; count=" - << ungrouped_offloading_objects.size(); - ungrouped_offloading_objects.clear(); - } - - for (int64_t i = static_cast(bucket_keys.size()); - i < bucket_backend_config_.bucket_keys_limit; ++i) { - if (it == offloading_objects.cend()) { - for (const auto& bucket_object : bucket_objects) { - ungrouped_offloading_objects.emplace(bucket_object.first, - bucket_object.second); - } - VLOG(1) << "Add offloading objects to ungrouped pool. " - << "Total ungrouped count: " - << ungrouped_offloading_objects.size(); - return {}; - } - - if (it->second > bucket_backend_config_.bucket_size_limit) { - LOG(ERROR) << "Object size exceeds bucket size limit: " - << "key=" << it->first - << ", object_size=" << it->second << ", limit=" - << bucket_backend_config_.bucket_size_limit; - ++it; - continue; - } - - auto is_exist_result = is_exist_func(it->first); - if (!is_exist_result) { - LOG(ERROR) << "Failed to check existence in storage backend: " - << "key=" << it->first - << ", error=" << is_exist_result.error(); - } - if (is_exist_result && is_exist_result.value()) { - ++it; - continue; - } - - if (bucket_data_size + it->second > - bucket_backend_config_.bucket_size_limit) { - break; - } - - bucket_data_size += it->second; - bucket_keys.push_back(it->first); - bucket_objects.emplace(it->first, it->second); - ++it; + std::vector bucket_keys; + int64_t bucket_data_size = 0; - if (bucket_data_size == bucket_backend_config_.bucket_size_limit) { - break; - } + auto flush_bucket = [&]() { + if (bucket_keys.empty()) { + return; } - - auto bucket_keys_count = static_cast(bucket_keys.size()); - residue_count -= bucket_keys_count; + const auto bucket_keys_count = static_cast(bucket_keys.size()); + grouped_count += bucket_keys_count; + grouped_keys.insert(grouped_keys.end(), bucket_keys.begin(), + bucket_keys.end()); buckets_keys.push_back(std::move(bucket_keys)); VLOG(1) << "Group objects with total object count: " << total_count << ", current bucket object count: " << bucket_keys_count << ", current bucket data size: " << bucket_data_size << ", grouped bucket count: " << buckets_keys.size() - << ", residue object count: " << residue_count; + << ", residue object count: " + << (total_count - grouped_count); + bucket_keys.clear(); + bucket_data_size = 0; + }; + + for (const auto& [key, size] : ungrouped_offloading_objects_) { + if (size > bucket_backend_config_.bucket_size_limit) { + LOG(ERROR) << "Object size exceeds bucket size limit: " + << "key=" << key << ", object_size=" << size + << ", limit=" + << bucket_backend_config_.bucket_size_limit; + continue; + } + + auto is_exist_result = is_exist_func(key); + if (!is_exist_result) { + LOG(ERROR) << "Failed to check existence in storage backend: " + << "key=" << key + << ", error=" << is_exist_result.error(); + } + if (is_exist_result && is_exist_result.value()) { + grouped_keys.push_back(key); + continue; + } + + if (!bucket_keys.empty() && + (static_cast(bucket_keys.size()) >= + bucket_backend_config_.bucket_keys_limit || + bucket_data_size + size > + bucket_backend_config_.bucket_size_limit)) { + flush_bucket(); + } + + bucket_data_size += size; + bucket_keys.push_back(key); + + if (static_cast(bucket_keys.size()) >= + bucket_backend_config_.bucket_keys_limit || + bucket_data_size == bucket_backend_config_.bucket_size_limit) { + flush_bucket(); + } } + if (!bucket_keys.empty() && + tail_idle_heartbeats_ >= + bucket_backend_config_.tail_flush_heartbeat_threshold) { + VLOG(1) << "Flush partial offload bucket after idle heartbeats: " + << tail_idle_heartbeats_ + << ", tail_key_count=" << bucket_keys.size(); + flush_bucket(); + tail_idle_heartbeats_ = 0; + } + for (const auto& key : grouped_keys) { + ungrouped_offloading_objects_.erase(key); + } + if (ungrouped_offloading_objects_.empty()) { + tail_idle_heartbeats_ = 0; + } return {}; } diff --git a/mooncake-store/tests/file_storage_test.cpp b/mooncake-store/tests/file_storage_test.cpp index 24b2f67243..3a5bfc211a 100644 --- a/mooncake-store/tests/file_storage_test.cpp +++ b/mooncake-store/tests/file_storage_test.cpp @@ -30,6 +30,7 @@ class FileStorageTest : public ::testing::Test { UnsetEnv("MOONCAKE_SCANMETA_ITERATOR_KEYS_LIMIT"); UnsetEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT"); UnsetEnv("MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES"); + UnsetEnv("MOONCAKE_OFFLOAD_BUCKET_TAIL_FLUSH_HEARTBEATS"); UnsetEnv("MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT"); UnsetEnv("MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES"); UnsetEnv("MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS"); @@ -182,13 +183,23 @@ TEST_F(FileStorageTest, GroupOffloadingKeysByBucket_bucket_keys_limit) { ASSERT_EQ(bucket_keys.size(), 10); } ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 5); + buckets_keys.clear(); + offloading_objects.clear(); ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( fileStorage, offloading_objects, buckets_keys)); - ASSERT_EQ(buckets_keys.size(), 4); - for (const auto& bucket_keys : buckets_keys) { - ASSERT_EQ(bucket_keys.size(), 10); - } + ASSERT_TRUE(buckets_keys.empty()); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 5); + + ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( + fileStorage, offloading_objects, buckets_keys)); + ASSERT_TRUE(buckets_keys.empty()); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 5); + + ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( + fileStorage, offloading_objects, buckets_keys)); + ASSERT_EQ(buckets_keys.size(), 1); + ASSERT_EQ(buckets_keys[0].size(), 5); ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 0); } @@ -209,13 +220,23 @@ TEST_F(FileStorageTest, GroupOffloadingKeysByBucket_bucket_size_limit) { ASSERT_EQ(bucket_keys.size(), 10); } ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 5); + buckets_keys.clear(); + offloading_objects.clear(); ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( fileStorage, offloading_objects, buckets_keys)); - ASSERT_EQ(buckets_keys.size(), 4); - for (const auto& bucket_keys : buckets_keys) { - ASSERT_EQ(bucket_keys.size(), 10); - } + ASSERT_TRUE(buckets_keys.empty()); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 5); + + ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( + fileStorage, offloading_objects, buckets_keys)); + ASSERT_TRUE(buckets_keys.empty()); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 5); + + ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( + fileStorage, offloading_objects, buckets_keys)); + ASSERT_EQ(buckets_keys.size(), 1); + ASSERT_EQ(buckets_keys[0].size(), 5); ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 0); } @@ -247,7 +268,7 @@ TEST_F(FileStorageTest, } TEST_F(FileStorageTest, - GroupOffloadingKeysByBucket_ungrouped_offloading_objects) { + GroupOffloadingKeysByBucket_flushes_tail_bucket) { std::unordered_map offloading_objects; for (size_t i = 0; i < 1; i++) { offloading_objects.emplace("test" + std::to_string(i), 1); @@ -258,14 +279,48 @@ TEST_F(FileStorageTest, FileStorage fileStorage(file_storage_config, nullptr, "localhost:9003"); ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( fileStorage, offloading_objects, buckets_keys)); + ASSERT_TRUE(buckets_keys.empty()); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 1); + + buckets_keys.clear(); offloading_objects.clear(); ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( fileStorage, offloading_objects, buckets_keys)); + ASSERT_TRUE(buckets_keys.empty()); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 1); + + ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( + fileStorage, offloading_objects, buckets_keys)); + ASSERT_TRUE(buckets_keys.empty()); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 1); + + ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( + fileStorage, offloading_objects, buckets_keys)); + ASSERT_EQ(buckets_keys.size(), 1); + ASSERT_EQ(buckets_keys[0].size(), 1); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 0); + for (size_t i = 0; i < 7; i++) { offloading_objects.emplace("test" + std::to_string(i), 1); } + buckets_keys.clear(); + ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( + fileStorage, offloading_objects, buckets_keys)); + ASSERT_TRUE(buckets_keys.empty()); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 7); + + offloading_objects.clear(); ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( fileStorage, offloading_objects, buckets_keys)); + ASSERT_TRUE(buckets_keys.empty()); + ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( + fileStorage, offloading_objects, buckets_keys)); + ASSERT_TRUE(buckets_keys.empty()); + ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( + fileStorage, offloading_objects, buckets_keys)); + ASSERT_EQ(buckets_keys.size(), 1); + ASSERT_EQ(buckets_keys[0].size(), 7); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 0); } TEST_F(FileStorageTest, DefaultValuesWhenNoEnvSet) { @@ -277,6 +332,7 @@ TEST_F(FileStorageTest, DefaultValuesWhenNoEnvSet) { EXPECT_EQ(config.scanmeta_iterator_keys_limit, 20000); EXPECT_EQ(bucket_backend_config.bucket_keys_limit, 500); EXPECT_EQ(bucket_backend_config.bucket_size_limit, 256 * 1024 * 1024); + EXPECT_EQ(bucket_backend_config.tail_flush_heartbeat_threshold, 3); EXPECT_EQ(config.total_keys_limit, 10'000'000); EXPECT_EQ(config.total_size_limit, 2ULL * 1024 * 1024 * 1024 * 1024); EXPECT_EQ(config.heartbeat_interval_seconds, 10u); @@ -542,4 +598,4 @@ TEST_F(FileStorageTest, NullSsdMetricDoesNotCrash) { // No crash = success. No metrics pointer, so nothing to verify. } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake