Skip to content

Fix offload bucket tail flush #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JerryLX wants to merge 1 commit into
base: supercache
Choose a base branch
from
Open

Fix offload bucket tail flush #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions mooncake-store/include/storage_backend.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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<void, ErrorCode> indicating operation status.
*/
tl::expected<void, ErrorCode> AllocateOffloadingBuckets(
Expand Down Expand Up @@ -980,6 +986,7 @@ class BucketStorageBackend : public StorageBackendInterface {
mutable Mutex offloading_mutex_;
std::unordered_map<std::string, int64_t> 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_;
Expand Down Expand Up @@ -1214,4 +1221,4 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface {
tl::expected<std::shared_ptr<StorageBackendInterface>, ErrorCode>
CreateStorageBackend(const FileStorageConfig& config);

} // namespace mooncake
} // namespace mooncake
168 changes: 97 additions & 71 deletions mooncake-store/src/storage_backend.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand All @@ -66,6 +71,10 @@ BucketBackendConfig BucketBackendConfig::FromEnvironment() {
config.bucket_size_limit = GetEnvOr<int64_t>(
"MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", config.bucket_size_limit);

config.tail_flush_heartbeat_threshold = GetEnvOr<int64_t>(
"MOONCAKE_OFFLOAD_BUCKET_TAIL_FLUSH_HEARTBEATS",
config.tail_flush_heartbeat_threshold);

config.max_total_size =
GetEnvOr<int64_t>("MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE",
GetEnvOr<int64_t>("MOONCAKE_BUCKET_MAX_TOTAL_SIZE",
Expand Down Expand Up @@ -1836,6 +1845,7 @@ tl::expected<void, ErrorCode> BucketStorageBackend::AllocateOffloadingBuckets(
void BucketStorageBackend::ClearUngroupedOffloadingObjects() {
MutexLocker locker(&offloading_mutex_);
ungrouped_offloading_objects_.clear();
tail_idle_heartbeats_ = 0;
}

size_t BucketStorageBackend::UngroupedOffloadingObjectsSize() const {
Expand All @@ -1847,92 +1857,108 @@ tl::expected<void, ErrorCode> BucketStorageBackend::GroupOffloadingKeysByBucket(
const std::unordered_map<std::string, int64_t>& offloading_objects,
std::vector<std::vector<std::string>>& 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<int64_t>(
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<int64_t>(ungrouped_offloading_objects_.size());
std::vector<std::string> grouped_keys;

auto is_exist_func =
[this](const std::string& key) -> tl::expected<bool, ErrorCode> {
return IsExist(key);
};

while (it != offloading_objects.cend()) {
std::vector<std::string> bucket_keys;
std::unordered_map<std::string, int64_t> 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<int64_t>(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<std::string> 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<int64_t>(bucket_keys.size());
residue_count -= bucket_keys_count;
const auto bucket_keys_count = static_cast<int64_t>(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<int64_t>(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<int64_t>(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 {};
}

Expand Down
76 changes: 66 additions & 10 deletions mooncake-store/tests/file_storage_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down Expand Up @@ -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);
}

Expand All @@ -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);
}

Expand Down Expand Up @@ -247,7 +268,7 @@ TEST_F(FileStorageTest,
}

TEST_F(FileStorageTest,
GroupOffloadingKeysByBucket_ungrouped_offloading_objects) {
GroupOffloadingKeysByBucket_flushes_tail_bucket) {
std::unordered_map<std::string, int64_t> offloading_objects;
for (size_t i = 0; i < 1; i++) {
offloading_objects.emplace("test" + std::to_string(i), 1);
Expand All @@ -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) {
Expand All @@ -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);
Expand Down Expand Up @@ -542,4 +598,4 @@ TEST_F(FileStorageTest, NullSsdMetricDoesNotCrash) {
// No crash = success. No metrics pointer, so nothing to verify.
}

} // namespace mooncake
} // namespace mooncake
Loading