Files
milvus/internal/querynodev2/segments/mock_loader.go
T
sijie-ni-0214 24db48a7e6 enhance: add QueryCoord autoscale load resource precheck (#51218)
## Summary

issue: #51217

Load requests introduce resource demand before QueryNode waterlines
reflect the new data. This PR adds optional QueryCoord resource
prechecks before LoadCollection/LoadPartitions configuration broadcasts
and publishes cumulative demand counters for external autoscaling
decisions.

Design: [Autoscale-Aware Load Resource Admission in
QueryCoord](https://github.com/sijie-ni-0214/milvus/blob/querycoord-autoscale-kernel/docs/design-docs/design_docs/20260721-querycoord-autoscale-load-resource-precheck.md).

## Admission behavior

- Estimate incremental final sealed-segment memory and local-disk demand
per resource group from the requested and current load configurations.
- With tiered eviction disabled, compare demand with memory availability
and each Worker's configured `LOCAL_STORAGE_SIZE` quota minus cached
local-directory usage after admission thresholds. Optionally admit
shortages within configured global autoscale upper bounds.
- With tiered eviction enabled, retain demand estimation and skip QN
capacity sampling, node availability checks, and global autoscale-limit
checks. Existing evictable caches make physical usage unsuitable for
this admission decision; Worker QN retains its loading-stage guard.
- Fail open on metadata, estimation, and capacity-metrics failures.
Capacity-metrics failures preserve successfully estimated demand;
metadata or estimation failures leave no demand to publish.
- Return a resource-insufficient status with `suggested_expand_percent`
when a complete capacity evaluation rejects the request. Proxy preserves
the status and ExtraInfo.
- Admission remains best-effort, without reservations or accounting for
concurrent accepted requests. QueryCoord does not subtract Worker
in-flight load reservations; the Worker guard remains authoritative. The
existing replica node-count prerequisite remains; node provisioning and
scale-from-zero are outside this PR.

## Load demand counters

Expose cumulative counters labeled by resource group:

- `milvus_querycoord_load_demand_memory_bytes`
- `milvus_querycoord_load_demand_disk_bytes`

Counters increase only when demand estimation succeeds, the request
proceeds, and its configuration broadcast succeeds. This includes tiered
loads and loads admitted through capacity-metrics fail-open. Disabled
prechecks, unchanged configurations, zero incremental demand, estimation
failures, explicit rejections, and failed broadcasts do not increase
counters. Load completion and release do not decrement them.

Series start at zero for existing resource groups and are removed on RG
deletion. The dashboard displays cumulative values in MiB; the control
plane computes interval increments, handles counter resets, and combines
demand with QN waterlines and provisioning state. These counters are
observations, not durable reservations. Internal handoff and balancing
tasks do not contribute to them.

## Resource estimation and sampling

- Share segment/index estimation in
`internal/util/segcore/loadresource`, retaining separate final-resource
and loading-peak APIs. Worker QN uses its CGo pool; QueryCoord estimates
sequentially within each request.
- Fetch recovery metadata, full segment metadata, and batched index
metadata. Support Storage V3 resource summaries and dropped-field
metadata. Filter indexes by loaded fields, retaining multiple JSON-path
indexes on one field.
- Apply the load-request collection mmap inheritance logic to cloned
expected/current schemas while preserving field overrides. The final
estimator applies cache ratios only with tiered eviction enabled.
Delegator-side Bloom/BM25 resources are outside the Worker estimate.
- Deduplicate QN resource sampling by node ID, use at most 16 concurrent
requests and one attempt per node, and bound the entire sampling round
to ten seconds. Reuse the snapshot for RG and global calculations.
Exclude embedded SQNs from global capacity and sampling; empty
eligible-node sets are unevaluable and fail open.
- Worker QNs report `LOCAL_STORAGE_SIZE` and cached local-directory
usage through `system_info`. QueryCoord uses these byte values for disk
availability, RG suggestions, and the global disk-capacity sum;
filesystem-wide disk metrics remain for monitoring. Shared filesystems
require per-QN quotas consistent with the shared capacity.

## Configuration

- `queryCoord.autoscale.precheckEnabled`: enable demand estimation and
precheck; default `false`.
- `queryCoord.autoscale.enabled`: allow admission within global
autoscale bounds when current capacity is insufficient; default `false`,
effective only with precheck enabled and tiered eviction disabled.
- `queryCoord.autoscale.maxMemoryLimit`: maximum global Worker QN
physical memory capacity after scaling, in GiB; default `0` means no
additional memory autoscale capacity.
- `queryCoord.autoscale.maxDiskLimit`: maximum sum of Worker QN
`LOCAL_STORAGE_SIZE` quotas after scaling, in GiB; default `0` means no
additional disk autoscale capacity. Both limits use `1 GiB = 1024^3
bytes`.

For non-tiered deployments, configure QueryCoord's memory high-water
level no higher than the Worker admission threshold. This relationship
is a rollout prerequisite and is not enforced by this PR.

## Validation

Local autoscale estimator and targeted QueryCoord tests passed,
including multiple-index filtering, collection mmap overrides, the
ten-second sampling timeout, empty-node fail-open, embedded-SQN
exclusion, tiered demand preservation, capacity-metrics failure
behavior, and local-storage quota accounting for QueryNodes sharing a
filesystem.

Both load callbacks are tested for precheck rejection, broadcast
failure, successful counter updates, tiered loads, and capacity-metrics
failures followed by successful or failed broadcasts. Tests assert
counters remain unchanged during Broadcast and increment only after it
succeeds. These are local unit/callback checks, not end-to-end external
autoscaler validation.

Local `make static-check` passed across core, pkg, client, and go_client
with zero issues, and segcore error-boundary checks passed. Targeted
QueryCoord and QueryNode tests, YAML consistency, and `make build-go`
passed. External autoscaler behavior remains subject to CI and rollout
validation.

---------

Signed-off-by: sijie-ni-0214 <[email protected]>
2026-09-29 01:47:36 +08:00

388 lines
12 KiB
Go

// Code generated by mockery v2.53.3. DO NOT EDIT.
package segments
import (
context "context"
commonpb "github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
pkoracle "github.com/milvus-io/milvus/internal/querynodev2/pkoracle"
storage "github.com/milvus-io/milvus/internal/storage"
querypb "github.com/milvus-io/milvus/pkg/v3/proto/querypb"
mock "github.com/stretchr/testify/mock"
)
// MockLoader is an autogenerated mock type for the Loader type
type MockLoader struct {
mock.Mock
}
type MockLoader_Expecter struct {
mock *mock.Mock
}
func (_m *MockLoader) EXPECT() *MockLoader_Expecter {
return &MockLoader_Expecter{mock: &_m.Mock}
}
// GetChunkManager provides a mock function with no fields
func (_m *MockLoader) GetChunkManager() storage.ChunkManager {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for GetChunkManager")
}
var r0 storage.ChunkManager
if rf, ok := ret.Get(0).(func() storage.ChunkManager); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(storage.ChunkManager)
}
}
return r0
}
// MockLoader_GetChunkManager_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetChunkManager'
type MockLoader_GetChunkManager_Call struct {
*mock.Call
}
// GetChunkManager is a helper method to define mock.On call
func (_e *MockLoader_Expecter) GetChunkManager() *MockLoader_GetChunkManager_Call {
return &MockLoader_GetChunkManager_Call{Call: _e.mock.On("GetChunkManager")}
}
func (_c *MockLoader_GetChunkManager_Call) Run(run func()) *MockLoader_GetChunkManager_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockLoader_GetChunkManager_Call) Return(_a0 storage.ChunkManager) *MockLoader_GetChunkManager_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockLoader_GetChunkManager_Call) RunAndReturn(run func() storage.ChunkManager) *MockLoader_GetChunkManager_Call {
_c.Call.Return(run)
return _c
}
// GetLocalDiskUsage provides a mock function with no fields
func (_m *MockLoader) GetLocalDiskUsage() (int64, error) {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for GetLocalDiskUsage")
}
var r0 int64
var r1 error
if rf, ok := ret.Get(0).(func() (int64, error)); ok {
return rf()
}
if rf, ok := ret.Get(0).(func() int64); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(int64)
}
if rf, ok := ret.Get(1).(func() error); ok {
r1 = rf()
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockLoader_GetLocalDiskUsage_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetLocalDiskUsage'
type MockLoader_GetLocalDiskUsage_Call struct {
*mock.Call
}
// GetLocalDiskUsage is a helper method to define mock.On call
func (_e *MockLoader_Expecter) GetLocalDiskUsage() *MockLoader_GetLocalDiskUsage_Call {
return &MockLoader_GetLocalDiskUsage_Call{Call: _e.mock.On("GetLocalDiskUsage")}
}
func (_c *MockLoader_GetLocalDiskUsage_Call) Run(run func()) *MockLoader_GetLocalDiskUsage_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockLoader_GetLocalDiskUsage_Call) Return(_a0 int64, _a1 error) *MockLoader_GetLocalDiskUsage_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockLoader_GetLocalDiskUsage_Call) RunAndReturn(run func() (int64, error)) *MockLoader_GetLocalDiskUsage_Call {
_c.Call.Return(run)
return _c
}
// Load provides a mock function with given fields: ctx, collectionID, segmentType, version, _a4
func (_m *MockLoader) Load(ctx context.Context, collectionID int64, segmentType commonpb.SegmentState, version int64, _a4 ...*querypb.SegmentLoadInfo) ([]Segment, error) {
_va := make([]interface{}, len(_a4))
for _i := range _a4 {
_va[_i] = _a4[_i]
}
var _ca []interface{}
_ca = append(_ca, ctx, collectionID, segmentType, version)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
if len(ret) == 0 {
panic("no return value specified for Load")
}
var r0 []Segment
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, int64, commonpb.SegmentState, int64, ...*querypb.SegmentLoadInfo) ([]Segment, error)); ok {
return rf(ctx, collectionID, segmentType, version, _a4...)
}
if rf, ok := ret.Get(0).(func(context.Context, int64, commonpb.SegmentState, int64, ...*querypb.SegmentLoadInfo) []Segment); ok {
r0 = rf(ctx, collectionID, segmentType, version, _a4...)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]Segment)
}
}
if rf, ok := ret.Get(1).(func(context.Context, int64, commonpb.SegmentState, int64, ...*querypb.SegmentLoadInfo) error); ok {
r1 = rf(ctx, collectionID, segmentType, version, _a4...)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockLoader_Load_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Load'
type MockLoader_Load_Call struct {
*mock.Call
}
// Load is a helper method to define mock.On call
// - ctx context.Context
// - collectionID int64
// - segmentType commonpb.SegmentState
// - version int64
// - _a4 ...*querypb.SegmentLoadInfo
func (_e *MockLoader_Expecter) Load(ctx interface{}, collectionID interface{}, segmentType interface{}, version interface{}, _a4 ...interface{}) *MockLoader_Load_Call {
return &MockLoader_Load_Call{Call: _e.mock.On("Load",
append([]interface{}{ctx, collectionID, segmentType, version}, _a4...)...)}
}
func (_c *MockLoader_Load_Call) Run(run func(ctx context.Context, collectionID int64, segmentType commonpb.SegmentState, version int64, _a4 ...*querypb.SegmentLoadInfo)) *MockLoader_Load_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]*querypb.SegmentLoadInfo, len(args)-4)
for i, a := range args[4:] {
if a != nil {
variadicArgs[i] = a.(*querypb.SegmentLoadInfo)
}
}
run(args[0].(context.Context), args[1].(int64), args[2].(commonpb.SegmentState), args[3].(int64), variadicArgs...)
})
return _c
}
func (_c *MockLoader_Load_Call) Return(_a0 []Segment, _a1 error) *MockLoader_Load_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockLoader_Load_Call) RunAndReturn(run func(context.Context, int64, commonpb.SegmentState, int64, ...*querypb.SegmentLoadInfo) ([]Segment, error)) *MockLoader_Load_Call {
_c.Call.Return(run)
return _c
}
// LoadBloomFilterSet provides a mock function with given fields: ctx, collectionID, infos
func (_m *MockLoader) LoadBloomFilterSet(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) ([]*pkoracle.BloomFilterSet, error) {
_va := make([]interface{}, len(infos))
for _i := range infos {
_va[_i] = infos[_i]
}
var _ca []interface{}
_ca = append(_ca, ctx, collectionID)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
if len(ret) == 0 {
panic("no return value specified for LoadBloomFilterSet")
}
var r0 []*pkoracle.BloomFilterSet
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, int64, ...*querypb.SegmentLoadInfo) ([]*pkoracle.BloomFilterSet, error)); ok {
return rf(ctx, collectionID, infos...)
}
if rf, ok := ret.Get(0).(func(context.Context, int64, ...*querypb.SegmentLoadInfo) []*pkoracle.BloomFilterSet); ok {
r0 = rf(ctx, collectionID, infos...)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*pkoracle.BloomFilterSet)
}
}
if rf, ok := ret.Get(1).(func(context.Context, int64, ...*querypb.SegmentLoadInfo) error); ok {
r1 = rf(ctx, collectionID, infos...)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockLoader_LoadBloomFilterSet_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadBloomFilterSet'
type MockLoader_LoadBloomFilterSet_Call struct {
*mock.Call
}
// LoadBloomFilterSet is a helper method to define mock.On call
// - ctx context.Context
// - collectionID int64
// - infos ...*querypb.SegmentLoadInfo
func (_e *MockLoader_Expecter) LoadBloomFilterSet(ctx interface{}, collectionID interface{}, infos ...interface{}) *MockLoader_LoadBloomFilterSet_Call {
return &MockLoader_LoadBloomFilterSet_Call{Call: _e.mock.On("LoadBloomFilterSet",
append([]interface{}{ctx, collectionID}, infos...)...)}
}
func (_c *MockLoader_LoadBloomFilterSet_Call) Run(run func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo)) *MockLoader_LoadBloomFilterSet_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]*querypb.SegmentLoadInfo, len(args)-2)
for i, a := range args[2:] {
if a != nil {
variadicArgs[i] = a.(*querypb.SegmentLoadInfo)
}
}
run(args[0].(context.Context), args[1].(int64), variadicArgs...)
})
return _c
}
func (_c *MockLoader_LoadBloomFilterSet_Call) Return(_a0 []*pkoracle.BloomFilterSet, _a1 error) *MockLoader_LoadBloomFilterSet_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockLoader_LoadBloomFilterSet_Call) RunAndReturn(run func(context.Context, int64, ...*querypb.SegmentLoadInfo) ([]*pkoracle.BloomFilterSet, error)) *MockLoader_LoadBloomFilterSet_Call {
_c.Call.Return(run)
return _c
}
// LoadDeltaLogs provides a mock function with given fields: ctx, segment, loadInfo
func (_m *MockLoader) LoadDeltaLogs(ctx context.Context, segment Segment, loadInfo *querypb.SegmentLoadInfo) error {
ret := _m.Called(ctx, segment, loadInfo)
if len(ret) == 0 {
panic("no return value specified for LoadDeltaLogs")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, Segment, *querypb.SegmentLoadInfo) error); ok {
r0 = rf(ctx, segment, loadInfo)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockLoader_LoadDeltaLogs_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadDeltaLogs'
type MockLoader_LoadDeltaLogs_Call struct {
*mock.Call
}
// LoadDeltaLogs is a helper method to define mock.On call
// - ctx context.Context
// - segment Segment
// - loadInfo *querypb.SegmentLoadInfo
func (_e *MockLoader_Expecter) LoadDeltaLogs(ctx interface{}, segment interface{}, loadInfo interface{}) *MockLoader_LoadDeltaLogs_Call {
return &MockLoader_LoadDeltaLogs_Call{Call: _e.mock.On("LoadDeltaLogs", ctx, segment, loadInfo)}
}
func (_c *MockLoader_LoadDeltaLogs_Call) Run(run func(ctx context.Context, segment Segment, loadInfo *querypb.SegmentLoadInfo)) *MockLoader_LoadDeltaLogs_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(Segment), args[2].(*querypb.SegmentLoadInfo))
})
return _c
}
func (_c *MockLoader_LoadDeltaLogs_Call) Return(_a0 error) *MockLoader_LoadDeltaLogs_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockLoader_LoadDeltaLogs_Call) RunAndReturn(run func(context.Context, Segment, *querypb.SegmentLoadInfo) error) *MockLoader_LoadDeltaLogs_Call {
_c.Call.Return(run)
return _c
}
// ReopenSegments provides a mock function with given fields: ctx, loadInfos
func (_m *MockLoader) ReopenSegments(ctx context.Context, loadInfos []*querypb.SegmentLoadInfo) error {
ret := _m.Called(ctx, loadInfos)
if len(ret) == 0 {
panic("no return value specified for ReopenSegments")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, []*querypb.SegmentLoadInfo) error); ok {
r0 = rf(ctx, loadInfos)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockLoader_ReopenSegments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ReopenSegments'
type MockLoader_ReopenSegments_Call struct {
*mock.Call
}
// ReopenSegments is a helper method to define mock.On call
// - ctx context.Context
// - loadInfos []*querypb.SegmentLoadInfo
func (_e *MockLoader_Expecter) ReopenSegments(ctx interface{}, loadInfos interface{}) *MockLoader_ReopenSegments_Call {
return &MockLoader_ReopenSegments_Call{Call: _e.mock.On("ReopenSegments", ctx, loadInfos)}
}
func (_c *MockLoader_ReopenSegments_Call) Run(run func(ctx context.Context, loadInfos []*querypb.SegmentLoadInfo)) *MockLoader_ReopenSegments_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].([]*querypb.SegmentLoadInfo))
})
return _c
}
func (_c *MockLoader_ReopenSegments_Call) Return(_a0 error) *MockLoader_ReopenSegments_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockLoader_ReopenSegments_Call) RunAndReturn(run func(context.Context, []*querypb.SegmentLoadInfo) error) *MockLoader_ReopenSegments_Call {
_c.Call.Return(run)
return _c
}
// NewMockLoader creates a new instance of MockLoader. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewMockLoader(t interface {
mock.TestingT
Cleanup(func())
}) *MockLoader {
mock := &MockLoader{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}