Skip to content

enhance: [2.6] add max loading size config for caching layer #121

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
sparknack wants to merge 2 commits into
base: 2.6
Choose a base branch
from
Open

enhance: [2.6] add max loading size config for caching layer #121

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e2596ac
enhance: add max loading memory limit
sparknack Aug 13, 2026
7cf7b7f
enhance: limit loading concurrency by overhead
sparknack Aug 13, 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
25 changes: 17 additions & 8 deletions include/cachinglayer/CacheSlot.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -427,17 +427,24 @@ class CacheSlot final : public std::enable_shared_from_this<CacheSlot<CellT>> {
// bonus cells should be empty if self_reserve_ is false.
auto bonus_cids = translator_->bonus_cells_to_be_loaded(loading_cids);

ResourceUsage essential_loaded_resource;
for (auto& cid : loading_cids) {
essential_loading_resource += translator_->estimated_byte_size_of_cell(cid).second;
const auto [loaded_resource, loading_resource] = translator_->estimated_byte_size_of_cell(cid);
essential_loaded_resource += loaded_resource;
essential_loading_resource += loading_resource;
}

ResourceUsage bonus_loaded_resource;
for (auto& cid : bonus_cids) {
bonus_loading_resource += translator_->estimated_byte_size_of_cell(cid).second;
const auto [loaded_resource, loading_resource] = translator_->estimated_byte_size_of_cell(cid);
bonus_loaded_resource += loaded_resource;
bonus_loading_resource += loading_resource;
}

auto resource_needed_for_loading = essential_loading_resource + bonus_loading_resource;
reservation_success =
SemiInlineGet(dlist_->ReserveLoadingResourceWithTimeout(resource_needed_for_loading, timeout, ctx));
auto loaded_resource = essential_loaded_resource + bonus_loaded_resource;

@buqian-zilliz buqian-zilliz Aug 14, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

[P1] Count discarded concurrent copies as loading overhead

loaded_resource includes bonus_loaded_resource for every reservation, although bonus CIDs are not claimed before I/O. CacheCell::set_cell explicitly permits concurrent bonus loads; when another RunLoad wins, this fetch's CellT is discarded without ChargeLoadedResource, so the discarded copy consumes its full allocation as transient loading memory. The accounting records only loading_estimate - loaded_estimate (possibly zero when estimates are equal), allowing duplicate prefetches to bypass the configured loading-overhead cap. Only subtract bytes for cells this reservation actually installs, or conservatively count raceable bonus copies as overhead.

reservation_success = SemiInlineGet(
dlist_->ReserveLoadingResourceWithTimeout(resource_needed_for_loading, loaded_resource, timeout, ctx));

if (!bonus_cids.empty()) {
// if the reservation failed, try to reserve only the essential loading resource
Expand All @@ -447,8 +454,9 @@ class CacheSlot final : public std::enable_shared_from_this<CacheSlot<CellT>> {
"essential "
"loading resource");
resource_needed_for_loading = essential_loading_resource;
reservation_success = SemiInlineGet(
dlist_->ReserveLoadingResourceWithTimeout(resource_needed_for_loading, timeout, ctx));
loaded_resource = essential_loaded_resource;
reservation_success = SemiInlineGet(dlist_->ReserveLoadingResourceWithTimeout(
resource_needed_for_loading, essential_loaded_resource, timeout, ctx));
} else {
// if the reservation succeeded, we can load the bonus cells
loading_cids.insert(loading_cids.end(), bonus_cids.begin(), bonus_cids.end());
Expand All @@ -475,9 +483,10 @@ class CacheSlot final : public std::enable_shared_from_this<CacheSlot<CellT>> {
monitor::cache_cell_loading_count(cell_data_type_, storage_type_).Increment(loading_cids.size());

// defer release resource_needed_for_loading
auto defer_release = folly::makeGuard([this, &resource_needed_for_loading, &loading_cids]() {
auto defer_release =
folly::makeGuard([this, &resource_needed_for_loading, &loaded_resource, &loading_cids]() {
try {
dlist_->ReleaseLoadingResource(resource_needed_for_loading);
dlist_->ReleaseLoadingResource(resource_needed_for_loading, loaded_resource);
monitor::cache_cell_loading_count(cell_data_type_, storage_type_).Decrement(loading_cids.size());
monitor::cache_loading_bytes(cell_data_type_, StorageType::MEMORY)
.Decrement(resource_needed_for_loading.memory_bytes);
Expand Down
25 changes: 25 additions & 0 deletions include/cachinglayer/Manager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,30 @@

namespace milvus::cachinglayer {

struct TieredStorageOptions {
CacheWarmupPolicies warmup_policies{};
CacheLimit cache_limit{};
bool storage_usage_tracking_enabled{false};
bool eviction_enabled{false};
EvictionConfig eviction_config{};
std::chrono::milliseconds loading_timeout{100000};
std::chrono::milliseconds warmup_loading_timeout{0};
uint32_t prefetch_pool_threads{0};
// Ratio of the effective memory limit.
double max_loading_mem_ratio{1.0};
};

class Manager {
public:
static Manager&
GetInstance();

// Must be called exactly once before any CacheSlot is created.
static void
ConfigureTieredStorage(const TieredStorageOptions& options);

// Compatibility overload for existing callers. Prefer TieredStorageOptions for new fields.
static void
ConfigureTieredStorage(CacheWarmupPolicies warmup_policies, CacheLimit cache_limit,
bool storage_usage_tracking_enabled, bool eviction_enabled, EvictionConfig eviction_config,
std::chrono::milliseconds loading_timeout,
Expand All @@ -42,6 +59,14 @@ class Manager {
UpdateConfig(std::chrono::milliseconds loading_timeout, std::chrono::milliseconds warmup_loading_timeout,
bool storage_usage_tracking_enabled, CacheWarmupPolicies warmup_policies);

static void
UpdateConfig(std::chrono::milliseconds loading_timeout, std::chrono::milliseconds warmup_loading_timeout,
bool storage_usage_tracking_enabled, CacheWarmupPolicies warmup_policies,
double max_loading_mem_ratio);

static void
UpdateMaxLoadingMemRatio(double max_loading_mem_ratio);

~Manager();

Manager(const Manager&) = delete;
Expand Down
34 changes: 33 additions & 1 deletion include/cachinglayer/TieredStorageConfig.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ class TieredStorageConfig {
bool storage_usage_tracking_enabled;
std::chrono::milliseconds loading_timeout;
std::chrono::milliseconds warmup_loading_timeout;
double max_loading_mem_ratio;
CacheWarmupPolicies warmup_policies;
};

Expand All @@ -48,7 +49,8 @@ class TieredStorageConfig {
[[nodiscard]] Snapshot
GetSnapshot() const {
std::shared_lock lock(mtx_);
return {storage_usage_tracking_enabled_, loading_timeout_, warmup_loading_timeout_, warmup_policies_};
return {storage_usage_tracking_enabled_, loading_timeout_, warmup_loading_timeout_, max_loading_mem_ratio_,
warmup_policies_};
}

// --- Individual readers (shared lock) ---
Expand All @@ -71,6 +73,12 @@ class TieredStorageConfig {
return warmup_loading_timeout_;
}

[[nodiscard]] double
max_loading_mem_ratio() const {
std::shared_lock lock(mtx_);
return max_loading_mem_ratio_;
}

[[nodiscard]] CacheWarmupPolicies
warmup_policies() const {
std::shared_lock lock(mtx_);
Expand All @@ -89,6 +97,20 @@ class TieredStorageConfig {
warmup_policies_ = warmup_policies;
}

void
UpdateAll(bool storage_usage_tracking_enabled, std::chrono::milliseconds loading_timeout,
std::chrono::milliseconds warmup_loading_timeout, CacheWarmupPolicies warmup_policies,
double max_loading_mem_ratio) {
AssertInfo(max_loading_mem_ratio >= 0 && max_loading_mem_ratio <= 1,
"[MCL] max loading memory ratio must be between 0 and 1, got {}", max_loading_mem_ratio);
std::unique_lock lock(mtx_);
storage_usage_tracking_enabled_ = storage_usage_tracking_enabled;
loading_timeout_ = loading_timeout;
warmup_loading_timeout_ = warmup_loading_timeout;
max_loading_mem_ratio_ = max_loading_mem_ratio;
warmup_policies_ = warmup_policies;
}

// --- Individual writers (exclusive lock) ---

void
Expand All @@ -109,6 +131,14 @@ class TieredStorageConfig {
warmup_loading_timeout_ = timeout;
}

void
SetMaxLoadingMemRatio(double max_loading_mem_ratio) {
AssertInfo(max_loading_mem_ratio >= 0 && max_loading_mem_ratio <= 1,
"[MCL] max loading memory ratio must be between 0 and 1, got {}", max_loading_mem_ratio);
std::unique_lock lock(mtx_);
max_loading_mem_ratio_ = max_loading_mem_ratio;
}

void
SetWarmupPolicies(CacheWarmupPolicies policies) {
std::unique_lock lock(mtx_);
Expand All @@ -122,6 +152,8 @@ class TieredStorageConfig {
bool storage_usage_tracking_enabled_{false};
std::chrono::milliseconds loading_timeout_{100000};
std::chrono::milliseconds warmup_loading_timeout_{0};
// Ratio of the effective memory limit.
double max_loading_mem_ratio_{1.0};
CacheWarmupPolicies warmup_policies_{};
};

Expand Down
4 changes: 4 additions & 0 deletions include/cachinglayer/Utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -338,6 +338,10 @@ getHostTotalMemory();
int64_t
getContainerMemLimit();

// Returns -1 when the effective memory limit cannot be determined.
int64_t
getMaxLoadingMemSize(double max_loading_mem_ratio);

// Returns unlimited if failed to get memory info, or if the platform is not supported.
SystemResourceInfo
getSystemMemoryInfo();
Expand Down
40 changes: 31 additions & 9 deletions include/cachinglayer/lrucache/DList.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,18 @@ namespace milvus::cachinglayer::internal {
class DList : public std::enable_shared_from_this<DList> {
public:
DList(bool eviction_enabled, ResourceUsage max_memory, ResourceUsage low_watermark, ResourceUsage high_watermark,
EvictionConfig eviction_config)
EvictionConfig eviction_config, int64_t max_loading_mem_size = -1)
: max_resource_limit_(max_memory),
low_watermark_(low_watermark),
high_watermark_(high_watermark),
max_loading_mem_size_(max_loading_mem_size),
eviction_config_(eviction_config),
next_request_id_(1) {
// if eviction is disabled, we don't need to initialize the event base and thread
AssertInfo(max_loading_mem_size >= -1, "[MCL] max loading memory size must be -1 or greater");

event_base_thread_ = std::make_unique<folly::ScopedEventBaseThread>("cache-eb");

// if eviction is disabled, we don't need to initialize the background eviction thread
if (!eviction_enabled) {
return;
}
Expand All @@ -58,9 +63,6 @@ class DList : public std::enable_shared_from_this<DList> {
monitor::cache_low_watermark_bytes(StorageType::MEMORY).Set(low_watermark.memory_bytes);
monitor::cache_low_watermark_bytes(StorageType::DISK).Set(low_watermark.file_bytes);

// Initialize event base and thread
event_base_thread_ = std::make_unique<folly::ScopedEventBaseThread>("cache-eb");

if (eviction_config_.background_eviction_enabled) {
LOG_INFO("[MCL] Starting periodic background eviction loop thread");
bg_eviction_thread_ = std::thread(&DList::evictionLoop, this);
Expand Down Expand Up @@ -112,11 +114,17 @@ class DList : public std::enable_shared_from_this<DList> {
void
UpdateHighWatermark(const ResourceUsage& new_high_watermark);

// Update max loading overhead memory size does not evict loaded cache; it may only wake waiting requests.
void
UpdateMaxLoadingMemSize(int64_t new_max_loading_mem_size);

// True if no nodes in the list.
bool
IsEmpty() const;

// Reserve loading resource with timeout, called before loading a cell.
// The two-size overload takes the estimated loaded size separately so the
// max loading limit can be applied to loading overhead only.
// When timeout > 0, the request will wait up to the specified duration.
// When timeout == 0, the request will fail immediately without entering the
// waiting queue (best-effort mode, used for warmup scenarios).
Expand All @@ -127,10 +135,17 @@ class DList : public std::enable_shared_from_this<DList> {
ReserveLoadingResourceWithTimeout(const ResourceUsage& size, std::chrono::milliseconds timeout,
OpContext* ctx = nullptr);

folly::SemiFuture<bool>
ReserveLoadingResourceWithTimeout(const ResourceUsage& loading_size, const ResourceUsage& loaded_size,
std::chrono::milliseconds timeout, OpContext* ctx = nullptr);

// Release resource used for loading, called after loading a cell.
void
ReleaseLoadingResource(const ResourceUsage& loading_size);

void
ReleaseLoadingResource(const ResourceUsage& loading_size, const ResourceUsage& loaded_size);

// Called when a cell is loaded.
void
ChargeLoadedResource(const ResourceUsage& size);
Expand Down Expand Up @@ -171,14 +186,15 @@ class DList : public std::enable_shared_from_this<DList> {
// Waiting request for timeout-based memory reservation
struct WaitingRequest {
ResourceUsage required_size;
ResourceUsage overhead_size;
std::chrono::steady_clock::time_point deadline;
folly::Promise<bool> promise;
uint64_t request_id;
std::optional<folly::CancellationCallback> cancel_cb{std::nullopt};

WaitingRequest(ResourceUsage size, std::chrono::steady_clock::time_point dl, folly::Promise<bool> p,
uint64_t id)
: required_size(size), deadline(dl), promise(std::move(p)), request_id(id) {
WaitingRequest(ResourceUsage size, ResourceUsage overhead, std::chrono::steady_clock::time_point dl,
folly::Promise<bool> p, uint64_t id)
: required_size(size), overhead_size(overhead), deadline(dl), promise(std::move(p)), request_id(id) {
}
};

Expand All @@ -199,7 +215,11 @@ class DList : public std::enable_shared_from_this<DList> {

// reserveResource without taking lock, must be called with lock held.
bool
reserveResourceInternal(const ResourceUsage& size);
reserveResourceInternal(const ResourceUsage& size, const ResourceUsage& overhead_size);

// Returns true when current loading overhead reaches the concurrent overhead memory limit.
bool
exceedMaxOverheadMemSize() const;

void
evictionLoop();
Expand Down Expand Up @@ -268,6 +288,8 @@ class DList : public std::enable_shared_from_this<DList> {
std::atomic<ResourceUsage> max_resource_limit_;
std::atomic<ResourceUsage> low_watermark_;
std::atomic<ResourceUsage> high_watermark_;
std::atomic<int64_t> max_loading_mem_size_;
std::atomic<ResourceUsage> total_loading_overhead_size_{};
const EvictionConfig eviction_config_;

std::thread bg_eviction_thread_;
Expand Down
Loading