milvus/internal/querycoordv2/task/mock_scheduler.go
wei liu b18de95817
enhance: Avoid assign too much segment/channels to new querynode (#34096) (#34245)
issue: #34095
pr: #34096

When a new query node comes online, the segment_checker,
channel_checker, and balance_checker simultaneously attempt to allocate
segments to it. If this occurs during the execution of a load task and
the distribution of the new query node hasn't been updated, the query
coordinator may mistakenly view the new query node as empty. As a
result, it assigns segments or channels to it, potentially overloading
the new query node with more segments or channels than expected.

This PR measures the workload of the executing tasks on the target query
node to prevent assigning an excessive number of segments to it.

---------

Signed-off-by: Wei Liu <wei.liu@zilliz.com>
2024-07-01 10:32:06 +08:00

483 lines
14 KiB
Go

// Code generated by mockery v2.32.4. DO NOT EDIT.
package task
import mock "github.com/stretchr/testify/mock"
// MockScheduler is an autogenerated mock type for the Scheduler type
type MockScheduler struct {
mock.Mock
}
type MockScheduler_Expecter struct {
mock *mock.Mock
}
func (_m *MockScheduler) EXPECT() *MockScheduler_Expecter {
return &MockScheduler_Expecter{mock: &_m.Mock}
}
// Add provides a mock function with given fields: task
func (_m *MockScheduler) Add(task Task) error {
ret := _m.Called(task)
var r0 error
if rf, ok := ret.Get(0).(func(Task) error); ok {
r0 = rf(task)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockScheduler_Add_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Add'
type MockScheduler_Add_Call struct {
*mock.Call
}
// Add is a helper method to define mock.On call
// - task Task
func (_e *MockScheduler_Expecter) Add(task interface{}) *MockScheduler_Add_Call {
return &MockScheduler_Add_Call{Call: _e.mock.On("Add", task)}
}
func (_c *MockScheduler_Add_Call) Run(run func(task Task)) *MockScheduler_Add_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(Task))
})
return _c
}
func (_c *MockScheduler_Add_Call) Return(_a0 error) *MockScheduler_Add_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockScheduler_Add_Call) RunAndReturn(run func(Task) error) *MockScheduler_Add_Call {
_c.Call.Return(run)
return _c
}
// AddExecutor provides a mock function with given fields: nodeID
func (_m *MockScheduler) AddExecutor(nodeID int64) {
_m.Called(nodeID)
}
// MockScheduler_AddExecutor_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AddExecutor'
type MockScheduler_AddExecutor_Call struct {
*mock.Call
}
// AddExecutor is a helper method to define mock.On call
// - nodeID int64
func (_e *MockScheduler_Expecter) AddExecutor(nodeID interface{}) *MockScheduler_AddExecutor_Call {
return &MockScheduler_AddExecutor_Call{Call: _e.mock.On("AddExecutor", nodeID)}
}
func (_c *MockScheduler_AddExecutor_Call) Run(run func(nodeID int64)) *MockScheduler_AddExecutor_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64))
})
return _c
}
func (_c *MockScheduler_AddExecutor_Call) Return() *MockScheduler_AddExecutor_Call {
_c.Call.Return()
return _c
}
func (_c *MockScheduler_AddExecutor_Call) RunAndReturn(run func(int64)) *MockScheduler_AddExecutor_Call {
_c.Call.Return(run)
return _c
}
// Dispatch provides a mock function with given fields: node
func (_m *MockScheduler) Dispatch(node int64) {
_m.Called(node)
}
// MockScheduler_Dispatch_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Dispatch'
type MockScheduler_Dispatch_Call struct {
*mock.Call
}
// Dispatch is a helper method to define mock.On call
// - node int64
func (_e *MockScheduler_Expecter) Dispatch(node interface{}) *MockScheduler_Dispatch_Call {
return &MockScheduler_Dispatch_Call{Call: _e.mock.On("Dispatch", node)}
}
func (_c *MockScheduler_Dispatch_Call) Run(run func(node int64)) *MockScheduler_Dispatch_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64))
})
return _c
}
func (_c *MockScheduler_Dispatch_Call) Return() *MockScheduler_Dispatch_Call {
_c.Call.Return()
return _c
}
func (_c *MockScheduler_Dispatch_Call) RunAndReturn(run func(int64)) *MockScheduler_Dispatch_Call {
_c.Call.Return(run)
return _c
}
// GetChannelTaskDelta provides a mock function with given fields: nodeID, collectionID
func (_m *MockScheduler) GetChannelTaskDelta(nodeID int64, collectionID int64) int {
ret := _m.Called(nodeID, collectionID)
var r0 int
if rf, ok := ret.Get(0).(func(int64, int64) int); ok {
r0 = rf(nodeID, collectionID)
} else {
r0 = ret.Get(0).(int)
}
return r0
}
// MockScheduler_GetChannelTaskDelta_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetChannelTaskDelta'
type MockScheduler_GetChannelTaskDelta_Call struct {
*mock.Call
}
// GetChannelTaskDelta is a helper method to define mock.On call
// - nodeID int64
// - collectionID int64
func (_e *MockScheduler_Expecter) GetChannelTaskDelta(nodeID interface{}, collectionID interface{}) *MockScheduler_GetChannelTaskDelta_Call {
return &MockScheduler_GetChannelTaskDelta_Call{Call: _e.mock.On("GetChannelTaskDelta", nodeID, collectionID)}
}
func (_c *MockScheduler_GetChannelTaskDelta_Call) Run(run func(nodeID int64, collectionID int64)) *MockScheduler_GetChannelTaskDelta_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64))
})
return _c
}
func (_c *MockScheduler_GetChannelTaskDelta_Call) Return(_a0 int) *MockScheduler_GetChannelTaskDelta_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockScheduler_GetChannelTaskDelta_Call) RunAndReturn(run func(int64, int64) int) *MockScheduler_GetChannelTaskDelta_Call {
_c.Call.Return(run)
return _c
}
// GetChannelTaskNum provides a mock function with given fields:
func (_m *MockScheduler) GetChannelTaskNum() int {
ret := _m.Called()
var r0 int
if rf, ok := ret.Get(0).(func() int); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(int)
}
return r0
}
// MockScheduler_GetChannelTaskNum_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetChannelTaskNum'
type MockScheduler_GetChannelTaskNum_Call struct {
*mock.Call
}
// GetChannelTaskNum is a helper method to define mock.On call
func (_e *MockScheduler_Expecter) GetChannelTaskNum() *MockScheduler_GetChannelTaskNum_Call {
return &MockScheduler_GetChannelTaskNum_Call{Call: _e.mock.On("GetChannelTaskNum")}
}
func (_c *MockScheduler_GetChannelTaskNum_Call) Run(run func()) *MockScheduler_GetChannelTaskNum_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockScheduler_GetChannelTaskNum_Call) Return(_a0 int) *MockScheduler_GetChannelTaskNum_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockScheduler_GetChannelTaskNum_Call) RunAndReturn(run func() int) *MockScheduler_GetChannelTaskNum_Call {
_c.Call.Return(run)
return _c
}
// GetExecutedFlag provides a mock function with given fields: nodeID
func (_m *MockScheduler) GetExecutedFlag(nodeID int64) <-chan struct{} {
ret := _m.Called(nodeID)
var r0 <-chan struct{}
if rf, ok := ret.Get(0).(func(int64) <-chan struct{}); ok {
r0 = rf(nodeID)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(<-chan struct{})
}
}
return r0
}
// MockScheduler_GetExecutedFlag_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetExecutedFlag'
type MockScheduler_GetExecutedFlag_Call struct {
*mock.Call
}
// GetExecutedFlag is a helper method to define mock.On call
// - nodeID int64
func (_e *MockScheduler_Expecter) GetExecutedFlag(nodeID interface{}) *MockScheduler_GetExecutedFlag_Call {
return &MockScheduler_GetExecutedFlag_Call{Call: _e.mock.On("GetExecutedFlag", nodeID)}
}
func (_c *MockScheduler_GetExecutedFlag_Call) Run(run func(nodeID int64)) *MockScheduler_GetExecutedFlag_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64))
})
return _c
}
func (_c *MockScheduler_GetExecutedFlag_Call) Return(_a0 <-chan struct{}) *MockScheduler_GetExecutedFlag_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockScheduler_GetExecutedFlag_Call) RunAndReturn(run func(int64) <-chan struct{}) *MockScheduler_GetExecutedFlag_Call {
_c.Call.Return(run)
return _c
}
// GetSegmentTaskDelta provides a mock function with given fields: nodeID, collectionID
func (_m *MockScheduler) GetSegmentTaskDelta(nodeID int64, collectionID int64) int {
ret := _m.Called(nodeID, collectionID)
var r0 int
if rf, ok := ret.Get(0).(func(int64, int64) int); ok {
r0 = rf(nodeID, collectionID)
} else {
r0 = ret.Get(0).(int)
}
return r0
}
// MockScheduler_GetSegmentTaskDelta_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSegmentTaskDelta'
type MockScheduler_GetSegmentTaskDelta_Call struct {
*mock.Call
}
// GetSegmentTaskDelta is a helper method to define mock.On call
// - nodeID int64
// - collectionID int64
func (_e *MockScheduler_Expecter) GetSegmentTaskDelta(nodeID interface{}, collectionID interface{}) *MockScheduler_GetSegmentTaskDelta_Call {
return &MockScheduler_GetSegmentTaskDelta_Call{Call: _e.mock.On("GetSegmentTaskDelta", nodeID, collectionID)}
}
func (_c *MockScheduler_GetSegmentTaskDelta_Call) Run(run func(nodeID int64, collectionID int64)) *MockScheduler_GetSegmentTaskDelta_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64))
})
return _c
}
func (_c *MockScheduler_GetSegmentTaskDelta_Call) Return(_a0 int) *MockScheduler_GetSegmentTaskDelta_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockScheduler_GetSegmentTaskDelta_Call) RunAndReturn(run func(int64, int64) int) *MockScheduler_GetSegmentTaskDelta_Call {
_c.Call.Return(run)
return _c
}
// GetSegmentTaskNum provides a mock function with given fields:
func (_m *MockScheduler) GetSegmentTaskNum() int {
ret := _m.Called()
var r0 int
if rf, ok := ret.Get(0).(func() int); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(int)
}
return r0
}
// MockScheduler_GetSegmentTaskNum_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSegmentTaskNum'
type MockScheduler_GetSegmentTaskNum_Call struct {
*mock.Call
}
// GetSegmentTaskNum is a helper method to define mock.On call
func (_e *MockScheduler_Expecter) GetSegmentTaskNum() *MockScheduler_GetSegmentTaskNum_Call {
return &MockScheduler_GetSegmentTaskNum_Call{Call: _e.mock.On("GetSegmentTaskNum")}
}
func (_c *MockScheduler_GetSegmentTaskNum_Call) Run(run func()) *MockScheduler_GetSegmentTaskNum_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockScheduler_GetSegmentTaskNum_Call) Return(_a0 int) *MockScheduler_GetSegmentTaskNum_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockScheduler_GetSegmentTaskNum_Call) RunAndReturn(run func() int) *MockScheduler_GetSegmentTaskNum_Call {
_c.Call.Return(run)
return _c
}
// RemoveByNode provides a mock function with given fields: node
func (_m *MockScheduler) RemoveByNode(node int64) {
_m.Called(node)
}
// MockScheduler_RemoveByNode_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveByNode'
type MockScheduler_RemoveByNode_Call struct {
*mock.Call
}
// RemoveByNode is a helper method to define mock.On call
// - node int64
func (_e *MockScheduler_Expecter) RemoveByNode(node interface{}) *MockScheduler_RemoveByNode_Call {
return &MockScheduler_RemoveByNode_Call{Call: _e.mock.On("RemoveByNode", node)}
}
func (_c *MockScheduler_RemoveByNode_Call) Run(run func(node int64)) *MockScheduler_RemoveByNode_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64))
})
return _c
}
func (_c *MockScheduler_RemoveByNode_Call) Return() *MockScheduler_RemoveByNode_Call {
_c.Call.Return()
return _c
}
func (_c *MockScheduler_RemoveByNode_Call) RunAndReturn(run func(int64)) *MockScheduler_RemoveByNode_Call {
_c.Call.Return(run)
return _c
}
// RemoveExecutor provides a mock function with given fields: nodeID
func (_m *MockScheduler) RemoveExecutor(nodeID int64) {
_m.Called(nodeID)
}
// MockScheduler_RemoveExecutor_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveExecutor'
type MockScheduler_RemoveExecutor_Call struct {
*mock.Call
}
// RemoveExecutor is a helper method to define mock.On call
// - nodeID int64
func (_e *MockScheduler_Expecter) RemoveExecutor(nodeID interface{}) *MockScheduler_RemoveExecutor_Call {
return &MockScheduler_RemoveExecutor_Call{Call: _e.mock.On("RemoveExecutor", nodeID)}
}
func (_c *MockScheduler_RemoveExecutor_Call) Run(run func(nodeID int64)) *MockScheduler_RemoveExecutor_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64))
})
return _c
}
func (_c *MockScheduler_RemoveExecutor_Call) Return() *MockScheduler_RemoveExecutor_Call {
_c.Call.Return()
return _c
}
func (_c *MockScheduler_RemoveExecutor_Call) RunAndReturn(run func(int64)) *MockScheduler_RemoveExecutor_Call {
_c.Call.Return(run)
return _c
}
// Start provides a mock function with given fields:
func (_m *MockScheduler) Start() {
_m.Called()
}
// MockScheduler_Start_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Start'
type MockScheduler_Start_Call struct {
*mock.Call
}
// Start is a helper method to define mock.On call
func (_e *MockScheduler_Expecter) Start() *MockScheduler_Start_Call {
return &MockScheduler_Start_Call{Call: _e.mock.On("Start")}
}
func (_c *MockScheduler_Start_Call) Run(run func()) *MockScheduler_Start_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockScheduler_Start_Call) Return() *MockScheduler_Start_Call {
_c.Call.Return()
return _c
}
func (_c *MockScheduler_Start_Call) RunAndReturn(run func()) *MockScheduler_Start_Call {
_c.Call.Return(run)
return _c
}
// Stop provides a mock function with given fields:
func (_m *MockScheduler) Stop() {
_m.Called()
}
// MockScheduler_Stop_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Stop'
type MockScheduler_Stop_Call struct {
*mock.Call
}
// Stop is a helper method to define mock.On call
func (_e *MockScheduler_Expecter) Stop() *MockScheduler_Stop_Call {
return &MockScheduler_Stop_Call{Call: _e.mock.On("Stop")}
}
func (_c *MockScheduler_Stop_Call) Run(run func()) *MockScheduler_Stop_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockScheduler_Stop_Call) Return() *MockScheduler_Stop_Call {
_c.Call.Return()
return _c
}
func (_c *MockScheduler_Stop_Call) RunAndReturn(run func()) *MockScheduler_Stop_Call {
_c.Call.Return(run)
return _c
}
// NewMockScheduler creates a new instance of MockScheduler. 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 NewMockScheduler(t interface {
mock.TestingT
Cleanup(func())
}) *MockScheduler {
mock := &MockScheduler{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}