milvus-io/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-10-12 02:19:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

enhance: extract proxy task model and scheduler (#53023)

Browse Source 
Related to #44761

Split the remaining "task & scheduler" item of issue #44761 into two new
sub-packages, completing the proxy component extraction plan.

- internal/proxy/taskmodel: the shared task contract (Task, DMLTask,
BaseTask, Condition/TaskCondition, TsoAllocator) plus the UniqueID,
Timestamp, VChan, PChan, PChanStatistics value types and the
SearchTaskName/QueryTaskName constants. This breaks cycle B, the
TASKS<->SCHEDULER dependency.

- internal/proxy/scheduler: the four task queues (DdTaskQueue,
DmTaskQueue, DqTaskQueue, DcTaskQueue) and the TaskScheduler loops,
previously task_scheduler.go. It depends only on taskmodel and pkg/v3.

The ~30 concrete task structs stay in the proxy root package and keep
their original names via a taskmodel_alias.go shim, so the diff is a
mechanical rename of the exported surface (getMetaCache->GetMetaCache,
setChannels/getChannels->SetChannels/GetChannels, Params.ProxyCfg->
paramtable.Get()) with no behavior change. Consumers (proxy.go, impl.go,
snapshot_impl.go, metrics_info.go, highlighter.go, task_delete.go) now
hold scheduler.TaskScheduler / scheduler.DmTaskQueue.

Signed-off-by: Congqi Xia <[email protected]>
This commit is contained in:
congqixia
2026-09-04 10:11:30 +08:00
committed by GitHub
parent 53a71ab6df
commit b889bb3c19
64 changed files with 2655 additions and 1260 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/proxy/alias_rbac_test.go
+5 -5
View File
@@ -345,7 +345,7 @@ func TestCreateAliasTask_ResolvesCollectionAlias(t *testing.T) {
task := &CreateAliasTask{
baseTask: baseTask{
metaCache: cache,
MetaCache: cache,
},
Condition: NewTaskCondition(ctx),
CreateAliasRequest: &milvuspb.CreateAliasRequest{
@@ -379,7 +379,7 @@ func TestAlterAliasTask_ResolvesCollectionAlias(t *testing.T) {
task := &AlterAliasTask{
baseTask: baseTask{
metaCache: cache,
MetaCache: cache,
},
Condition: NewTaskCondition(ctx),
AlterAliasRequest: &milvuspb.AlterAliasRequest{
@@ -416,7 +416,7 @@ func TestCreateAliasTask_ResolvesEvenWhenRBACFlagDisabled(t *testing.T) {
task := &CreateAliasTask{
baseTask: baseTask{
metaCache: cache,
MetaCache: cache,
},
Condition: NewTaskCondition(ctx),
CreateAliasRequest: &milvuspb.CreateAliasRequest{
@@ -450,7 +450,7 @@ func TestListAliasesTask_ResolvesCollectionAlias(t *testing.T) {
task := &ListAliasesTask{
baseTask: baseTask{
metaCache: cache,
MetaCache: cache,
},
Condition: NewTaskCondition(ctx),
ListAliasesRequest: &milvuspb.ListAliasesRequest{
@@ -478,7 +478,7 @@ func TestListAliasesTask_NoResolveWhenCollectionNameEmpty(t *testing.T) {
task := &ListAliasesTask{
baseTask: baseTask{
metaCache: cache,
MetaCache: cache,
},
Condition: NewTaskCondition(ctx),
ListAliasesRequest: &milvuspb.ListAliasesRequest{
internal/proxy/authentication_interceptor.go
+2 -2
View File
@@ -85,7 +85,7 @@ func GrpcAuthInterceptor(authFunc grpc_auth.AuthFunc) grpc.UnaryServerIntercepto
// as a readiness gate: until the proxy has published its meta cache (in
// Proxy.Init), all requests are rejected with ServiceUnavailable, mirroring the
// previous globalMetaCache == nil check that this PR's per-proxy cache removed.
func AuthenticationInterceptorWithMetaCache(getMetaCache func() Cache) grpc_auth.AuthFunc {
func AuthenticationInterceptorWithMetaCache(GetMetaCache func() Cache) grpc_auth.AuthFunc {
return func(ctx context.Context) (context.Context, error) {
// The keys within metadata.MD are normalized to lowercase.
// See: https://godoc.org/google.golang.org/grpc/metadata#New
@@ -93,7 +93,7 @@ func AuthenticationInterceptorWithMetaCache(getMetaCache func() Cache) grpc_auth
if !ok {
return nil, merr.WrapErrIoKeyNotFound("metadata", "auth check failure, due to occurs inner error: missing metadata")
}
if getMetaCache() == nil {
if GetMetaCache() == nil {
return nil, merr.WrapErrServiceUnavailable("internal: Milvus Proxy is not ready yet. please wait")
}
// check rpc call from sdk
internal/proxy/function_task.go
+1 -1
View File
@@ -102,7 +102,7 @@ func (t *alterCollectionFunctionTask) PreExecute(ctx context.Context) error {
if t.FunctionName != t.FunctionSchema.Name {
return merr.WrapErrParameterInvalidMsg("invalid function config, name not match")
}
coll, err := getCollectionInfo(ctx, t.getMetaCache(), t.GetDbName(), t.GetCollectionName())
coll, err := getCollectionInfo(ctx, t.GetMetaCache(), t.GetDbName(), t.GetCollectionName())
if err != nil {
mlog.Error(t.ctx, "AddCollectionTask, get collection info failed",
mlog.String("dbName", t.GetDbName()),
internal/proxy/function_task_test.go
+8 -8
View File
@@ -135,7 +135,7 @@ func (f *FunctionTaskSuite) TestAlterCollectionFunctionTaskPreExecute() {
}
cache := NewMockCache(f.T())
cache.EXPECT().GetCollectionID(ctx, req.DbName, req.CollectionName).Return(0, fmt.Errorf("Mock Error")).Maybe()
task.metaCache = cache
task.MetaCache = cache
err := task.PreExecute(ctx)
f.ErrorContains(err, "Mock Error")
@@ -162,7 +162,7 @@ func (f *FunctionTaskSuite) TestAlterCollectionFunctionTaskPreExecute() {
cache := NewMockCache(f.T())
cache.EXPECT().GetCollectionID(ctx, req.DbName, req.CollectionName).Return(int64(1), nil).Maybe()
cache.EXPECT().GetCollectionInfo(ctx, req.DbName, req.CollectionName, int64(1)).Return(nil, fmt.Errorf("Mock info error")).Maybe()
task.metaCache = cache
task.MetaCache = cache
err := task.PreExecute(ctx)
f.ErrorContains(err, "Mock info error")
@@ -193,7 +193,7 @@ func (f *FunctionTaskSuite) TestAlterCollectionFunctionTaskPreExecute() {
cache := NewMockCache(f.T())
cache.EXPECT().GetCollectionID(ctx, req.DbName, req.CollectionName).Return(int64(1), nil).Maybe()
cache.EXPECT().GetCollectionInfo(ctx, req.DbName, req.CollectionName, int64(1)).Return(nil, nil).Maybe()
task.metaCache = cache
task.MetaCache = cache
err := task.PreExecute(ctx)
f.ErrorContains(err, "not support alter BM25")
}
@@ -233,7 +233,7 @@ func (f *FunctionTaskSuite) TestAlterCollectionFunctionTaskPreExecute() {
}
cache.EXPECT().GetCollectionInfo(ctx, req.DbName, req.CollectionName, int64(1)).Return(coll, nil).Maybe()
task.metaCache = cache
task.MetaCache = cache
err := task.PreExecute(ctx)
f.ErrorContains(err, "not support alter BM25")
}
@@ -263,7 +263,7 @@ func (f *FunctionTaskSuite) TestAlterCollectionFunctionTaskPreExecute() {
cache := NewMockCache(f.T())
cache.EXPECT().GetCollectionID(ctx, req.DbName, req.CollectionName).Return(int64(1), nil).Maybe()
cache.EXPECT().GetCollectionInfo(ctx, req.DbName, req.CollectionName, int64(1)).Return(nil, nil).Maybe()
task.metaCache = cache
task.MetaCache = cache
err := task.PreExecute(ctx)
f.ErrorContains(err, "invalid function config, name not match")
}
@@ -306,7 +306,7 @@ func (f *FunctionTaskSuite) TestAlterCollectionFunctionTaskPreExecute() {
cache := NewMockCache(f.T())
cache.EXPECT().GetCollectionID(ctx, req.DbName, req.CollectionName).Return(int64(1), nil).Maybe()
cache.EXPECT().GetCollectionInfo(ctx, req.DbName, req.CollectionName, int64(1)).Return(coll, nil).Maybe()
task.metaCache = cache
task.MetaCache = cache
err := task.PreExecute(ctx)
f.ErrorContains(err, externalCollectionFunctionMutationUnsupportedMsg)
@@ -350,7 +350,7 @@ func (f *FunctionTaskSuite) TestAlterCollectionFunctionTaskPreExecute() {
cache.EXPECT().GetCollectionInfo(ctx, req.DbName, req.CollectionName, int64(1)).Return(coll, nil).Maybe()
m := mockey.Mock(validator.ValidateFunction).Return(nil).Build()
defer m.UnPatch()
task.metaCache = cache
task.MetaCache = cache
err := task.PreExecute(ctx)
f.NoError(err)
}
@@ -384,7 +384,7 @@ func (f *FunctionTaskSuite) TestAlterCollectionFunctionTaskPreExecute() {
cache := NewMockCache(f.T())
cache.EXPECT().GetCollectionID(ctx, req.DbName, req.CollectionName).Return(int64(1), nil).Maybe()
cache.EXPECT().GetCollectionInfo(ctx, req.DbName, req.CollectionName, int64(1)).Return(coll, nil).Maybe()
task.metaCache = cache
task.MetaCache = cache
err := task.PreExecute(ctx)
f.ErrorContains(err, "output fields cannot be altered")
}
internal/proxy/highlighter.go
+3 -2
View File
@@ -13,6 +13,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/proxy/scheduler"
"github.com/milvus-io/milvus/internal/proxy/shardclient"
"github.com/milvus-io/milvus/internal/util/function/highlight"
"github.com/milvus-io/milvus/internal/util/function/models"
@@ -320,7 +321,7 @@ type lexicalHighlightOperator struct {
tasks []*highlightTask
schema *schemaInfo
lbPolicy shardclient.LBPolicy
scheduler *taskScheduler
scheduler *scheduler.TaskScheduler
collectionName string
collectionID int64
@@ -466,7 +467,7 @@ func (op *lexicalHighlightOperator) run(ctx context.Context, span trace.Span, in
collectionID: op.collectionID,
dbName: op.dbName,
}
if err := op.scheduler.dqQueue.Enqueue(task); err != nil {
if err := op.scheduler.DqQueue.Enqueue(task); err != nil {
return nil, err
}
internal/proxy/impl.go
+137 -137
View File
File diff suppressed because it is too large Load Diff
internal/proxy/impl_test.go
+30 -28
View File
@@ -46,7 +46,9 @@ import (
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/mocks/distributed/mock_streaming"
"github.com/milvus-io/milvus/internal/proxy/channelmgr"
"github.com/milvus-io/milvus/internal/proxy/scheduler"
"github.com/milvus-io/milvus/internal/proxy/shardclient"
"github.com/milvus-io/milvus/internal/proxy/taskmodel"
"github.com/milvus-io/milvus/internal/util/dependency"
"github.com/milvus-io/milvus/internal/util/segcore"
"github.com/milvus-io/milvus/internal/util/sessionutil"
@@ -460,7 +462,7 @@ func TestProxyRenameCollection(t *testing.T) {
func TestProxyFunctionEdit(t *testing.T) {
mockey.PatchConvey("TestProxy_AddFunction", t, func() {
m1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
m1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
return nil
}).Build()
m2 := mockey.Mock((*TaskCondition).WaitToFinish).Return(nil).Build()
@@ -487,7 +489,7 @@ func TestProxyFunctionEdit(t *testing.T) {
})
mockey.PatchConvey("TestProxy_DropFunction", t, func() {
m1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
m1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
return nil
}).Build()
m2 := mockey.Mock((*TaskCondition).WaitToFinish).Return(nil).Build()
@@ -508,7 +510,7 @@ func TestProxyFunctionEdit(t *testing.T) {
})
mockey.PatchConvey("TestProxy_AlterFunction", t, func() {
m1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
m1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
return nil
}).Build()
m2 := mockey.Mock((*TaskCondition).WaitToFinish).Return(nil).Build()
@@ -553,7 +555,7 @@ func TestProxy_ResourceGroup(t *testing.T) {
qc.EXPECT().ShowLoadCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe()
tsoAllocatorIns := newMockTsoAllocator()
node.sched, err = newTaskScheduler(node.ctx, tsoAllocatorIns)
node.sched, err = scheduler.NewTaskScheduler(node.ctx, tsoAllocatorIns)
assert.NoError(t, err)
node.sched.Start()
defer node.sched.Close()
@@ -637,7 +639,7 @@ func TestProxy_InvalidResourceGroupName(t *testing.T) {
Status: merr.Success(),
}, nil).Maybe()
tsoAllocatorIns := newMockTsoAllocator()
node.sched, err = newTaskScheduler(node.ctx, tsoAllocatorIns)
node.sched, err = scheduler.NewTaskScheduler(node.ctx, tsoAllocatorIns)
assert.NoError(t, err)
node.sched.Start()
defer node.sched.Close()
@@ -696,7 +698,7 @@ func createTestProxy() *Proxy {
tso: newMockTimestampAllocatorInterface(),
}
node.sched, _ = newTaskScheduler(ctx, node.tsoAllocator)
node.sched, _ = scheduler.NewTaskScheduler(ctx, node.tsoAllocator)
node.sched.Start()
return node
@@ -1131,8 +1133,8 @@ func TestProxyCreateDatabase(t *testing.T) {
}
node.simpleLimiter = NewSimpleLimiter(0, 0)
node.UpdateStateCode(commonpb.StateCode_Healthy)
node.sched, err = newTaskScheduler(ctx, node.tsoAllocator)
node.sched.ddQueue.setMaxTaskNum(10)
node.sched, err = scheduler.NewTaskScheduler(ctx, node.tsoAllocator)
node.sched.DdQueue.SetMaxTaskNum(10)
assert.NoError(t, err)
err = node.sched.Start()
assert.NoError(t, err)
@@ -1186,8 +1188,8 @@ func TestProxyDropDatabase(t *testing.T) {
}
node.simpleLimiter = NewSimpleLimiter(0, 0)
node.UpdateStateCode(commonpb.StateCode_Healthy)
node.sched, err = newTaskScheduler(ctx, node.tsoAllocator)
node.sched.ddQueue.setMaxTaskNum(10)
node.sched, err = scheduler.NewTaskScheduler(ctx, node.tsoAllocator)
node.sched.DdQueue.SetMaxTaskNum(10)
assert.NoError(t, err)
err = node.sched.Start()
assert.NoError(t, err)
@@ -1248,8 +1250,8 @@ func TestProxyListDatabase(t *testing.T) {
}
node.simpleLimiter = NewSimpleLimiter(0, 0)
node.UpdateStateCode(commonpb.StateCode_Healthy)
node.sched, err = newTaskScheduler(ctx, node.tsoAllocator)
node.sched.ddQueue.setMaxTaskNum(10)
node.sched, err = scheduler.NewTaskScheduler(ctx, node.tsoAllocator)
node.sched.DdQueue.SetMaxTaskNum(10)
assert.NoError(t, err)
err = node.sched.Start()
assert.NoError(t, err)
@@ -1304,8 +1306,8 @@ func TestProxyAlterDatabase(t *testing.T) {
}
node.simpleLimiter = NewSimpleLimiter(0, 0)
node.UpdateStateCode(commonpb.StateCode_Healthy)
node.sched, err = newTaskScheduler(ctx, node.tsoAllocator)
node.sched.ddQueue.setMaxTaskNum(10)
node.sched, err = scheduler.NewTaskScheduler(ctx, node.tsoAllocator)
node.sched.DdQueue.SetMaxTaskNum(10)
assert.NoError(t, err)
err = node.sched.Start()
assert.NoError(t, err)
@@ -1357,8 +1359,8 @@ func TestProxyDescribeDatabase(t *testing.T) {
}
node.simpleLimiter = NewSimpleLimiter(0, 0)
node.UpdateStateCode(commonpb.StateCode_Healthy)
node.sched, err = newTaskScheduler(ctx, node.tsoAllocator)
node.sched.ddQueue.setMaxTaskNum(10)
node.sched, err = scheduler.NewTaskScheduler(ctx, node.tsoAllocator)
node.sched.DdQueue.SetMaxTaskNum(10)
assert.NoError(t, err)
err = node.sched.Start()
assert.NoError(t, err)
@@ -1610,7 +1612,7 @@ func TestProxy_Delete(t *testing.T) {
idAllocator, err := allocator.NewIDAllocator(ctx, rc, 0)
assert.NoError(t, err)
queue, err := newTaskScheduler(ctx, tsoAllocator)
queue, err := scheduler.NewTaskScheduler(ctx, tsoAllocator)
assert.NoError(t, err)
node := &Proxy{metaCache: cache, chMgr: chMgr, rowIDAllocator: idAllocator, sched: queue}
@@ -1641,7 +1643,7 @@ func TestProxy_ImportV2(t *testing.T) {
node.tsoAllocator = &timestampAllocator{
tso: newMockTimestampAllocatorInterface(),
}
scheduler, err := newTaskScheduler(ctx, node.tsoAllocator)
scheduler, err := scheduler.NewTaskScheduler(ctx, node.tsoAllocator)
assert.NoError(t, err)
node.sched = scheduler
err = node.sched.Start()
@@ -1929,7 +1931,7 @@ func TestRunAnalyzer(t *testing.T) {
p := &Proxy{}
tsoAllocatorIns := newMockTsoAllocator()
sched, err := newTaskScheduler(ctx, tsoAllocatorIns)
sched, err := scheduler.NewTaskScheduler(ctx, tsoAllocatorIns)
require.NoError(t, err)
sched.Start()
defer sched.Close()
@@ -3119,7 +3121,7 @@ func TestProxy_AddCollectionField_ExternalCollection(t *testing.T) {
}, nil).Build()
defer mockDescribe.UnPatch()
mockEnqueue := mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, queued task) error {
mockEnqueue := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(_ *scheduler.DdTaskQueue, queued taskmodel.Task) error {
_ = queued.OnEnqueue()
addTask := queued.(*addCollectionFieldTask)
assert.Equal(t, externalSchema, addTask.oldSchema)
@@ -3222,7 +3224,7 @@ func TestProxy_AddCollectionField_TextValidation(t *testing.T) {
}, nil).Build()
defer mockDescribe.UnPatch()
mockEnqueue := mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, queued task) error {
mockEnqueue := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(_ *scheduler.DdTaskQueue, queued taskmodel.Task) error {
require.NoError(t, queued.OnEnqueue())
addTask := queued.(*addCollectionFieldTask)
err := addTask.PreExecute(context.Background())
@@ -3269,7 +3271,7 @@ func TestProxy_AddCollectionField_DoesNotBlockOnSchemaVersion(t *testing.T) {
require.FailNow(t, "AddCollectionField should not query collection statistics")
return nil, errors.New("unexpected GetCollectionStatistics call")
}).Build()
mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(_ *scheduler.DdTaskQueue, t taskmodel.Task) error {
_ = t.OnEnqueue()
addTask := t.(*addCollectionFieldTask)
addTask.result = merr.Success()
@@ -3298,7 +3300,7 @@ func TestProxy_AddCollectionField_DoesNotBlockOnSchemaVersion(t *testing.T) {
mockey.Mock((*Proxy).GetCollectionStatistics).Return(&milvuspb.GetCollectionStatisticsResponse{
Status: merr.Success(),
}, nil).Build()
mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(_ *scheduler.DdTaskQueue, t taskmodel.Task) error {
_ = t.OnEnqueue()
addTask := t.(*addCollectionFieldTask)
addTask.result = merr.Success()
@@ -3620,7 +3622,7 @@ func TestProxy_AlterCollectionSchema(t *testing.T) {
require.FailNow(t, "AlterCollectionSchema should not query collection statistics")
return nil, errors.New("unexpected GetCollectionStatistics call")
}).Build()
mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(_ *scheduler.DdTaskQueue, t taskmodel.Task) error {
_ = t.OnEnqueue()
alterTask := t.(*alterCollectionSchemaTask)
alterTask.AlterCollectionSchemaResponse = &milvuspb.AlterCollectionSchemaResponse{AlterStatus: merr.Success()}
@@ -3649,7 +3651,7 @@ func TestProxy_AlterCollectionSchema(t *testing.T) {
mockey.Mock((*Proxy).GetCollectionStatistics).Return(&milvuspb.GetCollectionStatisticsResponse{
Status: merr.Success(),
}, nil).Build()
mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(_ *scheduler.DdTaskQueue, t taskmodel.Task) error {
_ = t.OnEnqueue()
alterTask := t.(*alterCollectionSchemaTask)
alterTask.AlterCollectionSchemaResponse = &milvuspb.AlterCollectionSchemaResponse{AlterStatus: merr.Success()}
@@ -3701,7 +3703,7 @@ func TestProxy_AlterCollectionSchema(t *testing.T) {
Schema: &schemapb.CollectionSchema{Name: "test_coll"},
}, nil).Build()
mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, _ task) error {
mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(_ *scheduler.DdTaskQueue, _ taskmodel.Task) error {
return errors.New("queue full")
}).Build()
@@ -3724,7 +3726,7 @@ func TestProxy_AlterCollectionSchema(t *testing.T) {
}, nil).Build()
// Call OnEnqueue so task.Base is initialized (BeginTs/EndTs are logged after Enqueue).
mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(_ *scheduler.DdTaskQueue, t taskmodel.Task) error {
_ = t.OnEnqueue()
return nil
}).Build()
@@ -3750,7 +3752,7 @@ func TestProxy_AlterCollectionSchema(t *testing.T) {
}, nil).Build()
// Call OnEnqueue so task.Base is initialized (BeginTs/EndTs are logged after Enqueue).
mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(_ *scheduler.DdTaskQueue, t taskmodel.Task) error {
_ = t.OnEnqueue()
return nil
}).Build()
internal/proxy/interface_def.go
-6
View File
@@ -24,12 +24,6 @@ import (
"github.com/milvus-io/milvus/pkg/v3/proto/rootcoordpb"
)
// use interface tsoAllocator to keep other components testable
// include: baseTaskQueue, taskScheduler
type tsoAllocator interface {
AllocOne(ctx context.Context) (Timestamp, error)
}
// use timestampAllocatorInterface to keep other components testable
//
//go:generate mockery --name=timestampAllocatorInterface --filename=mock_tso_test.go --outpkg=proxy --output=. --inpackage --structname=mockTimestampAllocator --with-expecter
internal/proxy/membership_filter_plan_size_test.go
+1 -1
View File
@@ -370,7 +370,7 @@ func TestQueryTaskMembershipFilterPlanSizeLimit(t *testing.T) {
Return(schema, nil).Maybe()
task := &queryTask{
baseTask: baseTask{
metaCache: cache,
MetaCache: cache,
},
ctx: ctx,
RetrieveRequest: &internalpb.RetrieveRequest{
internal/proxy/metrics_info.go
+1 -1
View File
@@ -76,7 +76,7 @@ func getQuotaMetrics(node *Proxy) (*metricsinfo.ProxyQuotaMetrics, error) {
return &metricsinfo.ProxyQuotaMetrics{
Hms: metricsinfo.HardwareMetrics{},
Rms: rms,
QueueMetrics: node.sched.getMetrics(),
QueueMetrics: node.sched.GetMetrics(),
}, nil
}
internal/proxy/metrics_info_test.go
+4 -95
View File
@@ -19,11 +19,11 @@ package proxy
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus/internal/proxy/scheduler"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
@@ -42,107 +42,16 @@ func TestProxy_metrics(t *testing.T) {
assert.NotNil(t, resp)
}
func createMockDDQueue() *ddTaskQueue {
q := &ddTaskQueue{
baseTaskQueue: newBaseTaskQueue(nil),
}
q.unissuedTasks.PushBack(&createCollectionTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
})
q.unissuedTasks.PushBack(&createPartitionTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
})
q.activeTasks[1] = &loadCollectionTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
}
return q
}
func createMockDMQueue() *dmTaskQueue {
q := &dmTaskQueue{
baseTaskQueue: newBaseTaskQueue(nil),
}
q.unissuedTasks.PushBack(&insertTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
})
q.unissuedTasks.PushBack(&upsertTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
})
q.activeTasks[1] = &deleteTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
}
return q
}
func createMockDQQueue() *dqTaskQueue {
q := &dqTaskQueue{
baseTaskQueue: newBaseTaskQueue(nil),
}
q.unissuedTasks.PushBack(&queryTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
})
q.unissuedTasks.PushBack(&searchTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
})
q.activeTasks[1] = &queryTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
}
return q
}
func createMockDCQueue() *ddTaskQueue {
q := &ddTaskQueue{
baseTaskQueue: newBaseTaskQueue(nil),
}
q.unissuedTasks.PushBack(&flushTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
})
q.unissuedTasks.PushBack(&flushTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
})
q.activeTasks[1] = &flushTask{
baseTask: baseTask{
onEnqueueTime: time.Now(),
},
}
return q
}
func getMockProxyRequestMetrics() *Proxy {
mixc := NewMixCoordMock()
defer mixc.Close()
sched, _ := scheduler.NewTaskScheduler(context.Background(), &mockTsoAllocator{})
proxy := &Proxy{
mixCoord: mixc,
session: &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{Address: funcutil.GenRandomStr()}},
sched: &taskScheduler{
ddQueue: createMockDDQueue(),
dmQueue: createMockDMQueue(),
dqQueue: createMockDQQueue(),
dcQueue: createMockDCQueue(),
},
sched: sched,
}
mixc.getMetricsFunc = func(ctx context.Context, request *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
internal/proxy/mock_test.go
-141
View File
@@ -23,13 +23,11 @@ import (
"google.golang.org/grpc"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/allocator"
"github.com/milvus-io/milvus/pkg/v3/mq/common"
"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
"github.com/milvus-io/milvus/pkg/v3/proto/rootcoordpb"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/testutils"
@@ -95,145 +93,6 @@ func newMockIDAllocatorInterface() allocator.Interface {
return &mockIDAllocatorInterface{}
}
type mockTask struct {
baseTask
*TaskCondition
id UniqueID
name string
tType commonpb.MsgType
ts Timestamp
}
func (m *mockTask) CanSkipAllocTimestamp() bool {
return false
}
func (m *mockTask) TraceCtx() context.Context {
return m.ctx
}
func (m *mockTask) ID() UniqueID {
return m.id
}
func (m *mockTask) SetID(uid UniqueID) {
m.id = uid
}
func (m *mockTask) Name() string {
return m.name
}
func (m *mockTask) Type() commonpb.MsgType {
return m.tType
}
func (m *mockTask) BeginTs() Timestamp {
return m.ts
}
func (m *mockTask) EndTs() Timestamp {
return m.ts
}
func (m *mockTask) SetTs(ts Timestamp) {
m.ts = ts
}
func (m *mockTask) OnEnqueue() error {
return nil
}
func (m *mockTask) PreExecute(ctx context.Context) error {
return nil
}
func (m *mockTask) Execute(ctx context.Context) error {
return nil
}
func (m *mockTask) PostExecute(ctx context.Context) error {
return nil
}
func newMockTask(ctx context.Context) *mockTask {
return &mockTask{
TaskCondition: NewTaskCondition(ctx),
id: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
name: funcutil.GenRandomStr(),
tType: commonpb.MsgType_Undefined,
ts: Timestamp(time.Now().Nanosecond()),
}
}
func newDefaultMockTask() *mockTask {
return newMockTask(context.Background())
}
type mockDdlTask struct {
*mockTask
}
func newMockDdlTask(ctx context.Context) *mockDdlTask {
return &mockDdlTask{
mockTask: newMockTask(ctx),
}
}
func newDefaultMockDdlTask() *mockDdlTask {
return newMockDdlTask(context.Background())
}
type mockDmlTask struct {
*mockTask
vchans []vChan
pchans []pChan
}
func (m *mockDmlTask) setChannels() error {
return nil
}
func (m *mockDmlTask) getChannels() []vChan {
return m.vchans
}
func newMockDmlTask(ctx context.Context) *mockDmlTask {
shardNum := 2
vchans := make([]vChan, 0, shardNum)
pchans := make([]pChan, 0, shardNum)
for i := 0; i < shardNum; i++ {
vchans = append(vchans, funcutil.GenRandomStr())
pchans = append(pchans, funcutil.GenRandomStr())
}
return &mockDmlTask{
mockTask: newMockTask(ctx),
vchans: vchans,
pchans: pchans,
}
}
func newDefaultMockDmlTask() *mockDmlTask {
return newMockDmlTask(context.Background())
}
type mockDqlTask struct {
*mockTask
}
func newMockDqlTask(ctx context.Context) *mockDqlTask {
return &mockDqlTask{
mockTask: newMockTask(ctx),
}
}
func newDefaultMockDqlTask() *mockDqlTask {
return newMockDqlTask(context.Background())
}
type simpleMockMsgStream struct {
msgChan chan *msgstream.ConsumeMsgPack
internal/proxy/privilege_interceptor.go
+3 -3
View File
@@ -60,7 +60,7 @@ func PrivilegeInterceptor(ctx context.Context, req interface{}) (context.Context
return PrivilegeInterceptorWithMetaCache(func() Cache { return nil })(ctx, req)
}
func PrivilegeInterceptorWithMetaCache(getMetaCache func() Cache) PrivilegeFunc {
func PrivilegeInterceptorWithMetaCache(GetMetaCache func() Cache) PrivilegeFunc {
return func(ctx context.Context, req interface{}) (context.Context, error) {
if !Params.CommonCfg.AuthorizationEnabled.GetAsBool() {
return ctx, nil
@@ -129,7 +129,7 @@ func PrivilegeInterceptorWithMetaCache(getMetaCache func() Cache) PrivilegeFunc
// Resolve alias to actual collection name for RBAC checks
if Params.ProxyCfg.ResolveAliasForPrivilege.GetAsBool() && objectType == commonpb.ObjectType_Collection.String() && objectNameIndex != 0 {
if objectName != util.AnyWord && objectName != "" {
if actualCollectionName, resolveErr := resolveCollectionAlias(ctx, getMetaCache(), dbName, objectName); resolveErr != nil {
if actualCollectionName, resolveErr := resolveCollectionAlias(ctx, GetMetaCache(), dbName, objectName); resolveErr != nil {
mlog.RatedWarn(ctx, rate.Limit(60), "failed to resolve collection alias for RBAC, using original name",
mlog.String("objectName", objectName), mlog.FieldDbName(dbName), mlog.Err(resolveErr))
} else {
@@ -157,7 +157,7 @@ func PrivilegeInterceptorWithMetaCache(getMetaCache func() Cache) PrivilegeFunc
resolvedNames = append(resolvedNames, name)
continue
}
if actualName, resolveErr := resolveCollectionAlias(ctx, getMetaCache(), dbName, name); resolveErr != nil {
if actualName, resolveErr := resolveCollectionAlias(ctx, GetMetaCache(), dbName, name); resolveErr != nil {
mlog.RatedWarn(ctx, rate.Limit(60), "failed to resolve collection alias for RBAC, using original name",
mlog.String("objectName", name), mlog.FieldDbName(dbName), mlog.Err(resolveErr))
resolvedNames = append(resolvedNames, name)
internal/proxy/proxy.go
+5 -4
View File
@@ -31,6 +31,7 @@ import (
"github.com/milvus-io/milvus/internal/allocator"
"github.com/milvus-io/milvus/internal/proxy/channelmgr"
"github.com/milvus-io/milvus/internal/proxy/connection"
"github.com/milvus-io/milvus/internal/proxy/scheduler"
"github.com/milvus-io/milvus/internal/proxy/shardclient"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/util/dependency"
@@ -89,7 +90,7 @@ type Proxy struct {
metaCache Cache
chMgr channelmgr.ChannelsMgr
sched *taskScheduler
sched *scheduler.TaskScheduler
rowIDAllocator *allocator.IDAllocator
tsoAllocator *timestampAllocator
@@ -176,7 +177,7 @@ func (node *Proxy) GetMetaCache() Cache {
// TooManyRequests. The REST layer probes it (via interface assertion, like
// GetMetaCache) to reject search/query before paying for body decoding.
func (node *Proxy) IsDQLQueueFull() bool {
return node.sched != nil && node.sched.dqQueue.isFull()
return node.sched != nil && node.sched.DqQueue.IsFull()
}
// Register registers proxy at etcd
@@ -287,7 +288,7 @@ func (node *Proxy) Init() error {
node.chMgr = chMgr
mlog.Debug(node.ctx, "create channels manager done", mlog.String("role", typeutil.ProxyRole))
node.sched, err = newTaskScheduler(node.ctx, node.tsoAllocator)
node.sched, err = scheduler.NewTaskScheduler(node.ctx, node.tsoAllocator)
if err != nil {
mlog.Warn(node.ctx, "failed to create task scheduler", mlog.String("role", typeutil.ProxyRole), mlog.Err(err))
return err
@@ -384,7 +385,7 @@ func (node *Proxy) Stop() error {
node.resourceManager.Close()
}
if metaCache := node.getMetaCache(); metaCache != nil {
if metaCache := node.GetMetaCache(); metaCache != nil {
metaCache.Close()
}
internal/proxy/proxy_test.go
+100 -25
View File
@@ -51,7 +51,9 @@ import (
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proxy/channelmgr"
"github.com/milvus-io/milvus/internal/proxy/privilege"
"github.com/milvus-io/milvus/internal/proxy/scheduler"
"github.com/milvus-io/milvus/internal/proxy/shardclient"
"github.com/milvus-io/milvus/internal/proxy/taskmodel"
"github.com/milvus-io/milvus/internal/util/componentutil"
"github.com/milvus-io/milvus/internal/util/dependency"
"github.com/milvus-io/milvus/internal/util/fileresource"
@@ -1203,7 +1205,7 @@ func TestProxy(t *testing.T) {
CollectionName: "alias",
})
_, err = proxy.getMetaCache().GetCollectionID(ctx, dbName, "alias")
_, err = proxy.GetMetaCache().GetCollectionID(ctx, dbName, "alias")
assert.Error(t, err)
})
@@ -1231,7 +1233,7 @@ func TestProxy(t *testing.T) {
})
t.Run("describe collection", func(t *testing.T) {
collectionID, err := proxy.getMetaCache().GetCollectionID(ctx, dbName, collectionName)
collectionID, err := proxy.GetMetaCache().GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
resp, err := proxy.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{
@@ -1415,7 +1417,7 @@ func TestProxy(t *testing.T) {
})
t.Run("show partitions", func(t *testing.T) {
collectionID, err := proxy.getMetaCache().GetCollectionID(ctx, dbName, collectionName)
collectionID, err := proxy.GetMetaCache().GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
resp, err := proxy.ShowPartitions(ctx, &milvuspb.ShowPartitionsRequest{
@@ -1983,7 +1985,7 @@ func TestProxy(t *testing.T) {
})
t.Run("get replicas", func(t *testing.T) {
collectionID, err := proxy.getMetaCache().GetCollectionID(ctx, dbName, collectionName)
collectionID, err := proxy.GetMetaCache().GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
resp, err := proxy.GetReplicas(ctx, &milvuspb.GetReplicasRequest{
@@ -2232,7 +2234,7 @@ func TestProxy(t *testing.T) {
})
t.Run("release collection", func(t *testing.T) {
_, err := proxy.getMetaCache().GetCollectionID(ctx, dbName, collectionName)
_, err := proxy.GetMetaCache().GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
resp, err := proxy.ReleaseCollection(ctx, &milvuspb.ReleaseCollectionRequest{
@@ -2260,7 +2262,7 @@ func TestProxy(t *testing.T) {
pLoaded := true
t.Run("load partitions", func(t *testing.T) {
collectionID, err := proxy.getMetaCache().GetCollectionID(ctx, dbName, collectionName)
collectionID, err := proxy.GetMetaCache().GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
resp, err := proxy.LoadPartitions(ctx, &milvuspb.LoadPartitionsRequest{
@@ -2310,7 +2312,7 @@ func TestProxy(t *testing.T) {
assert.True(t, pLoaded)
t.Run("show in-memory partitions", func(t *testing.T) {
collectionID, err := proxy.getMetaCache().GetCollectionID(ctx, dbName, collectionName)
collectionID, err := proxy.GetMetaCache().GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
resp, err := proxy.ShowPartitions(ctx, &milvuspb.ShowPartitionsRequest{
@@ -2557,7 +2559,7 @@ func TestProxy(t *testing.T) {
})
t.Run("show in-memory partitions after release partition", func(t *testing.T) {
collectionID, err := proxy.getMetaCache().GetCollectionID(ctx, dbName, collectionName)
collectionID, err := proxy.GetMetaCache().GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
resp, err := proxy.ShowPartitions(ctx, &milvuspb.ShowPartitionsRequest{
@@ -2607,7 +2609,7 @@ func TestProxy(t *testing.T) {
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
_, err = proxy.getMetaCache().GetPartitionID(ctx, dbName, collectionName, partitionName)
_, err = proxy.GetMetaCache().GetPartitionID(ctx, dbName, collectionName, partitionName)
assert.Error(t, err)
// drop non-exist partition -> fail
@@ -2646,7 +2648,7 @@ func TestProxy(t *testing.T) {
})
t.Run("show partitions after drop partition", func(t *testing.T) {
collectionID, err := proxy.getMetaCache().GetCollectionID(ctx, dbName, collectionName)
collectionID, err := proxy.GetMetaCache().GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
resp, err := proxy.ShowPartitions(ctx, &milvuspb.ShowPartitionsRequest{
@@ -2687,7 +2689,7 @@ func TestProxy(t *testing.T) {
})
t.Run("truncate collection", func(t *testing.T) {
_, err := proxy.getMetaCache().GetCollectionID(ctx, dbName, collectionName)
_, err := proxy.GetMetaCache().GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
resp, err := proxy.TruncateCollection(ctx, &milvuspb.TruncateCollectionRequest{
@@ -2710,7 +2712,7 @@ func TestProxy(t *testing.T) {
})
t.Run("drop collection", func(t *testing.T) {
_, err := proxy.getMetaCache().GetCollectionID(ctx, dbName, collectionName)
_, err := proxy.GetMetaCache().GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
resp, err := proxy.DropCollection(ctx, &milvuspb.DropCollectionRequest{
@@ -2732,7 +2734,7 @@ func TestProxy(t *testing.T) {
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
_, err = proxy.getMetaCache().GetCollectionID(ctx, dbName, collectionName)
_, err = proxy.GetMetaCache().GetCollectionID(ctx, dbName, collectionName)
assert.Error(t, err)
resp, err = proxy.InvalidateCollectionMetaCache(ctx, &proxypb.InvalidateCollMetaCacheRequest{
@@ -2744,7 +2746,7 @@ func TestProxy(t *testing.T) {
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
hasDatabase := proxy.getMetaCache().HasDatabase(ctx, dbName)
hasDatabase := proxy.GetMetaCache().HasDatabase(ctx, dbName)
assert.False(t, hasDatabase)
})
@@ -3218,8 +3220,8 @@ func TestProxy(t *testing.T) {
// queue full
ddParallel := proxy.sched.ddQueue.getMaxTaskNum()
proxy.sched.ddQueue.setMaxTaskNum(0)
ddParallel := proxy.sched.DdQueue.GetMaxTaskNum()
proxy.sched.DdQueue.SetMaxTaskNum(0)
t.Run("CreateCollection fail, dd queue full", func(t *testing.T) {
resp, err := proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{})
@@ -3364,7 +3366,7 @@ func TestProxy(t *testing.T) {
})
t.Run("Flush fail, dc queue full", func(t *testing.T) {
proxy.sched.dcQueue.setMaxTaskNum(0)
proxy.sched.DcQueue.SetMaxTaskNum(0)
resp, err := proxy.Flush(ctx, &milvuspb.FlushRequest{})
assert.NoError(t, err)
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
@@ -3400,10 +3402,10 @@ func TestProxy(t *testing.T) {
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
})
proxy.sched.ddQueue.setMaxTaskNum(ddParallel)
proxy.sched.DdQueue.SetMaxTaskNum(ddParallel)
dmParallelism := proxy.sched.dmQueue.getMaxTaskNum()
proxy.sched.dmQueue.setMaxTaskNum(0)
dmParallelism := proxy.sched.DmQueue.GetMaxTaskNum()
proxy.sched.DmQueue.SetMaxTaskNum(0)
t.Run("Insert fail, dm queue full", func(t *testing.T) {
resp, err := proxy.Insert(ctx, &milvuspb.InsertRequest{})
@@ -3423,10 +3425,10 @@ func TestProxy(t *testing.T) {
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
})
proxy.sched.dmQueue.setMaxTaskNum(dmParallelism)
proxy.sched.DmQueue.SetMaxTaskNum(dmParallelism)
dqParallelism := proxy.sched.dqQueue.getMaxTaskNum()
proxy.sched.dqQueue.setMaxTaskNum(0)
dqParallelism := proxy.sched.DqQueue.GetMaxTaskNum()
proxy.sched.DqQueue.SetMaxTaskNum(0)
t.Run("Search fail, dq queue full", func(t *testing.T) {
resp, err := proxy.Search(ctx, &milvuspb.SearchRequest{})
@@ -3440,7 +3442,7 @@ func TestProxy(t *testing.T) {
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
})
proxy.sched.dqQueue.setMaxTaskNum(dqParallelism)
proxy.sched.DqQueue.SetMaxTaskNum(dqParallelism)
// timeout
@@ -4522,7 +4524,7 @@ func TestProxy_Import(t *testing.T) {
proxy.tsoAllocator = &timestampAllocator{
tso: newMockTimestampAllocatorInterface(),
}
scheduler, err := newTaskScheduler(ctx, proxy.tsoAllocator)
scheduler, err := scheduler.NewTaskScheduler(ctx, proxy.tsoAllocator)
assert.NoError(t, err)
proxy.sched = scheduler
err = proxy.sched.Start()
@@ -4831,3 +4833,76 @@ func (c CheckExtension) Report(info any) int {
c.reportChecker(info)
return 0
}
type dqlFullMockTask struct {
taskmodel.BaseTask
*taskmodel.TaskCondition
id taskmodel.UniqueID
ts taskmodel.Timestamp
name string
}
func (t *dqlFullMockTask) TraceCtx() context.Context {
return context.Background()
}
func (t *dqlFullMockTask) ID() taskmodel.UniqueID {
return t.id
}
func (t *dqlFullMockTask) SetID(uid taskmodel.UniqueID) {
t.id = uid
}
func (t *dqlFullMockTask) Name() string {
return t.name
}
func (t *dqlFullMockTask) Type() commonpb.MsgType {
return commonpb.MsgType_Retrieve
}
func (t *dqlFullMockTask) BeginTs() taskmodel.Timestamp {
return t.ts
}
func (t *dqlFullMockTask) EndTs() taskmodel.Timestamp {
return t.ts
}
func (t *dqlFullMockTask) SetTs(ts taskmodel.Timestamp) {
t.ts = ts
}
func (t *dqlFullMockTask) OnEnqueue() error {
return nil
}
func (t *dqlFullMockTask) PreExecute(ctx context.Context) error {
return nil
}
func (t *dqlFullMockTask) Execute(ctx context.Context) error {
return nil
}
func (t *dqlFullMockTask) PostExecute(ctx context.Context) error {
return nil
}
func TestProxy_IsDQLQueueFull(t *testing.T) {
node := &Proxy{}
assert.False(t, node.IsDQLQueueFull())
sched, err := scheduler.NewTaskScheduler(context.Background(), newMockTsoAllocator())
assert.NoError(t, err)
node.sched = sched
sched.DqQueue.SetMaxTaskNum(1)
assert.False(t, node.IsDQLQueueFull())
assert.NoError(t, sched.DqQueue.Enqueue(&dqlFullMockTask{
BaseTask: taskmodel.BaseTask{},
TaskCondition: taskmodel.NewTaskCondition(context.Background()),
name: "query",
}))
assert.True(t, node.IsDQLQueueFull())
}
internal/proxy/rate_limit_interceptor.go
+3 -3
View File
@@ -41,15 +41,15 @@ func RateLimitInterceptor(limiter types.Limiter) grpc.UnaryServerInterceptor {
}
// RateLimitInterceptorWithMetaCache returns a new unary server interceptor that performs
// request rate limiting. getMetaCache is resolved per request so the interceptor observes
// request rate limiting. GetMetaCache is resolved per request so the interceptor observes
// the meta cache initialized by the proxy after startup instead of a construction-time snapshot.
func RateLimitInterceptorWithMetaCache(getMetaCache func() Cache, limiter types.Limiter) grpc.UnaryServerInterceptor {
func RateLimitInterceptorWithMetaCache(GetMetaCache func() Cache, limiter types.Limiter) grpc.UnaryServerInterceptor {
return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
request, ok := req.(proto.Message)
if !ok {
return nil, merr.WrapErrParameterInvalidMsg("wrong req format when check limiter")
}
dbID, collectionIDToPartIDs, rt, n, err := GetRequestInfo(ctx, getMetaCache(), request)
dbID, collectionIDToPartIDs, rt, n, err := GetRequestInfo(ctx, GetMetaCache(), request)
if err != nil {
mlog.Warn(context.TODO(), "failed to get request info", mlog.Err(err))
return handler(ctx, req)
internal/proxy/scheduler/OWNERS
+7
View File
@@ -0,0 +1,7 @@
reviewers:
- czs007
- bigsheeper
- congqixia
approvers:
- maintainers
internal/proxy/scheduler/README.md
+80
View File
@@ -0,0 +1,80 @@
# Scheduler Package
The `scheduler` package owns the proxy's task queues and the scheduling loops
that drive task execution. It was extracted verbatim from the proxy root
package's `task_scheduler.go` (issue #44761) and now depends only on the shared
task model (`taskmodel`) plus `pkg/v3`.
## Overview
The proxy has four task queues, each backed by its own scheduler loop:
- **DdQueue** (definition) — DDL tasks such as create/drop collection, alias,
index, database, resource-group operations.
- **DmQueue** (manipulation) — DML tasks: insert/delete/upsert.
- **DqQueue** (query) — DQL tasks: search/query.
- **DcQueue** (control) — data-control operations such as flush.
Each task is enqueued by the proxy's RPC handlers and, when picked up by its
loop, runs `PreExecute -> Execute -> PostExecute` under a bounded worker pool.
## Responsibilities
1. **Queues** (`DdTaskQueue`, `DmTaskQueue`, `DqTaskQueue`, `DcTaskQueue`) —
each maintains an unissued list and an active map, plus the enqueue/dequeue
and task-lookup primitives.
2. **`TaskScheduler`** — owns the four queues and their loops (`definitionLoop`,
`controlLoop`, `manipulationLoop`, `queryLoop`), and exposes `Start`/`Close`.
3. **TSO + ID allocation** — `Enqueue` allocates a timestamp (or an ID from the
meta cache for tasks that skip timestamp allocation) before a task is
unissued.
4. **DML channel statistics** — `DmTaskQueue` tracks per-physical-channel
min/max timestamps for DML tasks (`commitPChanStats`/`popPChanStats`).
5. **Metrics** — `GetMetrics` reports per-queue pending/executing task counts
and timing, consumed by the proxy's quota/system-info metrics.
## Architecture
```
┌───────────────────────────────────────────────────────────────┐
│ TaskScheduler │
│ │
│ DdQueue ──► definitionLoop ──► processTask(Pre/Exec/Post) │
│ DmQueue ──► manipulationLoop ──► processTask │
│ DqQueue ──► queryLoop ──► processTask │
│ DcQueue ──► controlLoop ──► processTask │
│ │
│ queues hold: unissued list · active map · TSO allocator │
└───────────────────────────────────────────────────────────────┘
```
### Key types
```go
func NewTaskScheduler(ctx context.Context, tsoAllocator taskmodel.TsoAllocator,
opts ...SchedOpt) (*TaskScheduler, error)
func (s *TaskScheduler) Start() error
func (s *TaskScheduler) Close()
func (s *TaskScheduler) GetMetrics() []metricsinfo.TaskQueueMetrics
func (s *TaskScheduler) ClearDQLQueue(taskType string, reason string) ClearTaskQueueResult
```
The `TaskScheduler` exposes its queues as fields (`DdQueue`, `DmQueue`,
`DqQueue`, `DcQueue`) so the proxy's RPC handlers can call `Enqueue` directly.
## Dependency rule
`scheduler` imports `taskmodel` (for `Task`/`DMLTask`/`TsoAllocator` and the
channel/timestamp types) and `pkg/v3` (`conc`, `metrics`, `metricsinfo`, `merr`,
`paramtable`, `tsoutil`, `typeutil`). It has no `internal/*` imports and never
imports the proxy root package, so the one-way `proxy -> scheduler -> taskmodel`
edge stays acyclic.
## Related components
- **taskmodel** (`internal/proxy/taskmodel/`): the interfaces this package
schedules against.
- **proxy root** (`internal/proxy/`): constructs the scheduler in `Proxy.Init`
and enqueues concrete tasks from `impl.go` / `snapshot_impl.go`; consumes
`GetMetrics` from `metrics_info.go`.
internal/proxy/scheduler/mock_cache_test.go
+1137
View File
File diff suppressed because it is too large Load Diff
internal/proxy/scheduler/mock_task_test.go
+193
View File
@@ -0,0 +1,193 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package scheduler
import (
"context"
"sync"
"time"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus/internal/proxy/metacache"
"github.com/milvus-io/milvus/internal/proxy/taskmodel"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/uniquegenerator"
)
type mockTsoAllocator struct {
mu sync.Mutex
logicPart uint32
}
func (tso *mockTsoAllocator) AllocOne(ctx context.Context) (taskmodel.Timestamp, error) {
tso.mu.Lock()
defer tso.mu.Unlock()
tso.logicPart++
physical := uint64(time.Now().UnixMilli())
return (physical << 18) + uint64(tso.logicPart), nil
}
func newMockTsoAllocator() taskmodel.TsoAllocator {
return &mockTsoAllocator{}
}
type mockTask struct {
taskmodel.BaseTask
*taskmodel.TaskCondition
id taskmodel.UniqueID
name string
tType commonpb.MsgType
ts taskmodel.Timestamp
skipAllocTS bool
}
func (m *mockTask) CanSkipAllocTimestamp() bool {
return m.skipAllocTS
}
func (m *mockTask) TraceCtx() context.Context {
return m.Ctx()
}
func (m *mockTask) ID() taskmodel.UniqueID {
return m.id
}
func (m *mockTask) SetID(uid taskmodel.UniqueID) {
m.id = uid
}
func (m *mockTask) Name() string {
return m.name
}
func (m *mockTask) Type() commonpb.MsgType {
return m.tType
}
func (m *mockTask) BeginTs() taskmodel.Timestamp {
return m.ts
}
func (m *mockTask) EndTs() taskmodel.Timestamp {
return m.ts
}
func (m *mockTask) SetTs(ts taskmodel.Timestamp) {
m.ts = ts
}
func (m *mockTask) OnEnqueue() error {
return nil
}
func (m *mockTask) PreExecute(ctx context.Context) error {
return nil
}
func (m *mockTask) Execute(ctx context.Context) error {
return nil
}
func (m *mockTask) PostExecute(ctx context.Context) error {
return nil
}
func newMockTask(ctx context.Context) *mockTask {
return &mockTask{
TaskCondition: taskmodel.NewTaskCondition(ctx),
id: taskmodel.UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
name: funcutil.GenRandomStr(),
tType: commonpb.MsgType_Undefined,
ts: taskmodel.Timestamp(time.Now().Nanosecond()),
}
}
func newDefaultMockTask() *mockTask {
return newMockTask(context.Background())
}
func newSkipAllocMockTask(metaCache metacache.Cache) *mockTask {
t := newMockTask(context.Background())
t.skipAllocTS = true
t.MetaCache = metaCache
return t
}
type mockDdlTask struct {
*mockTask
}
func newMockDdlTask(ctx context.Context) *mockDdlTask {
return &mockDdlTask{
mockTask: newMockTask(ctx),
}
}
func newDefaultMockDdlTask() *mockDdlTask {
return newMockDdlTask(context.Background())
}
type mockDmlTask struct {
*mockTask
vchans []taskmodel.VChan
pchans []taskmodel.PChan
}
func (m *mockDmlTask) SetChannels() error {
return nil
}
func (m *mockDmlTask) GetChannels() []taskmodel.PChan {
return m.pchans
}
func newMockDmlTask(ctx context.Context) *mockDmlTask {
shardNum := 2
vchans := make([]taskmodel.VChan, 0, shardNum)
pchans := make([]taskmodel.PChan, 0, shardNum)
for i := 0; i < shardNum; i++ {
vchans = append(vchans, funcutil.GenRandomStr())
pchans = append(pchans, funcutil.GenRandomStr())
}
return &mockDmlTask{
mockTask: newMockTask(ctx),
vchans: vchans,
pchans: pchans,
}
}
func newDefaultMockDmlTask() *mockDmlTask {
return newMockDmlTask(context.Background())
}
type mockDqlTask struct {
*mockTask
}
func newMockDqlTask(ctx context.Context) *mockDqlTask {
return &mockDqlTask{
mockTask: newMockTask(ctx),
}
}
func newDefaultMockDqlTask() *mockDqlTask {
return newMockDqlTask(context.Background())
}
internal/proxy/task_scheduler.go → internal/proxy/scheduler/task_scheduler.go
+169 -168
View File
@@ -14,7 +14,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package proxy
package scheduler
import (
"container/list"
@@ -28,6 +28,7 @@ import (
"github.com/cockroachdb/errors"
"go.opentelemetry.io/otel"
"github.com/milvus-io/milvus/internal/proxy/taskmodel"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/util/conc"
@@ -38,28 +39,28 @@ import (
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
type taskQueue interface {
type TaskQueue interface {
utChan() <-chan int
utEmpty() bool
utFull() bool
addUnissuedTask(t task) error
FrontUnissuedTask() task
PopUnissuedTask() task
AddActiveTask(t task)
PopActiveTask(taskID UniqueID) task
getTaskByReqID(reqID UniqueID) task
Enqueue(t task) error
setMaxTaskNum(num int64)
getMaxTaskNum() int64
addUnissuedTask(t taskmodel.Task) error
FrontUnissuedTask() taskmodel.Task
PopUnissuedTask() taskmodel.Task
AddActiveTask(t taskmodel.Task)
PopActiveTask(taskID taskmodel.UniqueID) taskmodel.Task
getTaskByReqID(reqID taskmodel.UniqueID) taskmodel.Task
Enqueue(t taskmodel.Task) error
SetMaxTaskNum(num int64)
GetMaxTaskNum() int64
}
// make sure baseTaskQueue implements taskQueue.
var _ taskQueue = (*baseTaskQueue)(nil)
// make sure BaseTaskQueue implements TaskQueue.
var _ TaskQueue = (*BaseTaskQueue)(nil)
// baseTaskQueue implements taskQueue.
type baseTaskQueue struct {
// BaseTaskQueue implements TaskQueue.
type BaseTaskQueue struct {
unissuedTasks *list.List
activeTasks map[UniqueID]task
activeTasks map[taskmodel.UniqueID]taskmodel.Task
utLock sync.RWMutex
atLock sync.RWMutex
@@ -69,37 +70,37 @@ type baseTaskQueue struct {
utBufChan chan int // to block scheduler
tsoAllocatorIns tsoAllocator
tsoAllocatorIns taskmodel.TsoAllocator
}
func (queue *baseTaskQueue) utChan() <-chan int {
func (queue *BaseTaskQueue) utChan() <-chan int {
return queue.utBufChan
}
func (queue *baseTaskQueue) utEmpty() bool {
func (queue *BaseTaskQueue) utEmpty() bool {
queue.utLock.RLock()
defer queue.utLock.RUnlock()
return queue.unissuedTasks.Len() == 0
}
func (queue *baseTaskQueue) utFull() bool {
return int64(queue.unissuedTasks.Len()) >= queue.getMaxTaskNum()
func (queue *BaseTaskQueue) utFull() bool {
return int64(queue.unissuedTasks.Len()) >= queue.GetMaxTaskNum()
}
// isFull is the lock-acquiring counterpart of utFull; utFull assumes the
// IsFull is the lock-acquiring counterpart of utFull; utFull assumes the
// caller already holds utLock.
func (queue *baseTaskQueue) isFull() bool {
func (queue *BaseTaskQueue) IsFull() bool {
queue.utLock.RLock()
defer queue.utLock.RUnlock()
return queue.utFull()
}
func (queue *baseTaskQueue) addUnissuedTask(t task) error {
func (queue *BaseTaskQueue) addUnissuedTask(t taskmodel.Task) error {
queue.utLock.Lock()
defer queue.utLock.Unlock()
if queue.utFull() {
return merr.WrapErrTooManyRequests(int32(queue.getMaxTaskNum()))
return merr.WrapErrTooManyRequests(int32(queue.GetMaxTaskNum()))
}
queue.unissuedTasks.PushBack(t)
// utBufChan is an edge-triggered, capacity-1 notifier: a pending token
@@ -112,7 +113,7 @@ func (queue *baseTaskQueue) addUnissuedTask(t task) error {
return nil
}
func (queue *baseTaskQueue) FrontUnissuedTask() task {
func (queue *BaseTaskQueue) FrontUnissuedTask() taskmodel.Task {
queue.utLock.RLock()
defer queue.utLock.RUnlock()
@@ -120,10 +121,10 @@ func (queue *baseTaskQueue) FrontUnissuedTask() task {
return nil
}
return queue.unissuedTasks.Front().Value.(task)
return queue.unissuedTasks.Front().Value.(taskmodel.Task)
}
func (queue *baseTaskQueue) PopUnissuedTask() task {
func (queue *BaseTaskQueue) PopUnissuedTask() taskmodel.Task {
queue.utLock.Lock()
defer queue.utLock.Unlock()
@@ -134,17 +135,17 @@ func (queue *baseTaskQueue) PopUnissuedTask() task {
ft := queue.unissuedTasks.Front()
queue.unissuedTasks.Remove(ft)
return ft.Value.(task)
return ft.Value.(taskmodel.Task)
}
func (queue *baseTaskQueue) popUnissuedTasks(filter func(task) bool) []task {
func (queue *BaseTaskQueue) popUnissuedTasks(filter func(taskmodel.Task) bool) []taskmodel.Task {
queue.utLock.Lock()
defer queue.utLock.Unlock()
removed := make([]task, 0)
removed := make([]taskmodel.Task, 0)
for e := queue.unissuedTasks.Front(); e != nil; {
next := e.Next()
t := e.Value.(task)
t := e.Value.(taskmodel.Task)
if filter == nil || filter(t) {
queue.unissuedTasks.Remove(e)
removed = append(removed, t)
@@ -154,7 +155,7 @@ func (queue *baseTaskQueue) popUnissuedTasks(filter func(task) bool) []task {
return removed
}
func (queue *baseTaskQueue) AddActiveTask(t task) {
func (queue *BaseTaskQueue) AddActiveTask(t taskmodel.Task) {
queue.atLock.Lock()
defer queue.atLock.Unlock()
tID := t.ID()
@@ -167,7 +168,7 @@ func (queue *baseTaskQueue) AddActiveTask(t task) {
t.SetExecutingTime()
}
func (queue *baseTaskQueue) PopActiveTask(taskID UniqueID) task {
func (queue *BaseTaskQueue) PopActiveTask(taskID taskmodel.UniqueID) taskmodel.Task {
queue.atLock.Lock()
defer queue.atLock.Unlock()
t, ok := queue.activeTasks[taskID]
@@ -179,12 +180,12 @@ func (queue *baseTaskQueue) PopActiveTask(taskID UniqueID) task {
return t
}
func (queue *baseTaskQueue) getTaskByReqID(reqID UniqueID) task {
func (queue *BaseTaskQueue) getTaskByReqID(reqID taskmodel.UniqueID) taskmodel.Task {
queue.utLock.RLock()
for e := queue.unissuedTasks.Front(); e != nil; e = e.Next() {
if e.Value.(task).ID() == reqID {
if e.Value.(taskmodel.Task).ID() == reqID {
queue.utLock.RUnlock()
return e.Value.(task)
return e.Value.(taskmodel.Task)
}
}
queue.utLock.RUnlock()
@@ -198,7 +199,7 @@ func (queue *baseTaskQueue) getTaskByReqID(reqID UniqueID) task {
return nil
}
func (queue *baseTaskQueue) Enqueue(t task) error {
func (queue *BaseTaskQueue) Enqueue(t taskmodel.Task) error {
err := t.OnEnqueue()
if err != nil {
return err
@@ -212,14 +213,14 @@ func (queue *baseTaskQueue) Enqueue(t task) error {
full := queue.utFull()
queue.utLock.RUnlock()
if full {
return merr.WrapErrTooManyRequests(int32(queue.getMaxTaskNum()))
return merr.WrapErrTooManyRequests(int32(queue.GetMaxTaskNum()))
}
var ts Timestamp
var id UniqueID
var ts taskmodel.Timestamp
var id taskmodel.UniqueID
if t.CanSkipAllocTimestamp() {
ts = tsoutil.ComposeTS(time.Now().UnixMilli(), 0)
id, err = t.getMetaCache().AllocID(t.TraceCtx())
id, err = t.GetMetaCache().AllocID(t.TraceCtx())
if err != nil {
return err
}
@@ -229,7 +230,7 @@ func (queue *baseTaskQueue) Enqueue(t task) error {
return err
}
// we always use same msg id and ts for now.
id = UniqueID(ts)
id = taskmodel.UniqueID(ts)
}
t.SetTs(ts)
t.SetID(id)
@@ -238,39 +239,39 @@ func (queue *baseTaskQueue) Enqueue(t task) error {
return queue.addUnissuedTask(t)
}
func (queue *baseTaskQueue) setMaxTaskNum(num int64) {
func (queue *BaseTaskQueue) SetMaxTaskNum(num int64) {
queue.maxTaskNumMtx.Lock()
defer queue.maxTaskNumMtx.Unlock()
queue.maxTaskNum = num
}
func (queue *baseTaskQueue) getMaxTaskNum() int64 {
func (queue *BaseTaskQueue) GetMaxTaskNum() int64 {
queue.maxTaskNumMtx.RLock()
defer queue.maxTaskNumMtx.RUnlock()
return queue.maxTaskNum
}
func newBaseTaskQueue(tsoAllocatorIns tsoAllocator) *baseTaskQueue {
return &baseTaskQueue{
func newBaseTaskQueue(tsoAllocatorIns taskmodel.TsoAllocator) *BaseTaskQueue {
return &BaseTaskQueue{
unissuedTasks: list.New(),
activeTasks: make(map[UniqueID]task),
activeTasks: make(map[taskmodel.UniqueID]taskmodel.Task),
utLock: sync.RWMutex{},
atLock: sync.RWMutex{},
maxTaskNum: Params.ProxyCfg.MaxTaskNum.GetAsInt64(),
maxTaskNum: paramtable.Get().ProxyCfg.MaxTaskNum.GetAsInt64(),
utBufChan: make(chan int, 1),
tsoAllocatorIns: tsoAllocatorIns,
}
}
// ddTaskQueue represents queue for DDL task such as createCollection/createPartition/dropCollection/dropPartition/hasCollection/hasPartition
type ddTaskQueue struct {
*baseTaskQueue
// DdTaskQueue represents queue for DDL task such as createCollection/createPartition/dropCollection/dropPartition/hasCollection/hasPartition
type DdTaskQueue struct {
*BaseTaskQueue
lock sync.Mutex
}
func (queue *ddTaskQueue) updateMetrics() {
func (queue *DdTaskQueue) updateMetrics() {
queue.utLock.RLock()
unissuedTasksNum := queue.unissuedTasks.Len()
queue.utLock.RUnlock()
@@ -283,19 +284,19 @@ func (queue *ddTaskQueue) updateMetrics() {
}
type pChanStatInfo struct {
pChanStatistics
tsSet map[Timestamp]struct{}
taskmodel.PChanStatistics
tsSet map[taskmodel.Timestamp]struct{}
}
// dmTaskQueue represents queue for DML task such as insert/delete/upsert
type dmTaskQueue struct {
*baseTaskQueue
// DmTaskQueue represents queue for DML task such as insert/delete/upsert
type DmTaskQueue struct {
*BaseTaskQueue
statsLock sync.RWMutex
pChanStatisticsInfos map[pChan]*pChanStatInfo
pChanStatisticsInfos map[taskmodel.PChan]*pChanStatInfo
}
func (queue *dmTaskQueue) updateMetrics() {
func (queue *DmTaskQueue) updateMetrics() {
queue.utLock.RLock()
unissuedTasksNum := queue.unissuedTasks.Len()
queue.utLock.RUnlock()
@@ -307,14 +308,14 @@ func (queue *dmTaskQueue) updateMetrics() {
metrics.ProxyQueueTaskNum.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), "dml", metrics.InProgressIndexTaskLabel).Set(float64(activateTaskNum))
}
func (queue *dmTaskQueue) Enqueue(t task) error {
func (queue *DmTaskQueue) Enqueue(t taskmodel.Task) error {
// This statsLock has two functions:
// 1) Protect member pChanStatisticsInfos
// 2) Serialize the timestamp allocation for dml tasks
// 1. set the current pChannels for this dmTask
dmt := t.(dmlTask)
err := dmt.setChannels()
dmt := t.(taskmodel.DMLTask)
err := dmt.SetChannels()
if err != nil {
mlog.Warn(t.TraceCtx(), "setChannels failed when Enqueue", mlog.FieldTaskID(t.ID()), mlog.Err(err))
return err
@@ -323,12 +324,12 @@ func (queue *dmTaskQueue) Enqueue(t task) error {
// 2. enqueue dml task
queue.statsLock.Lock()
defer queue.statsLock.Unlock()
err = queue.baseTaskQueue.Enqueue(t)
err = queue.BaseTaskQueue.Enqueue(t)
if err != nil {
return err
}
// 3. commit will use pChannels got previously when preAdding and will definitely succeed
pChannels := dmt.getChannels()
pChannels := dmt.GetChannels()
queue.commitPChanStats(dmt, pChannels)
// there's indeed a possibility that the collection info cache was expired after preAddPChanStats
// but considering root coord knows everything about meta modification, invalid stats appended after the meta changed
@@ -336,7 +337,7 @@ func (queue *dmTaskQueue) Enqueue(t task) error {
return nil
}
func (queue *dmTaskQueue) PopActiveTask(taskID UniqueID) task {
func (queue *DmTaskQueue) PopActiveTask(taskID taskmodel.UniqueID) taskmodel.Task {
queue.atLock.Lock()
defer queue.atLock.Unlock()
t, ok := queue.activeTasks[taskID]
@@ -345,7 +346,7 @@ func (queue *dmTaskQueue) PopActiveTask(taskID UniqueID) task {
defer queue.statsLock.Unlock()
delete(queue.activeTasks, taskID)
mlog.Debug(t.TraceCtx(), "Proxy dmTaskQueue popPChanStats", mlog.FieldTaskID(t.ID()))
mlog.Debug(t.TraceCtx(), "Proxy DmTaskQueue popPChanStats", mlog.FieldTaskID(t.ID()))
queue.popPChanStats(t)
} else {
mlog.Warn(context.TODO(), "Proxy task not in active task list!", mlog.FieldTaskID(taskID))
@@ -353,15 +354,15 @@ func (queue *dmTaskQueue) PopActiveTask(taskID UniqueID) task {
return t
}
func (queue *dmTaskQueue) commitPChanStats(dmt dmlTask, pChannels []pChan) {
func (queue *DmTaskQueue) commitPChanStats(dmt taskmodel.DMLTask, pChannels []taskmodel.PChan) {
// 1. prepare new stat for all pChannels
newStats := make(map[pChan]pChanStatistics)
newStats := make(map[taskmodel.PChan]taskmodel.PChanStatistics)
beginTs := dmt.BeginTs()
endTs := dmt.EndTs()
for _, channel := range pChannels {
newStats[channel] = pChanStatistics{
minTs: beginTs,
maxTs: endTs,
newStats[channel] = taskmodel.PChanStatistics{
MinTs: beginTs,
MaxTs: endTs,
}
}
// 2. update stats for all pChannels
@@ -369,26 +370,26 @@ func (queue *dmTaskQueue) commitPChanStats(dmt dmlTask, pChannels []pChan) {
currentStat, ok := queue.pChanStatisticsInfos[cName]
if !ok {
currentStat = &pChanStatInfo{
pChanStatistics: newStat,
tsSet: map[Timestamp]struct{}{
newStat.minTs: {},
PChanStatistics: newStat,
tsSet: map[taskmodel.Timestamp]struct{}{
newStat.MinTs: {},
},
}
queue.pChanStatisticsInfos[cName] = currentStat
} else {
if currentStat.minTs > newStat.minTs {
currentStat.minTs = newStat.minTs
if currentStat.MinTs > newStat.MinTs {
currentStat.MinTs = newStat.MinTs
}
if currentStat.maxTs < newStat.maxTs {
currentStat.maxTs = newStat.maxTs
if currentStat.MaxTs < newStat.MaxTs {
currentStat.MaxTs = newStat.MaxTs
}
currentStat.tsSet[newStat.minTs] = struct{}{}
currentStat.tsSet[newStat.MinTs] = struct{}{}
}
}
}
func (queue *dmTaskQueue) popPChanStats(t task) {
channels := t.(dmlTask).getChannels()
func (queue *DmTaskQueue) popPChanStats(t taskmodel.Task) {
channels := t.(taskmodel.DMLTask).GetChannels()
taskTs := t.BeginTs()
for _, cName := range channels {
info, ok := queue.pChanStatisticsInfos[cName]
@@ -397,48 +398,48 @@ func (queue *dmTaskQueue) popPChanStats(t task) {
if len(info.tsSet) <= 0 {
delete(queue.pChanStatisticsInfos, cName)
} else {
newMinTs := info.maxTs
newMinTs := info.MaxTs
for ts := range info.tsSet {
if newMinTs > ts {
newMinTs = ts
}
}
info.minTs = newMinTs
info.MinTs = newMinTs
}
}
}
}
func (queue *dmTaskQueue) getPChanStatsInfo() (map[pChan]*pChanStatistics, error) {
ret := make(map[pChan]*pChanStatistics)
func (queue *DmTaskQueue) getPChanStatsInfo() (map[taskmodel.PChan]*taskmodel.PChanStatistics, error) {
ret := make(map[taskmodel.PChan]*taskmodel.PChanStatistics)
queue.statsLock.RLock()
defer queue.statsLock.RUnlock()
for cName, info := range queue.pChanStatisticsInfos {
ret[cName] = &pChanStatistics{
minTs: info.minTs,
maxTs: info.maxTs,
ret[cName] = &taskmodel.PChanStatistics{
MinTs: info.MinTs,
MaxTs: info.MaxTs,
}
}
return ret, nil
}
// dqTaskQueue represents queue for DQL task such as search/query
type dqTaskQueue struct {
*baseTaskQueue
// DqTaskQueue represents queue for DQL task such as search/query
type DqTaskQueue struct {
*BaseTaskQueue
}
type clearTaskQueueResult struct {
queuedCleared int64
type ClearTaskQueueResult struct {
QueuedCleared int64
}
func isDQLTaskMatched(t task, taskType string) bool {
func isDQLTaskMatched(t taskmodel.Task, taskType string) bool {
switch taskType {
case "", "all":
return true
case "search":
return t.Name() == SearchTaskName
return t.Name() == taskmodel.SearchTaskName
case "query":
return t.Name() == QueryTaskName
return t.Name() == taskmodel.QueryTaskName
default:
return false
}
@@ -451,13 +452,13 @@ func clearTaskQueueError(reason string) error {
return errors.Wrap(context.Canceled, fmt.Sprintf("read task queue cleared by admin: %s", reason))
}
func (queue *dqTaskQueue) clearQueuedTasks(taskType string, reason string) clearTaskQueueResult {
removed := queue.popUnissuedTasks(func(t task) bool {
func (queue *DqTaskQueue) clearQueuedTasks(taskType string, reason string) ClearTaskQueueResult {
removed := queue.popUnissuedTasks(func(t taskmodel.Task) bool {
return isDQLTaskMatched(t, taskType)
})
if len(removed) == 0 {
queue.updateMetrics()
return clearTaskQueueResult{}
return ClearTaskQueueResult{}
}
clearErr := clearTaskQueueError(reason)
@@ -465,10 +466,10 @@ func (queue *dqTaskQueue) clearQueuedTasks(taskType string, reason string) clear
task.Notify(clearErr)
}
queue.updateMetrics()
return clearTaskQueueResult{queuedCleared: int64(len(removed))}
return ClearTaskQueueResult{QueuedCleared: int64(len(removed))}
}
func (queue *dqTaskQueue) updateMetrics() {
func (queue *DqTaskQueue) updateMetrics() {
queue.utLock.RLock()
unissuedTasksNum := queue.unissuedTasks.Len()
queue.utLock.RUnlock()
@@ -480,61 +481,61 @@ func (queue *dqTaskQueue) updateMetrics() {
metrics.ProxyQueueTaskNum.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), "dql", metrics.InProgressIndexTaskLabel).Set(float64(activateTaskNum))
}
func (queue *ddTaskQueue) Enqueue(t task) error {
func (queue *DdTaskQueue) Enqueue(t taskmodel.Task) error {
queue.lock.Lock()
defer queue.lock.Unlock()
return queue.baseTaskQueue.Enqueue(t)
return queue.BaseTaskQueue.Enqueue(t)
}
func newDdTaskQueue(tsoAllocatorIns tsoAllocator) *ddTaskQueue {
return &ddTaskQueue{
baseTaskQueue: newBaseTaskQueue(tsoAllocatorIns),
func newDdTaskQueue(tsoAllocatorIns taskmodel.TsoAllocator) *DdTaskQueue {
return &DdTaskQueue{
BaseTaskQueue: newBaseTaskQueue(tsoAllocatorIns),
}
}
func newDmTaskQueue(tsoAllocatorIns tsoAllocator) *dmTaskQueue {
return &dmTaskQueue{
baseTaskQueue: newBaseTaskQueue(tsoAllocatorIns),
pChanStatisticsInfos: make(map[pChan]*pChanStatInfo),
func newDmTaskQueue(tsoAllocatorIns taskmodel.TsoAllocator) *DmTaskQueue {
return &DmTaskQueue{
BaseTaskQueue: newBaseTaskQueue(tsoAllocatorIns),
pChanStatisticsInfos: make(map[taskmodel.PChan]*pChanStatInfo),
}
}
func newDqTaskQueue(tsoAllocatorIns tsoAllocator) *dqTaskQueue {
return &dqTaskQueue{
baseTaskQueue: newBaseTaskQueue(tsoAllocatorIns),
func newDqTaskQueue(tsoAllocatorIns taskmodel.TsoAllocator) *DqTaskQueue {
return &DqTaskQueue{
BaseTaskQueue: newBaseTaskQueue(tsoAllocatorIns),
}
}
// taskScheduler schedules the gRPC tasks.
type taskScheduler struct {
ddQueue *ddTaskQueue
dmQueue *dmTaskQueue
dqQueue *dqTaskQueue
// TaskScheduler schedules the gRPC tasks.
type TaskScheduler struct {
DdQueue *DdTaskQueue
DmQueue *DmTaskQueue
DqQueue *DqTaskQueue
// data control queue, use for such as flush operation, which control the data status
dcQueue *ddTaskQueue
DcQueue *DdTaskQueue
wg sync.WaitGroup
ctx context.Context
cancel context.CancelFunc
}
type schedOpt func(*taskScheduler)
type SchedOpt func(*TaskScheduler)
func newTaskScheduler(ctx context.Context,
tsoAllocatorIns tsoAllocator,
opts ...schedOpt,
) (*taskScheduler, error) {
func NewTaskScheduler(ctx context.Context,
tsoAllocatorIns taskmodel.TsoAllocator,
opts ...SchedOpt,
) (*TaskScheduler, error) {
ctx1, cancel := context.WithCancel(ctx)
s := &taskScheduler{
s := &TaskScheduler{
ctx: ctx1,
cancel: cancel,
}
s.ddQueue = newDdTaskQueue(tsoAllocatorIns)
s.dmQueue = newDmTaskQueue(tsoAllocatorIns)
s.dqQueue = newDqTaskQueue(tsoAllocatorIns)
s.DdQueue = newDdTaskQueue(tsoAllocatorIns)
s.DmQueue = newDmTaskQueue(tsoAllocatorIns)
s.DqQueue = newDqTaskQueue(tsoAllocatorIns)
s.dcQueue = newDdTaskQueue(tsoAllocatorIns)
s.DcQueue = newDdTaskQueue(tsoAllocatorIns)
for _, opt := range opts {
opt(s)
@@ -543,27 +544,27 @@ func newTaskScheduler(ctx context.Context,
return s, nil
}
func (sched *taskScheduler) scheduleDdTask() task {
return sched.ddQueue.PopUnissuedTask()
func (sched *TaskScheduler) scheduleDdTask() taskmodel.Task {
return sched.DdQueue.PopUnissuedTask()
}
func (sched *taskScheduler) scheduleDcTask() task {
return sched.dcQueue.PopUnissuedTask()
func (sched *TaskScheduler) scheduleDcTask() taskmodel.Task {
return sched.DcQueue.PopUnissuedTask()
}
func (sched *taskScheduler) scheduleDmTask() task {
return sched.dmQueue.PopUnissuedTask()
func (sched *TaskScheduler) scheduleDmTask() taskmodel.Task {
return sched.DmQueue.PopUnissuedTask()
}
func (sched *taskScheduler) scheduleDqTask() task {
return sched.dqQueue.PopUnissuedTask()
func (sched *TaskScheduler) scheduleDqTask() taskmodel.Task {
return sched.DqQueue.PopUnissuedTask()
}
func (sched *taskScheduler) clearDQLQueue(taskType string, reason string) clearTaskQueueResult {
return sched.dqQueue.clearQueuedTasks(taskType, reason)
func (sched *TaskScheduler) ClearDQLQueue(taskType string, reason string) ClearTaskQueueResult {
return sched.DqQueue.clearQueuedTasks(taskType, reason)
}
func (sched *taskScheduler) processTask(t task, q taskQueue) {
func (sched *TaskScheduler) processTask(t taskmodel.Task, q TaskQueue) {
ctx, span := otel.Tracer(typeutil.ProxyRole).Start(t.TraceCtx(), t.Name())
defer span.End()
@@ -610,7 +611,7 @@ func (sched *taskScheduler) processTask(t task, q taskQueue) {
}
// definitionLoop schedules the ddl tasks.
func (sched *taskScheduler) definitionLoop() {
func (sched *TaskScheduler) definitionLoop() {
defer sched.wg.Done()
pool := conc.NewPool[struct{}](paramtable.Get().ProxyCfg.DDLConcurrency.GetAsInt(), conc.WithExpiryDuration(time.Minute))
@@ -619,21 +620,21 @@ func (sched *taskScheduler) definitionLoop() {
select {
case <-sched.ctx.Done():
return
case <-sched.ddQueue.utChan():
case <-sched.DdQueue.utChan():
for t := sched.scheduleDdTask(); t != nil; t = sched.scheduleDdTask() {
task := t
pool.Submit(func() (struct{}, error) {
sched.processTask(task, sched.ddQueue)
sched.processTask(task, sched.DdQueue)
return struct{}{}, nil
})
}
sched.ddQueue.updateMetrics()
sched.DdQueue.updateMetrics()
}
}
}
// controlLoop schedule the data control operation, such as flush
func (sched *taskScheduler) controlLoop() {
func (sched *TaskScheduler) controlLoop() {
defer sched.wg.Done()
pool := conc.NewPool[struct{}](paramtable.Get().ProxyCfg.DCLConcurrency.GetAsInt(), conc.WithExpiryDuration(time.Minute))
@@ -642,20 +643,20 @@ func (sched *taskScheduler) controlLoop() {
select {
case <-sched.ctx.Done():
return
case <-sched.dcQueue.utChan():
case <-sched.DcQueue.utChan():
for t := sched.scheduleDcTask(); t != nil; t = sched.scheduleDcTask() {
task := t
pool.Submit(func() (struct{}, error) {
sched.processTask(task, sched.dcQueue)
sched.processTask(task, sched.DcQueue)
return struct{}{}, nil
})
}
sched.dcQueue.updateMetrics()
sched.DcQueue.updateMetrics()
}
}
}
func (sched *taskScheduler) manipulationLoop() {
func (sched *TaskScheduler) manipulationLoop() {
defer sched.wg.Done()
pool := conc.NewPool[struct{}](paramtable.Get().ProxyCfg.MaxTaskNum.GetAsInt())
defer pool.Release()
@@ -663,20 +664,20 @@ func (sched *taskScheduler) manipulationLoop() {
select {
case <-sched.ctx.Done():
return
case <-sched.dmQueue.utChan():
case <-sched.DmQueue.utChan():
for t := sched.scheduleDmTask(); t != nil; t = sched.scheduleDmTask() {
task := t
pool.Submit(func() (struct{}, error) {
sched.processTask(task, sched.dmQueue)
sched.processTask(task, sched.DmQueue)
return struct{}{}, nil
})
}
sched.dmQueue.updateMetrics()
sched.DmQueue.updateMetrics()
}
}
}
func (sched *taskScheduler) queryLoop() {
func (sched *TaskScheduler) queryLoop() {
defer sched.wg.Done()
poolSize := paramtable.Get().ProxyCfg.MaxTaskNum.GetAsInt()
@@ -689,7 +690,7 @@ func (sched *taskScheduler) queryLoop() {
select {
case <-sched.ctx.Done():
return
case <-sched.dqQueue.utChan():
case <-sched.DqQueue.utChan():
for t := sched.scheduleDqTask(); t != nil; t = sched.scheduleDqTask() {
task := t
p := pool
@@ -698,16 +699,16 @@ func (sched *taskScheduler) queryLoop() {
p = subTaskPool
}
p.Submit(func() (struct{}, error) {
sched.processTask(task, sched.dqQueue)
sched.processTask(task, sched.DqQueue)
return struct{}{}, nil
})
}
sched.dqQueue.updateMetrics()
sched.DqQueue.updateMetrics()
}
}
}
func (sched *taskScheduler) Start() error {
func (sched *TaskScheduler) Start() error {
sched.wg.Add(1)
go sched.definitionLoop()
@@ -723,16 +724,16 @@ func (sched *taskScheduler) Start() error {
return nil
}
func (sched *taskScheduler) Close() {
func (sched *TaskScheduler) Close() {
sched.cancel()
sched.wg.Wait()
}
func (sched *taskScheduler) getPChanStatistics() (map[pChan]*pChanStatistics, error) {
return sched.dmQueue.getPChanStatsInfo()
func (sched *TaskScheduler) GetPChanStatistics() (map[taskmodel.PChan]*taskmodel.PChanStatistics, error) {
return sched.DmQueue.getPChanStatsInfo()
}
func (sched *taskScheduler) getTaskQueueMetrics(queue *baseTaskQueue, queueType string) metricsinfo.TaskQueueMetrics {
func (sched *TaskScheduler) getTaskQueueMetrics(queue *BaseTaskQueue, queueType string) metricsinfo.TaskQueueMetrics {
pendingTaskStats := make(map[string]*TaskStatsTracker, 0)
executingTaskStats := make(map[string]*TaskStatsTracker, 0)
queue.atLock.RLock()
@@ -765,7 +766,7 @@ func (sched *taskScheduler) getTaskQueueMetrics(queue *baseTaskQueue, queueType
utNum := queue.unissuedTasks.Len()
for e := queue.unissuedTasks.Front(); e != nil; e = e.Next() {
task := e.Value.(task)
task := e.Value.(taskmodel.Task)
taskType := task.Name()
queueTimeMs := task.GetDurationInQueue().Milliseconds()
@@ -830,11 +831,11 @@ func (t *TaskStatsTracker) AvgQueueTime() int64 {
return t.TotalQueueTime / t.Count
}
func (sched *taskScheduler) getMetrics() []metricsinfo.TaskQueueMetrics {
dmlQueueMetrics := sched.getTaskQueueMetrics(sched.dmQueue.baseTaskQueue, "dml")
ddlQueueMetrics := sched.getTaskQueueMetrics(sched.ddQueue.baseTaskQueue, "ddl")
dqlQueueMetrics := sched.getTaskQueueMetrics(sched.dqQueue.baseTaskQueue, "dql")
dcQueueMetrics := sched.getTaskQueueMetrics(sched.dcQueue.baseTaskQueue, "dc")
func (sched *TaskScheduler) GetMetrics() []metricsinfo.TaskQueueMetrics {
dmlQueueMetrics := sched.getTaskQueueMetrics(sched.DmQueue.BaseTaskQueue, "dml")
ddlQueueMetrics := sched.getTaskQueueMetrics(sched.DdQueue.BaseTaskQueue, "ddl")
dqlQueueMetrics := sched.getTaskQueueMetrics(sched.DqQueue.BaseTaskQueue, "dql")
dcQueueMetrics := sched.getTaskQueueMetrics(sched.DcQueue.BaseTaskQueue, "dc")
return []metricsinfo.TaskQueueMetrics{
dmlQueueMetrics,
ddlQueueMetrics,
internal/proxy/task_scheduler_test.go → internal/proxy/scheduler/task_scheduler_test.go
+54 -134
View File
@@ -14,7 +14,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package proxy
package scheduler
import (
"context"
@@ -29,47 +29,29 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus/internal/proxy/channelmgr"
"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
"github.com/milvus-io/milvus/internal/proxy/taskmodel"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
)
func TestBaseTaskQueue_isFull(t *testing.T) {
queue := newBaseTaskQueue(newMockTsoAllocator())
queue.setMaxTaskNum(2)
queue.SetMaxTaskNum(2)
assert.False(t, queue.isFull())
assert.False(t, queue.IsFull())
assert.NoError(t, queue.Enqueue(newDefaultMockTask()))
assert.False(t, queue.isFull())
assert.False(t, queue.IsFull())
assert.NoError(t, queue.Enqueue(newDefaultMockTask()))
assert.True(t, queue.isFull())
assert.True(t, queue.IsFull())
queue.PopUnissuedTask()
assert.False(t, queue.isFull())
}
func TestProxy_IsDQLQueueFull(t *testing.T) {
node := &Proxy{}
assert.False(t, node.IsDQLQueueFull())
sched, err := newTaskScheduler(context.Background(), newMockTsoAllocator())
assert.NoError(t, err)
node.sched = sched
sched.dqQueue.setMaxTaskNum(1)
assert.False(t, node.IsDQLQueueFull())
assert.NoError(t, sched.dqQueue.Enqueue(newDefaultMockDqlTask()))
assert.True(t, node.IsDQLQueueFull())
assert.False(t, queue.IsFull())
}
func TestBaseTaskQueue(t *testing.T) {
var err error
var unissuedTask task
var activeTask task
var unissuedTask taskmodel.Task
var activeTask taskmodel.Task
tsoAllocatorIns := newMockTsoAllocator()
queue := newBaseTaskQueue(tsoAllocatorIns)
@@ -132,8 +114,8 @@ func TestBaseTaskQueue(t *testing.T) {
assert.NotNil(t, activeTask)
// test utFull
queue.setMaxTaskNum(10) // utBufChan is an edge-triggered notifier; capacity is tracked by utFull
for i := 0; i < int(queue.getMaxTaskNum()); i++ {
queue.SetMaxTaskNum(10) // utBufChan is an edge-triggered notifier; capacity is tracked by utFull
for i := 0; i < int(queue.GetMaxTaskNum()); i++ {
err = queue.Enqueue(newDefaultMockTask())
assert.NoError(t, err)
}
@@ -144,8 +126,8 @@ func TestBaseTaskQueue(t *testing.T) {
func TestDdTaskQueue(t *testing.T) {
var err error
var unissuedTask task
var activeTask task
var unissuedTask taskmodel.Task
var activeTask taskmodel.Task
tsoAllocatorIns := newMockTsoAllocator()
queue := newDdTaskQueue(tsoAllocatorIns)
@@ -208,8 +190,8 @@ func TestDdTaskQueue(t *testing.T) {
assert.NotNil(t, activeTask)
// test utFull
queue.setMaxTaskNum(10) // utBufChan is an edge-triggered notifier; capacity is tracked by utFull
for i := 0; i < int(queue.getMaxTaskNum()); i++ {
queue.SetMaxTaskNum(10) // utBufChan is an edge-triggered notifier; capacity is tracked by utFull
for i := 0; i < int(queue.GetMaxTaskNum()); i++ {
err = queue.Enqueue(newDefaultMockDdlTask())
assert.NoError(t, err)
}
@@ -221,8 +203,8 @@ func TestDdTaskQueue(t *testing.T) {
// test the logic of queue
func TestDmTaskQueue_Basic(t *testing.T) {
var err error
var unissuedTask task
var activeTask task
var unissuedTask taskmodel.Task
var activeTask taskmodel.Task
tsoAllocatorIns := newMockTsoAllocator()
queue := newDmTaskQueue(tsoAllocatorIns)
@@ -284,8 +266,8 @@ func TestDmTaskQueue_Basic(t *testing.T) {
assert.NotNil(t, activeTask)
// test utFull
queue.setMaxTaskNum(10) // utBufChan is an edge-triggered notifier; capacity is tracked by utFull
for i := 0; i < int(queue.getMaxTaskNum()); i++ {
queue.SetMaxTaskNum(10) // utBufChan is an edge-triggered notifier; capacity is tracked by utFull
for i := 0; i < int(queue.GetMaxTaskNum()); i++ {
err = queue.Enqueue(newDefaultMockDmlTask())
assert.NoError(t, err)
}
@@ -297,7 +279,7 @@ func TestDmTaskQueue_Basic(t *testing.T) {
// test the timestamp statistics
func TestDmTaskQueue_TimestampStatistics(t *testing.T) {
var err error
var unissuedTask task
var unissuedTask taskmodel.Task
tsoAllocatorIns := newMockTsoAllocator()
queue := newDmTaskQueue(tsoAllocatorIns)
@@ -315,8 +297,8 @@ func TestDmTaskQueue_TimestampStatistics(t *testing.T) {
unissuedTask = queue.FrontUnissuedTask()
assert.NotNil(t, unissuedTask)
for _, stat := range stats {
assert.Equal(t, unissuedTask.BeginTs(), stat.minTs)
assert.Equal(t, unissuedTask.EndTs(), stat.maxTs)
assert.Equal(t, unissuedTask.BeginTs(), stat.MinTs)
assert.Equal(t, unissuedTask.EndTs(), stat.MaxTs)
}
unissuedTask = queue.PopUnissuedTask()
@@ -355,7 +337,7 @@ func TestDmTaskQueue_TimestampStatistics2(t *testing.T) {
continue
}
utTask := queue.PopUnissuedTask()
go func(ut task) {
go func(ut taskmodel.Task) {
defer workerWg.Done()
assert.NotNil(t, ut)
queue.AddActiveTask(ut)
@@ -370,7 +352,7 @@ func TestDmTaskQueue_TimestampStatistics2(t *testing.T) {
workerWg.Wait()
}()
var currPChanStats map[pChan]*pChanStatistics
var currPChanStats map[taskmodel.PChan]*taskmodel.PChanStatistics
var wgSchedule sync.WaitGroup
scheduleCtx, scheduleCancel := context.WithCancel(context.TODO())
schedule := func() {
@@ -391,12 +373,12 @@ func TestDmTaskQueue_TimestampStatistics2(t *testing.T) {
for p, stat := range stats {
curInfo, ok := currPChanStats[p]
if ok {
fmt.Println("stat.minTs", stat.minTs, " ", "curInfo.minTs:", curInfo.minTs)
fmt.Println("stat.maxTs", stat.maxTs, " ", "curInfo.minTs:", curInfo.maxTs)
assert.True(t, stat.minTs >= curInfo.minTs)
curInfo.minTs = stat.minTs
assert.True(t, stat.maxTs >= curInfo.maxTs)
curInfo.maxTs = stat.maxTs
fmt.Println("stat.MinTs", stat.MinTs, " ", "curInfo.MinTs:", curInfo.MinTs)
fmt.Println("stat.MaxTs", stat.MaxTs, " ", "curInfo.MinTs:", curInfo.MaxTs)
assert.True(t, stat.MinTs >= curInfo.MinTs)
curInfo.MinTs = stat.MinTs
assert.True(t, stat.MaxTs >= curInfo.MaxTs)
curInfo.MaxTs = stat.MaxTs
}
}
}
@@ -443,8 +425,8 @@ func TestDmTaskQueue_TimestampStatistics2(t *testing.T) {
func TestDqTaskQueue(t *testing.T) {
var err error
var unissuedTask task
var activeTask task
var unissuedTask taskmodel.Task
var activeTask taskmodel.Task
tsoAllocatorIns := newMockTsoAllocator()
queue := newDqTaskQueue(tsoAllocatorIns)
@@ -507,8 +489,8 @@ func TestDqTaskQueue(t *testing.T) {
assert.NotNil(t, activeTask)
// test utFull
queue.setMaxTaskNum(10) // utBufChan is an edge-triggered notifier; capacity is tracked by utFull
for i := 0; i < int(queue.getMaxTaskNum()); i++ {
queue.SetMaxTaskNum(10) // utBufChan is an edge-triggered notifier; capacity is tracked by utFull
for i := 0; i < int(queue.GetMaxTaskNum()); i++ {
err = queue.Enqueue(newDefaultMockDqlTask())
assert.NoError(t, err)
}
@@ -523,7 +505,7 @@ func TestTaskScheduler(t *testing.T) {
ctx := context.Background()
tsoAllocatorIns := newMockTsoAllocator()
sched, err := newTaskScheduler(ctx, tsoAllocatorIns)
sched, err := NewTaskScheduler(ctx, tsoAllocatorIns)
assert.NoError(t, err)
assert.NotNil(t, sched)
@@ -531,7 +513,7 @@ func TestTaskScheduler(t *testing.T) {
assert.NoError(t, err)
defer sched.Close()
stats, err := sched.getPChanStatistics()
stats, err := sched.GetPChanStatistics()
assert.NoError(t, err)
assert.Equal(t, 0, len(stats))
@@ -550,7 +532,7 @@ func TestTaskScheduler(t *testing.T) {
go func() {
defer wg.Done()
err := sched.ddQueue.Enqueue(newDefaultMockDdlTask())
err := sched.DdQueue.Enqueue(newDefaultMockDdlTask())
assert.NoError(t, err)
}()
}
@@ -565,7 +547,7 @@ func TestTaskScheduler(t *testing.T) {
go func() {
defer wg.Done()
err := sched.dmQueue.Enqueue(newDefaultMockDmlTask())
err := sched.DmQueue.Enqueue(newDefaultMockDmlTask())
assert.NoError(t, err)
}()
}
@@ -580,7 +562,7 @@ func TestTaskScheduler(t *testing.T) {
go func() {
defer wg.Done()
err := sched.dqQueue.Enqueue(newDefaultMockDqlTask())
err := sched.DqQueue.Enqueue(newDefaultMockDqlTask())
assert.NoError(t, err)
}()
}
@@ -590,40 +572,18 @@ func TestTaskScheduler(t *testing.T) {
}
func TestTaskScheduler_concurrentPushAndPop(t *testing.T) {
collectionID := UniqueID(0)
collectionName := "col-0"
channels := []pChan{"mock-chan-0", "mock-chan-1"}
cache := NewMockCache(t)
cache.On("GetCollectionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(collectionID, nil)
tsoAllocatorIns := newMockTsoAllocator()
scheduler, err := newTaskScheduler(context.Background(), tsoAllocatorIns)
scheduler, err := NewTaskScheduler(context.Background(), tsoAllocatorIns)
assert.NoError(t, err)
run := func(wg *sync.WaitGroup) {
defer wg.Done()
chMgr := channelmgr.NewMockChannelsMgr(t)
chMgr.EXPECT().GetChannels(mock.Anything).Return(channels, nil)
it := &insertTask{
baseTask: baseTask{metaCache: cache},
ctx: context.Background(),
insertMsg: &msgstream.InsertMsg{
InsertRequest: &msgpb.InsertRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
},
},
chMgr: chMgr,
}
err := scheduler.dmQueue.Enqueue(it)
it := newDefaultMockDmlTask()
err := scheduler.DmQueue.Enqueue(it)
assert.NoError(t, err)
task := scheduler.scheduleDmTask()
scheduler.dmQueue.AddActiveTask(task)
chMgr.EXPECT().GetChannels(mock.Anything).Return(nil, errors.New("mock err"))
scheduler.dmQueue.PopActiveTask(task.ID()) // assert no panic
scheduler.DmQueue.AddActiveTask(task)
scheduler.DmQueue.PopActiveTask(task.ID()) // assert no panic
}
wg := &sync.WaitGroup{}
@@ -635,9 +595,6 @@ func TestTaskScheduler_concurrentPushAndPop(t *testing.T) {
}
func TestTaskScheduler_SkipAllocTimestamp(t *testing.T) {
dbName := "test_query"
collName := "test_skip_alloc_timestamp"
collID := UniqueID(111)
mockMetaCache := NewMockCache(t)
tsoAllocatorIns := newMockTsoAllocator()
@@ -647,63 +604,26 @@ func TestTaskScheduler_SkipAllocTimestamp(t *testing.T) {
assert.True(t, queue.utEmpty())
assert.False(t, queue.utFull())
mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, nil)
mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
&collectionInfo{
CollID: collID,
ConsistencyLevel: commonpb.ConsistencyLevel_Eventually,
}, nil)
mockMetaCache.EXPECT().AllocID(mock.Anything).Return(1, nil).Twice()
t.Run("query", func(t *testing.T) {
qt := &queryTask{
baseTask: baseTask{metaCache: mockMetaCache},
RetrieveRequest: &internalpb.RetrieveRequest{
QueryLabel: "query",
Base: &commonpb.MsgBase{},
},
request: &milvuspb.QueryRequest{
DbName: dbName,
CollectionName: collName,
UseDefaultConsistency: true,
},
}
qt := newSkipAllocMockTask(mockMetaCache)
qt.name = "query"
err := queue.Enqueue(qt)
assert.NoError(t, err)
})
t.Run("search", func(t *testing.T) {
st := &searchTask{
baseTask: baseTask{metaCache: mockMetaCache},
SearchRequest: &internalpb.SearchRequest{
Base: &commonpb.MsgBase{},
},
request: &milvuspb.SearchRequest{
DbName: dbName,
CollectionName: collName,
UseDefaultConsistency: true,
},
}
st := newSkipAllocMockTask(mockMetaCache)
st.name = "search"
err := queue.Enqueue(st)
assert.NoError(t, err)
})
mockMetaCache.EXPECT().AllocID(mock.Anything).Return(0, errors.New("mock error")).Once()
t.Run("failed", func(t *testing.T) {
st := &searchTask{
baseTask: baseTask{metaCache: mockMetaCache},
SearchRequest: &internalpb.SearchRequest{
Base: &commonpb.MsgBase{},
},
request: &milvuspb.SearchRequest{
DbName: dbName,
CollectionName: collName,
UseDefaultConsistency: true,
},
}
st := newSkipAllocMockTask(mockMetaCache)
st.name = "search"
err := queue.Enqueue(st)
assert.Error(t, err)
})
@@ -715,7 +635,7 @@ type blockingTsoAllocator struct {
calls atomic.Int64
}
func (b *blockingTsoAllocator) AllocOne(ctx context.Context) (Timestamp, error) {
func (b *blockingTsoAllocator) AllocOne(ctx context.Context) (taskmodel.Timestamp, error) {
b.calls.Add(1)
<-ctx.Done()
return 0, ctx.Err()
@@ -727,7 +647,7 @@ func (b *blockingTsoAllocator) AllocOne(ctx context.Context) (Timestamp, error)
func TestBaseTaskQueue_EnqueueFastFailBeforeAlloc(t *testing.T) {
tsoAllocatorIns := newMockTsoAllocator()
queue := newBaseTaskQueue(tsoAllocatorIns)
queue.setMaxTaskNum(2)
queue.SetMaxTaskNum(2)
// Fill queue to capacity with the non-blocking mock allocator.
for i := 0; i < 2; i++ {
@@ -758,7 +678,7 @@ func TestBaseTaskQueue_EnqueueFastFailBeforeAlloc(t *testing.T) {
// triggered notifier: many enqueues produce at most one pending token.
func TestBaseTaskQueue_NotifierCoalesces(t *testing.T) {
queue := newBaseTaskQueue(newMockTsoAllocator())
queue.setMaxTaskNum(100)
queue.SetMaxTaskNum(100)
for i := 0; i < 10; i++ {
assert.NoError(t, queue.Enqueue(newDefaultMockTask()))
@@ -773,7 +693,7 @@ func TestBaseTaskQueue_NotifierCoalesces(t *testing.T) {
// notifier send must not block Enqueue.
func TestBaseTaskQueue_EnqueueNotifierNonBlocking(t *testing.T) {
queue := newBaseTaskQueue(newMockTsoAllocator())
queue.setMaxTaskNum(1024)
queue.SetMaxTaskNum(1024)
done := make(chan struct{})
go func() {
internal/proxy/search_pipeline.go
+1 -1
View File
@@ -1336,7 +1336,7 @@ func (op *requeryOperator) requery(ctx context.Context, span trace.Span, ids *sc
}
qt := &queryTask{
baseTask: baseTask{
metaCache: op.node.(*Proxy).getMetaCache(),
MetaCache: op.node.(*Proxy).GetMetaCache(),
},
ctx: op.traceCtx,
Condition: NewTaskCondition(op.traceCtx),
internal/proxy/search_pipeline_test.go
+3 -2
View File
@@ -35,6 +35,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proxy/scheduler"
"github.com/milvus-io/milvus/internal/proxy/search_agg"
"github.com/milvus-io/milvus/internal/proxy/shardclient"
"github.com/milvus-io/milvus/internal/util/function/chain"
@@ -1692,7 +1693,7 @@ func (s *SearchPipelineSuite) TestHighlightOp() {
proxy.tsoAllocator = &timestampAllocator{
tso: newMockTimestampAllocatorInterface(),
}
sched, err := newTaskScheduler(ctx, proxy.tsoAllocator)
sched, err := scheduler.NewTaskScheduler(ctx, proxy.tsoAllocator)
s.Require().NoError(err)
err = sched.Start()
@@ -1815,7 +1816,7 @@ func (s *SearchPipelineSuite) TestLexicalHighlightOpNullableStringKeepsEmptyHigh
proxy.tsoAllocator = &timestampAllocator{
tso: newMockTimestampAllocatorInterface(),
}
sched, err := newTaskScheduler(ctx, proxy.tsoAllocator)
sched, err := scheduler.NewTaskScheduler(ctx, proxy.tsoAllocator)
s.Require().NoError(err)
err = sched.Start()
internal/proxy/service_provider.go
+4 -4
View File
@@ -278,7 +278,7 @@ func (node *CachedProxyServiceProvider) DescribeCollection(ctx context.Context,
// discarded the returned collectionInfo and forced a second describe on a
// rolling-upgrade old RootCoord response that omitted DbName (such entries
// are deliberately returned uncached because their database is unknown).
c, err = node.getMetaCache().GetCollectionInfo(ctx, request.DbName, "", collectionID)
c, err = node.GetMetaCache().GetCollectionInfo(ctx, request.DbName, "", collectionID)
if err != nil {
resp.Status = describeCollectionErrorStatus(err, request.DbName, collectionName)
return resp, nil
@@ -295,12 +295,12 @@ func (node *CachedProxyServiceProvider) DescribeCollection(ctx context.Context,
// Resolve the id and complete entry from the name only when the caller did
// not provide an id. The id-only path already has the complete entry above.
if !resolvedNameByID {
collectionID, err = node.getMetaCache().GetCollectionID(ctx, request.DbName, collectionName)
collectionID, err = node.GetMetaCache().GetCollectionID(ctx, request.DbName, collectionName)
if err != nil {
resp.Status = describeCollectionErrorStatus(err, request.DbName, collectionName)
return resp, nil
}
c, err = node.getMetaCache().GetCollectionInfo(ctx, request.DbName, collectionName, collectionID)
c, err = node.GetMetaCache().GetCollectionInfo(ctx, request.DbName, collectionName, collectionID)
if err != nil {
resp.Status = describeCollectionErrorStatus(err, request.DbName, collectionName)
return resp, nil
@@ -372,7 +372,7 @@ func (node *RemoteProxyServiceProvider) DescribeCollection(ctx context.Context,
log.Debug(ctx, "DescribeCollection received")
if err := node.sched.ddQueue.Enqueue(dct); err != nil {
if err := node.sched.DdQueue.Enqueue(dct); err != nil {
log.Warn(ctx, "DescribeCollection failed to enqueue",
mlog.Err(err))
internal/proxy/snapshot_impl.go
+18 -18
View File
@@ -49,14 +49,14 @@ func (node *Proxy) CreateSnapshot(ctx context.Context, req *milvuspb.CreateSnaps
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
t := &createSnapshotTask{
baseTask: baseTask{metaCache: node.getMetaCache()},
baseTask: baseTask{MetaCache: node.GetMetaCache()},
req: req,
ctx: ctx,
Condition: NewTaskCondition(ctx),
mixCoord: node.mixCoord,
}
err := node.sched.ddQueue.Enqueue(t)
err := node.sched.DdQueue.Enqueue(t)
if err != nil {
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
log.Warn(ctx, "CreateSnapshot failed to Enqueue",
@@ -91,14 +91,14 @@ func (node *Proxy) DropSnapshot(ctx context.Context, req *milvuspb.DropSnapshotR
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
t := &dropSnapshotTask{
baseTask: baseTask{metaCache: node.getMetaCache()},
baseTask: baseTask{MetaCache: node.GetMetaCache()},
req: req,
ctx: ctx,
Condition: NewTaskCondition(ctx),
mixCoord: node.mixCoord,
}
err := node.sched.ddQueue.Enqueue(t)
err := node.sched.DdQueue.Enqueue(t)
if err != nil {
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
log.Warn(ctx, "DropSnapshot failed to Enqueue",
@@ -133,14 +133,14 @@ func (node *Proxy) DescribeSnapshot(ctx context.Context, req *milvuspb.DescribeS
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
t := &describeSnapshotTask{
baseTask: baseTask{metaCache: node.getMetaCache()},
baseTask: baseTask{MetaCache: node.GetMetaCache()},
req: req,
ctx: ctx,
Condition: NewTaskCondition(ctx),
mixCoord: node.mixCoord,
}
err := node.sched.ddQueue.Enqueue(t)
err := node.sched.DdQueue.Enqueue(t)
if err != nil {
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
log.Warn(ctx, "DescribeSnapshot failed to Enqueue",
@@ -178,14 +178,14 @@ func (node *Proxy) ListSnapshots(ctx context.Context, req *milvuspb.ListSnapshot
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
t := &listSnapshotsTask{
baseTask: baseTask{metaCache: node.getMetaCache()},
baseTask: baseTask{MetaCache: node.GetMetaCache()},
req: req,
ctx: ctx,
Condition: NewTaskCondition(ctx),
mixCoord: node.mixCoord,
}
err := node.sched.ddQueue.Enqueue(t)
err := node.sched.DdQueue.Enqueue(t)
if err != nil {
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
log.Warn(ctx, "ListSnapshots failed to Enqueue",
@@ -300,7 +300,7 @@ func (node *Proxy) ExportSnapshot(ctx context.Context, req *milvuspb.ExportSnaps
return &milvuspb.ExportSnapshotResponse{Status: merr.Status(err)}, nil
}
collectionID, err := node.getMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName())
collectionID, err := node.GetMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName())
if err != nil {
failStatus, failCause := failMetricLabel(err)
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, failStatus, failCause, req.GetDbName(), req.GetCollectionName()).Inc()
@@ -454,14 +454,14 @@ func (node *Proxy) RestoreSnapshot(ctx context.Context, req *milvuspb.RestoreSna
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
t := &restoreSnapshotTask{
baseTask: baseTask{metaCache: node.getMetaCache()},
baseTask: baseTask{MetaCache: node.GetMetaCache()},
req: req,
ctx: ctx,
Condition: NewTaskCondition(ctx),
mixCoord: node.mixCoord,
}
err := node.sched.ddQueue.Enqueue(t)
err := node.sched.DdQueue.Enqueue(t)
if err != nil {
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
log.Warn(ctx, "RestoreSnapshot failed to Enqueue",
@@ -496,14 +496,14 @@ func (node *Proxy) GetRestoreSnapshotState(ctx context.Context, req *milvuspb.Ge
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel, metrics.CauseNA, "", "").Inc()
t := &getRestoreSnapshotStateTask{
baseTask: baseTask{metaCache: node.getMetaCache()},
baseTask: baseTask{MetaCache: node.GetMetaCache()},
req: req,
ctx: ctx,
Condition: NewTaskCondition(ctx),
mixCoord: node.mixCoord,
}
err := node.sched.ddQueue.Enqueue(t)
err := node.sched.DdQueue.Enqueue(t)
if err != nil {
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel, metrics.CauseNA, "", "").Inc()
log.Warn(ctx, "GetRestoreSnapshotState failed to Enqueue",
@@ -542,14 +542,14 @@ func (node *Proxy) ListRestoreSnapshotJobs(ctx context.Context, req *milvuspb.Li
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
t := &listRestoreSnapshotJobsTask{
baseTask: baseTask{metaCache: node.getMetaCache()},
baseTask: baseTask{MetaCache: node.GetMetaCache()},
req: req,
ctx: ctx,
Condition: NewTaskCondition(ctx),
mixCoord: node.mixCoord,
}
err := node.sched.ddQueue.Enqueue(t)
err := node.sched.DdQueue.Enqueue(t)
if err != nil {
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
log.Warn(ctx, "ListRestoreSnapshotJobs failed to Enqueue",
@@ -590,14 +590,14 @@ func (node *Proxy) PinSnapshotData(ctx context.Context, req *milvuspb.PinSnapsho
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
t := &pinSnapshotDataTask{
baseTask: baseTask{metaCache: node.getMetaCache()},
baseTask: baseTask{MetaCache: node.GetMetaCache()},
req: req,
ctx: ctx,
Condition: NewTaskCondition(ctx),
mixCoord: node.mixCoord,
}
if err := node.sched.ddQueue.Enqueue(t); err != nil {
if err := node.sched.DdQueue.Enqueue(t); err != nil {
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel, metrics.CauseNA, req.GetDbName(), req.GetCollectionName()).Inc()
log.Warn(ctx, "PinSnapshotData failed to Enqueue",
mlog.Err(err))
@@ -641,7 +641,7 @@ func (node *Proxy) UnpinSnapshotData(ctx context.Context, req *milvuspb.UnpinSna
mixCoord: node.mixCoord,
}
if err := node.sched.ddQueue.Enqueue(t); err != nil {
if err := node.sched.DdQueue.Enqueue(t); err != nil {
metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel, metrics.CauseNA, "", "").Inc()
log.Warn(ctx, "UnpinSnapshotData failed to Enqueue",
mlog.Err(err))
internal/proxy/snapshot_impl_test.go
+61 -59
View File
@@ -28,6 +28,8 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus/internal/proxy/scheduler"
"github.com/milvus-io/milvus/internal/proxy/taskmodel"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
@@ -35,8 +37,8 @@ import (
func TestProxy_CreateSnapshot_Success(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -47,7 +49,7 @@ func TestProxy_CreateSnapshot_Success(t *testing.T) {
}
// Mock successful enqueue
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
// Set task result to simulate successful execution
if cst, ok := t.(*createSnapshotTask); ok {
cst.result = merr.Success()
@@ -69,8 +71,8 @@ func TestProxy_CreateSnapshot_Success(t *testing.T) {
func TestProxy_CreateSnapshot_EnqueueFailure(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -83,7 +85,7 @@ func TestProxy_CreateSnapshot_EnqueueFailure(t *testing.T) {
expectedError := errors.New("enqueue failed")
// Mock enqueue failure
mock := mockey.Mock((*ddTaskQueue).Enqueue).Return(expectedError).Build()
mock := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).Return(expectedError).Build()
defer mock.UnPatch()
result, err := proxy.CreateSnapshot(context.Background(), req)
@@ -96,8 +98,8 @@ func TestProxy_CreateSnapshot_EnqueueFailure(t *testing.T) {
func TestProxy_CreateSnapshot_WaitToFinishFailure(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -110,7 +112,7 @@ func TestProxy_CreateSnapshot_WaitToFinishFailure(t *testing.T) {
expectedError := errors.New("task execution failed")
// Mock successful enqueue but failed task execution
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).Return(nil).Build()
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).Return(nil).Build()
defer mock1.UnPatch()
mock2 := mockey.Mock((*TaskCondition).WaitToFinish).Return(expectedError).Build()
defer mock2.UnPatch()
@@ -125,8 +127,8 @@ func TestProxy_CreateSnapshot_WaitToFinishFailure(t *testing.T) {
func TestProxy_DropSnapshot_Success(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -135,7 +137,7 @@ func TestProxy_DropSnapshot_Success(t *testing.T) {
}
// Mock successful enqueue and task completion
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
if dst, ok := t.(*dropSnapshotTask); ok {
dst.result = merr.Success()
}
@@ -155,8 +157,8 @@ func TestProxy_DropSnapshot_Success(t *testing.T) {
func TestProxy_DropSnapshot_EnqueueFailure(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -166,7 +168,7 @@ func TestProxy_DropSnapshot_EnqueueFailure(t *testing.T) {
expectedError := errors.New("enqueue failed")
mock := mockey.Mock((*ddTaskQueue).Enqueue).Return(expectedError).Build()
mock := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).Return(expectedError).Build()
defer mock.UnPatch()
result, err := proxy.DropSnapshot(context.Background(), req)
@@ -179,8 +181,8 @@ func TestProxy_DropSnapshot_EnqueueFailure(t *testing.T) {
func TestProxy_DropSnapshot_WaitToFinishFailure(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -190,7 +192,7 @@ func TestProxy_DropSnapshot_WaitToFinishFailure(t *testing.T) {
expectedError := errors.New("task execution failed")
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).Return(nil).Build()
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).Return(nil).Build()
defer mock1.UnPatch()
mock2 := mockey.Mock((*TaskCondition).WaitToFinish).Return(expectedError).Build()
defer mock2.UnPatch()
@@ -205,8 +207,8 @@ func TestProxy_DropSnapshot_WaitToFinishFailure(t *testing.T) {
func TestProxy_DescribeSnapshot_Success(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -215,7 +217,7 @@ func TestProxy_DescribeSnapshot_Success(t *testing.T) {
}
// Mock successful enqueue and task completion
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
if dst, ok := t.(*describeSnapshotTask); ok {
dst.result = &milvuspb.DescribeSnapshotResponse{
Status: merr.Success(),
@@ -241,8 +243,8 @@ func TestProxy_DescribeSnapshot_Success(t *testing.T) {
func TestProxy_DescribeSnapshot_EnqueueFailure(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -252,7 +254,7 @@ func TestProxy_DescribeSnapshot_EnqueueFailure(t *testing.T) {
expectedError := errors.New("enqueue failed")
mock := mockey.Mock((*ddTaskQueue).Enqueue).Return(expectedError).Build()
mock := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).Return(expectedError).Build()
defer mock.UnPatch()
result, err := proxy.DescribeSnapshot(context.Background(), req)
@@ -265,8 +267,8 @@ func TestProxy_DescribeSnapshot_EnqueueFailure(t *testing.T) {
func TestProxy_DescribeSnapshot_WaitToFinishFailure(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -276,7 +278,7 @@ func TestProxy_DescribeSnapshot_WaitToFinishFailure(t *testing.T) {
expectedError := errors.New("task execution failed")
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).Return(nil).Build()
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).Return(nil).Build()
defer mock1.UnPatch()
mock2 := mockey.Mock((*TaskCondition).WaitToFinish).Return(expectedError).Build()
defer mock2.UnPatch()
@@ -291,8 +293,8 @@ func TestProxy_DescribeSnapshot_WaitToFinishFailure(t *testing.T) {
func TestProxy_ListSnapshots_Success(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -301,7 +303,7 @@ func TestProxy_ListSnapshots_Success(t *testing.T) {
}
// Mock successful enqueue and task completion
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
if lst, ok := t.(*listSnapshotsTask); ok {
lst.result = &milvuspb.ListSnapshotsResponse{
Status: merr.Success(),
@@ -330,8 +332,8 @@ func TestProxy_ListSnapshots_Success(t *testing.T) {
func TestProxy_ListSnapshots_EnqueueFailure(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -341,7 +343,7 @@ func TestProxy_ListSnapshots_EnqueueFailure(t *testing.T) {
expectedError := errors.New("enqueue failed")
mock := mockey.Mock((*ddTaskQueue).Enqueue).Return(expectedError).Build()
mock := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).Return(expectedError).Build()
defer mock.UnPatch()
result, err := proxy.ListSnapshots(context.Background(), req)
@@ -354,8 +356,8 @@ func TestProxy_ListSnapshots_EnqueueFailure(t *testing.T) {
func TestProxy_ListSnapshots_WaitToFinishFailure(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -365,7 +367,7 @@ func TestProxy_ListSnapshots_WaitToFinishFailure(t *testing.T) {
expectedError := errors.New("task execution failed")
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).Return(nil).Build()
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).Return(nil).Build()
defer mock1.UnPatch()
mock2 := mockey.Mock((*TaskCondition).WaitToFinish).Return(expectedError).Build()
defer mock2.UnPatch()
@@ -380,8 +382,8 @@ func TestProxy_ListSnapshots_WaitToFinishFailure(t *testing.T) {
func TestProxy_RestoreSnapshot_Success(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -392,7 +394,7 @@ func TestProxy_RestoreSnapshot_Success(t *testing.T) {
}
// Mock successful enqueue and task completion
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
if rst, ok := t.(*restoreSnapshotTask); ok {
rst.result = &milvuspb.RestoreSnapshotResponse{
Status: merr.Success(),
@@ -414,8 +416,8 @@ func TestProxy_RestoreSnapshot_Success(t *testing.T) {
func TestProxy_RestoreSnapshot_EnqueueFailure(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -426,7 +428,7 @@ func TestProxy_RestoreSnapshot_EnqueueFailure(t *testing.T) {
expectedError := errors.New("enqueue failed")
mock := mockey.Mock((*ddTaskQueue).Enqueue).Return(expectedError).Build()
mock := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).Return(expectedError).Build()
defer mock.UnPatch()
result, err := proxy.RestoreSnapshot(context.Background(), req)
@@ -439,8 +441,8 @@ func TestProxy_RestoreSnapshot_EnqueueFailure(t *testing.T) {
func TestProxy_RestoreSnapshot_WaitToFinishFailure(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -451,7 +453,7 @@ func TestProxy_RestoreSnapshot_WaitToFinishFailure(t *testing.T) {
expectedError := errors.New("task execution failed")
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).Return(nil).Build()
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).Return(nil).Build()
defer mock1.UnPatch()
mock2 := mockey.Mock((*TaskCondition).WaitToFinish).Return(expectedError).Build()
defer mock2.UnPatch()
@@ -658,8 +660,8 @@ func TestSnapshotRequestOption_RBACGuardrails(t *testing.T) {
// Test task creation and basic properties
func TestProxy_CreateSnapshot_TaskCreation(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -672,7 +674,7 @@ func TestProxy_CreateSnapshot_TaskCreation(t *testing.T) {
var enqueuedTask *createSnapshotTask
// Capture the task that gets enqueued
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
if cst, ok := t.(*createSnapshotTask); ok {
enqueuedTask = cst
cst.result = merr.Success()
@@ -697,8 +699,8 @@ func TestProxy_CreateSnapshot_TaskCreation(t *testing.T) {
// Test edge cases
func TestProxy_CreateSnapshot_NilRequest(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -710,8 +712,8 @@ func TestProxy_CreateSnapshot_NilRequest(t *testing.T) {
func TestProxy_DropSnapshot_EmptySnapshotName(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -719,7 +721,7 @@ func TestProxy_DropSnapshot_EmptySnapshotName(t *testing.T) {
Name: "", // Empty snapshot name
}
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
if dst, ok := t.(*dropSnapshotTask); ok {
dst.result = merr.Success()
}
@@ -741,8 +743,8 @@ func TestProxy_DropSnapshot_EmptySnapshotName(t *testing.T) {
// Test metrics and observability
func TestProxy_CreateSnapshot_MetricsRecording(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -753,7 +755,7 @@ func TestProxy_CreateSnapshot_MetricsRecording(t *testing.T) {
}
// Mock successful execution
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
if cst, ok := t.(*createSnapshotTask); ok {
cst.result = merr.Success()
}
@@ -776,8 +778,8 @@ func TestProxy_CreateSnapshot_MetricsRecording(t *testing.T) {
// Test context handling
func TestProxy_CreateSnapshot_ContextCancellation(t *testing.T) {
proxy := &Proxy{
sched: &taskScheduler{
ddQueue: &ddTaskQueue{},
sched: &scheduler.TaskScheduler{
DdQueue: &scheduler.DdTaskQueue{},
},
}
@@ -792,7 +794,7 @@ func TestProxy_CreateSnapshot_ContextCancellation(t *testing.T) {
cancel()
// Mock successful enqueue
mock1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
mock1 := mockey.Mock((*scheduler.DdTaskQueue).Enqueue).To(func(t taskmodel.Task) error {
if cst, ok := t.(*createSnapshotTask); ok {
cst.result = merr.Success()
}
internal/proxy/task.go
+35 -103
View File
@@ -21,7 +21,6 @@ import (
"fmt"
"math"
"strconv"
"time"
"github.com/cockroachdb/errors"
"google.golang.org/protobuf/proto"
@@ -37,7 +36,6 @@ import (
"github.com/milvus-io/milvus/internal/util/schemautil"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
@@ -146,72 +144,6 @@ const (
NormScoreKey = "norm_score"
)
type task interface {
TraceCtx() context.Context
ID() UniqueID // return ReqID
SetID(uid UniqueID) // set ReqID
Name() string
Type() commonpb.MsgType
BeginTs() Timestamp
EndTs() Timestamp
SetTs(ts Timestamp)
OnEnqueue() error
PreExecute(ctx context.Context) error
Execute(ctx context.Context) error
PostExecute(ctx context.Context) error
WaitToFinish() error
Notify(err error)
CanSkipAllocTimestamp() bool
getMetaCache() Cache
SetOnEnqueueTime()
GetDurationInQueue() time.Duration
IsSubTask() bool
SetExecutingTime()
GetDurationInExecuting() time.Duration
}
type baseTask struct {
onEnqueueTime time.Time
executingTime time.Time
metaCache Cache
}
func (bt *baseTask) getMetaCache() Cache {
return bt.metaCache
}
func (bt *baseTask) CanSkipAllocTimestamp() bool {
return false
}
func (bt *baseTask) SetOnEnqueueTime() {
bt.onEnqueueTime = time.Now()
}
func (bt *baseTask) GetDurationInQueue() time.Duration {
return time.Since(bt.onEnqueueTime)
}
func (bt *baseTask) IsSubTask() bool {
return false
}
func (bt *baseTask) SetExecutingTime() {
bt.executingTime = time.Now()
}
func (bt *baseTask) GetDurationInExecuting() time.Duration {
return time.Since(bt.executingTime)
}
type dmlTask interface {
task
setChannels() error
getChannels() []pChan
}
type BaseInsertTask = msgstream.InsertMsg
func validateTextStorageV3Enabled(schema *schemapb.CollectionSchema) error {
if err := typeutil.ValidateTextRequiresStorageV3(schema, Params.CommonCfg.UseLoonFFI.GetAsBool()); err != nil {
return merr.WrapErrParameterInvalidMsg("%s", err.Error())
@@ -1509,7 +1441,7 @@ func (t *dropCollectionTask) PreExecute(ctx context.Context) error {
// No need to check collection name
// Validation shall be preformed in `CreateCollection`
// also permit drop collection one with bad collection name
_, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.GetCollectionName())
_, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.GetCollectionName())
if err != nil {
if errors.Is(err, merr.ErrCollectionNotFound) || errors.Is(err, merr.ErrDatabaseNotFound) {
// make dropping collection idempotent.
@@ -1595,7 +1527,7 @@ func (t *truncateCollectionTask) PreExecute(ctx context.Context) error {
// putting the collection in an inconsistent state from which the next
// load/search would fail. Reject up front; users who want to reset the
// view should use RefreshExternalCollection or DropCollection.
collSchema, err := t.getMetaCache().GetCollectionSchema(ctx, t.GetDbName(), t.GetCollectionName())
collSchema, err := t.GetMetaCache().GetCollectionSchema(ctx, t.GetDbName(), t.GetCollectionName())
if err != nil {
return err
}
@@ -1677,7 +1609,7 @@ func (t *hasCollectionTask) Execute(ctx context.Context) error {
t.result = &milvuspb.BoolResponse{
Status: merr.Success(),
}
_, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.GetCollectionName())
_, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.GetCollectionName())
// error other than
if err != nil && !errors.Is(err, merr.ErrCollectionNotFound) {
t.result.Status = merr.Status(err)
@@ -1880,7 +1812,7 @@ func (t *showCollectionsTask) Execute(ctx context.Context) error {
}
collectionIDs := make([]UniqueID, 0)
for _, collectionName := range t.CollectionNames {
collectionID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), collectionName)
collectionID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), collectionName)
if err != nil {
mlog.Debug(ctx, "Failed to get collection id.", mlog.String("collectionName", collectionName),
mlog.Int64("requestID", t.Base.MsgID), mlog.String("requestType", "showCollections"))
@@ -1937,7 +1869,7 @@ func (t *showCollectionsTask) Execute(ctx context.Context) error {
mlog.Int64("requestID", t.Base.MsgID), mlog.String("requestType", "showCollections"))
continue
}
collectionInfo, err := t.getMetaCache().GetCollectionInfo(ctx, t.GetDbName(), collectionName, id)
collectionInfo, err := t.GetMetaCache().GetCollectionInfo(ctx, t.GetDbName(), collectionName, id)
if err != nil {
mlog.Debug(ctx, "Failed to get collection info.", mlog.String("collectionName", collectionName),
mlog.Int64("requestID", t.Base.MsgID), mlog.String("requestType", "showCollections"))
@@ -2197,11 +2129,11 @@ func (t *alterCollectionTask) PreExecute(ctx context.Context) error {
return err
}
collSchema, err := t.getMetaCache().GetCollectionSchema(ctx, t.GetDbName(), t.CollectionName)
collSchema, err := t.GetMetaCache().GetCollectionSchema(ctx, t.GetDbName(), t.CollectionName)
if err != nil {
return err
}
collectionID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
collectionID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
if err != nil {
return err
}
@@ -2306,11 +2238,11 @@ func (t *alterCollectionTask) PreExecute(ctx context.Context) error {
}
}
isPartitionKeyMode, err := isPartitionKeyMode(ctx, t.getMetaCache(), t.GetDbName(), t.CollectionName)
isPartitionKeyMode, err := isPartitionKeyMode(ctx, t.GetMetaCache(), t.GetDbName(), t.CollectionName)
if err != nil {
return err
}
collBasicInfo, err := t.getMetaCache().GetCollectionInfo(t.ctx, t.GetDbName(), t.CollectionName, t.CollectionID)
collBasicInfo, err := t.GetMetaCache().GetCollectionInfo(t.ctx, t.GetDbName(), t.CollectionName, t.CollectionID)
if err != nil {
return err
}
@@ -2344,7 +2276,7 @@ func (t *alterCollectionTask) PreExecute(ctx context.Context) error {
// If partition key isolation or query_mode changed, check for existing vector index.
// Changing these properties requires dropping the vector index first.
if isoChanged || queryModeChanged {
if vecField, err := checkVectorIndexExist(ctx, t.getMetaCache(), t.GetDbName(), t.CollectionName, t.CollectionID, t.mixCoord); err != nil {
if vecField, err := checkVectorIndexExist(ctx, t.GetMetaCache(), t.GetDbName(), t.CollectionName, t.CollectionID, t.mixCoord); err != nil {
return err
} else if vecField != "" {
if isoChanged {
@@ -2533,13 +2465,13 @@ func validateAlterAnalyzerFieldParam(collSchema *schemapb.CollectionSchema, fiel
}
func (t *alterCollectionFieldTask) PreExecute(ctx context.Context) error {
collSchema, err := t.getMetaCache().GetCollectionSchema(ctx, t.GetDbName(), t.CollectionName)
collSchema, err := t.GetMetaCache().GetCollectionSchema(ctx, t.GetDbName(), t.CollectionName)
if err != nil {
return err
}
isCollectionLoadedFn := func() (bool, error) {
collectionID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
collectionID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
if err != nil {
return false, err
}
@@ -2735,7 +2667,7 @@ func (t *createPartitionTask) PreExecute(ctx context.Context) error {
}
// Check partition key mode
collSchema, err := t.getMetaCache().GetCollectionSchema(ctx, t.GetDbName(), collName)
collSchema, err := t.GetMetaCache().GetCollectionSchema(ctx, t.GetDbName(), collName)
if err != nil {
return err
}
@@ -2759,7 +2691,7 @@ func (t *createPartitionTask) Execute(ctx context.Context) (err error) {
if err := merr.CheckRPCCall(t.result, err); err != nil {
return err
}
collectionID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.GetCollectionName())
collectionID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.GetCollectionName())
if err != nil {
t.result = merr.Status(err)
return err
@@ -2834,7 +2766,7 @@ func (t *dropPartitionTask) PreExecute(ctx context.Context) error {
}
// Check partition key mode
collSchema, err := t.getMetaCache().GetCollectionSchema(ctx, t.GetDbName(), collName)
collSchema, err := t.GetMetaCache().GetCollectionSchema(ctx, t.GetDbName(), collName)
if err != nil {
return err
}
@@ -2846,11 +2778,11 @@ func (t *dropPartitionTask) PreExecute(ctx context.Context) error {
return err
}
collID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.GetCollectionName())
collID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.GetCollectionName())
if err != nil {
return err
}
partID, err := t.getMetaCache().GetPartitionID(ctx, t.GetDbName(), t.GetCollectionName(), t.GetPartitionName())
partID, err := t.GetMetaCache().GetPartitionID(ctx, t.GetDbName(), t.GetCollectionName(), t.GetPartitionName())
if err != nil {
if errors.Is(merr.ErrPartitionNotFound, err) || errors.Is(merr.ErrCollectionNotFound, err) || errors.Is(merr.ErrDatabaseNotFound, err) {
return nil
@@ -3030,7 +2962,7 @@ func (t *showPartitionsTask) Execute(ctx context.Context) error {
if t.GetType() == milvuspb.ShowType_InMemory {
collectionName := t.CollectionName
collectionID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), collectionName)
collectionID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), collectionName)
if err != nil {
mlog.Debug(ctx, "Failed to get collection id.", mlog.String("collectionName", collectionName),
mlog.Int64("requestID", t.Base.MsgID), mlog.String("requestType", "showPartitions"))
@@ -3043,7 +2975,7 @@ func (t *showPartitionsTask) Execute(ctx context.Context) error {
}
partitionIDs := make([]UniqueID, 0)
for _, partitionName := range t.PartitionNames {
partitionID, err := t.getMetaCache().GetPartitionID(ctx, t.GetDbName(), collectionName, partitionName)
partitionID, err := t.GetMetaCache().GetPartitionID(ctx, t.GetDbName(), collectionName, partitionName)
if err != nil {
mlog.Debug(ctx, "Failed to get partition id.", mlog.String("partitionName", partitionName),
mlog.Int64("requestID", t.Base.MsgID), mlog.String("requestType", "showPartitions"))
@@ -3080,7 +3012,7 @@ func (t *showPartitionsTask) Execute(ctx context.Context) error {
mlog.Int64("requestID", t.Base.MsgID), mlog.String("requestType", "showPartitions"))
return merr.WrapErrParameterInvalidMsg("failed to show partitions")
}
partitionInfo, err := t.getMetaCache().GetPartitionInfo(ctx, t.GetDbName(), collectionName, partitionName)
partitionInfo, err := t.GetMetaCache().GetPartitionInfo(ctx, t.GetDbName(), collectionName, partitionName)
if err != nil {
mlog.Debug(ctx, "Failed to get partition id.", mlog.String("partitionName", partitionName),
mlog.Int64("requestID", t.Base.MsgID), mlog.String("requestType", "showPartitions"))
@@ -3180,7 +3112,7 @@ func (t *loadCollectionTask) GetLoadPriority() commonpb.LoadPriority {
}
func (t *loadCollectionTask) Execute(ctx context.Context) (err error) {
collID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
collID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
log := mlog.With(
mlog.String("role", typeutil.ProxyRole),
@@ -3192,7 +3124,7 @@ func (t *loadCollectionTask) Execute(ctx context.Context) (err error) {
}
t.collectionID = collID
collSchema, err := t.getMetaCache().GetCollectionSchema(ctx, t.GetDbName(), t.CollectionName)
collSchema, err := t.GetMetaCache().GetCollectionSchema(ctx, t.GetDbName(), t.CollectionName)
if err != nil {
return err
}
@@ -3268,7 +3200,7 @@ func (t *loadCollectionTask) Execute(ctx context.Context) (err error) {
}
func (t *loadCollectionTask) PostExecute(ctx context.Context) error {
collID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
collID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
mlog.Debug(ctx, "loadCollectionTask PostExecute",
mlog.String("role", typeutil.ProxyRole),
mlog.Int64("collectionID", collID))
@@ -3341,7 +3273,7 @@ func (t *releaseCollectionTask) PreExecute(ctx context.Context) error {
}
func (t *releaseCollectionTask) Execute(ctx context.Context) (err error) {
collID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
collID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
if err != nil {
return err
}
@@ -3426,7 +3358,7 @@ func (t *loadPartitionsTask) PreExecute(ctx context.Context) error {
return err
}
partitionKeyMode, err := isPartitionKeyMode(ctx, t.getMetaCache(), t.GetDbName(), collName)
partitionKeyMode, err := isPartitionKeyMode(ctx, t.GetMetaCache(), t.GetDbName(), collName)
if err != nil {
return err
}
@@ -3448,12 +3380,12 @@ func (t *loadPartitionsTask) GetLoadPriority() commonpb.LoadPriority {
func (t *loadPartitionsTask) Execute(ctx context.Context) error {
var partitionIDs []int64
collID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
collID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
if err != nil {
return err
}
t.collectionID = collID
collSchema, err := t.getMetaCache().GetCollectionSchema(ctx, t.GetDbName(), t.CollectionName)
collSchema, err := t.GetMetaCache().GetCollectionSchema(ctx, t.GetDbName(), t.CollectionName)
if err != nil {
return err
}
@@ -3504,7 +3436,7 @@ func (t *loadPartitionsTask) Execute(ctx context.Context) error {
}
for _, partitionName := range t.PartitionNames {
partitionID, err := t.getMetaCache().GetPartitionID(ctx, t.GetDbName(), t.CollectionName, partitionName)
partitionID, err := t.GetMetaCache().GetPartitionID(ctx, t.GetDbName(), t.CollectionName, partitionName)
if err != nil {
return err
}
@@ -3603,7 +3535,7 @@ func (t *releasePartitionsTask) PreExecute(ctx context.Context) error {
return err
}
partitionKeyMode, err := isPartitionKeyMode(ctx, t.getMetaCache(), t.GetDbName(), collName)
partitionKeyMode, err := isPartitionKeyMode(ctx, t.GetMetaCache(), t.GetDbName(), collName)
if err != nil {
return err
}
@@ -3616,13 +3548,13 @@ func (t *releasePartitionsTask) PreExecute(ctx context.Context) error {
func (t *releasePartitionsTask) Execute(ctx context.Context) (err error) {
var partitionIDs []int64
collID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
collID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
if err != nil {
return err
}
t.collectionID = collID
for _, partitionName := range t.PartitionNames {
partitionID, err := t.getMetaCache().GetPartitionID(ctx, t.GetDbName(), t.CollectionName, partitionName)
partitionID, err := t.GetMetaCache().GetPartitionID(ctx, t.GetDbName(), t.CollectionName, partitionName)
if err != nil {
return err
}
@@ -3909,7 +3841,7 @@ func (t *DescribeResourceGroupTask) Execute(ctx context.Context) error {
getCollectionName := func(collections map[int64]int32) (map[string]int32, error) {
ret := make(map[string]int32)
for key, value := range collections {
name, err := t.getMetaCache().GetCollectionName(ctx, "", key)
name, err := t.GetMetaCache().GetCollectionName(ctx, "", key)
if err != nil {
mlog.Warn(ctx, "failed to get collection name",
mlog.Int64("collectionID", key),
@@ -4088,7 +4020,7 @@ func (t *TransferReplicaTask) PreExecute(ctx context.Context) error {
func (t *TransferReplicaTask) Execute(ctx context.Context) error {
var err error
collID, err := t.getMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
collID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
if err != nil {
return err
}
@@ -4226,14 +4158,14 @@ func (t *RunAnalyzerTask) OnEnqueue() error {
func (t *RunAnalyzerTask) PreExecute(ctx context.Context) error {
t.dbName = t.GetDbName()
collID, err := t.getMetaCache().GetCollectionID(ctx, t.dbName, t.GetCollectionName())
collID, err := t.GetMetaCache().GetCollectionID(ctx, t.dbName, t.GetCollectionName())
if err != nil { // err is not nil if collection not exists
return err
}
t.collectionID = collID
schema, err := t.getMetaCache().GetCollectionSchema(ctx, t.dbName, t.GetCollectionName())
schema, err := t.GetMetaCache().GetCollectionSchema(ctx, t.dbName, t.GetCollectionName())
if err != nil { // err is not nil if collection not exists
return err
}
internal/proxy/task_alias.go
+3 -3
View File
@@ -104,7 +104,7 @@ func (t *CreateAliasTask) PreExecute(ctx context.Context) error {
// This is needed for correctness: rootcoord's CheckIfAliasCreatable only checks mt.names,
// so passing an alias as CollectionName would fail.
dbName := t.GetDbName()
if resolved, err := resolveCollectionAlias(ctx, t.getMetaCache(), dbName, collName); err == nil {
if resolved, err := resolveCollectionAlias(ctx, t.GetMetaCache(), dbName, collName); err == nil {
t.CollectionName = resolved
}
return nil
@@ -269,7 +269,7 @@ func (t *AlterAliasTask) PreExecute(ctx context.Context) error {
// This is needed for correctness: rootcoord's CheckIfAliasAlterable only checks mt.names,
// so passing an alias as CollectionName would fail.
dbName := t.GetDbName()
if resolved, err := resolveCollectionAlias(ctx, t.getMetaCache(), dbName, collName); err == nil {
if resolved, err := resolveCollectionAlias(ctx, t.GetMetaCache(), dbName, collName); err == nil {
t.CollectionName = resolved
}
return nil
@@ -419,7 +419,7 @@ func (a *ListAliasesTask) PreExecute(ctx context.Context) error {
// Resolve CollectionName in case the user passed an alias instead of the real name.
// rootcoord filters by collection name in mt.names, so an alias would return no results.
dbName := a.GetDbName()
if resolved, err := resolveCollectionAlias(ctx, a.getMetaCache(), dbName, a.GetCollectionName()); err == nil {
if resolved, err := resolveCollectionAlias(ctx, a.GetMetaCache(), dbName, a.GetCollectionName()); err == nil {
a.CollectionName = resolved
}
}
internal/proxy/task_batch_update_manifest.go
+1 -1
View File
@@ -89,7 +89,7 @@ func (bt *batchUpdateManifestTask) PreExecute(ctx context.Context) error {
return merr.WrapErrParameterInvalidMsg("items is empty")
}
collectionID, err := bt.getMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName())
collectionID, err := bt.GetMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName())
if err != nil {
return err
}
internal/proxy/task_database.go
+1 -1
View File
@@ -152,7 +152,7 @@ func (ddt *dropDatabaseTask) Execute(ctx context.Context) error {
// Local best-effort cleanup on the issuing proxy; the authoritative
// eviction is the DropDatabase broadcast handled in
// InvalidateCollectionMetaCache.
ddt.getMetaCache().RemoveDatabase(ctx, ddt.DbName)
ddt.GetMetaCache().RemoveDatabase(ctx, ddt.DbName)
}
return err
}
internal/proxy/task_database_test.go
+1 -1
View File
@@ -80,7 +80,7 @@ func TestDropDatabaseTask(t *testing.T) {
).Maybe()
task := &dropDatabaseTask{
baseTask: baseTask{
metaCache: cache,
MetaCache: cache,
},
Condition: NewTaskCondition(ctx),
DropDatabaseRequest: &milvuspb.DropDatabaseRequest{
internal/proxy/task_delete.go
+17 -16
View File
@@ -17,6 +17,7 @@ import (
"github.com/milvus-io/milvus/internal/allocator"
"github.com/milvus-io/milvus/internal/parser/planparserv2"
"github.com/milvus-io/milvus/internal/proxy/channelmgr"
"github.com/milvus-io/milvus/internal/proxy/scheduler"
"github.com/milvus-io/milvus/internal/proxy/shardclient"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/util/exprutil"
@@ -111,8 +112,8 @@ func (dt *deleteTask) OnEnqueue() error {
return nil
}
func (dt *deleteTask) setChannels() error {
collID, err := dt.getMetaCache().GetCollectionID(dt.ctx, dt.req.GetDbName(), dt.req.GetCollectionName())
func (dt *deleteTask) SetChannels() error {
collID, err := dt.GetMetaCache().GetCollectionID(dt.ctx, dt.req.GetDbName(), dt.req.GetCollectionName())
if err != nil {
return err
}
@@ -124,7 +125,7 @@ func (dt *deleteTask) setChannels() error {
return nil
}
func (dt *deleteTask) getChannels() []pChan {
func (dt *deleteTask) GetChannels() []pChan {
return dt.pChannels
}
@@ -132,7 +133,7 @@ func (dt *deleteTask) PreExecute(ctx context.Context) error {
if dt.req.Namespace == nil {
return nil
}
schema, err := dt.getMetaCache().GetCollectionSchema(ctx, dt.req.GetDbName(), dt.req.GetCollectionName())
schema, err := dt.GetMetaCache().GetCollectionSchema(ctx, dt.req.GetDbName(), dt.req.GetCollectionName())
if err != nil {
return err
}
@@ -311,7 +312,7 @@ type deleteRunner struct {
count atomic.Int64
// task queue
queue *dmTaskQueue
queue *scheduler.DmTaskQueue
allQueryCnt atomic.Int64
sessionTS atomic.Uint64
@@ -320,7 +321,7 @@ type deleteRunner struct {
scannedTotalBytes atomic.Int64
}
func (dr *deleteRunner) getMetaCache() Cache {
func (dr *deleteRunner) GetMetaCache() Cache {
return dr.metaCache
}
@@ -336,18 +337,18 @@ func (dr *deleteRunner) Init(ctx context.Context) error {
return ErrWithLog(log, "Invalid collection name", err)
}
db, err := dr.getMetaCache().GetDatabaseInfo(ctx, dr.req.GetDbName())
db, err := dr.GetMetaCache().GetDatabaseInfo(ctx, dr.req.GetDbName())
if err != nil {
return err
}
dr.dbID = db.DBID
dr.collectionID, err = dr.getMetaCache().GetCollectionID(ctx, dr.req.GetDbName(), collName)
dr.collectionID, err = dr.GetMetaCache().GetCollectionID(ctx, dr.req.GetDbName(), collName)
if err != nil {
return ErrWithLog(log, "Failed to get collection id", err)
}
dr.schema, err = dr.getMetaCache().GetCollectionSchema(ctx, dr.req.GetDbName(), collName)
dr.schema, err = dr.GetMetaCache().GetCollectionSchema(ctx, dr.req.GetDbName(), collName)
if err != nil {
return ErrWithLog(log, "Failed to get collection schema", err)
}
@@ -362,7 +363,7 @@ func (dr *deleteRunner) Init(ctx context.Context) error {
dr.req.PartitionName = partitionName
}
colInfo, err := dr.getMetaCache().GetCollectionInfo(ctx, dr.req.GetDbName(), collName, dr.collectionID)
colInfo, err := dr.GetMetaCache().GetCollectionInfo(ctx, dr.req.GetDbName(), collName, dr.collectionID)
if err != nil {
return ErrWithLog(log, "Failed to get collection info", err)
}
@@ -398,11 +399,11 @@ func (dr *deleteRunner) Init(ctx context.Context) error {
if len(partName) > 0 {
return merr.WrapErrParameterInvalidMsg("not support manually specifying the partition names if namespace is used")
}
hashedPartitionNames, err := assignNamespacePartitionKey(ctx, dr.getMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), dr.req.Namespace)
hashedPartitionNames, err := assignNamespacePartitionKey(ctx, dr.GetMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), dr.req.Namespace)
if err != nil {
return err
}
dr.partitionIDs, err = getPartitionIDs(ctx, dr.getMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), hashedPartitionNames)
dr.partitionIDs, err = getPartitionIDs(ctx, dr.GetMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), hashedPartitionNames)
if err != nil {
return err
}
@@ -415,11 +416,11 @@ func (dr *deleteRunner) Init(ctx context.Context) error {
return err
}
partitionKeys := exprutil.ParseKeys(expr, exprutil.PartitionKey)
hashedPartitionNames, err := assignPartitionKeys(ctx, dr.getMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), partitionKeys)
hashedPartitionNames, err := assignPartitionKeys(ctx, dr.GetMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), partitionKeys)
if err != nil {
return err
}
dr.partitionIDs, err = getPartitionIDs(ctx, dr.getMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), hashedPartitionNames)
dr.partitionIDs, err = getPartitionIDs(ctx, dr.GetMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), hashedPartitionNames)
if err != nil {
return err
}
@@ -430,7 +431,7 @@ func (dr *deleteRunner) Init(ctx context.Context) error {
}
// dynamic validation
partID, err := dr.getMetaCache().GetPartitionID(ctx, dr.req.GetDbName(), collName, partName)
partID, err := dr.GetMetaCache().GetPartitionID(ctx, dr.req.GetDbName(), collName, partName)
if err != nil {
return ErrWithLog(log, "Failed to get partition id", err)
}
@@ -475,7 +476,7 @@ func (dr *deleteRunner) Run(ctx context.Context) error {
func (dr *deleteRunner) produce(ctx context.Context, primaryKeys *schemapb.IDs, partitionID UniqueID) (*deleteTask, error) {
dt := &deleteTask{
baseTask: baseTask{metaCache: dr.getMetaCache()},
baseTask: baseTask{MetaCache: dr.GetMetaCache()},
ctx: ctx,
Condition: NewTaskCondition(ctx),
req: dr.req,
internal/proxy/task_delete_streaming.go
+1 -1
View File
@@ -30,7 +30,7 @@ func (dt *deleteTask) Execute(ctx context.Context) (err error) {
var collectionSchema *schemapb.CollectionSchema
if dt.req.Namespace != nil || hookutil.IsClusterEncryptionEnabled() {
schema, err := dt.getMetaCache().GetCollectionSchema(ctx, dt.req.GetDbName(), dt.req.GetCollectionName())
schema, err := dt.GetMetaCache().GetCollectionSchema(ctx, dt.req.GetDbName(), dt.req.GetCollectionName())
if err != nil {
mlog.Warn(ctx, "get collection schema from meta cache failed", mlog.String("collectionName", dt.req.GetCollectionName()), mlog.Err(err))
return err
internal/proxy/task_delete_test.go
+12 -11
View File
@@ -21,6 +21,7 @@ import (
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/parser/planparserv2"
"github.com/milvus-io/milvus/internal/proxy/channelmgr"
"github.com/milvus-io/milvus/internal/proxy/scheduler"
"github.com/milvus-io/milvus/internal/proxy/shardclient"
"github.com/milvus-io/milvus/internal/util/streamrpc"
"github.com/milvus-io/milvus/pkg/v3/common"
@@ -127,16 +128,16 @@ func TestDeleteTask_GetChannels(t *testing.T) {
chMgr := channelmgr.NewMockChannelsMgr(t)
chMgr.EXPECT().GetChannels(mock.Anything).Return(channels, nil)
dt := deleteTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: context.Background(),
req: &milvuspb.DeleteRequest{
CollectionName: collectionName,
},
chMgr: chMgr,
}
err := dt.setChannels()
err := dt.SetChannels()
assert.NoError(t, err)
resChannels := dt.getChannels()
resChannels := dt.GetChannels()
assert.ElementsMatch(t, channels, resChannels)
assert.ElementsMatch(t, channels, dt.pChannels)
}
@@ -778,7 +779,7 @@ func TestDeleteRunner_Run(t *testing.T) {
tsoAllocator := &mockTsoAllocator{}
idAllocator := &mockIDAllocatorInterface{}
queue, err := newTaskScheduler(ctx, tsoAllocator)
queue, err := scheduler.NewTaskScheduler(ctx, tsoAllocator)
assert.NoError(t, err)
queue.Start()
defer queue.Close()
@@ -824,7 +825,7 @@ func TestDeleteRunner_Run(t *testing.T) {
vChannels: channels,
tsoAllocatorIns: tsoAllocator,
idAllocator: idAllocator,
queue: queue.dmQueue,
queue: queue.DmQueue,
lb: lb,
result: &milvuspb.MutationResult{
Status: merr.Success(),
@@ -858,7 +859,7 @@ func TestDeleteRunner_Run(t *testing.T) {
idAllocator: idAllocator,
tsoAllocatorIns: tsoAllocator,
metaCache: metaCache,
queue: queue.dmQueue,
queue: queue.DmQueue,
chMgr: mockMgr,
schema: schema,
collectionID: collectionID,
@@ -927,7 +928,7 @@ func TestDeleteRunner_Run(t *testing.T) {
require.NoError(t, err)
dr := deleteRunner{
queue: queue.dmQueue,
queue: queue.DmQueue,
metaCache: metaCache,
chMgr: mockMgr,
schema: schema,
@@ -995,7 +996,7 @@ func TestDeleteRunner_Run(t *testing.T) {
dr := deleteRunner{
chMgr: mockMgr,
queue: queue.dmQueue,
queue: queue.DmQueue,
metaCache: metaCache,
schema: schema,
collectionID: collectionID,
@@ -1059,7 +1060,7 @@ func TestDeleteRunner_Run(t *testing.T) {
dr := deleteRunner{
chMgr: mockMgr,
queue: queue.dmQueue,
queue: queue.DmQueue,
metaCache: metaCache,
schema: schema,
collectionID: collectionID,
@@ -1123,7 +1124,7 @@ func TestDeleteRunner_Run(t *testing.T) {
require.NoError(t, err)
dr := deleteRunner{
queue: queue.dmQueue,
queue: queue.DmQueue,
metaCache: metaCache,
chMgr: mockMgr,
schema: schema,
@@ -1196,7 +1197,7 @@ func TestDeleteRunner_Run(t *testing.T) {
dr := deleteRunner{
metaCache: mockCache,
queue: queue.dmQueue,
queue: queue.DmQueue,
chMgr: mockMgr,
schema: schema,
collectionID: collectionID,
internal/proxy/task_flush_all_test.go
+2 -3
View File
@@ -26,6 +26,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proxy/taskmodel"
"github.com/milvus-io/milvus/pkg/v3/util/uniquegenerator"
)
@@ -147,7 +148,7 @@ func TestFlushAllTaskPostExecute(t *testing.T) {
func TestFlushAllTaskImplementsTaskInterface(t *testing.T) {
// Verify that flushAllTask implements the task interface
var _ task = (*flushAllTask)(nil)
var _ taskmodel.Task = (*flushAllTask)(nil)
task, mixCoord, _ := createTestFlushAllTask(t)
defer mixCoord.AssertExpectations(t)
@@ -213,7 +214,6 @@ func TestFlushAllTaskBaseTaskMethods(t *testing.T) {
// Test SetOnEnqueueTime
task.SetOnEnqueueTime()
assert.False(t, task.onEnqueueTime.IsZero())
// Test GetDurationInQueue
time.Sleep(1 * time.Millisecond)
@@ -225,7 +225,6 @@ func TestFlushAllTaskBaseTaskMethods(t *testing.T) {
// Test SetExecutingTime
task.SetExecutingTime()
assert.False(t, task.executingTime.IsZero())
// Test GetDurationInExecuting
time.Sleep(1 * time.Millisecond)
internal/proxy/task_flush_streaming.go
+1 -1
View File
@@ -44,7 +44,7 @@ func (t *flushTask) Execute(ctx context.Context) error {
mlog.Info(ctx, "flushTaskByStreamingService.Execute", mlog.Int("collectionNum", len(t.CollectionNames)), mlog.Uint64("flushTs", flushTs))
timeOfSeal, _ := tsoutil.ParseTS(flushTs)
for _, collName := range t.CollectionNames {
collID, err := t.getMetaCache().GetCollectionID(t.ctx, t.DbName, collName)
collID, err := t.GetMetaCache().GetCollectionID(t.ctx, t.DbName, collName)
if err != nil {
return err
}
internal/proxy/task_import.go
+9 -9
View File
@@ -91,12 +91,12 @@ func (it *importTask) OnEnqueue() error {
func (it *importTask) PreExecute(ctx context.Context) error {
req := it.req
node := it.node
collectionID, err := it.getMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName())
collectionID, err := it.GetMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName())
if err != nil {
return err
}
it.collectionID = collectionID
schema, err := it.getMetaCache().GetCollectionSchema(ctx, req.GetDbName(), req.GetCollectionName())
schema, err := it.GetMetaCache().GetCollectionSchema(ctx, req.GetDbName(), req.GetCollectionName())
if err != nil {
return err
}
@@ -124,7 +124,7 @@ func (it *importTask) PreExecute(ctx context.Context) error {
return merr.WrapErrParameterMissingMsg("partition not specified")
}
// Currently, Backup tool call import must with a partition name, each time restore a partition
partitionID, err := it.getMetaCache().GetPartitionID(ctx, req.GetDbName(), req.GetCollectionName(), req.GetPartitionName())
partitionID, err := it.GetMetaCache().GetPartitionID(ctx, req.GetDbName(), req.GetCollectionName(), req.GetPartitionName())
if err != nil {
return err
}
@@ -133,7 +133,7 @@ func (it *importTask) PreExecute(ctx context.Context) error {
if req.GetPartitionName() == "" {
partitionIDs = []UniqueID{common.AllPartitionsID}
} else {
partitionID, err := it.getMetaCache().GetPartitionID(ctx, req.GetDbName(), req.GetCollectionName(), req.PartitionName)
partitionID, err := it.GetMetaCache().GetPartitionID(ctx, req.GetDbName(), req.GetCollectionName(), req.PartitionName)
if err != nil {
return err
}
@@ -156,7 +156,7 @@ func (it *importTask) PreExecute(ctx context.Context) error {
if req.GetPartitionName() != "" {
return merr.WrapErrImportFailed("not allow to set partition name for collection with partition key")
}
partitions, err := it.getMetaCache().GetPartitions(ctx, req.GetDbName(), req.GetCollectionName())
partitions, err := it.GetMetaCache().GetPartitions(ctx, req.GetDbName(), req.GetCollectionName())
if err != nil {
return err
}
@@ -168,7 +168,7 @@ func (it *importTask) PreExecute(ctx context.Context) error {
if req.GetPartitionName() == "" {
req.PartitionName = Params.CommonCfg.DefaultPartitionName.GetValue()
}
partitionID, err := it.getMetaCache().GetPartitionID(ctx, req.GetDbName(), req.GetCollectionName(), req.PartitionName)
partitionID, err := it.GetMetaCache().GetPartitionID(ctx, req.GetDbName(), req.GetCollectionName(), req.PartitionName)
if err != nil {
return err
}
@@ -199,13 +199,13 @@ func (it *importTask) PreExecute(ctx context.Context) error {
return nil
}
func (it *importTask) setChannels() error {
func (it *importTask) SetChannels() error {
// import task only send message by broadcast, which didn't affect the time tick rule at proxy.
// so we don't need to set channels here.
return nil
}
func (it *importTask) getChannels() []pChan {
func (it *importTask) GetChannels() []pChan {
return nil
}
@@ -214,7 +214,7 @@ func (it *importTask) Execute(ctx context.Context) error {
// DataCoord will handle validation, broadcasting, and job creation
// Get database ID from database name
dbInfo, err := it.getMetaCache().GetDatabaseInfo(ctx, it.req.GetDbName())
dbInfo, err := it.GetMetaCache().GetDatabaseInfo(ctx, it.req.GetDbName())
if err != nil {
mlog.Warn(ctx, "failed to get database info", mlog.FieldDbName(it.req.GetDbName()), mlog.Err(err))
return err
internal/proxy/task_import_test.go
+9 -9
View File
@@ -64,7 +64,7 @@ func (s *ImportTaskSuite) TestExecute_GetDatabaseInfoFailsReturnsError() {
},
resp: &internalpb.ImportResponse{},
}
task.metaCache = mockCache
task.MetaCache = mockCache
err := task.Execute(ctx)
@@ -104,7 +104,7 @@ func (s *ImportTaskSuite) TestExecute_ImportV2RPCFailsReturnsError() {
},
resp: &internalpb.ImportResponse{},
}
task.metaCache = mockCache
task.MetaCache = mockCache
err := task.Execute(ctx)
@@ -146,7 +146,7 @@ func (s *ImportTaskSuite) TestExecute_ImportV2ReturnsErrorStatusReturnsError() {
},
resp: &internalpb.ImportResponse{},
}
task.metaCache = mockCache
task.MetaCache = mockCache
err := task.Execute(ctx)
@@ -190,7 +190,7 @@ func (s *ImportTaskSuite) TestExecute_SuccessSetsJobID() {
},
resp: resp,
}
task.metaCache = mockCache
task.MetaCache = mockCache
err := task.Execute(ctx)
@@ -241,7 +241,7 @@ func (s *ImportTaskSuite) TestExecute_PassesCorrectRequestParameters() {
},
resp: &internalpb.ImportResponse{},
}
task.metaCache = mockCache
task.MetaCache = mockCache
err := task.Execute(ctx)
@@ -320,12 +320,12 @@ func (s *ImportTaskSuite) TestTaskBasicMethods() {
func (s *ImportTaskSuite) TestSetChannels_ReturnsNil() {
task := &importTask{}
s.NoError(task.setChannels())
s.NoError(task.SetChannels())
}
func (s *ImportTaskSuite) TestGetChannels_ReturnsNil() {
task := &importTask{}
s.Nil(task.getChannels())
s.Nil(task.GetChannels())
}
// --------------------------------
@@ -373,7 +373,7 @@ func (s *ImportTaskSuite) TestExecute_DataTimestampIsAlwaysZero() {
},
resp: &internalpb.ImportResponse{},
}
task.metaCache = mockCache
task.MetaCache = mockCache
task.Execute(ctx)
@@ -401,7 +401,7 @@ func (s *ImportTaskSuite) TestPreExecute_GetCollectionIDFailsReturnsError() {
CollectionName: "test_collection",
},
}
task.metaCache = mockCache
task.MetaCache = mockCache
err := task.PreExecute(ctx)
internal/proxy/task_index.go
+11 -11
View File
@@ -513,7 +513,7 @@ func (cit *createIndexTask) parseIndexParams(ctx context.Context) error {
}
func (cit *createIndexTask) getIndexedFieldAndFunction(ctx context.Context) error {
schema, err := cit.getMetaCache().GetCollectionSchema(ctx, cit.req.GetDbName(), cit.req.GetCollectionName())
schema, err := cit.GetMetaCache().GetCollectionSchema(ctx, cit.req.GetDbName(), cit.req.GetCollectionName())
if err != nil {
mlog.Error(ctx, "failed to get collection schema", mlog.Err(err))
return merr.Wrap(err, "failed to get collection schema")
@@ -585,13 +585,13 @@ func checkTrain(ctx context.Context, field *schemapb.FieldSchema, indexParams ma
func (cit *createIndexTask) PreExecute(ctx context.Context) error {
collName := cit.req.GetCollectionName()
collID, err := cit.getMetaCache().GetCollectionID(ctx, cit.req.GetDbName(), collName)
collID, err := cit.GetMetaCache().GetCollectionID(ctx, cit.req.GetDbName(), collName)
if err != nil {
return err
}
cit.collectionID = collID
collInfo, err := cit.getMetaCache().GetCollectionInfo(ctx, cit.req.GetDbName(), cit.req.GetCollectionName(), cit.collectionID)
collInfo, err := cit.GetMetaCache().GetCollectionInfo(ctx, cit.req.GetDbName(), cit.req.GetCollectionName(), cit.collectionID)
if err != nil {
return err
}
@@ -723,7 +723,7 @@ func (t *alterIndexTask) PreExecute(ctx context.Context) error {
collName := t.req.GetCollectionName()
collection, err := t.getMetaCache().GetCollectionID(ctx, t.req.GetDbName(), collName)
collection, err := t.GetMetaCache().GetCollectionID(ctx, t.req.GetDbName(), collName)
if err != nil {
return err
}
@@ -833,7 +833,7 @@ func (dit *describeIndexTask) PreExecute(ctx context.Context) error {
return err
}
collID, err := dit.getMetaCache().GetCollectionID(ctx, dit.GetDbName(), dit.CollectionName)
collID, err := dit.GetMetaCache().GetCollectionID(ctx, dit.GetDbName(), dit.CollectionName)
if err != nil {
return err
}
@@ -842,7 +842,7 @@ func (dit *describeIndexTask) PreExecute(ctx context.Context) error {
}
func (dit *describeIndexTask) Execute(ctx context.Context) error {
schema, err := dit.getMetaCache().GetCollectionSchema(ctx, dit.GetDbName(), dit.GetCollectionName())
schema, err := dit.GetMetaCache().GetCollectionSchema(ctx, dit.GetDbName(), dit.GetCollectionName())
if err != nil {
mlog.Error(ctx, "failed to get collection schema", mlog.Err(err))
return merr.Wrap(err, "failed to get collection schema")
@@ -973,7 +973,7 @@ func (dit *getIndexStatisticsTask) PreExecute(ctx context.Context) error {
return err
}
collID, err := dit.getMetaCache().GetCollectionID(ctx, dit.GetDbName(), dit.CollectionName)
collID, err := dit.GetMetaCache().GetCollectionID(ctx, dit.GetDbName(), dit.CollectionName)
if err != nil {
return err
}
@@ -982,7 +982,7 @@ func (dit *getIndexStatisticsTask) PreExecute(ctx context.Context) error {
}
func (dit *getIndexStatisticsTask) Execute(ctx context.Context) error {
schema, err := dit.getMetaCache().GetCollectionSchema(ctx, dit.GetDbName(), dit.GetCollectionName())
schema, err := dit.GetMetaCache().GetCollectionSchema(ctx, dit.GetDbName(), dit.GetCollectionName())
if err != nil {
mlog.Error(ctx, "failed to get collection schema", mlog.String("collection_name", dit.GetCollectionName()), mlog.Err(err))
return merr.Wrap(err, "failed to get collection schema")
@@ -1081,7 +1081,7 @@ func (dit *dropIndexTask) OnEnqueue() error {
}
func (dit *dropIndexTask) PreExecute(ctx context.Context) error {
collID, err := dit.getMetaCache().GetCollectionID(ctx, dit.GetDbName(), dit.CollectionName)
collID, err := dit.GetMetaCache().GetCollectionID(ctx, dit.GetDbName(), dit.CollectionName)
if err != nil {
return err
}
@@ -1177,7 +1177,7 @@ func (gibpt *getIndexBuildProgressTask) PreExecute(ctx context.Context) error {
func (gibpt *getIndexBuildProgressTask) Execute(ctx context.Context) error {
collectionName := gibpt.CollectionName
collectionID, err := gibpt.getMetaCache().GetCollectionID(ctx, gibpt.GetDbName(), collectionName)
collectionID, err := gibpt.GetMetaCache().GetCollectionID(ctx, gibpt.GetDbName(), collectionName)
if err != nil { // err is not nil if collection not exists
return err
}
@@ -1264,7 +1264,7 @@ func (gist *getIndexStateTask) PreExecute(ctx context.Context) error {
}
func (gist *getIndexStateTask) Execute(ctx context.Context) error {
collectionID, err := gist.getMetaCache().GetCollectionID(ctx, gist.GetDbName(), gist.CollectionName)
collectionID, err := gist.GetMetaCache().GetCollectionID(ctx, gist.GetDbName(), gist.CollectionName)
if err != nil {
return err
}
internal/proxy/task_index_test.go
+5 -5
View File
@@ -72,7 +72,7 @@ func TestGetIndexStateTask_Execute(t *testing.T) {
assert.NoError(t, err)
gist := &getIndexStateTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
GetIndexStateRequest: &milvuspb.GetIndexStateRequest{
Base: &commonpb.MsgBase{},
DbName: dbName,
@@ -110,7 +110,7 @@ func TestDropIndexTask_PreExecute(t *testing.T) {
mock.AnythingOfType("string"),
).Return(collectionID, nil)
dit := dropIndexTask{
baseTask: baseTask{metaCache: mockCache},
baseTask: baseTask{MetaCache: mockCache},
ctx: ctx,
DropIndexRequest: &milvuspb.DropIndexRequest{
Base: &commonpb.MsgBase{
@@ -141,12 +141,12 @@ func TestDropIndexTask_PreExecute(t *testing.T) {
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(0), errors.New("error"))
dit.metaCache = mockCache
dit.MetaCache = mockCache
err := dit.PreExecute(ctx)
assert.Error(t, err)
})
dit.metaCache = mockCache
dit.MetaCache = mockCache
mockCache.On("GetCollectionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
@@ -223,7 +223,7 @@ func TestCreateIndexTask_PreExecute(t *testing.T) {
mock.AnythingOfType("int64"),
).Return(&collectionInfo{}, nil)
cit := createIndexTask{
baseTask: baseTask{metaCache: mockCache},
baseTask: baseTask{MetaCache: mockCache},
ctx: ctx,
req: &milvuspb.CreateIndexRequest{
Base: &commonpb.MsgBase{
internal/proxy/task_insert.go
+8 -8
View File
@@ -76,8 +76,8 @@ func (it *insertTask) EndTs() Timestamp {
return it.insertMsg.EndTimestamp
}
func (it *insertTask) setChannels() error {
collID, err := it.getMetaCache().GetCollectionID(it.ctx, it.insertMsg.GetDbName(), it.insertMsg.CollectionName)
func (it *insertTask) SetChannels() error {
collID, err := it.GetMetaCache().GetCollectionID(it.ctx, it.insertMsg.GetDbName(), it.insertMsg.CollectionName)
if err != nil {
return err
}
@@ -89,7 +89,7 @@ func (it *insertTask) setChannels() error {
return nil
}
func (it *insertTask) getChannels() []pChan {
func (it *insertTask) GetChannels() []pChan {
return it.pChannels
}
@@ -143,14 +143,14 @@ func (it *insertTask) PreExecute(ctx context.Context) error {
return err
}
collID, err := it.getMetaCache().GetCollectionID(context.Background(), it.insertMsg.GetDbName(), collectionName)
collID, err := it.GetMetaCache().GetCollectionID(context.Background(), it.insertMsg.GetDbName(), collectionName)
if err != nil {
log.Warn(ctx, "fail to get collection id", mlog.Err(err))
return err
}
it.collectionID = collID
colInfo, err := it.getMetaCache().GetCollectionInfo(ctx, it.insertMsg.GetDbName(), collectionName, collID)
colInfo, err := it.GetMetaCache().GetCollectionInfo(ctx, it.insertMsg.GetDbName(), collectionName, collID)
if err != nil {
log.Warn(ctx, "fail to get collection info", mlog.Err(err))
return err
@@ -168,7 +168,7 @@ func (it *insertTask) PreExecute(ctx context.Context) error {
}
}
schema, err := it.getMetaCache().GetCollectionSchema(ctx, it.insertMsg.GetDbName(), collectionName)
schema, err := it.GetMetaCache().GetCollectionSchema(ctx, it.insertMsg.GetDbName(), collectionName)
if err != nil {
log.Warn(ctx, "get collection schema from global meta cache failed", mlog.String("collectionName", collectionName), mlog.Err(err))
return err
@@ -272,7 +272,7 @@ func (it *insertTask) PreExecute(ctx context.Context) error {
return err
}
partitionKeyMode, err := isPartitionKeyMode(ctx, it.getMetaCache(), it.insertMsg.GetDbName(), collectionName)
partitionKeyMode, err := isPartitionKeyMode(ctx, it.GetMetaCache(), it.insertMsg.GetDbName(), collectionName)
if err != nil {
log.Warn(ctx, "check partition key mode failed", mlog.String("collectionName", collectionName), mlog.Err(err))
return err
@@ -289,7 +289,7 @@ func (it *insertTask) PreExecute(ctx context.Context) error {
// insert to _default partition
partitionTag := it.insertMsg.GetPartitionName()
if len(partitionTag) <= 0 {
pinfo, err := it.getMetaCache().GetPartitionInfo(ctx, it.insertMsg.GetDbName(), collectionName, "")
pinfo, err := it.GetMetaCache().GetPartitionInfo(ctx, it.insertMsg.GetDbName(), collectionName, "")
if err != nil {
log.Warn(ctx, "get partition info failed", mlog.String("collectionName", collectionName), mlog.Err(err))
return err
internal/proxy/task_insert_streaming.go
+3 -3
View File
@@ -38,7 +38,7 @@ func (it *insertTask) Execute(ctx context.Context) error {
tr := timerecord.NewTimeRecorder(fmt.Sprintf("proxy execute insert streaming %d", it.ID()))
collectionName := it.insertMsg.CollectionName
collID, err := it.getMetaCache().GetCollectionID(it.ctx, it.insertMsg.GetDbName(), collectionName)
collID, err := it.GetMetaCache().GetCollectionID(it.ctx, it.insertMsg.GetDbName(), collectionName)
if err != nil {
mlog.Warn(ctx, "fail to get collection id", mlog.Err(err))
return err
@@ -69,9 +69,9 @@ func (it *insertTask) Execute(ctx context.Context) error {
// start to repack insert data
var msgs []message.MutableMessage
if it.partitionKeys == nil {
msgs, err = repackInsertDataForStreamingService(it.TraceCtx(), it.getMetaCache(), channelNames, it.insertMsg, it.result, ez, it.schemaVersion, nil)
msgs, err = repackInsertDataForStreamingService(it.TraceCtx(), it.GetMetaCache(), channelNames, it.insertMsg, it.result, ez, it.schemaVersion, nil)
} else {
msgs, err = repackInsertDataWithPartitionKeyForStreamingService(it.TraceCtx(), it.getMetaCache(), channelNames, it.insertMsg, it.result, it.partitionKeys, ez, it.schema, it.schemaVersion, nil)
msgs, err = repackInsertDataWithPartitionKeyForStreamingService(it.TraceCtx(), it.GetMetaCache(), channelNames, it.insertMsg, it.result, it.partitionKeys, ez, it.schema, it.schemaVersion, nil)
}
if err != nil {
mlog.Warn(ctx, "assign segmentID and repack insert data failed", mlog.Err(err))
internal/proxy/task_insert_test.go
+13 -13
View File
@@ -107,7 +107,7 @@ func TestInsertTaskPreExecuteTextRequiresStorageV3(t *testing.T) {
},
},
}
task.metaCache = cache
task.MetaCache = cache
err := task.PreExecute(context.Background())
assert.Error(t, err)
@@ -361,7 +361,7 @@ func TestInsertTask_CheckAligned(t *testing.T) {
}
func TestInsertTask(t *testing.T) {
t.Run("test getChannels", func(t *testing.T) {
t.Run("test GetChannels", func(t *testing.T) {
collectionID := UniqueID(0)
collectionName := "col-0"
channels := []pChan{"mock-chan-0", "mock-chan-1"}
@@ -374,7 +374,7 @@ func TestInsertTask(t *testing.T) {
chMgr := channelmgr.NewMockChannelsMgr(t)
chMgr.EXPECT().GetChannels(mock.Anything).Return(channels, nil)
it := insertTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: context.Background(),
insertMsg: &msgstream.InsertMsg{
InsertRequest: &msgpb.InsertRequest{
@@ -383,9 +383,9 @@ func TestInsertTask(t *testing.T) {
},
chMgr: chMgr,
}
err := it.setChannels()
err := it.SetChannels()
assert.NoError(t, err)
resChannels := it.getChannels()
resChannels := it.GetChannels()
assert.ElementsMatch(t, channels, resChannels)
assert.ElementsMatch(t, channels, it.pChannels)
})
@@ -462,7 +462,7 @@ func TestMaxInsertSize(t *testing.T) {
cache.On("GetCollectionInfo", mock.Anything, dbName, collectionName, UniqueID(100)).Return(&collectionInfo{Schema: schema}, nil)
cache.On("GetCollectionSchema", mock.Anything, dbName, collectionName).Return(schema, nil)
it := insertTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: context.Background(),
insertMsg: &msgstream.InsertMsg{
InsertRequest: &msgpb.InsertRequest{
@@ -622,7 +622,7 @@ func TestInsertTask_KeepUserPK_WhenAllowInsertAutoIDTrue(t *testing.T) {
mock.Anything,
mock.Anything,
).Return(&collectionInfo{Schema: info}, nil)
task.metaCache = cache
task.MetaCache = cache
err = task.PreExecute(context.Background())
assert.NoError(t, err)
@@ -768,7 +768,7 @@ func TestInsertTask_Function(t *testing.T) {
mock.Anything,
mock.Anything,
).Return(&collectionInfo{Schema: info}, nil)
task.metaCache = cache
task.MetaCache = cache
err = task.PreExecute(ctx)
assert.NoError(t, err)
}
@@ -798,7 +798,7 @@ func TestInsertTaskForSchemaMismatch(t *testing.T) {
},
}),
}, nil)
it.metaCache = mockCache
it.MetaCache = mockCache
err := it.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrCollectionSchemaMismatch)
@@ -868,7 +868,7 @@ func TestInsertTask_Namespace(t *testing.T) {
Version: msgpb.InsertDataVersion_ColumnBased,
},
},
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
schema: schemaWithNamespaceEnabled,
idAllocator: idAllocator,
}
@@ -887,7 +887,7 @@ func TestInsertTask_Namespace(t *testing.T) {
},
},
idAllocator: idAllocator,
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
}
err = it.PreExecute(context.Background())
assert.Error(t, err)
@@ -915,7 +915,7 @@ func TestInsertTask_Namespace(t *testing.T) {
Version: msgpb.InsertDataVersion_ColumnBased,
},
},
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
schema: schemaWithNamespaceDisabled,
idAllocator: idAllocator,
}
@@ -935,7 +935,7 @@ func TestInsertTask_Namespace(t *testing.T) {
},
},
idAllocator: idAllocator,
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
}
err = it.PreExecute(context.Background())
assert.Error(t, err)
internal/proxy/task_query.go
+10 -11
View File
@@ -48,7 +48,6 @@ const (
const (
RetrieveTaskName = "RetrieveTask"
QueryTaskName = "QueryTask"
)
type queryTask struct {
@@ -643,14 +642,14 @@ func (t *queryTask) CanSkipAllocTimestamp() bool {
}
consistencyLevel = t.request.GetConsistencyLevel()
} else {
collID, err := t.getMetaCache().GetCollectionID(context.Background(), t.request.GetDbName(), t.request.GetCollectionName())
collID, err := t.GetMetaCache().GetCollectionID(context.Background(), t.request.GetDbName(), t.request.GetCollectionName())
if err != nil { // err is not nil if collection not exists
mlog.Warn(t.ctx, "query task get collectionID failed, can't skip alloc timestamp",
mlog.String("collectionName", t.request.GetCollectionName()), mlog.Err(err))
return false
}
collectionInfo, err2 := t.getMetaCache().GetCollectionInfo(context.Background(), t.request.GetDbName(), t.request.GetCollectionName(), collID)
collectionInfo, err2 := t.GetMetaCache().GetCollectionInfo(context.Background(), t.request.GetDbName(), t.request.GetCollectionName(), collID)
if err2 != nil {
mlog.Warn(t.ctx, "query task get collection info failed, can't skip alloc timestamp",
mlog.String("collectionName", t.request.GetCollectionName()), mlog.Err(err))
@@ -678,14 +677,14 @@ func (t *queryTask) PreExecute(ctx context.Context) error {
}
log.Debug(ctx, "Validate collectionName.")
collID, err := t.getMetaCache().GetCollectionID(ctx, t.request.GetDbName(), collectionName)
collID, err := t.GetMetaCache().GetCollectionID(ctx, t.request.GetDbName(), collectionName)
if err != nil {
log.Warn(ctx, "Failed to get collection id.", mlog.String("collectionName", collectionName), mlog.Err(err))
return err
}
t.CollectionID = collID
colInfo, err := t.getMetaCache().GetCollectionInfo(ctx, t.request.GetDbName(), collectionName, t.CollectionID)
colInfo, err := t.GetMetaCache().GetCollectionInfo(ctx, t.request.GetDbName(), collectionName, t.CollectionID)
if err != nil {
log.Warn(ctx, "Failed to get collection info.", mlog.String("collectionName", collectionName),
mlog.Int64("collectionID", t.CollectionID), mlog.Err(err))
@@ -697,7 +696,7 @@ func (t *queryTask) PreExecute(ctx context.Context) error {
}
log.Debug(ctx, "Get collection ID by name", mlog.Int64("collectionID", t.CollectionID))
schema, err := t.getMetaCache().GetCollectionSchema(ctx, t.request.GetDbName(), t.collectionName)
schema, err := t.GetMetaCache().GetCollectionSchema(ctx, t.request.GetDbName(), t.collectionName)
if err != nil {
log.Warn(ctx, "get collection schema failed", mlog.Err(err))
return err
@@ -714,7 +713,7 @@ func (t *queryTask) PreExecute(ctx context.Context) error {
t.request.PartitionNames = partitionNames
}
t.partitionKeyMode, err = isPartitionKeyMode(ctx, t.getMetaCache(), t.request.GetDbName(), collectionName)
t.partitionKeyMode, err = isPartitionKeyMode(ctx, t.GetMetaCache(), t.request.GetDbName(), collectionName)
if err != nil {
log.Warn(ctx, "check partition key mode failed", mlog.Int64("collectionID", t.CollectionID), mlog.Err(err))
return err
@@ -809,7 +808,7 @@ func (t *queryTask) PreExecute(ctx context.Context) error {
if !t.reQuery {
partitionNames := t.request.GetPartitionNames()
if namespacePartitionKeyMode(t.schema.CollectionSchema) && t.request.Namespace != nil {
hashedPartitionNames, err := assignNamespacePartitionKey(ctx, t.getMetaCache(), t.request.GetDbName(), t.request.CollectionName, t.request.Namespace)
hashedPartitionNames, err := assignNamespacePartitionKey(ctx, t.GetMetaCache(), t.request.GetDbName(), t.request.CollectionName, t.request.Namespace)
if err != nil {
return err
}
@@ -821,14 +820,14 @@ func (t *queryTask) PreExecute(ctx context.Context) error {
return err
}
partitionKeys := exprutil.ParseKeys(expr, exprutil.PartitionKey)
hashedPartitionNames, err := assignPartitionKeys(ctx, t.getMetaCache(), t.request.GetDbName(), t.request.CollectionName, partitionKeys)
hashedPartitionNames, err := assignPartitionKeys(ctx, t.GetMetaCache(), t.request.GetDbName(), t.request.CollectionName, partitionKeys)
if err != nil {
return err
}
partitionNames = append(partitionNames, hashedPartitionNames...)
}
t.PartitionIDs, err = getPartitionIDs(ctx, t.getMetaCache(), t.request.GetDbName(), t.request.CollectionName, partitionNames)
t.PartitionIDs, err = getPartitionIDs(ctx, t.GetMetaCache(), t.request.GetDbName(), t.request.CollectionName, partitionNames)
if err != nil {
return err
}
@@ -851,7 +850,7 @@ func (t *queryTask) PreExecute(ctx context.Context) error {
t.Username = username
}
collectionInfo, err2 := t.getMetaCache().GetCollectionInfo(ctx, t.request.GetDbName(), collectionName, t.CollectionID)
collectionInfo, err2 := t.GetMetaCache().GetCollectionInfo(ctx, t.request.GetDbName(), collectionName, t.CollectionID)
if err2 != nil {
log.Warn(ctx, "Proxy::queryTask::PreExecute failed to GetCollectionInfo from cache",
mlog.String("collectionName", collectionName), mlog.Int64("collectionID", t.CollectionID),
internal/proxy/task_query_namespace_test.go
+2 -2
View File
@@ -84,7 +84,7 @@ func TestQueryTask_PlanNamespace_AfterPreExecute(t *testing.T) {
}).Build()
task := &queryTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(context.Background()),
RetrieveRequest: &internalpb.RetrieveRequest{QueryLabel: "query", Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Retrieve}},
ctx: context.Background(),
@@ -149,7 +149,7 @@ func TestQueryTask_NamespaceSetsPartitionIDs(t *testing.T) {
for _, ns := range namespaces {
namespace := ns
task := &queryTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(context.Background()),
RetrieveRequest: &internalpb.RetrieveRequest{QueryLabel: "query", Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Retrieve}},
ctx: context.Background(),
internal/proxy/task_query_test.go
+11 -11
View File
@@ -131,7 +131,7 @@ func TestQueryTask_all(t *testing.T) {
t.Run("test query task parameters", func(t *testing.T) {
task := &queryTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
RetrieveRequest: &internalpb.RetrieveRequest{
QueryLabel: "query",
@@ -279,7 +279,7 @@ func TestQueryTask_all(t *testing.T) {
t.Run("test query for iterator", func(t *testing.T) {
qt := &queryTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
RetrieveRequest: &internalpb.RetrieveRequest{
QueryLabel: "query",
@@ -332,7 +332,7 @@ func TestQueryTask_all(t *testing.T) {
// next page query task
qt = &queryTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
RetrieveRequest: &internalpb.RetrieveRequest{
QueryLabel: "query",
@@ -401,7 +401,7 @@ func TestQueryTask_all(t *testing.T) {
})
}
task := &queryTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
RetrieveRequest: &internalpb.RetrieveRequest{
QueryLabel: "query",
@@ -501,7 +501,7 @@ func TestQueryTask_all(t *testing.T) {
t.Run("test order by without limit", func(t *testing.T) {
// ORDER BY without explicit limit should fail with an error
task := &queryTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
RetrieveRequest: &internalpb.RetrieveRequest{
Base: &commonpb.MsgBase{
@@ -544,7 +544,7 @@ func TestQueryTask_all(t *testing.T) {
t.Run("test order by with limit succeeds", func(t *testing.T) {
// ORDER BY with explicit limit should pass the validation
task := &queryTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
RetrieveRequest: &internalpb.RetrieveRequest{
Base: &commonpb.MsgBase{
@@ -1781,7 +1781,7 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
t.Run("default consistency level", func(t *testing.T) {
qt := &queryTask{
baseTask: baseTask{metaCache: mockMetaCache},
baseTask: baseTask{MetaCache: mockMetaCache},
request: &milvuspb.QueryRequest{
Base: nil,
DbName: dbName,
@@ -1824,7 +1824,7 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
}, nil).Times(3)
qt := &queryTask{
baseTask: baseTask{metaCache: mockMetaCache},
baseTask: baseTask{MetaCache: mockMetaCache},
request: &milvuspb.QueryRequest{
Base: nil,
DbName: dbName,
@@ -1848,7 +1848,7 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
t.Run("legacy_guarantee_ts", func(t *testing.T) {
qt := &queryTask{
baseTask: baseTask{metaCache: mockMetaCache},
baseTask: baseTask{MetaCache: mockMetaCache},
request: &milvuspb.QueryRequest{
Base: nil,
DbName: dbName,
@@ -1877,7 +1877,7 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
nil, errors.New("mock error")).Once()
qt := &queryTask{
baseTask: baseTask{metaCache: mockMetaCache},
baseTask: baseTask{MetaCache: mockMetaCache},
request: &milvuspb.QueryRequest{
Base: nil,
DbName: dbName,
@@ -1902,7 +1902,7 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
assert.False(t, skip)
qt2 := &queryTask{
baseTask: baseTask{metaCache: mockMetaCache},
baseTask: baseTask{MetaCache: mockMetaCache},
request: &milvuspb.QueryRequest{
Base: nil,
DbName: dbName,
internal/proxy/task_search.go
+13 -13
View File
@@ -24,6 +24,7 @@ import (
"github.com/milvus-io/milvus/internal/proxy/fieldvalidator"
"github.com/milvus-io/milvus/internal/proxy/search_agg"
"github.com/milvus-io/milvus/internal/proxy/shardclient"
"github.com/milvus-io/milvus/internal/proxy/taskmodel"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/util/exprutil"
"github.com/milvus-io/milvus/internal/util/function/embedding"
@@ -48,7 +49,6 @@ import (
)
const (
SearchTaskName = "SearchTask"
SearchLevelKey = "level"
// requeryThreshold is the estimated threshold for the size of the search results.
@@ -140,14 +140,14 @@ func (t *searchTask) CanSkipAllocTimestamp() bool {
}
consistencyLevel = t.request.GetConsistencyLevel()
} else {
collID, err := t.getMetaCache().GetCollectionID(context.Background(), t.request.GetDbName(), t.request.GetCollectionName())
collID, err := t.GetMetaCache().GetCollectionID(context.Background(), t.request.GetDbName(), t.request.GetCollectionName())
if err != nil { // err is not nil if collection not exists
mlog.Warn(t.ctx, "search task get collectionID failed, can't skip alloc timestamp",
mlog.String("collectionName", t.request.GetCollectionName()), mlog.Err(err))
return false
}
collectionInfo, err2 := t.getMetaCache().GetCollectionInfo(context.Background(), t.request.GetDbName(), t.request.GetCollectionName(), collID)
collectionInfo, err2 := t.GetMetaCache().GetCollectionInfo(context.Background(), t.request.GetDbName(), t.request.GetCollectionName(), collID)
if err2 != nil {
mlog.Warn(t.ctx, "search task get collection info failed, can't skip alloc timestamp",
mlog.String("collectionName", t.request.GetCollectionName()), mlog.Err(err))
@@ -168,7 +168,7 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
collectionName := t.request.CollectionName
t.collectionName = collectionName
collID, err := t.getMetaCache().GetCollectionID(ctx, t.request.GetDbName(), collectionName)
collID, err := t.GetMetaCache().GetCollectionID(ctx, t.request.GetDbName(), collectionName)
if err != nil { // err is not nil if collection not exists
return err
}
@@ -176,7 +176,7 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
t.DbID = 0 // todo
t.CollectionID = collID
log := mlog.With(mlog.Int64("collID", collID), mlog.String("collName", collectionName))
t.schema, err = t.getMetaCache().GetCollectionSchema(ctx, t.request.GetDbName(), collectionName)
t.schema, err = t.GetMetaCache().GetCollectionSchema(ctx, t.request.GetDbName(), collectionName)
if err != nil {
log.Warn(ctx, "get collection schema failed", mlog.Err(err))
return err
@@ -185,7 +185,7 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
return err
}
collectionInfo, err2 := t.getMetaCache().GetCollectionInfo(ctx, t.request.GetDbName(), collectionName, t.CollectionID)
collectionInfo, err2 := t.GetMetaCache().GetCollectionInfo(ctx, t.request.GetDbName(), collectionName, t.CollectionID)
if err2 != nil {
log.Warn(ctx, "Proxy::searchTask::PreExecute failed to GetCollectionInfo from cache",
mlog.String("collectionName", collectionName), mlog.Int64("collectionID", t.CollectionID), mlog.Err(err2))
@@ -194,7 +194,7 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
t.largeTopKEnabled = collectionInfo.QueryMode == common.QueryModeLargeTopK
t.partitionKeyIsolation = collectionInfo.PartitionKeyIsolation
t.partitionKeyMode, err = isPartitionKeyMode(ctx, t.getMetaCache(), t.request.GetDbName(), collectionName)
t.partitionKeyMode, err = isPartitionKeyMode(ctx, t.GetMetaCache(), t.request.GetDbName(), collectionName)
if err != nil {
log.Warn(ctx, "is partition key mode failed", mlog.Err(err))
return err
@@ -216,7 +216,7 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
if !t.partitionKeyMode && len(t.request.GetPartitionNames()) > 0 {
// translate partition name to partition ids. Use regex-pattern to match partition name.
t.PartitionIDs, err = getPartitionIDs(ctx, t.getMetaCache(), t.request.GetDbName(), collectionName, t.request.GetPartitionNames())
t.PartitionIDs, err = getPartitionIDs(ctx, t.GetMetaCache(), t.request.GetDbName(), collectionName, t.request.GetPartitionNames())
if err != nil {
log.Warn(ctx, "failed to get partition ids", mlog.Err(err))
return err
@@ -1217,13 +1217,13 @@ func (t *searchTask) tryGeneratePlan(
func (t *searchTask) tryParsePartitionIDsFromPlan(plan *planpb.PlanNode) ([]int64, error) {
if namespacePartitionKeyMode(t.schema.CollectionSchema) && t.request.Namespace != nil {
hashedPartitionNames, err := assignNamespacePartitionKey(t.ctx, t.getMetaCache(), t.request.GetDbName(), t.collectionName, t.request.Namespace)
hashedPartitionNames, err := assignNamespacePartitionKey(t.ctx, t.GetMetaCache(), t.request.GetDbName(), t.collectionName, t.request.Namespace)
if err != nil {
mlog.Warn(t.ctx, "failed to assign namespace partition key", mlog.Err(err))
return nil, err
}
if len(hashedPartitionNames) > 0 {
PartitionIDs, err2 := getPartitionIDs(t.ctx, t.getMetaCache(), t.request.GetDbName(), t.collectionName, hashedPartitionNames)
PartitionIDs, err2 := getPartitionIDs(t.ctx, t.GetMetaCache(), t.request.GetDbName(), t.collectionName, hashedPartitionNames)
if err2 != nil {
mlog.Warn(t.ctx, "failed to get namespace partition ids", mlog.Err(err2))
return nil, err2
@@ -1239,7 +1239,7 @@ func (t *searchTask) tryParsePartitionIDsFromPlan(plan *planpb.PlanNode) ([]int6
return nil, err
}
partitionKeys := exprutil.ParseKeys(expr, exprutil.PartitionKey)
hashedPartitionNames, err := assignPartitionKeys(t.ctx, t.getMetaCache(), t.request.GetDbName(), t.collectionName, partitionKeys)
hashedPartitionNames, err := assignPartitionKeys(t.ctx, t.GetMetaCache(), t.request.GetDbName(), t.collectionName, partitionKeys)
if err != nil {
mlog.Warn(t.ctx, "failed to assign partition keys", mlog.Err(err))
return nil, err
@@ -1247,7 +1247,7 @@ func (t *searchTask) tryParsePartitionIDsFromPlan(plan *planpb.PlanNode) ([]int6
if len(hashedPartitionNames) > 0 {
// translate partition name to partition ids. Use regex-pattern to match partition name.
PartitionIDs, err2 := getPartitionIDs(t.ctx, t.getMetaCache(), t.request.GetDbName(), t.collectionName, hashedPartitionNames)
PartitionIDs, err2 := getPartitionIDs(t.ctx, t.GetMetaCache(), t.request.GetDbName(), t.collectionName, hashedPartitionNames)
if err2 != nil {
mlog.Warn(t.ctx, "failed to get partition ids", mlog.Err(err2))
return nil, err2
@@ -1628,7 +1628,7 @@ func (t *searchTask) SetID(uid UniqueID) {
}
func (t *searchTask) Name() string {
return SearchTaskName
return taskmodel.SearchTaskName
}
func (t *searchTask) Type() commonpb.MsgType {
internal/proxy/task_search_namespace_test.go
+2 -2
View File
@@ -52,7 +52,7 @@ func TestSearchTask_PlanNamespace_AfterPreExecute(t *testing.T) {
// Build task
task := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(context.Background()),
SearchRequest: &internalpb.SearchRequest{Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Search}},
ctx: context.Background(),
@@ -105,7 +105,7 @@ func TestSearchTask_NamespaceSetsPartitionIDs(t *testing.T) {
for _, ns := range namespaces {
namespace := ns
task := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(context.Background()),
SearchRequest: &internalpb.SearchRequest{Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Search}},
ctx: context.Background(),
internal/proxy/task_search_test.go
+25 -24
View File
@@ -40,6 +40,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proxy/scheduler"
"github.com/milvus-io/milvus/internal/proxy/search_agg"
"github.com/milvus-io/milvus/internal/proxy/shardclient"
"github.com/milvus-io/milvus/internal/types"
@@ -99,7 +100,7 @@ func TestSearchTaskPreExecuteTextRequiresStorageV3(t *testing.T) {
cache.EXPECT().GetCollectionSchema(mock.Anything, dbName, collectionName).Return(schema, nil)
task := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: context.Background(),
SearchRequest: &internalpb.SearchRequest{},
request: &milvuspb.SearchRequest{
@@ -153,7 +154,7 @@ func TestSearchTaskPreExecuteUsesTimezoneForTimestamptzFilter(t *testing.T) {
Value: "Asia/Shanghai",
})
task := &searchTask{
baseTask: baseTask{metaCache: &MetaCache{}},
baseTask: baseTask{MetaCache: &MetaCache{}},
Condition: NewTaskCondition(context.Background()),
SearchRequest: &internalpb.SearchRequest{Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Search}},
ctx: context.Background(),
@@ -195,7 +196,7 @@ func TestSearchTask_PostExecute(t *testing.T) {
getSearchTask := func(t *testing.T, collName string) *searchTask {
task := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
collectionName: collName,
SearchRequest: &internalpb.SearchRequest{
@@ -268,7 +269,7 @@ func TestSearchTask_PostExecute(t *testing.T) {
}
qt := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
SearchRequest: &internalpb.SearchRequest{
Base: &commonpb.MsgBase{
@@ -362,7 +363,7 @@ func TestSearchTask_PostExecute(t *testing.T) {
},
}
task := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
collectionName: collName,
SearchRequest: &internalpb.SearchRequest{
@@ -474,7 +475,7 @@ func TestSearchTask_PostExecute(t *testing.T) {
},
}
task := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
collectionName: collName,
SearchRequest: &internalpb.SearchRequest{
@@ -903,7 +904,7 @@ func TestSearchTask_PreExecute(t *testing.T) {
getSearchTask := func(t *testing.T, collName string) *searchTask {
task := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
collectionName: collName,
SearchRequest: &internalpb.SearchRequest{},
@@ -921,7 +922,7 @@ func TestSearchTask_PreExecute(t *testing.T) {
getSearchTaskWithNq := func(t *testing.T, collName string, nq int64) *searchTask {
task := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
collectionName: collName,
SearchRequest: &internalpb.SearchRequest{},
@@ -952,7 +953,7 @@ func TestSearchTask_PreExecute(t *testing.T) {
},
}
task := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
collectionName: collName,
SearchRequest: &internalpb.SearchRequest{},
@@ -1517,7 +1518,7 @@ func TestSearchTask_WithFunctions(t *testing.T) {
}
task := &searchTask{
baseTask: baseTask{metaCache: mockCache},
baseTask: baseTask{MetaCache: mockCache},
ctx: ctx,
collectionName: collectionName,
SearchRequest: &internalpb.SearchRequest{
@@ -1602,7 +1603,7 @@ func TestSearchTask_WithFunctions(t *testing.T) {
subReqs = append(subReqs, subReq)
}
task := &searchTask{
baseTask: baseTask{metaCache: mockCache},
baseTask: baseTask{MetaCache: mockCache},
ctx: ctx,
collectionName: collectionName,
SearchRequest: &internalpb.SearchRequest{
@@ -3460,7 +3461,7 @@ func TestSearchTask_ErrExecute(t *testing.T) {
// test begins
task := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
SearchRequest: &internalpb.SearchRequest{
Base: &commonpb.MsgBase{
@@ -4583,7 +4584,7 @@ func TestSearchTask_Requery(t *testing.T) {
node.tsoAllocator = &timestampAllocator{
tso: newMockTimestampAllocatorInterface(),
}
scheduler, err := newTaskScheduler(ctx, node.tsoAllocator)
scheduler, err := scheduler.NewTaskScheduler(ctx, node.tsoAllocator)
assert.NoError(t, err)
node.sched = scheduler
err = node.sched.Start()
@@ -4676,7 +4677,7 @@ func TestSearchTask_Requery(t *testing.T) {
outputFields := []string{pkField, vecField}
qt := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
SearchRequest: &internalpb.SearchRequest{
Base: &commonpb.MsgBase{
@@ -4721,7 +4722,7 @@ func TestSearchTask_Requery(t *testing.T) {
node := mocks.NewMockProxy(t)
qt := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
SearchRequest: &internalpb.SearchRequest{
Base: &commonpb.MsgBase{
@@ -4754,7 +4755,7 @@ func TestSearchTask_Requery(t *testing.T) {
node.lbPolicy = lb
qt := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
SearchRequest: &internalpb.SearchRequest{
Base: &commonpb.MsgBase{
@@ -4800,7 +4801,7 @@ func TestSearchTask_Requery(t *testing.T) {
}
qt := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
SearchRequest: &internalpb.SearchRequest{
Base: &commonpb.MsgBase{
@@ -4972,7 +4973,7 @@ func TestSearchTask_CanSkipAllocTimestamp(t *testing.T) {
t.Run("default consistency level", func(t *testing.T) {
st := &searchTask{
baseTask: baseTask{metaCache: mockMetaCache},
baseTask: baseTask{MetaCache: mockMetaCache},
request: &milvuspb.SearchRequest{
Base: nil,
DbName: dbName,
@@ -5015,7 +5016,7 @@ func TestSearchTask_CanSkipAllocTimestamp(t *testing.T) {
}, nil).Times(3)
st := &searchTask{
baseTask: baseTask{metaCache: mockMetaCache},
baseTask: baseTask{MetaCache: mockMetaCache},
request: &milvuspb.SearchRequest{
Base: nil,
DbName: dbName,
@@ -5039,7 +5040,7 @@ func TestSearchTask_CanSkipAllocTimestamp(t *testing.T) {
t.Run("legacy_guarantee_ts", func(t *testing.T) {
st := &searchTask{
baseTask: baseTask{metaCache: mockMetaCache},
baseTask: baseTask{MetaCache: mockMetaCache},
request: &milvuspb.SearchRequest{
Base: nil,
DbName: dbName,
@@ -5068,7 +5069,7 @@ func TestSearchTask_CanSkipAllocTimestamp(t *testing.T) {
nil, errors.New("mock error")).Once()
st := &searchTask{
baseTask: baseTask{metaCache: mockMetaCache},
baseTask: baseTask{MetaCache: mockMetaCache},
request: &milvuspb.SearchRequest{
Base: nil,
DbName: dbName,
@@ -5093,7 +5094,7 @@ func TestSearchTask_CanSkipAllocTimestamp(t *testing.T) {
assert.False(t, skip)
st2 := &searchTask{
baseTask: baseTask{metaCache: mockMetaCache},
baseTask: baseTask{MetaCache: mockMetaCache},
request: &milvuspb.SearchRequest{
Base: nil,
DbName: dbName,
@@ -5148,7 +5149,7 @@ func (s *MaterializedViewTestSuite) SetupTest() {
func (s *MaterializedViewTestSuite) getSearchTask() *searchTask {
task := &searchTask{
baseTask: baseTask{metaCache: s.mockMetaCache},
baseTask: baseTask{MetaCache: s.mockMetaCache},
ctx: s.ctx,
collectionName: s.colName,
SearchRequest: &internalpb.SearchRequest{},
@@ -5179,7 +5180,7 @@ func (s *MaterializedViewTestSuite) getHybridSearchTask(dsl string) *searchTask
}
task := &searchTask{
baseTask: baseTask{metaCache: s.mockMetaCache},
baseTask: baseTask{MetaCache: s.mockMetaCache},
ctx: s.ctx,
collectionName: s.colName,
SearchRequest: &internalpb.SearchRequest{},
internal/proxy/task_snapshot.go
+13 -13
View File
@@ -123,7 +123,7 @@ func (cst *createSnapshotTask) PreExecute(ctx context.Context) error {
return merr.WrapErrParameterInvalidMsg("compaction_protection_seconds must not exceed %d", maxCompactionProtectionSeconds)
}
collectionID, err := cst.getMetaCache().GetCollectionID(ctx, cst.req.GetDbName(), cst.req.GetCollectionName())
collectionID, err := cst.GetMetaCache().GetCollectionID(ctx, cst.req.GetDbName(), cst.req.GetCollectionName())
if err != nil {
return err
}
@@ -220,7 +220,7 @@ func (dst *dropSnapshotTask) PreExecute(ctx context.Context) error {
if dst.req.GetCollectionName() == "" {
return merr.WrapErrParameterMissingMsg("collection_name is required for drop snapshot")
}
collectionID, err := dst.getMetaCache().GetCollectionID(ctx, dst.req.GetDbName(), dst.req.GetCollectionName())
collectionID, err := dst.GetMetaCache().GetCollectionID(ctx, dst.req.GetDbName(), dst.req.GetCollectionName())
if err != nil {
return err
}
@@ -315,7 +315,7 @@ func (dst *describeSnapshotTask) PreExecute(ctx context.Context) error {
if dst.req.GetCollectionName() == "" {
return merr.WrapErrParameterMissingMsg("collection_name is required for describe snapshot")
}
collectionID, err := dst.getMetaCache().GetCollectionID(ctx, dst.req.GetDbName(), dst.req.GetCollectionName())
collectionID, err := dst.GetMetaCache().GetCollectionID(ctx, dst.req.GetDbName(), dst.req.GetCollectionName())
if err != nil {
return err
}
@@ -346,7 +346,7 @@ func (dst *describeSnapshotTask) Execute(ctx context.Context) error {
snapshotInfo := result.GetSnapshotInfo()
collectionName, err := dst.getMetaCache().GetCollectionName(ctx, "", snapshotInfo.GetCollectionId())
collectionName, err := dst.GetMetaCache().GetCollectionName(ctx, "", snapshotInfo.GetCollectionId())
if err != nil {
mlog.Warn(ctx, "DescribeSnapshot fail to get collection name",
mlog.Err(err))
@@ -354,7 +354,7 @@ func (dst *describeSnapshotTask) Execute(ctx context.Context) error {
}
var partitionNames []string
for _, partitionID := range snapshotInfo.GetPartitionIds() {
partitionName, err := dst.getMetaCache().GetPartitionName(ctx, "", collectionName, partitionID)
partitionName, err := dst.GetMetaCache().GetPartitionName(ctx, "", collectionName, partitionID)
if err != nil {
mlog.Warn(ctx, "DescribeSnapshot fail to get partition name",
mlog.FieldPartitionID(partitionID),
@@ -436,7 +436,7 @@ func (lst *listSnapshotsTask) OnEnqueue() error {
func (lst *listSnapshotsTask) PreExecute(ctx context.Context) error {
// Resolve database ID for db-level filtering
if lst.req.GetDbName() != "" {
dbInfo, err := lst.getMetaCache().GetDatabaseInfo(ctx, lst.req.GetDbName())
dbInfo, err := lst.GetMetaCache().GetDatabaseInfo(ctx, lst.req.GetDbName())
if err != nil {
return err
}
@@ -447,7 +447,7 @@ func (lst *listSnapshotsTask) PreExecute(ctx context.Context) error {
return merr.WrapErrParameterMissingMsg("collection_name is required for ListSnapshots")
}
collectionID, err := lst.getMetaCache().GetCollectionID(ctx, lst.req.GetDbName(), lst.req.GetCollectionName())
collectionID, err := lst.GetMetaCache().GetCollectionID(ctx, lst.req.GetDbName(), lst.req.GetCollectionName())
if err != nil {
return err
}
@@ -560,7 +560,7 @@ func (rst *restoreSnapshotTask) PreExecute(ctx context.Context) error {
}
// Resolve source collection ID for per-collection snapshot lookup (RPC call)
collectionID, err := rst.getMetaCache().GetCollectionID(ctx, rst.req.GetDbName(), rst.req.GetCollectionName())
collectionID, err := rst.GetMetaCache().GetCollectionID(ctx, rst.req.GetDbName(), rst.req.GetCollectionName())
if err != nil {
return err
}
@@ -682,7 +682,7 @@ func (grst *getRestoreSnapshotStateTask) Execute(ctx context.Context) error {
info := result.GetInfo()
var milvusInfo *milvuspb.RestoreSnapshotInfo
if info != nil {
dbName, collectionName := resolveCollectionNames(ctx, grst.getMetaCache(), info.GetCollectionId())
dbName, collectionName := resolveCollectionNames(ctx, grst.GetMetaCache(), info.GetCollectionId())
milvusInfo = &milvuspb.RestoreSnapshotInfo{
JobId: info.GetJobId(),
SnapshotName: info.GetSnapshotName(),
@@ -764,7 +764,7 @@ func (lrst *listRestoreSnapshotJobsTask) OnEnqueue() error {
func (lrst *listRestoreSnapshotJobsTask) PreExecute(ctx context.Context) error {
// Resolve database ID for db-level filtering
if lrst.req.GetDbName() != "" {
dbInfo, err := lrst.getMetaCache().GetDatabaseInfo(ctx, lrst.req.GetDbName())
dbInfo, err := lrst.GetMetaCache().GetDatabaseInfo(ctx, lrst.req.GetDbName())
if err != nil {
return err
}
@@ -772,7 +772,7 @@ func (lrst *listRestoreSnapshotJobsTask) PreExecute(ctx context.Context) error {
}
if lrst.req.GetCollectionName() != "" {
collectionID, err := lrst.getMetaCache().GetCollectionID(ctx, lrst.req.GetDbName(), lrst.req.GetCollectionName())
collectionID, err := lrst.GetMetaCache().GetCollectionID(ctx, lrst.req.GetDbName(), lrst.req.GetCollectionName())
if err != nil {
return err
}
@@ -805,7 +805,7 @@ func (lrst *listRestoreSnapshotJobsTask) Execute(ctx context.Context) error {
jobs := result.GetJobs()
milvusJobs := make([]*milvuspb.RestoreSnapshotInfo, 0, len(jobs))
for _, job := range jobs {
dbName, collectionName := resolveCollectionNames(ctx, lrst.getMetaCache(), job.GetCollectionId())
dbName, collectionName := resolveCollectionNames(ctx, lrst.GetMetaCache(), job.GetCollectionId())
milvusJobs = append(milvusJobs, &milvuspb.RestoreSnapshotInfo{
JobId: job.GetJobId(),
SnapshotName: job.GetSnapshotName(),
@@ -904,7 +904,7 @@ func (pst *pinSnapshotDataTask) PreExecute(ctx context.Context) error {
return merr.WrapErrParameterInvalidMsg("ttl_seconds exceeds maximum of %d (30 days)", maxPinTTLSeconds)
}
collectionID, err := pst.getMetaCache().GetCollectionID(ctx, pst.req.GetDbName(), pst.req.GetCollectionName())
collectionID, err := pst.GetMetaCache().GetCollectionID(ctx, pst.req.GetDbName(), pst.req.GetCollectionName())
if err != nil {
return err
}
internal/proxy/task_snapshot_test.go
+29 -29
View File
@@ -75,7 +75,7 @@ func TestCreateSnapshotTask_OnEnqueue_BaseAlreadyExists(t *testing.T) {
func TestCreateSnapshotTask_PreExecute_Success(t *testing.T) {
cache := &MetaCache{}
task := &createSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.CreateSnapshotRequest{
Name: "test_snapshot",
DbName: "default",
@@ -96,7 +96,7 @@ func TestCreateSnapshotTask_PreExecute_Success(t *testing.T) {
func TestCreateSnapshotTask_PreExecute_CollectionNotFound(t *testing.T) {
cache := &MetaCache{}
task := &createSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.CreateSnapshotRequest{
Name: "test_snapshot",
DbName: "default",
@@ -148,7 +148,7 @@ func TestCreateSnapshotTask_PreExecute_ProtectionExceedsMax(t *testing.T) {
func TestCreateSnapshotTask_PreExecute_ProtectionZero(t *testing.T) {
cache := &MetaCache{}
task := &createSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.CreateSnapshotRequest{
Name: "test_snapshot",
DbName: "default",
@@ -167,7 +167,7 @@ func TestCreateSnapshotTask_PreExecute_ProtectionZero(t *testing.T) {
func TestCreateSnapshotTask_PreExecute_ProtectionValid(t *testing.T) {
cache := &MetaCache{}
task := &createSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.CreateSnapshotRequest{
Name: "test_snapshot",
DbName: "default",
@@ -285,7 +285,7 @@ func TestDropSnapshotTask_OnEnqueue_Success(t *testing.T) {
func TestDropSnapshotTask_PreExecute(t *testing.T) {
cache := &MetaCache{}
task := &dropSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.DropSnapshotRequest{
Name: "test_snapshot",
DbName: "default",
@@ -318,7 +318,7 @@ func TestDropSnapshotTask_PreExecute_MissingCollectionName(t *testing.T) {
func TestDropSnapshotTask_PreExecute_CollectionNotFound(t *testing.T) {
cache := &MetaCache{}
task := &dropSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.DropSnapshotRequest{
Name: "test_snapshot",
DbName: "default",
@@ -395,7 +395,7 @@ func TestDescribeSnapshotTask_Execute_Success(t *testing.T) {
mockMixCoord := NewMixCoordMock()
cache := &MetaCache{}
task := &describeSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.DescribeSnapshotRequest{
Name: "test_snapshot",
},
@@ -468,7 +468,7 @@ func TestDescribeSnapshotTask_Execute_CollectionNameResolutionError(t *testing.T
mockMixCoord := NewMixCoordMock()
cache := &MetaCache{}
task := &describeSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.DescribeSnapshotRequest{
Name: "test_snapshot",
},
@@ -503,7 +503,7 @@ func TestDescribeSnapshotTask_Execute_PartitionNameResolutionError(t *testing.T)
mockMixCoord := NewMixCoordMock()
cache := &MetaCache{}
task := &describeSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.DescribeSnapshotRequest{
Name: "test_snapshot",
},
@@ -561,7 +561,7 @@ func TestListSnapshotsTask_OnEnqueue_Success(t *testing.T) {
func TestListSnapshotsTask_PreExecute_Success(t *testing.T) {
cache := &MetaCache{}
task := &listSnapshotsTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.ListSnapshotsRequest{
DbName: "default",
CollectionName: "test_collection",
@@ -702,7 +702,7 @@ func TestCreateSnapshotTask_FullLifecycle(t *testing.T) {
mockMixCoord := NewMixCoordMock()
cache := &MetaCache{}
task := &createSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.CreateSnapshotRequest{
DbName: "default",
CollectionName: "test_collection",
@@ -748,7 +748,7 @@ func TestCreateSnapshotTask_FullLifecycle(t *testing.T) {
func TestCreateSnapshotTask_EmptyPartitionNames(t *testing.T) {
cache := &MetaCache{}
task := &createSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.CreateSnapshotRequest{
Name: "test_snapshot",
DbName: "default",
@@ -920,7 +920,7 @@ func TestSnapshotTasks_PreExecute_ValidNames(t *testing.T) {
t.Run(tc.name, func(t *testing.T) {
cache := &MetaCache{}
task := &createSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.CreateSnapshotRequest{
Name: tc.snapshotName,
DbName: "default",
@@ -945,7 +945,7 @@ func TestSnapshotTasks_PreExecute_ValidNames(t *testing.T) {
t.Run(tc.name, func(t *testing.T) {
cache := &MetaCache{}
task := &dropSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.DropSnapshotRequest{
Name: tc.snapshotName,
DbName: "default",
@@ -969,7 +969,7 @@ func TestSnapshotTasks_PreExecute_ValidNames(t *testing.T) {
t.Run(tc.name, func(t *testing.T) {
cache := &MetaCache{}
task := &describeSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.DescribeSnapshotRequest{
Name: tc.snapshotName,
DbName: "default",
@@ -1004,7 +1004,7 @@ func TestSnapshotTasks_PreExecute_ValidNames(t *testing.T) {
t.Run(tc.name, func(t *testing.T) {
cache := &MetaCache{}
task := &restoreSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.RestoreSnapshotRequest{
Name: tc.snapshotName,
CollectionName: tc.collectionName,
@@ -1028,7 +1028,7 @@ func TestSnapshotTasks_PreExecute_ValidNames(t *testing.T) {
func TestListSnapshotsTask_PreExecute_EmptyCollectionName(t *testing.T) {
cache := &MetaCache{}
task := &listSnapshotsTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.ListSnapshotsRequest{
DbName: "default",
CollectionName: "", // Empty collection name should be rejected
@@ -1116,7 +1116,7 @@ func TestRestoreSnapshotTask_PreExecute_MissingTargetCollectionName(t *testing.T
func TestRestoreSnapshotTask_PreExecute_GetCollectionIDError(t *testing.T) {
cache := &MetaCache{}
task := &restoreSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.RestoreSnapshotRequest{
Name: "valid_snapshot",
CollectionName: "source_collection",
@@ -1154,7 +1154,7 @@ func TestDescribeSnapshotTask_PreExecute_MissingCollectionName(t *testing.T) {
func TestDescribeSnapshotTask_PreExecute_GetCollectionIDError(t *testing.T) {
cache := &MetaCache{}
task := &describeSnapshotTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.DescribeSnapshotRequest{
Name: "valid_snapshot",
DbName: "default",
@@ -1177,7 +1177,7 @@ func TestDescribeSnapshotTask_PreExecute_GetCollectionIDError(t *testing.T) {
func TestListSnapshotsTask_PreExecute_GetDatabaseInfoError(t *testing.T) {
cache := &MetaCache{}
task := &listSnapshotsTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.ListSnapshotsRequest{
DbName: "nonexistent_db",
},
@@ -1196,7 +1196,7 @@ func TestListSnapshotsTask_PreExecute_GetDatabaseInfoError(t *testing.T) {
func TestListSnapshotsTask_PreExecute_GetCollectionIDError(t *testing.T) {
cache := &MetaCache{}
task := &listSnapshotsTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.ListSnapshotsRequest{
DbName: "default",
CollectionName: "nonexistent_collection",
@@ -1249,7 +1249,7 @@ func TestGetRestoreSnapshotStateTask_Execute_Success(t *testing.T) {
mockMixCoord := NewMixCoordMock()
cache := &MetaCache{}
task := &getRestoreSnapshotStateTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.GetRestoreSnapshotStateRequest{
JobId: 1,
},
@@ -1394,7 +1394,7 @@ func TestListRestoreSnapshotJobsTask_OnEnqueue(t *testing.T) {
func TestListRestoreSnapshotJobsTask_PreExecute_DbNameError(t *testing.T) {
cache := &MetaCache{}
task := &listRestoreSnapshotJobsTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.ListRestoreSnapshotJobsRequest{
DbName: "nonexistent_db",
},
@@ -1413,7 +1413,7 @@ func TestListRestoreSnapshotJobsTask_PreExecute_DbNameError(t *testing.T) {
func TestListRestoreSnapshotJobsTask_PreExecute_CollectionNameError(t *testing.T) {
cache := &MetaCache{}
task := &listRestoreSnapshotJobsTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.ListRestoreSnapshotJobsRequest{
DbName: "default",
CollectionName: "nonexistent_collection",
@@ -1436,7 +1436,7 @@ func TestListRestoreSnapshotJobsTask_PreExecute_CollectionNameError(t *testing.T
func TestListRestoreSnapshotJobsTask_PreExecute_BothDbAndCollectionSuccess(t *testing.T) {
cache := &MetaCache{}
task := &listRestoreSnapshotJobsTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.ListRestoreSnapshotJobsRequest{
DbName: "default",
CollectionName: "test_collection",
@@ -1460,7 +1460,7 @@ func TestListRestoreSnapshotJobsTask_PreExecute_BothDbAndCollectionSuccess(t *te
func TestListRestoreSnapshotJobsTask_PreExecute_OnlyDbName(t *testing.T) {
cache := &MetaCache{}
task := &listRestoreSnapshotJobsTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.ListRestoreSnapshotJobsRequest{
DbName: "default",
},
@@ -1493,7 +1493,7 @@ func TestListRestoreSnapshotJobsTask_Execute_Success(t *testing.T) {
mockMixCoord := NewMixCoordMock()
cache := &MetaCache{}
task := &listRestoreSnapshotJobsTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.ListRestoreSnapshotJobsRequest{},
mixCoord: mockMixCoord,
collectionID: 0,
@@ -1680,7 +1680,7 @@ func TestListRestoreSnapshotJobsTask_TaskInterface(t *testing.T) {
func TestPinSnapshotDataTask_PreExecute_Success(t *testing.T) {
cache := &MetaCache{}
task := &pinSnapshotDataTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.PinSnapshotDataRequest{
Name: "test_snapshot",
DbName: "default",
@@ -1757,7 +1757,7 @@ func TestPinSnapshotDataTask_PreExecute_TTLExceedsMax(t *testing.T) {
func TestPinSnapshotDataTask_PreExecute_TTLAtMaxBoundary(t *testing.T) {
cache := &MetaCache{}
task := &pinSnapshotDataTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
req: &milvuspb.PinSnapshotDataRequest{
Name: "test_snapshot",
DbName: "default",
internal/proxy/task_statistic.go
+6 -6
View File
@@ -110,11 +110,11 @@ func (g *getStatisticsTask) PreExecute(ctx context.Context) error {
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetStatistics-PreExecute")
defer sp.End()
collID, err := g.getMetaCache().GetCollectionID(ctx, g.request.GetDbName(), g.collectionName)
collID, err := g.GetMetaCache().GetCollectionID(ctx, g.request.GetDbName(), g.collectionName)
if err != nil { // err is not nil if collection not exists
return err
}
partIDs, err := getPartitionIDs(ctx, g.getMetaCache(), g.request.GetDbName(), g.collectionName, g.partitionNames)
partIDs, err := getPartitionIDs(ctx, g.GetMetaCache(), g.request.GetDbName(), g.collectionName, g.partitionNames)
if err != nil { // err is not nil if partition not exists
return err
}
@@ -131,7 +131,7 @@ func (g *getStatisticsTask) PreExecute(ctx context.Context) error {
}
// check if collection/partitions are loaded into query node
loaded, unloaded, err := checkFullLoaded(ctx, g.getMetaCache(), g.mixc, g.request.GetDbName(), g.collectionName, g.CollectionID, partIDs)
loaded, unloaded, err := checkFullLoaded(ctx, g.GetMetaCache(), g.mixc, g.request.GetDbName(), g.collectionName, g.CollectionID, partIDs)
log := mlog.With(
mlog.String("collectionName", g.collectionName),
mlog.Int64("collectionID", g.CollectionID),
@@ -643,7 +643,7 @@ func (g *getCollectionStatisticsTask) PreExecute(ctx context.Context) error {
}
func (g *getCollectionStatisticsTask) Execute(ctx context.Context) error {
collID, err := g.getMetaCache().GetCollectionID(ctx, g.GetDbName(), g.CollectionName)
collID, err := g.GetMetaCache().GetCollectionID(ctx, g.GetDbName(), g.CollectionName)
if err != nil {
return err
}
@@ -726,12 +726,12 @@ func (g *getPartitionStatisticsTask) PreExecute(ctx context.Context) error {
}
func (g *getPartitionStatisticsTask) Execute(ctx context.Context) error {
collID, err := g.getMetaCache().GetCollectionID(ctx, g.GetDbName(), g.CollectionName)
collID, err := g.GetMetaCache().GetCollectionID(ctx, g.GetDbName(), g.CollectionName)
if err != nil {
return err
}
g.collectionID = collID
partitionID, err := g.getMetaCache().GetPartitionID(ctx, g.GetDbName(), g.CollectionName, g.PartitionName)
partitionID, err := g.GetMetaCache().GetPartitionID(ctx, g.GetDbName(), g.CollectionName, g.PartitionName)
if err != nil {
return err
}
internal/proxy/task_statistic_test.go
+1 -1
View File
@@ -175,7 +175,7 @@ func (s *StatisticTaskSuite) TestStatisticTask_Timeout() {
func (s *StatisticTaskSuite) getStatisticsTask(ctx context.Context) *getStatisticsTask {
return &getStatisticsTask{
baseTask: baseTask{metaCache: s.metaCache},
baseTask: baseTask{MetaCache: s.metaCache},
Condition: NewTaskCondition(ctx),
ctx: ctx,
collectionName: s.collectionName,
internal/proxy/task_test.go
+66 -66
View File
@@ -571,7 +571,7 @@ func TestAlterCollection_AllowInsertAutoID_Validation(t *testing.T) {
assert.NoError(t, err)
task := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_AlterCollectionField},
DbName: dbName,
@@ -591,7 +591,7 @@ func TestAlterCollection_AllowInsertAutoID_Validation(t *testing.T) {
assert.NoError(t, err)
task := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_AlterCollectionField},
DbName: dbName,
@@ -2230,7 +2230,7 @@ func TestHasCollectionTask(t *testing.T) {
// CreateCollection
task := &hasCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
HasCollectionRequest: &milvuspb.HasCollectionRequest{
Base: &commonpb.MsgBase{
@@ -2769,7 +2769,7 @@ func TestCreatePartitionTask(t *testing.T) {
mixCoord: rc,
result: nil,
}
task.metaCache = mockCache
task.MetaCache = mockCache
task.OnEnqueue()
task.PreExecute(ctx)
@@ -2841,7 +2841,7 @@ func TestDropPartitionTask(t *testing.T) {
mixCoord: mixc,
result: nil,
}
task.metaCache = mockCache
task.MetaCache = mockCache
task.OnEnqueue()
task.PreExecute(ctx)
@@ -2875,7 +2875,7 @@ func TestDropPartitionTask(t *testing.T) {
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(mustNewSchemaInfo(&schemapb.CollectionSchema{}), nil)
task.metaCache = mockCache
task.MetaCache = mockCache
task.PartitionName = "partition1"
err = task.PreExecute(ctx)
assert.Error(t, err)
@@ -2902,7 +2902,7 @@ func TestDropPartitionTask(t *testing.T) {
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(mustNewSchemaInfo(&schemapb.CollectionSchema{}), nil)
task.metaCache = mockCache
task.MetaCache = mockCache
err = task.PreExecute(ctx)
assert.NoError(t, err)
})
@@ -2928,7 +2928,7 @@ func TestDropPartitionTask(t *testing.T) {
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(mustNewSchemaInfo(&schemapb.CollectionSchema{}), nil)
task.metaCache = mockCache
task.MetaCache = mockCache
err = task.PreExecute(ctx)
assert.Error(t, err)
})
@@ -3115,7 +3115,7 @@ func TestTask_Int64PrimaryKey(t *testing.T) {
t.Run("insert", func(t *testing.T) {
hash := testutils.GenerateHashKeys(nb)
task := &insertTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
insertMsg: &BaseInsertTask{
BaseMsg: msgstream.BaseMsg{
HashValues: hash,
@@ -3166,7 +3166,7 @@ func TestTask_Int64PrimaryKey(t *testing.T) {
t.Run("simple delete", func(t *testing.T) {
task := &deleteTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
req: &milvuspb.DeleteRequest{
CollectionName: collectionName,
@@ -3337,7 +3337,7 @@ func TestTask_VarCharPrimaryKey(t *testing.T) {
t.Run("insert", func(t *testing.T) {
hash := testutils.GenerateHashKeys(nb)
task := &insertTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
insertMsg: &BaseInsertTask{
BaseMsg: msgstream.BaseMsg{
HashValues: hash,
@@ -3391,7 +3391,7 @@ func TestTask_VarCharPrimaryKey(t *testing.T) {
t.Run("upsert", func(t *testing.T) {
hash := testutils.GenerateHashKeys(nb)
task := &upsertTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
upsertMsg: &msgstream.UpsertMsg{
InsertMsg: &BaseInsertTask{
BaseMsg: msgstream.BaseMsg{
@@ -3474,7 +3474,7 @@ func TestTask_VarCharPrimaryKey(t *testing.T) {
t.Run("simple delete", func(t *testing.T) {
task := &deleteTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
req: &milvuspb.DeleteRequest{
CollectionName: collectionName,
@@ -3563,7 +3563,7 @@ func Test_createIndexTask_getIndexedFieldAndFunction(t *testing.T) {
},
}), nil)
cit.metaCache = cache
cit.MetaCache = cache
err := cit.getIndexedFieldAndFunction(context.Background())
assert.NoError(t, err)
assert.Equal(t, fieldName, cit.fieldSchema.GetName())
@@ -3576,7 +3576,7 @@ func Test_createIndexTask_getIndexedFieldAndFunction(t *testing.T) {
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(nil, errors.New("mock"))
cit.metaCache = cache
cit.MetaCache = cache
err := cit.getIndexedFieldAndFunction(context.Background())
assert.Error(t, err)
})
@@ -3595,7 +3595,7 @@ func Test_createIndexTask_getIndexedFieldAndFunction(t *testing.T) {
otherField,
},
}), nil)
cit.metaCache = cache
cit.MetaCache = cache
err := cit.getIndexedFieldAndFunction(context.Background())
assert.Error(t, err)
})
@@ -3762,7 +3762,7 @@ func Test_createIndexTask_PreExecute(t *testing.T) {
mock.AnythingOfType("string"),
mock.AnythingOfType("int64"),
).Return(&collectionInfo{}, nil)
cit.metaCache = cache
cit.MetaCache = cache
cit.req.ExtraParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
@@ -3787,7 +3787,7 @@ func Test_createIndexTask_PreExecute(t *testing.T) {
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(0), errors.New("mock"))
cit.metaCache = cache
cit.MetaCache = cache
assert.Error(t, cit.PreExecute(context.Background()))
})
@@ -3804,7 +3804,7 @@ func Test_createIndexTask_PreExecute(t *testing.T) {
mock.AnythingOfType("string"),
mock.AnythingOfType("int64"),
).Return(&collectionInfo{}, nil)
cit.metaCache = cache
cit.MetaCache = cache
for i := 0; i < 256; i++ {
cit.req.IndexName += "a"
@@ -3835,7 +3835,7 @@ func Test_dropCollectionTask_PreExecute(t *testing.T) {
CollectionName: "valid", // invalid
}}
dct.metaCache = &MetaCache{}
dct.MetaCache = &MetaCache{}
mockGetCollectionID := mockey.Mock((*MetaCache).GetCollectionID).To(func(ctx context.Context, dbName, collectionName string) (UniqueID, error) {
return 1, nil
}).Build()
@@ -3929,7 +3929,7 @@ func Test_truncateCollectionTask_PreExecute(t *testing.T) {
tct := &truncateCollectionTask{TruncateCollectionRequest: &milvuspb.TruncateCollectionRequest{
Base: &commonpb.MsgBase{}, DbName: dbName, CollectionName: regularName,
}}
tct.metaCache = cache
tct.MetaCache = cache
assert.NoError(t, tct.PreExecute(ctx))
})
@@ -3947,7 +3947,7 @@ func Test_truncateCollectionTask_PreExecute(t *testing.T) {
tct := &truncateCollectionTask{TruncateCollectionRequest: &milvuspb.TruncateCollectionRequest{
Base: &commonpb.MsgBase{}, DbName: dbName, CollectionName: externalName,
}}
tct.metaCache = cache
tct.MetaCache = cache
err := tct.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "external collection")
@@ -4037,7 +4037,7 @@ func Test_loadCollectionTask_Execute(t *testing.T) {
result: nil,
collectionID: 0,
}
lct.metaCache = cache
lct.MetaCache = cache
t.Run("indexcoord describe index error", func(t *testing.T) {
err := lct.Execute(ctx)
@@ -4146,7 +4146,7 @@ func TestLoadCollectionTaskExecuteTextRequiresStorageV3(t *testing.T) {
},
ctx: context.Background(),
}
task.metaCache = cache
task.MetaCache = cache
err := task.Execute(context.Background())
assert.Error(t, err)
@@ -4196,7 +4196,7 @@ func Test_loadPartitionTask_Execute(t *testing.T) {
result: nil,
collectionID: 0,
}
lpt.metaCache = cache
lpt.MetaCache = cache
t.Run("indexcoord describe index error", func(t *testing.T) {
err := lpt.Execute(ctx)
@@ -4387,7 +4387,7 @@ func TestTransferReplicaTask(t *testing.T) {
ctx: ctx,
mixCoord: rc,
}
task.metaCache = cache
task.MetaCache = cache
task.OnEnqueue()
task.PreExecute(ctx)
@@ -4486,7 +4486,7 @@ func TestDescribeResourceGroupTask(t *testing.T) {
ctx: ctx,
mixCoord: rc,
}
task.metaCache = cache
task.MetaCache = cache
task.OnEnqueue()
task.PreExecute(ctx)
@@ -4539,7 +4539,7 @@ func TestDescribeResourceGroupTaskFailed(t *testing.T) {
ctx: ctx,
mixCoord: rc,
}
task.metaCache = cache
task.MetaCache = cache
task.OnEnqueue()
task.PreExecute(ctx)
@@ -4649,7 +4649,7 @@ func TestCreateCollectionTaskWithPartitionKey(t *testing.T) {
result: nil,
schema: nil,
}
task.metaCache = mockCache
task.MetaCache = mockCache
t.Run("PreExecute", func(t *testing.T) {
defer Params.Reset(Params.RootCoordCfg.MaxPartitionNum.Key)
@@ -4774,7 +4774,7 @@ func TestCreateCollectionTaskWithPartitionKey(t *testing.T) {
assert.Equal(t, task.GetNumPartitions(), int64(len(partitionNames)))
createPartitionTask := &createPartitionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
CreatePartitionRequest: &milvuspb.CreatePartitionRequest{
Base: &commonpb.MsgBase{
@@ -4792,7 +4792,7 @@ func TestCreateCollectionTaskWithPartitionKey(t *testing.T) {
assert.Error(t, err)
dropPartitionTask := &dropPartitionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
DropPartitionRequest: &milvuspb.DropPartitionRequest{
Base: &commonpb.MsgBase{
@@ -4810,7 +4810,7 @@ func TestCreateCollectionTaskWithPartitionKey(t *testing.T) {
assert.Error(t, err)
loadPartitionTask := &loadPartitionsTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
LoadPartitionsRequest: &milvuspb.LoadPartitionsRequest{
Base: &commonpb.MsgBase{
@@ -4827,7 +4827,7 @@ func TestCreateCollectionTaskWithPartitionKey(t *testing.T) {
assert.Error(t, err)
releasePartitionsTask := &releasePartitionsTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
ReleasePartitionsRequest: &milvuspb.ReleasePartitionsRequest{
Base: &commonpb.MsgBase{
@@ -4923,7 +4923,7 @@ func TestPartitionKey(t *testing.T) {
t.Run("Insert", func(t *testing.T) {
it := &insertTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
insertMsg: &BaseInsertTask{
BaseMsg: msgstream.BaseMsg{},
InsertRequest: &msgpb.InsertRequest{
@@ -4973,7 +4973,7 @@ func TestPartitionKey(t *testing.T) {
t.Run("Upsert", func(t *testing.T) {
hash := testutils.GenerateHashKeys(nb)
ut := &upsertTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
Condition: NewTaskCondition(ctx),
baseMsg: msgstream.BaseMsg{
@@ -5012,7 +5012,7 @@ func TestPartitionKey(t *testing.T) {
t.Run("delete", func(t *testing.T) {
dt := &deleteTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
req: &milvuspb.DeleteRequest{
CollectionName: collectionName,
@@ -5036,7 +5036,7 @@ func TestPartitionKey(t *testing.T) {
t.Run("search", func(t *testing.T) {
searchTask := &searchTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
SearchRequest: &internalpb.SearchRequest{
Base: &commonpb.MsgBase{},
@@ -5057,7 +5057,7 @@ func TestPartitionKey(t *testing.T) {
t.Run("query", func(t *testing.T) {
queryTask := &queryTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
RetrieveRequest: &internalpb.RetrieveRequest{
QueryLabel: "query",
@@ -5144,7 +5144,7 @@ func TestDefaultPartition(t *testing.T) {
t.Run("Insert", func(t *testing.T) {
it := &insertTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
insertMsg: &BaseInsertTask{
BaseMsg: msgstream.BaseMsg{},
InsertRequest: &msgpb.InsertRequest{
@@ -5190,7 +5190,7 @@ func TestDefaultPartition(t *testing.T) {
t.Run("Upsert", func(t *testing.T) {
hash := testutils.GenerateHashKeys(nb)
ut := &upsertTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: ctx,
Condition: NewTaskCondition(ctx),
baseMsg: msgstream.BaseMsg{
@@ -5225,7 +5225,7 @@ func TestDefaultPartition(t *testing.T) {
t.Run("delete", func(t *testing.T) {
dt := &deleteTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
Condition: NewTaskCondition(ctx),
req: &milvuspb.DeleteRequest{
CollectionName: collectionName,
@@ -5557,7 +5557,7 @@ func TestAlterCollectionCheckLoaded(t *testing.T) {
}
task := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
@@ -5604,7 +5604,7 @@ func TestAlterCollectionTaskValidateTTLAndTTLField(t *testing.T) {
col := "alter_ttl_seconds_then_field_" + funcutil.GenRandomStr()
createCollectionWithProps(col, []*commonpb.KeyValuePair{{Key: common.CollectionTTLConfigKey, Value: "3600"}})
task := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
@@ -5621,7 +5621,7 @@ func TestAlterCollectionTaskValidateTTLAndTTLField(t *testing.T) {
col := "alter_ttl_field_then_seconds_" + funcutil.GenRandomStr()
createCollectionWithProps(col, []*commonpb.KeyValuePair{{Key: common.CollectionTTLFieldKey, Value: "ttl"}})
task := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
@@ -5638,7 +5638,7 @@ func TestAlterCollectionTaskValidateTTLAndTTLField(t *testing.T) {
col := "alter_ttl_field_invalid_" + funcutil.GenRandomStr()
createCollectionWithProps(col, nil)
task := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
@@ -5654,7 +5654,7 @@ func TestAlterCollectionTaskValidateTTLAndTTLField(t *testing.T) {
col := "alter_ttl_field_invalid_type_" + funcutil.GenRandomStr()
createCollectionWithProps(col, nil)
task := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
@@ -5670,7 +5670,7 @@ func TestAlterCollectionTaskValidateTTLAndTTLField(t *testing.T) {
col := "alter_delete_ttl_seconds_" + funcutil.GenRandomStr()
createCollectionWithProps(col, []*commonpb.KeyValuePair{{Key: common.CollectionTTLConfigKey, Value: "3600"}})
task := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
@@ -5686,7 +5686,7 @@ func TestAlterCollectionTaskValidateTTLAndTTLField(t *testing.T) {
col := "alter_delete_ttl_field_" + funcutil.GenRandomStr()
createCollectionWithProps(col, []*commonpb.KeyValuePair{{Key: common.CollectionTTLFieldKey, Value: "ttl"}})
task := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
@@ -5740,7 +5740,7 @@ func TestAlterCollectionTaskValidateDescription(t *testing.T) {
runAlter := func(collectionName string, properties []*commonpb.KeyValuePair) error {
task := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
@@ -5914,7 +5914,7 @@ func TestTaskPartitionKeyIsolation(t *testing.T) {
}
return &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
@@ -5983,7 +5983,7 @@ func TestTaskPartitionKeyIsolation(t *testing.T) {
colName := collectionName + "AlterNoIso"
createIsoCollection(colName, true, false, true)
alterTask := alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
@@ -6028,7 +6028,7 @@ func TestTaskPartitionKeyIsolation(t *testing.T) {
createIsoCollection(colName, true, true, false)
mockVectorIndexForCollection(t, ctx, qc, colName)
alterTask := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
@@ -6045,7 +6045,7 @@ func TestTaskPartitionKeyIsolation(t *testing.T) {
colName := collectionName + "DeleteIsoNoIdx"
createIsoCollection(colName, true, true, false)
alterTask := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
@@ -6064,7 +6064,7 @@ func TestTaskPartitionKeyIsolation(t *testing.T) {
createIsoCollection(colName, true, true, false)
mockVectorIndexForCollection(t, ctx, qc, colName)
alterTask := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
@@ -6120,7 +6120,7 @@ func TestAlterCollectionQueryMode(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createCollection(colName, false)
alterTask := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
@@ -6136,7 +6136,7 @@ func TestAlterCollectionQueryMode(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createCollection(colName, true)
alterTask := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
@@ -6153,7 +6153,7 @@ func TestAlterCollectionQueryMode(t *testing.T) {
createCollection(colName, false)
mockVectorIndexForCollection(t, ctx, qc, colName)
alterTask := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
@@ -6171,7 +6171,7 @@ func TestAlterCollectionQueryMode(t *testing.T) {
createCollection(colName, true)
mockVectorIndexForCollection(t, ctx, qc, colName)
alterTask := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
@@ -6188,7 +6188,7 @@ func TestAlterCollectionQueryMode(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createCollection(colName, true)
alterTask := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
@@ -6205,7 +6205,7 @@ func TestAlterCollectionQueryMode(t *testing.T) {
createCollection(colName, true)
mockVectorIndexForCollection(t, ctx, qc, colName)
alterTask := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
@@ -6352,7 +6352,7 @@ func TestAlterCollectionFieldCheckLoaded(t *testing.T) {
// update property "mmap.enabled" but the collection is loaded
task := &alterCollectionFieldTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionFieldRequest: &milvuspb.AlterCollectionFieldRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
@@ -6365,7 +6365,7 @@ func TestAlterCollectionFieldCheckLoaded(t *testing.T) {
// delete property "mmap.enabled" but the collection is loaded
task = &alterCollectionFieldTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionFieldRequest: &milvuspb.AlterCollectionFieldRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
@@ -6688,7 +6688,7 @@ func TestAlterCollectionField(t *testing.T) {
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
task := &alterCollectionFieldTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionFieldRequest: &milvuspb.AlterCollectionFieldRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
@@ -6717,7 +6717,7 @@ func TestAlterCollectionField(t *testing.T) {
defer paramtable.Get().Reset(paramtable.Get().ProxyCfg.MaxArrayCapacity.Key)
task := &alterCollectionFieldTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionFieldRequest: &milvuspb.AlterCollectionFieldRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
@@ -6853,7 +6853,7 @@ func TestCreateCollectionTaskWithStructArrayField(t *testing.T) {
result: nil,
schema: nil,
}
task.metaCache = mockCache
task.MetaCache = mockCache
t.Run("create collection with struct array field", func(t *testing.T) {
err := task.OnEnqueue()
@@ -7169,7 +7169,7 @@ func TestAlterCollection_AllowInsertAutoID_AutoIDFalse(t *testing.T) {
qc.CreateCollection(ctx, createColReq)
task := &alterCollectionTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
internal/proxy/task_upsert.go
+14 -14
View File
@@ -134,22 +134,22 @@ func (it *upsertTask) refreshMutationResultCounts() {
func (it *upsertTask) getPChanStats() (map[pChan]pChanStatistics, error) {
ret := make(map[pChan]pChanStatistics)
channels := it.getChannels()
channels := it.GetChannels()
beginTs := it.BeginTs()
endTs := it.EndTs()
for _, channel := range channels {
ret[channel] = pChanStatistics{
minTs: beginTs,
maxTs: endTs,
MinTs: beginTs,
MaxTs: endTs,
}
}
return ret, nil
}
func (it *upsertTask) setChannels() error {
collID, err := it.getMetaCache().GetCollectionID(it.ctx, it.req.GetDbName(), it.req.CollectionName)
func (it *upsertTask) SetChannels() error {
collID, err := it.GetMetaCache().GetCollectionID(it.ctx, it.req.GetDbName(), it.req.CollectionName)
if err != nil {
return err
}
@@ -161,7 +161,7 @@ func (it *upsertTask) setChannels() error {
return nil
}
func (it *upsertTask) getChannels() []pChan {
func (it *upsertTask) GetChannels() []pChan {
return it.pChannels
}
@@ -217,7 +217,7 @@ func retrieveByPKs(ctx context.Context, t *upsertTask, ids *schemapb.IDs, output
log.Warn(ctx, "Invalid partition name", mlog.String("partitionName", partName), mlog.Err(err))
return nil, segcore.StorageCost{}, err
}
partID, err := t.getMetaCache().GetPartitionID(ctx, t.req.GetDbName(), t.req.GetCollectionName(), partName)
partID, err := t.GetMetaCache().GetPartitionID(ctx, t.req.GetDbName(), t.req.GetCollectionName(), partName)
if err != nil {
log.Warn(ctx, "Failed to get partition id", mlog.String("partitionName", partName), mlog.Err(err))
return nil, segcore.StorageCost{}, err
@@ -230,7 +230,7 @@ func retrieveByPKs(ctx context.Context, t *upsertTask, ids *schemapb.IDs, output
plan.Namespace = namespaceForPlan(t.schema.CollectionSchema, t.req.Namespace)
qt := &queryTask{
baseTask: baseTask{
metaCache: t.getMetaCache(),
MetaCache: t.GetMetaCache(),
},
ctx: t.ctx,
Condition: NewTaskCondition(t.ctx),
@@ -1555,7 +1555,7 @@ func (it *upsertTask) deletePreExecute(ctx context.Context) error {
log.Warn(ctx, "Invalid partition name", mlog.String("partitionName", partName), mlog.Err(err))
return err
}
partID, err := it.getMetaCache().GetPartitionID(ctx, it.req.GetDbName(), collName, partName)
partID, err := it.GetMetaCache().GetPartitionID(ctx, it.req.GetDbName(), collName, partName)
if err != nil {
log.Warn(ctx, "Failed to get partition id", mlog.String("collectionName", collName), mlog.String("partitionName", partName), mlog.Err(err))
return err
@@ -1590,14 +1590,14 @@ func (it *upsertTask) PreExecute(ctx context.Context) error {
}
// check collection exists
collID, err := it.getMetaCache().GetCollectionID(context.Background(), it.req.GetDbName(), collectionName)
collID, err := it.GetMetaCache().GetCollectionID(context.Background(), it.req.GetDbName(), collectionName)
if err != nil {
log.Warn(ctx, "fail to get collection id", mlog.Err(err))
return err
}
it.collectionID = collID
colInfo, err := it.getMetaCache().GetCollectionInfo(ctx, it.req.GetDbName(), collectionName, collID)
colInfo, err := it.GetMetaCache().GetCollectionInfo(ctx, it.req.GetDbName(), collectionName, collID)
if err != nil {
log.Warn(ctx, "fail to get collection info", mlog.Err(err))
return err
@@ -1614,7 +1614,7 @@ func (it *upsertTask) PreExecute(ctx context.Context) error {
}
}
schema, err := it.getMetaCache().GetCollectionSchema(ctx, it.req.GetDbName(), collectionName)
schema, err := it.GetMetaCache().GetCollectionSchema(ctx, it.req.GetDbName(), collectionName)
if err != nil {
log.Warn(ctx, "Failed to get collection schema",
mlog.String("collectionName", collectionName),
@@ -1647,7 +1647,7 @@ func (it *upsertTask) PreExecute(ctx context.Context) error {
it.req.PartitionName = partitionName
}
it.partitionKeyMode, err = isPartitionKeyMode(ctx, it.getMetaCache(), it.req.GetDbName(), collectionName)
it.partitionKeyMode, err = isPartitionKeyMode(ctx, it.GetMetaCache(), it.req.GetDbName(), collectionName)
if err != nil {
log.Warn(ctx, "check partition key mode failed",
mlog.String("collectionName", collectionName),
@@ -1663,7 +1663,7 @@ func (it *upsertTask) PreExecute(ctx context.Context) error {
// insert to _default partition
partitionTag := it.req.GetPartitionName()
if len(partitionTag) <= 0 {
pinfo, err := it.getMetaCache().GetPartitionInfo(ctx, it.req.GetDbName(), collectionName, "")
pinfo, err := it.GetMetaCache().GetPartitionInfo(ctx, it.req.GetDbName(), collectionName, "")
if err != nil {
log.Warn(ctx, "get partition info failed", mlog.String("collectionName", collectionName), mlog.Err(err))
return err
internal/proxy/task_upsert_streaming.go
+3 -3
View File
@@ -201,7 +201,7 @@ func (ut *upsertTask) packInsertMessage(ctx context.Context, ez *message.CipherC
defer tr.Elapse("insert execute done when insertExecute")
collectionName := ut.upsertMsg.InsertMsg.CollectionName
collID, err := ut.getMetaCache().GetCollectionID(ctx, ut.req.GetDbName(), collectionName)
collID, err := ut.GetMetaCache().GetCollectionID(ctx, ut.req.GetDbName(), collectionName)
if err != nil {
return nil, err
}
@@ -231,9 +231,9 @@ func (ut *upsertTask) packInsertMessage(ctx context.Context, ez *message.CipherC
// start to repack insert data
var msgs []message.MutableMessage
if ut.partitionKeys == nil {
msgs, err = repackInsertDataForStreamingService(ut.TraceCtx(), ut.getMetaCache(), channelNames, ut.upsertMsg.InsertMsg, ut.result, ez, ut.schemaVersion, ut.partialUpdateCASGroups)
msgs, err = repackInsertDataForStreamingService(ut.TraceCtx(), ut.GetMetaCache(), channelNames, ut.upsertMsg.InsertMsg, ut.result, ez, ut.schemaVersion, ut.partialUpdateCASGroups)
} else {
msgs, err = repackInsertDataWithPartitionKeyForStreamingService(ut.TraceCtx(), ut.getMetaCache(), channelNames, ut.upsertMsg.InsertMsg, ut.result, ut.partitionKeys, ez, ut.schema.CollectionSchema, ut.schemaVersion, ut.partialUpdateCASGroups)
msgs, err = repackInsertDataWithPartitionKeyForStreamingService(ut.TraceCtx(), ut.GetMetaCache(), channelNames, ut.upsertMsg.InsertMsg, ut.result, ut.partitionKeys, ez, ut.schema.CollectionSchema, ut.schemaVersion, ut.partialUpdateCASGroups)
}
if err != nil {
log.Warn(ctx, "assign segmentID and repack insert data failed", mlog.Err(err))
internal/proxy/task_upsert_test.go
+10 -10
View File
@@ -324,7 +324,7 @@ func TestUpsertTask_CheckAligned(t *testing.T) {
}
func TestUpsertTask(t *testing.T) {
t.Run("test getChannels", func(t *testing.T) {
t.Run("test GetChannels", func(t *testing.T) {
collectionID := UniqueID(0)
collectionName := "col-0"
channels := []pChan{"mock-chan-0", "mock-chan-1"}
@@ -338,16 +338,16 @@ func TestUpsertTask(t *testing.T) {
chMgr := channelmgr.NewMockChannelsMgr(t)
chMgr.EXPECT().GetChannels(mock.Anything).Return(channels, nil)
ut := upsertTask{
baseTask: baseTask{metaCache: cache},
baseTask: baseTask{MetaCache: cache},
ctx: context.Background(),
req: &milvuspb.UpsertRequest{
CollectionName: collectionName,
},
chMgr: chMgr,
}
err := ut.setChannels()
err := ut.SetChannels()
assert.NoError(t, err)
resChannels := ut.getChannels()
resChannels := ut.GetChannels()
assert.ElementsMatch(t, channels, resChannels)
assert.ElementsMatch(t, channels, ut.pChannels)
})
@@ -456,7 +456,7 @@ func TestUpsertTask_Function(t *testing.T) {
defer idAllocator.Close()
assert.NoError(t, err)
task := upsertTask{
baseTask: baseTask{metaCache: &MetaCache{}},
baseTask: baseTask{MetaCache: &MetaCache{}},
ctx: context.Background(),
req: &milvuspb.UpsertRequest{
CollectionName: collectionName,
@@ -501,7 +501,7 @@ func TestUpsertTaskForSchemaMismatch(t *testing.T) {
t.Run("schema ts mismatch", func(t *testing.T) {
ut := upsertTask{
baseTask: baseTask{metaCache: mockCache},
baseTask: baseTask{MetaCache: mockCache},
ctx: ctx,
req: &milvuspb.UpsertRequest{
CollectionName: "col-0",
@@ -530,7 +530,7 @@ func createTestUpdateTask() *upsertTask {
mcClient := &grpcmixcoordclient.Client{}
upsertTask := &upsertTask{
baseTask: baseTask{metaCache: &MetaCache{}},
baseTask: baseTask{MetaCache: &MetaCache{}},
Condition: NewTaskCondition(context.Background()),
req: &milvuspb.UpsertRequest{
Base: commonpbutil.NewMsgBase(
@@ -1424,7 +1424,7 @@ func TestInsertTaskExecuteSelectsPartitionRouting(t *testing.T) {
t.Cleanup(func() { streaming.SetWALForTest(oldWAL) })
task := &insertTask{
baseTask: baseTask{metaCache: &MetaCache{}},
baseTask: baseTask{MetaCache: &MetaCache{}},
ctx: context.Background(),
insertMsg: &msgstream.InsertMsg{InsertRequest: &msgpb.InsertRequest{
Base: &commonpb.MsgBase{},
@@ -4148,7 +4148,7 @@ func TestUpsertTask_queryPreExecute_EmptyDataArray(t *testing.T) {
defer idAllocator.Close()
task := &upsertTask{
baseTask: baseTask{metaCache: &MetaCache{}},
baseTask: baseTask{MetaCache: &MetaCache{}},
ctx: ctx,
schema: schema,
req: &milvuspb.UpsertRequest{
@@ -4265,7 +4265,7 @@ func TestUpsertTask_queryPreExecute_EmptyDataArray(t *testing.T) {
defer idAllocator.Close()
task := &upsertTask{
baseTask: baseTask{metaCache: &MetaCache{}},
baseTask: baseTask{MetaCache: &MetaCache{}},
ctx: ctx,
schema: schema,
req: &milvuspb.UpsertRequest{
internal/proxy/taskmodel/OWNERS
+8
View File
@@ -0,0 +1,8 @@
reviewers:
- weiliu1031
- SpadeA-Tang
- czs007
- congqixia
approvers:
- maintainers
internal/proxy/taskmodel/README.md
+107
View File
@@ -0,0 +1,107 @@
# TaskModel Package
The `taskmodel` package is the shared task model layer of the Milvus proxy. It
defines the interfaces and value types that decouple the task scheduler from
the concrete task implementations that remain in the `internal/proxy` root
package.
This package was extracted from the proxy root package (issue #44761): the
`task` interface, `baseTask`, the `Condition`/`TaskCondition` primitives, the
TSO allocator interface, and the channel/timestamp value types moved here
verbatim. The concrete task structs (search/query/insert/delete/upsert and the
DDL tasks) still live in the root package and implement these interfaces.
## Overview
The proxy's scheduler consumes the `Task` / `DMLTask` interfaces, so it never
needs to know about the concrete task types. That one-way dependency
(`scheduler -> taskmodel`) is what allowed the scheduler to be extracted into
its own package without dragging every task implementation along.
## Responsibilities
1. **`Task`** — the contract every proxy task implements (lifecycle:
`PreExecute`/`Execute`/`PostExecute`, timing bookkeeping, `GetMetaCache`,
`WaitToFinish`/`Notify`).
2. **`DMLTask`** — the `Task` variant for insert/delete/upsert tasks, which
resolve their physical channels (`SetChannels`/`GetChannels`) before
enqueueing.
3. **`BaseTask`** — embedded by concrete tasks to provide the shared meta-cache
accessor and queue/execute timing fields.
4. **`Condition` / `TaskCondition`** — the notification primitive tasks embed to
implement `WaitToFinish`/`Notify`/`Ctx`.
5. **`TsoAllocator`** — the timestamp-allocation interface implemented by the
proxy's timestamp allocator and consumed by the scheduler.
6. **Value types** — `UniqueID`, `Timestamp`, `VChan`, `PChan`,
`PChanStatistics`, and the `BaseInsertTask` alias.
## Architecture
```
┌────────────────────────────────────────────────────────────┐
│ taskmodel │
│ │
│ Task ◄───────────── DMLTask │
│ ▲ ▲ │
│ │ embeds │ embeds │
│ BaseTask ── Condition / TaskCondition │
│ │
│ TsoAllocator UniqueID / Timestamp / VChan / PChan │
│ PChanStatistics BaseInsertTask │
└────────────────────────────────────────────────────────────┘
```
### Key types
```go
type Task interface {
TraceCtx() context.Context
ID() UniqueID
SetID(uid UniqueID)
Name() string
Type() commonpb.MsgType
BeginTs() Timestamp
EndTs() Timestamp
SetTs(ts Timestamp)
OnEnqueue() error
PreExecute(ctx context.Context) error
Execute(ctx context.Context) error
PostExecute(ctx context.Context) error
WaitToFinish() error
Notify(err error)
CanSkipAllocTimestamp() bool
GetMetaCache() metacache.Cache
SetOnEnqueueTime()
GetDurationInQueue() time.Duration
IsSubTask() bool
SetExecutingTime()
GetDurationInExecuting() time.Duration
}
type DMLTask interface {
Task
SetChannels() error
GetChannels() []PChan
}
type TsoAllocator interface {
AllocOne(ctx context.Context) (Timestamp, error)
}
```
## Dependency rule
`taskmodel` imports only `metacache` (for the `Cache` type), `msgstream` (for
the `BaseInsertTask` alias), the proto packages, and `pkg/v3`. It never imports
the proxy root package, so concrete tasks can implement these interfaces
without introducing an import cycle.
## Related components
- **scheduler** (`internal/proxy/scheduler/`): the sole consumer of `Task` /
`DMLTask` / `TsoAllocator`.
- **proxy root** (`internal/proxy/`): keeps the ~30 concrete task structs,
which embed `BaseTask` and `Condition` and return the shared task-name
constants.
- **metacache** (`internal/proxy/metacache/`): provides the `Cache` interface
exposed through `Task.GetMetaCache()`.
internal/proxy/condition.go → internal/proxy/taskmodel/condition.go
+1 -1
View File
@@ -14,7 +14,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package proxy
package taskmodel
import (
"context"
internal/proxy/taskmodel/task.go
+104
View File
@@ -0,0 +1,104 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package taskmodel
import (
"context"
"time"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus/internal/proxy/metacache"
"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
)
const (
SearchTaskName = "SearchTask"
QueryTaskName = "QueryTask"
)
// Task is the interface implemented by every proxy task. It is consumed by the
// scheduler and implemented by the concrete task structs in the proxy package.
type Task interface {
TraceCtx() context.Context
ID() UniqueID // return ReqID
SetID(uid UniqueID) // set ReqID
Name() string
Type() commonpb.MsgType
BeginTs() Timestamp
EndTs() Timestamp
SetTs(ts Timestamp)
OnEnqueue() error
PreExecute(ctx context.Context) error
Execute(ctx context.Context) error
PostExecute(ctx context.Context) error
WaitToFinish() error
Notify(err error)
CanSkipAllocTimestamp() bool
GetMetaCache() metacache.Cache
SetOnEnqueueTime()
GetDurationInQueue() time.Duration
IsSubTask() bool
SetExecutingTime()
GetDurationInExecuting() time.Duration
}
// BaseTask is embedded by concrete task structs to provide the common
// meta-cache accessor and queue/execute timing bookkeeping.
type BaseTask struct {
onEnqueueTime time.Time
executingTime time.Time
MetaCache metacache.Cache
}
func (bt *BaseTask) GetMetaCache() metacache.Cache {
return bt.MetaCache
}
func (bt *BaseTask) CanSkipAllocTimestamp() bool {
return false
}
func (bt *BaseTask) SetOnEnqueueTime() {
bt.onEnqueueTime = time.Now()
}
func (bt *BaseTask) GetDurationInQueue() time.Duration {
return time.Since(bt.onEnqueueTime)
}
func (bt *BaseTask) IsSubTask() bool {
return false
}
func (bt *BaseTask) SetExecutingTime() {
bt.executingTime = time.Now()
}
func (bt *BaseTask) GetDurationInExecuting() time.Duration {
return time.Since(bt.executingTime)
}
// DMLTask is the interface implemented by DML tasks (insert/delete/upsert)
// whose physical channels are resolved before enqueueing.
type DMLTask interface {
Task
SetChannels() error
GetChannels() []PChan
}
// BaseInsertTask is an alias of msgstream.InsertMsg.
type BaseInsertTask = msgstream.InsertMsg
internal/proxy/taskmodel/types.go
+47
View File
@@ -0,0 +1,47 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package taskmodel
import (
"context"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
// UniqueID is an alias of typeutil.UniqueID.
type UniqueID = typeutil.UniqueID
// Timestamp is an alias of typeutil.Timestamp.
type Timestamp = typeutil.Timestamp
// VChan shortcuts for virtual channel.
type VChan = string
// PChan shortcuts for physical channel.
type PChan = string
// PChanStatistics holds the min/max timestamp of a physical channel's DML.
type PChanStatistics struct {
MinTs Timestamp
MaxTs Timestamp
}
// TsoAllocator allocates timestamps. It is the interface implemented by the
// proxy's timestamp allocator and consumed by the task scheduler.
type TsoAllocator interface {
AllocOne(ctx context.Context) (Timestamp, error)
}
internal/proxy/type_def.go → internal/proxy/taskmodel_alias.go
+21 -7
View File
@@ -16,13 +16,27 @@
package proxy
// vChan shortcuts for virtual channel.
type vChan = string
import (
"context"
// pChan shortcuts for physical channel.
type pChan = string
"github.com/milvus-io/milvus/internal/proxy/taskmodel"
)
type pChanStatistics struct {
minTs Timestamp
maxTs Timestamp
// Aliases to the extracted task model, kept so the concrete task structs that
// remain in this package can keep their original (unexported) type names.
type (
baseTask = taskmodel.BaseTask
dmlTask = taskmodel.DMLTask
tsoAllocator = taskmodel.TsoAllocator
Condition = taskmodel.Condition
TaskCondition = taskmodel.TaskCondition
pChan = taskmodel.PChan
vChan = taskmodel.VChan
pChanStatistics = taskmodel.PChanStatistics
BaseInsertTask = taskmodel.BaseInsertTask
)
func NewTaskCondition(ctx context.Context) *taskmodel.TaskCondition {
return taskmodel.NewTaskCondition(ctx)
}