Skip to content

[Store][POC] Add TTL-bounded group retention leases #2835

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

Draft
ishandhanani wants to merge 5 commits into
base: main
Choose a base branch
from
Draft

[Store][POC] Add TTL-bounded group retention leases #2835

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
619f900
feat(store): retain object groups with server-owned leases
ishandhanani Jul 9, 2026
c72ad37
fix(store): retain groups before members arrive
ishandhanani Jul 9, 2026
1b9ffd5
feat(store): log group retention outcomes
ishandhanani Jul 9, 2026
56974e3
feat(store): bound group retention admission
ishandhanani Jul 9, 2026
f9b03d4
style(store): format retention POC
ishandhanani Jul 10, 2026
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
16 changes: 16 additions & 0 deletions mooncake-integration/store/store_py.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2289,6 +2289,22 @@ PYBIND11_MODULE(store, m) {
py::arg("keys"),
"Check if multiple objects exist. Returns list of results: 1 if "
"exists, 0 if not exists, -1 if error")
.def(
"retain_groups",
[](MooncakeStorePyWrapper &self,
const std::vector<std::string> &group_ids, uint64_t ttl_ms) {
if (!self.real_client_) {
return std::vector<int>(
group_ids.size(),
static_cast<int>(ErrorCode::INVALID_PARAMS));
}
Comment on lines +2296 to +2300

@gemini-code-assist gemini-code-assist Bot Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Directly casting ErrorCode::INVALID_PARAMS to int via static_cast<int> might bypass the standard error-to-integer mapping (such as the toInt() helper used in RealClient), potentially returning a positive integer instead of a negative error code as documented.

To ensure consistency with the rest of the client and adhere to the 'negative error code on failure' contract, we should use a negative value or a proper conversion helper. If ErrorCode values are positive, we should negate it or use the appropriate mapping function.

Suggested change
if (!self.real_client_) {
return std::vector<int>(
group_ids.size(),
static_cast<int>(ErrorCode::INVALID_PARAMS));
}
if (!self.real_client_) {
return std::vector<int>(
group_ids.size(),
-static_cast<int>(ErrorCode::INVALID_PARAMS));
}

py::gil_scoped_release release;
return self.real_client_->retainGroups(group_ids, ttl_ms);
},
py::arg("group_ids"), py::arg("ttl_ms"),
"Retain current and future objects in each group for the requested "
"TTL. Returns 1 if accepted, 0 if admission is full, and a "
"negative error code on failure.")
.def("close",
[](MooncakeStorePyWrapper &self) {
if (!self.store_) return 0;
Expand Down
3 changes: 3 additions & 0 deletions mooncake-store/include/client_service.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -363,6 +363,9 @@ class Client {
std::vector<tl::expected<bool, ErrorCode>> BatchIsExist(
const std::vector<std::string>& keys);

std::vector<tl::expected<bool, ErrorCode>> RetainGroups(
const std::vector<std::string>& group_ids, uint64_t ttl_ms);

/**
* @brief Create a copy task to copy an object's replicas to target segments
* @param key Object key
Expand Down
3 changes: 3 additions & 0 deletions mooncake-store/include/master_client.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,9 @@ class MasterClient {
[[nodiscard]] std::vector<tl::expected<bool, ErrorCode>> BatchExistKey(
const std::vector<std::string>& object_keys);

[[nodiscard]] std::vector<tl::expected<bool, ErrorCode>> RetainGroups(
const std::vector<std::string>& group_ids, uint64_t ttl_ms);

/**
* @brief Calculate Store-observed cache reuse metrics
* @param object_keys None
Expand Down
16 changes: 16 additions & 0 deletions mooncake-store/include/master_config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -694,6 +694,8 @@ class MasterServiceConfigBuilder {
private:
uint64_t default_kv_lease_ttl_ = DEFAULT_DEFAULT_KV_LEASE_TTL;
uint64_t default_kv_soft_pin_ttl_ = DEFAULT_KV_SOFT_PIN_TTL_MS;
size_t max_retained_groups_ = DEFAULT_MAX_RETAINED_GROUPS;
uint64_t max_group_retention_ttl_ms_ = DEFAULT_MAX_GROUP_RETENTION_TTL_MS;
bool allow_evict_soft_pinned_objects_ =
DEFAULT_ALLOW_EVICT_SOFT_PINNED_OBJECTS;
double eviction_ratio_ = DEFAULT_EVICTION_RATIO;
Expand Down Expand Up @@ -760,6 +762,16 @@ class MasterServiceConfigBuilder {
return *this;
}

MasterServiceConfigBuilder& set_max_retained_groups(size_t max_groups) {
max_retained_groups_ = max_groups;
return *this;
}

MasterServiceConfigBuilder& set_max_group_retention_ttl_ms(uint64_t ttl) {
max_group_retention_ttl_ms_ = ttl;
return *this;
}

MasterServiceConfigBuilder& set_allow_evict_soft_pinned_objects(
bool allow) {
allow_evict_soft_pinned_objects_ = allow;
Expand Down Expand Up @@ -1032,6 +1044,8 @@ class MasterServiceConfig {
public:
uint64_t default_kv_lease_ttl = DEFAULT_DEFAULT_KV_LEASE_TTL;
uint64_t default_kv_soft_pin_ttl = DEFAULT_KV_SOFT_PIN_TTL_MS;
size_t max_retained_groups = DEFAULT_MAX_RETAINED_GROUPS;
uint64_t max_group_retention_ttl_ms = DEFAULT_MAX_GROUP_RETENTION_TTL_MS;
bool allow_evict_soft_pinned_objects =
DEFAULT_ALLOW_EVICT_SOFT_PINNED_OBJECTS;
double eviction_ratio = DEFAULT_EVICTION_RATIO;
Expand Down Expand Up @@ -1202,6 +1216,8 @@ inline MasterServiceConfig MasterServiceConfigBuilder::build() const {
MasterServiceConfig config;
config.default_kv_lease_ttl = default_kv_lease_ttl_;
config.default_kv_soft_pin_ttl = default_kv_soft_pin_ttl_;
config.max_retained_groups = max_retained_groups_;
config.max_group_retention_ttl_ms = max_group_retention_ttl_ms_;
config.allow_evict_soft_pinned_objects = allow_evict_soft_pinned_objects_;
config.eviction_ratio = eviction_ratio_;
config.eviction_high_watermark_ratio = eviction_high_watermark_ratio_;
Expand Down
13 changes: 12 additions & 1 deletion mooncake-store/include/master_service.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,10 @@ class MasterService {
std::vector<tl::expected<bool, ErrorCode>> BatchExistKey(
const std::vector<std::string>& keys, const std::string& tenant_id);

std::vector<tl::expected<bool, ErrorCode>> RetainGroups(
const std::vector<std::string>& group_ids, uint64_t ttl_ms,
const std::string& tenant_id);

/**
* @brief Fetch all keys for a single tenant.
* @return ErrorCode::OK if exists
Expand Down Expand Up @@ -1218,11 +1222,14 @@ class MasterService {

std::unordered_map<std::string, std::unordered_set<std::string>>
group_members; // group_id → set of keys
std::unordered_map<std::string, std::chrono::system_clock::time_point>
group_retention_deadlines;

bool Empty() const {
return metadata.empty() && processing_keys.empty() &&
replication_tasks.empty() && offloading_tasks.empty() &&
promotion_tasks.empty() && group_members.empty();
promotion_tasks.empty() && group_members.empty() &&
group_retention_deadlines.empty();
}
};

Expand Down Expand Up @@ -1397,6 +1404,7 @@ class MasterService {
const std::string& tenant_id,
const std::string& key,
const std::string& group_id);
void PruneExpiredGroupRetentions();
std::unordered_map<std::string, ObjectMetadata>::iterator EraseMetadata(
TenantState& tenant_state,
std::unordered_map<std::string, ObjectMetadata>::iterator it,
Expand Down Expand Up @@ -1538,6 +1546,9 @@ class MasterService {
// Lease related members
const uint64_t default_kv_lease_ttl_; // in milliseconds
const uint64_t default_kv_soft_pin_ttl_; // in milliseconds
const size_t max_retained_groups_;
const uint64_t max_group_retention_ttl_ms_;
std::atomic<size_t> retained_group_count_{0};
const bool allow_evict_soft_pinned_objects_;

// Eviction related members
Expand Down
6 changes: 6 additions & 0 deletions mooncake-store/include/real_client.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -314,6 +314,9 @@ class RealClient : public PyClient {
*/
std::vector<int> batchIsExist(const std::vector<std::string> &keys);

std::vector<int> retainGroups(const std::vector<std::string> &group_ids,
uint64_t ttl_ms);

/**
* @brief Get the size of an object
* @param key Key of the object
Expand Down Expand Up @@ -655,6 +658,9 @@ class RealClient : public PyClient {
std::vector<tl::expected<bool, ErrorCode>> batchIsExist_internal(
const std::vector<std::string> &keys);

std::vector<tl::expected<bool, ErrorCode>> retainGroups_internal(
const std::vector<std::string> &group_ids, uint64_t ttl_ms);

tl::expected<int64_t, ErrorCode> getSize_internal(const std::string &key);

std::shared_ptr<BufferHandle> get_buffer_internal(
Expand Down
4 changes: 4 additions & 0 deletions mooncake-store/include/rpc_service.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,10 @@ class WrappedMasterService {
const std::vector<std::string>& keys,
const std::string& tenant_id = "default");

std::vector<tl::expected<bool, ErrorCode>> RetainGroups(
const std::vector<std::string>& group_ids, uint64_t ttl_ms,
const std::string& tenant_id = "default");

tl::expected<
std::unordered_map<UUID, std::vector<std::string>, boost::hash<UUID>>,
ErrorCode>
Expand Down
3 changes: 3 additions & 0 deletions mooncake-store/include/types.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,9 @@ static constexpr uint64_t DEFAULT_DEFAULT_KV_LEASE_TTL =
5000; // in milliseconds
static constexpr uint64_t DEFAULT_KV_SOFT_PIN_TTL_MS =
30 * 60 * 1000; // 30 minutes
static constexpr size_t DEFAULT_MAX_RETAINED_GROUPS = 65536;
static constexpr uint64_t DEFAULT_MAX_GROUP_RETENTION_TTL_MS =
24 * 60 * 60 * 1000; // 24 hours
static constexpr bool DEFAULT_ALLOW_EVICT_SOFT_PINNED_OBJECTS = true;
static constexpr double DEFAULT_EVICTION_RATIO = 0.05;
static constexpr double DEFAULT_EVICTION_HIGH_WATERMARK_RATIO = 0.95;
Expand Down
10 changes: 10 additions & 0 deletions mooncake-store/src/client_service.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3013,6 +3013,16 @@ std::vector<tl::expected<bool, ErrorCode>> Client::BatchIsExist(
return response;
}

std::vector<tl::expected<bool, ErrorCode>> Client::RetainGroups(
const std::vector<std::string>& group_ids, uint64_t ttl_ms) {
auto response = master_client_.RetainGroups(group_ids, ttl_ms);
if (response.size() != group_ids.size()) {
return std::vector<tl::expected<bool, ErrorCode>>(
group_ids.size(), tl::unexpected(ErrorCode::RPC_FAIL));
}
return response;
}

void* Client::GetBaseAddr() { return transfer_engine_->getBaseAddr(); }

tl::expected<void, ErrorCode> Client::MountLocalDiskSegment(
Expand Down
11 changes: 11 additions & 0 deletions mooncake-store/src/master_client.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,11 @@ struct RpcNameTraits<&WrappedMasterService::BatchExistKey> {
static constexpr const char* value = "BatchExistKey";
};

template <>
struct RpcNameTraits<&WrappedMasterService::RetainGroups> {
static constexpr const char* value = "RetainGroups";
};

template <>
struct RpcNameTraits<&WrappedMasterService::GetReplicaList> {
static constexpr const char* value = "GetReplicaList";
Expand Down Expand Up @@ -469,6 +474,12 @@ std::vector<tl::expected<bool, ErrorCode>> MasterClient::BatchExistKey(
return result;
}

std::vector<tl::expected<bool, ErrorCode>> MasterClient::RetainGroups(
const std::vector<std::string>& group_ids, uint64_t ttl_ms) {
return invoke_batch_rpc<&WrappedMasterService::RetainGroups, bool>(
group_ids.size(), group_ids, ttl_ms, tenant_id_);
}

tl::expected<MasterMetricManager::CacheHitStatDict, ErrorCode>
MasterClient::CalcCacheStats() {
return invoke_rpc<&WrappedMasterService::CalcCacheStats,
Expand Down
Loading
Loading