diff --git a/mooncake-store/src/file_storage.cpp b/mooncake-store/src/file_storage.cpp index 7500059f53..0fcd113210 100644 --- a/mooncake-store/src/file_storage.cpp +++ b/mooncake-store/src/file_storage.cpp @@ -324,21 +324,23 @@ tl::expected FileStorage::Init() { return init_storage_backend_result; } auto enable_offloading_result = IsEnableOffloading(); + bool initial_enable_offloading = false; if (enable_offloading_result.has_value()) { + initial_enable_offloading = enable_offloading_result.value(); LOG(INFO) << "IsEnableOffloading result: " - << (enable_offloading_result.value() ? "true" : "false"); + << (initial_enable_offloading ? "true" : "false"); + } else if (enable_offloading_result.error() == + ErrorCode::UNABLE_OFFLOADING) { + LOG(INFO) << "Offloading is not ready before metadata scan; mounting " + "with offloading disabled"; } else { - LOG(INFO) << "IsEnableOffloading result: error: " - << enable_offloading_result.error(); - } - if (!enable_offloading_result) { LOG(ERROR) << "Failed to get enable persist result, error : " << enable_offloading_result.error(); return tl::make_unexpected(enable_offloading_result.error()); } { MutexLocker locker(&offloading_mutex_); - enable_offloading_ = enable_offloading_result.value(); + enable_offloading_ = initial_enable_offloading; auto mount_file_storage_result = client_->MountLocalDiskSegment(enable_offloading_); if (!mount_file_storage_result) { @@ -382,6 +384,20 @@ tl::expected FileStorage::Init() { return scan_meta_result; } + // Some backends need the initial scan to reconstruct quota accounting. + // Refresh the local state before the first heartbeat publishes it to the + // master. Until this point such backends mount conservatively as disabled. + auto refreshed_offloading_result = IsEnableOffloading(); + if (!refreshed_offloading_result) { + LOG(ERROR) << "Failed to refresh offloading state after metadata scan: " + << refreshed_offloading_result.error(); + return tl::make_unexpected(refreshed_offloading_result.error()); + } + { + MutexLocker locker(&offloading_mutex_); + enable_offloading_ = refreshed_offloading_result.value(); + } + heartbeat_running_.store(true); heartbeat_thread_ = std::thread([this]() { LOG(INFO) << "Starting periodic task with interval: " @@ -736,8 +752,13 @@ tl::expected FileStorage::RunDiskWatermarkEviction() { tl::expected FileStorage::IsEnableOffloading() { auto is_enable_offloading_result = storage_backend_->IsEnableOffloading(); if (!is_enable_offloading_result) { - LOG(ERROR) << "Failed to get enabling offload: " - << is_enable_offloading_result.error(); + if (is_enable_offloading_result.error() == + ErrorCode::UNABLE_OFFLOADING) { + VLOG(1) << "Offloading is not ready"; + } else { + LOG(ERROR) << "Failed to get enabling offload: " + << is_enable_offloading_result.error(); + } return tl::make_unexpected(is_enable_offloading_result.error()); } diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 3b666fade2..593e089829 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -1342,6 +1342,11 @@ tl::expected StorageBackendAdaptor::BatchOffload( if (!enable_offloading_res) { return tl::make_unexpected(enable_offloading_res.error()); } + if (!enable_offloading_res.value()) { + // Once metadata has been scanned, false means the configured + // non-evicting quota has been exhausted. + return tl::make_unexpected(ErrorCode::KEYS_ULTRA_LIMIT); + } std::vector metadatas; std::vector keys; metadatas.reserve(batch_object.size()); @@ -1462,8 +1467,8 @@ tl::expected StorageBackendAdaptor::IsEnableOffloading() { } if (!meta_scanned_.load(std::memory_order_acquire)) { - LOG(ERROR) << "Metadata has not been loaded yet"; - return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + VLOG(1) << "Metadata has not been loaded yet; offloading is not ready"; + return tl::make_unexpected(ErrorCode::UNABLE_OFFLOADING); } MutexLocker lock(&mutex_); diff --git a/mooncake-store/tests/file_storage_test.cpp b/mooncake-store/tests/file_storage_test.cpp index 16413486df..0c150bf901 100644 --- a/mooncake-store/tests/file_storage_test.cpp +++ b/mooncake-store/tests/file_storage_test.cpp @@ -36,6 +36,7 @@ class FileStorageTest : public ::testing::Test { UnsetEnv("MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT"); UnsetEnv("MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES"); UnsetEnv("MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS"); + UnsetEnv("MOONCAKE_OFFLOAD_ENABLE_EVICTION"); UnsetEnv("MOONCAKE_OFFLOAD_ENABLE_DISK_WATERMARK_EVICTION"); UnsetEnv("MOONCAKE_OFFLOAD_DISK_EVICTION_HIGH_WATERMARK_RATIO"); UnsetEnv("MOONCAKE_OFFLOAD_DISK_EVICTION_LOW_WATERMARK_RATIO"); @@ -406,6 +407,39 @@ TEST_F(FileStorageTest, ReadDiskWatermarkConfigFromEnv) { EXPECT_DOUBLE_EQ(invalid_config.disk_eviction_low_watermark_ratio, 0.80); } +TEST_F(FileStorageTest, InitFilePerKeyWithEvictionDisabled) { + SetEnv("MOONCAKE_OFFLOAD_ENABLE_EVICTION", "false"); + + std::filesystem::path master_root = + std::filesystem::path(data_path) / "no_eviction_master"; + std::filesystem::create_directories(master_root); + testing::InProcMaster master; + auto master_config = InProcMasterConfigBuilder() + .set_enable_offload(true) + .set_root_fs_dir(master_root.string()) + .build(); + ASSERT_TRUE(master.Start(master_config)); + + std::string local_rpc_addr = + "127.0.0.1:" + std::to_string(getFreeTcpPort()); + auto client = Client::Create(local_rpc_addr, master.metadata_url(), "tcp", + std::nullopt, master.master_address()); + ASSERT_TRUE(client); + + FileStorageConfig config = FileStorageConfig::FromEnvironment(); + config.storage_backend_type = StorageBackendType::kFilePerKey; + config.storage_filepath = data_path + "/no_eviction_storage"; + config.local_buffer_size = 4 * 1024 * 1024; + config.heartbeat_interval_seconds = 1; + fs::create_directories(config.storage_filepath); + + FileStorage file_storage(config, client.value(), local_rpc_addr); + ASSERT_TRUE(file_storage.Init()); + auto enable_result = FileStorageIsEnableOffloading(file_storage); + ASSERT_TRUE(enable_result); + EXPECT_TRUE(enable_result.value()); +} + TEST_F(FileStorageTest, HeartbeatRunsDiskWatermarkEvictionWithoutOffloadWork) { std::filesystem::path master_root = std::filesystem::path(data_path) / "heartbeat_master"; diff --git a/mooncake-store/tests/storage_backend_test.cpp b/mooncake-store/tests/storage_backend_test.cpp index b87272358a..4a7fc9d222 100644 --- a/mooncake-store/tests/storage_backend_test.cpp +++ b/mooncake-store/tests/storage_backend_test.cpp @@ -774,12 +774,22 @@ TEST_F(StorageBackendTest, AdaptorScanMetaAndIsEnableOffloading) { ASSERT_TRUE(adaptor.Init()); auto enable_before = adaptor.IsEnableOffloading(); - EXPECT_FALSE(enable_before); - EXPECT_EQ(enable_before.error(), ErrorCode::INTERNAL_ERROR); + ASSERT_FALSE(enable_before); + EXPECT_EQ(enable_before.error(), ErrorCode::UNABLE_OFFLOADING); - // New behavior: must call ScanMeta once before BatchOffload when eviction - // is disabled, otherwise meta_scanned_ is false and BatchOffload is - // rejected. + std::string pre_scan_value = "not written"; + std::unordered_map> pre_scan_batch = { + {"pre_scan_key", {{pre_scan_value.data(), pre_scan_value.size()}}}, + }; + auto pre_scan_offload = adaptor.BatchOffload( + pre_scan_batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_FALSE(pre_scan_offload); + EXPECT_EQ(pre_scan_offload.error(), ErrorCode::UNABLE_OFFLOADING); + + // ScanMeta reconstructs quota accounting before writes are accepted when + // eviction is disabled. auto scan_init_res = adaptor.ScanMeta( [](const std::vector&, std::vector&) { return ErrorCode::OK; }); @@ -858,6 +868,17 @@ TEST_F(StorageBackendTest, AdaptorScanMetaAndIsEnableOffloading) { auto enable_strict = strict_adaptor.IsEnableOffloading(); ASSERT_TRUE(enable_strict); EXPECT_FALSE(enable_strict.value()); + + std::string strict_value = "capacity limited"; + std::unordered_map> strict_batch = { + {"strict_key", {{strict_value.data(), strict_value.size()}}}, + }; + auto strict_offload = strict_adaptor.BatchOffload( + strict_batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_FALSE(strict_offload); + EXPECT_EQ(strict_offload.error(), ErrorCode::KEYS_ULTRA_LIMIT); } TEST_F(StorageBackendTest, AdaptorScanMetaAndBatchLoadAcrossRestart) {