[test] add e2e test cases for aggressive location pruning (#116)

[test] add e2e test cases for aggressive location pruning

[data_storage] introducing the dummy storage backend

The dummy storage backend is a local filesystem based storage backend
for testing purposes only.

The e2e tests for aggressive location pruning need a storage backend
with MightExist() actually implemented, but currently no storage backend
can satisfy the need. This is where the dummy backend comes into play:
it can implement MightExist() without worrying about the read/write
path overhead.

The dummy backend reuses the LocalStorageSpec for StorageConfig in the
proto definition, so "local" should be used instead of "dummy" when
constructing requests like AddStorageRequest.

Backward-compatibility is naturally handled so no new instance behavior
version needs to be introduced:

v0-data: key_count
v1-data: key_count, usage_data[_,]
v2-data: key_count, usage_data[_,dummy]

During recovery:
----

|    | v0-data | v1-data        | v2-data                |
| -- | ------- | ----------     | ----------             |
| v0 | v0,ok   | v0,ok(ign)     | v0,ok(ign)             |
| v1 | v0,ok   | v1,ok          | err,ok(can't rollback) |
| v2 | v0,ok   | v1,ok(dummy=0) | v2,ok                  |

Author: oldsharp

integration_test/reclaimer/BUILD

@@ -1,5 +1,19 @@
 package(default_visibility = ["//integration_test/reclaimer:__subpackages__"])
 
+py_test(
+    name = "location_pruning_test",
+    srcs = ["location_pruning_test.py"],
+    tags = ["no-remote-exec"],
+    data = [
+        "//kv_cache_manager:kv_cache_manager_bin",
+    ],
+    deps = [
+        "//integration_test/testlib:test_base",
+        "//integration_test/admin_service:http_interface_test",
+        "//integration_test/meta_service:http_interface_test",
+    ],
+)
+
 py_test(
     name = "reclaiming_test",
     srcs = ["reclaiming_test.py"],

integration_test/reclaimer/location_pruning_test.py

@@ -4,13 +4,15 @@ cc_library(
     name = "data_storage",
     srcs = [
         "data_storage_manager.cc",
+        "dummy_backend.cc",
         "hf3fs_backend.cc",
         "mooncake_backend.cc",
         "nfs_backend.cc",
     ],
     hdrs = [
         "data_storage_backend.h",
         "data_storage_manager.h",
+        "dummy_backend.h",
         "hf3fs_backend.h",
         "mooncake_backend.h",
         "nfs_backend.h",

kv_cache_manager/data_storage/BUILD

@@ -5,6 +5,7 @@
 #include <utility>
 
 #include "kv_cache_manager/common/logger.h"
+#include "kv_cache_manager/data_storage/dummy_backend.h"
 #include "kv_cache_manager/data_storage/hf3fs_backend.h"
 #include "kv_cache_manager/data_storage/mooncake_backend.h"
 #include "kv_cache_manager/data_storage/nfs_backend.h"
@@ -169,6 +170,8 @@ std::shared_ptr<DataStorageBackend> DataStorageManager::CreateStorageBackend(con
         return std::make_shared<TairMempoolBackend>(metrics_registry_);
     case DataStorageType::DATA_STORAGE_TYPE_NFS:
         return std::make_shared<NfsBackend>(metrics_registry_);
+    case DataStorageType::DATA_STORAGE_TYPE_DUMMY:
+        return std::make_shared<DummyBackend>(metrics_registry_);
     default:
         return nullptr;
     }
@@ -260,4 +263,4 @@ std::vector<ErrorCode> DataStorageManager::UnLock(const std::string &unique_name
     auto storage_backend = iter->second;
     return storage_backend->UnLock(storage_uris);
 }
-} // namespace kv_cache_manager
+} // namespace kv_cache_manager

kv_cache_manager/data_storage/data_storage_manager.cc

@@ -0,0 +1,169 @@
+#include "kv_cache_manager/data_storage/dummy_backend.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <filesystem>
+#include <fstream>
+#include <functional>
+#include <iterator>
+#include <memory>
+#include <string>
+#include <system_error>
+#include <utility>
+#include <vector>
+
+#include "kv_cache_manager/common/error_code.h"
+#include "kv_cache_manager/common/hash/hash.h"
+#include "kv_cache_manager/common/logger.h"
+#include "kv_cache_manager/common/string_util.h"
+#include "kv_cache_manager/data_storage/data_storage_uri.h"
+#include "kv_cache_manager/metrics/metrics_registry.h"
+
+namespace kv_cache_manager {
+
+DummyBackend::DummyBackend(std::shared_ptr<MetricsRegistry> metrics_registry)
+    : DataStorageBackend(std::move(metrics_registry)) {}
+
+DataStorageType DummyBackend::GetType() { return DataStorageType::DATA_STORAGE_TYPE_DUMMY; }
+
+bool DummyBackend::Available() { return IsOpen() && IsAvailable(); }
+
+double DummyBackend::GetStorageUsageRatio(const std::string &trace_id) const { return 0.0; }
+
+ErrorCode DummyBackend::DoOpen(const StorageConfig &storage_config, const std::string &trace_id) {
+    if (const auto cfg = std::dynamic_pointer_cast<DummyStorageSpec>(storage_config.storage_spec())) {
+        spec_ = *cfg;
+    } else {
+        KVCM_LOG_WARN("unexpected config type, storage config: [%s]", storage_config.ToString().c_str());
+        return ErrorCode::EC_ERROR;
+    }
+    if (spec_.root_path().empty()) {
+        KVCM_LOG_WARN("open dummy backend failed, root_path is empty");
+        return ErrorCode::EC_ERROR;
+    }
+    base_path_ = std::filesystem::path(spec_.root_path());
+    std::error_code ec;
+    std::filesystem::create_directories(base_path_, ec);
+    if (ec) {
+        KVCM_LOG_WARN("open dummy backend failed, cannot create root_path [%s], msg: [%s]",
+                      base_path_.string().c_str(),
+                      ec.message().c_str());
+        return ErrorCode::EC_ERROR;
+    }
+    KVCM_LOG_INFO("open dummy backend success, config: [%s]", spec_.ToString().c_str());
+    SetOpen(true);
+    SetAvailable(true);
+    return ErrorCode::EC_OK;
+}
+
+ErrorCode DummyBackend::Close() {
+    KVCM_LOG_INFO("close dummy backend");
+    SetOpen(false);
+    SetAvailable(false);
+    return ErrorCode::EC_OK;
+}
+
+std::vector<std::pair<ErrorCode, DataStorageUri>> DummyBackend::Create(const std::vector<std::string> &keys,
+                                                                       const std::size_t size_per_key,
+                                                                       const std::string &trace_id,
+                                                                       const std::function<void()> cb) {
+    std::vector<std::pair<ErrorCode, DataStorageUri>> results;
+    std::vector<std::vector<std::string>> batches;
+
+    auto batch_size = spec_.key_count_per_file();
+    batch_size = batch_size <= 0 ? 1 : batch_size;
+    using diff_t = std::vector<std::string>::difference_type;
+    for (std::size_t start = 0; start < keys.size(); start += batch_size) {
+        batches.emplace_back(std::next(keys.begin(), static_cast<diff_t>(start)),
+                             std::next(keys.begin(), static_cast<diff_t>(std::min(start + batch_size, keys.size()))));
+    }
+
+    for (auto &batch : batches) {
+        DataStorageUri storage_uri;
+        storage_uri.SetProtocol(ToString(GetType()));
+
+        if (batch.size() > 1) {
+            std::string combine_key = StringUtil::Join(batch, "|");
+            std::string hash_str = StringUtil::Uint64ToHex(Hash64(combine_key.c_str(), combine_key.size(), 42));
+            storage_uri.SetPath(base_path_ / (batch[0] + "_" + hash_str));
+        } else {
+            storage_uri.SetPath(base_path_ / batch[0]);
+        }
+
+        storage_uri.SetParam("size", std::to_string(size_per_key));
+
+        for (std::size_t i = 0; i != batch.size(); ++i) {
+            if (batch_size > 1) {
+                storage_uri.SetParam("blkid", std::to_string(i));
+            }
+            results.emplace_back(ErrorCode::EC_OK, storage_uri);
+        }
+    }
+
+    if (cb) {
+        cb();
+    }
+
+    return results;
+}
+
+std::vector<ErrorCode> DummyBackend::Delete(const std::vector<DataStorageUri> &storage_uris,
+                                            const std::string &trace_id,
+                                            const std::function<void()> cb) {
+    std::vector<ErrorCode> results;
+    for (auto &uri : storage_uris) {
+        std::filesystem::path file_path = uri.GetPath();
+        std::error_code ec;
+        const bool removed = std::filesystem::remove(file_path, ec);
+        if (ec) {
+            KVCM_LOG_ERROR(
+                "failed to delete file, path: [%s], msg: [%s]", file_path.string().c_str(), ec.message().c_str());
+            results.push_back(ErrorCode::EC_ERROR);
+            continue;
+        }
+        if (!removed) {
+            KVCM_LOG_WARN("file not exist, path: [%s]", file_path.string().c_str());
+        }
+        results.push_back(ErrorCode::EC_OK);
+    }
+
+    if (cb) {
+        cb();
+    }
+
+    return results;
+}
+
+std::vector<bool> DummyBackend::Exist(const std::vector<DataStorageUri> &storage_uris) {
+    std::vector<bool> results;
+    for (auto &uri : storage_uris) {
+        std::error_code ec;
+        const bool res = std::filesystem::exists(uri.GetPath(), ec);
+        if (ec) {
+            KVCM_LOG_ERROR("std::filesystem::exists call failed, err code: [%d], err msg: [%s], path: [%s]",
+                           ec.value(),
+                           ec.message().c_str(),
+                           uri.GetPath().c_str());
+            results.push_back(false);
+            continue;
+        }
+        results.push_back(res);
+    }
+    return results;
+}
+
+std::vector<bool> DummyBackend::MightExist(const std::vector<DataStorageUri> &storage_uris) {
+    return Exist(storage_uris);
+}
+
+std::vector<ErrorCode> DummyBackend::Lock(const std::vector<DataStorageUri> &storage_uris) {
+    std::vector<ErrorCode> results(storage_uris.size(), ErrorCode::EC_OK);
+    return results;
+}
+
+std::vector<ErrorCode> DummyBackend::UnLock(const std::vector<DataStorageUri> &storage_uris) {
+    std::vector<ErrorCode> results(storage_uris.size(), ErrorCode::EC_OK);
+    return results;
+}
+
+} // namespace kv_cache_manager

kv_cache_manager/data_storage/dummy_backend.cc

@@ -0,0 +1,49 @@
+#pragma once
+
+#include <cstddef>
+#include <filesystem>
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "kv_cache_manager/common/error_code.h"
+#include "kv_cache_manager/data_storage/data_storage_backend.h"
+#include "kv_cache_manager/data_storage/data_storage_uri.h"
+#include "kv_cache_manager/data_storage/storage_config.h"
+
+namespace kv_cache_manager {
+
+class MetricsRegistry;
+
+class DummyBackend : public DataStorageBackend {
+public:
+    DummyBackend() = delete;
+    explicit DummyBackend(std::shared_ptr<MetricsRegistry> metrics_registry);
+    ~DummyBackend() override = default;
+
+    DataStorageType GetType() override;
+    bool Available() override;
+    [[nodiscard]] double GetStorageUsageRatio(const std::string &trace_id) const override;
+
+    ErrorCode DoOpen(const StorageConfig &storage_config, const std::string &trace_id) override;
+    ErrorCode Close() override;
+
+    std::vector<std::pair<ErrorCode, DataStorageUri>> Create(const std::vector<std::string> &keys,
+                                                             std::size_t size_per_key,
+                                                             const std::string &trace_id,
+                                                             std::function<void()> cb) override;
+    std::vector<ErrorCode> Delete(const std::vector<DataStorageUri> &storage_uris,
+                                  const std::string &trace_id,
+                                  std::function<void()> cb) override;
+    std::vector<bool> Exist(const std::vector<DataStorageUri> &storage_uris) override;
+    std::vector<bool> MightExist(const std::vector<DataStorageUri> &storage_uris) override;
+    std::vector<ErrorCode> Lock(const std::vector<DataStorageUri> &storage_uris) override;
+    std::vector<ErrorCode> UnLock(const std::vector<DataStorageUri> &storage_uris) override;
+
+private:
+    DummyStorageSpec spec_;
+    std::filesystem::path base_path_;
+};
+
+} // namespace kv_cache_manager

kv_cache_manager/data_storage/dummy_backend.h

@@ -160,6 +160,8 @@ std::string ToString(const DataStorageType &type) {
         return "pace";
     case DataStorageType::DATA_STORAGE_TYPE_NFS:
         return "file";
+    case DataStorageType::DATA_STORAGE_TYPE_DUMMY:
+        return "dummy";
     default:
         return "unrecognized";
     }
@@ -176,6 +178,8 @@ DataStorageType ToDataStorageType(const std::string &type) {
         return DataStorageType::DATA_STORAGE_TYPE_TAIR_MEMPOOL;
     } else if (type == "file") {
         return DataStorageType::DATA_STORAGE_TYPE_NFS;
+    } else if (type == "dummy") {
+        return DataStorageType::DATA_STORAGE_TYPE_DUMMY;
     } else {
         return DataStorageType::DATA_STORAGE_TYPE_UNKNOWN;
     }
@@ -271,6 +275,38 @@ void NfsStorageSpec::ToRapidWriter(rapidjson::Writer<rapidjson::StringBuffer> &w
     Put(writer, "key_count_per_file", key_count_per_file_);
 }
 
+// DummyStorageSpec
+std::string DummyStorageSpec::ToString() const {
+    std::ostringstream oss;
+    oss << "root_path: " << root_path_;
+    oss << " , key_count_per_file: " << key_count_per_file_;
+    return oss.str();
+}
+
+bool DummyStorageSpec::ValidateRequiredFields(std::string &invalid_fields) const {
+    bool valid = true;
+    std::string local_invalid_fields;
+    if (root_path_.empty()) {
+        valid = false;
+        local_invalid_fields += "{root_path}";
+    }
+    if (!valid) {
+        invalid_fields += "{DummyStorageSpec: " + local_invalid_fields + "}";
+    }
+    return valid;
+}
+
+bool DummyStorageSpec::FromRapidValue(const rapidjson::Value &rapid_value) {
+    KVCM_JSON_GET_MACRO(rapid_value, "root_path", root_path_);
+    KVCM_JSON_GET_DEFAULT_MACRO(rapid_value, "key_count_per_file", key_count_per_file_, 1);
+    return true;
+}
+
+void DummyStorageSpec::ToRapidWriter(rapidjson::Writer<rapidjson::StringBuffer> &writer) const noexcept {
+    Put(writer, "root_path", root_path_);
+    Put(writer, "key_count_per_file", key_count_per_file_);
+}
+
 bool StorageConfig::FromRapidValue(const rapidjson::Value &rapid_value) {
     std::string type_str;
     KVCM_JSON_GET_MACRO(rapid_value, "type", type_str);
@@ -297,6 +333,10 @@ bool StorageConfig::FromRapidValue(const rapidjson::Value &rapid_value) {
         auto tmp = std::make_shared<TairMemPoolStorageSpec>();
         KVCM_JSON_GET_MACRO(rapid_value, "storage_spec", tmp);
         storage_spec_ = tmp;
+    } else if (type_ == DataStorageType::DATA_STORAGE_TYPE_DUMMY) {
+        auto tmp = std::make_shared<DummyStorageSpec>();
+        KVCM_JSON_GET_MACRO(rapid_value, "storage_spec", tmp);
+        storage_spec_ = tmp;
     } else {
         storage_spec_ = nullptr; // 对未知或未支持类型，设为空
     }

kv_cache_manager/data_storage/storage_config.cc

@@ -14,6 +14,7 @@ enum class DataStorageType : uint8_t {
     DATA_STORAGE_TYPE_TAIR_MEMPOOL = 3,
     DATA_STORAGE_TYPE_NFS = 4,
     DATA_STORAGE_TYPE_VCNS_HF3FS = 5,
+    DATA_STORAGE_TYPE_DUMMY = 6,
     COUNT, // as sentinel
 };
 
@@ -173,6 +174,23 @@ class NfsStorageSpec : public StorageSpec {
     int32_t key_count_per_file_ = 0;
 };
 
+class DummyStorageSpec : public StorageSpec {
+public:
+    bool FromRapidValue(const rapidjson::Value &rapid_value) override;
+    void ToRapidWriter(rapidjson::Writer<rapidjson::StringBuffer> &writer) const noexcept override;
+    bool ValidateRequiredFields(std::string &invalid_fields) const override;
+
+    std::string ToString() const override;
+    const std::string &root_path() const { return root_path_; }
+    void set_root_path(const std::string &root_path) { root_path_ = root_path; }
+    int32_t key_count_per_file() const { return key_count_per_file_; }
+    void set_key_count_per_file(int32_t value) { key_count_per_file_ = value; }
+
+private:
+    std::string root_path_;
+    int32_t key_count_per_file_ = 0;
+};
+
 class StorageConfig : public Jsonizable {
 public:
     StorageConfig() = default;

kv_cache_manager/data_storage/storage_config.h

@@ -39,6 +39,19 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "DummyBackendTest",
+    srcs = [
+        "dummy_backend_test.cc",
+    ],
+    copts = ["-fno-access-control"],
+    data = [],
+    deps = [
+        "//kv_cache_manager/common:unittest",
+        "//kv_cache_manager/data_storage",
+    ],
+)
+
 cc_test(
     name = "StorageConfigTest",
     srcs = [