From b889bb3c19a5a6d2b85a94253fed1d5ddf52cf4f Mon Sep 17 00:00:00 2001 From: congqixia Date: Fri, 4 Sep 2026 10:11:30 +0800 Subject: [PATCH] enhance: extract proxy task model and scheduler (#53023) 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 --- internal/proxy/alias_rbac_test.go | 10 +- internal/proxy/authentication_interceptor.go | 4 +- internal/proxy/function_task.go | 2 +- internal/proxy/function_task_test.go | 16 +- internal/proxy/highlighter.go | 5 +- internal/proxy/impl.go | 274 ++-- internal/proxy/impl_test.go | 58 +- internal/proxy/interface_def.go | 6 - .../proxy/membership_filter_plan_size_test.go | 2 +- internal/proxy/metrics_info.go | 2 +- internal/proxy/metrics_info_test.go | 99 +- internal/proxy/mock_test.go | 141 -- internal/proxy/privilege_interceptor.go | 6 +- internal/proxy/proxy.go | 9 +- internal/proxy/proxy_test.go | 125 +- internal/proxy/rate_limit_interceptor.go | 6 +- internal/proxy/scheduler/OWNERS | 7 + internal/proxy/scheduler/README.md | 80 ++ internal/proxy/scheduler/mock_cache_test.go | 1137 +++++++++++++++++ internal/proxy/scheduler/mock_task_test.go | 193 +++ .../proxy/{ => scheduler}/task_scheduler.go | 337 ++--- .../{ => scheduler}/task_scheduler_test.go | 188 +-- internal/proxy/search_pipeline.go | 2 +- internal/proxy/search_pipeline_test.go | 5 +- internal/proxy/service_provider.go | 8 +- internal/proxy/snapshot_impl.go | 36 +- internal/proxy/snapshot_impl_test.go | 120 +- internal/proxy/task.go | 138 +- internal/proxy/task_alias.go | 6 +- internal/proxy/task_batch_update_manifest.go | 2 +- internal/proxy/task_database.go | 2 +- internal/proxy/task_database_test.go | 2 +- internal/proxy/task_delete.go | 33 +- internal/proxy/task_delete_streaming.go | 2 +- internal/proxy/task_delete_test.go | 23 +- internal/proxy/task_flush_all_test.go | 5 +- internal/proxy/task_flush_streaming.go | 2 +- internal/proxy/task_import.go | 18 +- internal/proxy/task_import_test.go | 18 +- internal/proxy/task_index.go | 22 +- internal/proxy/task_index_test.go | 10 +- internal/proxy/task_insert.go | 16 +- internal/proxy/task_insert_streaming.go | 6 +- internal/proxy/task_insert_test.go | 26 +- internal/proxy/task_query.go | 21 +- internal/proxy/task_query_namespace_test.go | 4 +- internal/proxy/task_query_test.go | 22 +- internal/proxy/task_search.go | 26 +- internal/proxy/task_search_namespace_test.go | 4 +- internal/proxy/task_search_test.go | 49 +- internal/proxy/task_snapshot.go | 26 +- internal/proxy/task_snapshot_test.go | 58 +- internal/proxy/task_statistic.go | 12 +- internal/proxy/task_statistic_test.go | 2 +- internal/proxy/task_test.go | 132 +- internal/proxy/task_upsert.go | 28 +- internal/proxy/task_upsert_streaming.go | 6 +- internal/proxy/task_upsert_test.go | 20 +- internal/proxy/taskmodel/OWNERS | 8 + internal/proxy/taskmodel/README.md | 107 ++ internal/proxy/{ => taskmodel}/condition.go | 2 +- internal/proxy/taskmodel/task.go | 104 ++ internal/proxy/taskmodel/types.go | 47 + .../proxy/{type_def.go => taskmodel_alias.go} | 28 +- 64 files changed, 2655 insertions(+), 1260 deletions(-) create mode 100644 internal/proxy/scheduler/OWNERS create mode 100644 internal/proxy/scheduler/README.md create mode 100644 internal/proxy/scheduler/mock_cache_test.go create mode 100644 internal/proxy/scheduler/mock_task_test.go rename internal/proxy/{ => scheduler}/task_scheduler.go (66%) rename internal/proxy/{ => scheduler}/task_scheduler_test.go (76%) create mode 100644 internal/proxy/taskmodel/OWNERS create mode 100644 internal/proxy/taskmodel/README.md rename internal/proxy/{ => taskmodel}/condition.go (99%) create mode 100644 internal/proxy/taskmodel/task.go create mode 100644 internal/proxy/taskmodel/types.go rename internal/proxy/{type_def.go => taskmodel_alias.go} (53%) diff --git a/internal/proxy/alias_rbac_test.go b/internal/proxy/alias_rbac_test.go index 94e39074c0..30c3ba2cbb 100644 --- a/internal/proxy/alias_rbac_test.go +++ b/internal/proxy/alias_rbac_test.go @@ -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{ diff --git a/internal/proxy/authentication_interceptor.go b/internal/proxy/authentication_interceptor.go index 82514cfd83..32e5f84dae 100644 --- a/internal/proxy/authentication_interceptor.go +++ b/internal/proxy/authentication_interceptor.go @@ -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 diff --git a/internal/proxy/function_task.go b/internal/proxy/function_task.go index dc04c3d831..15cabce663 100644 --- a/internal/proxy/function_task.go +++ b/internal/proxy/function_task.go @@ -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()), diff --git a/internal/proxy/function_task_test.go b/internal/proxy/function_task_test.go index eb1ac70845..f1c30517bd 100644 --- a/internal/proxy/function_task_test.go +++ b/internal/proxy/function_task_test.go @@ -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") } diff --git a/internal/proxy/highlighter.go b/internal/proxy/highlighter.go index 5be347008e..41ce46d0f5 100644 --- a/internal/proxy/highlighter.go +++ b/internal/proxy/highlighter.go @@ -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 } diff --git a/internal/proxy/impl.go b/internal/proxy/impl.go index a9872f46b9..ad246bc43b 100644 --- a/internal/proxy/impl.go +++ b/internal/proxy/impl.go @@ -173,12 +173,12 @@ func (node *Proxy) InvalidateCollectionMetaCache(ctx context.Context, request *p } } - if node.getMetaCache() != nil { + if node.GetMetaCache() != nil { switch msgType { case commonpb.MsgType_DropCollection, commonpb.MsgType_RenameCollection, commonpb.MsgType_DropAlias, commonpb.MsgType_AlterAlias, commonpb.MsgType_CreateAlias: aliasOperation := msgType == commonpb.MsgType_CreateAlias || msgType == commonpb.MsgType_AlterAlias || msgType == commonpb.MsgType_DropAlias - aliasName = node.getMetaCache().InvalidateCollectionMeta( + aliasName = node.GetMetaCache().InvalidateCollectionMeta( ctx, request.GetDbName(), collectionName, collectionID, aliasOperation) deprecateShardCaches(append(aliasName, collectionName)...) if aliasOperation && collectionName != "" { @@ -197,14 +197,14 @@ func (node *Proxy) InvalidateCollectionMetaCache(ctx context.Context, request *p // resolution is race-free -- nothing re-points an alias // concurrently with its drop), so scanning there would // reintroduce a permanent O(N) stall on the normal path. - node.getMetaCache().RemoveAliasHolders(ctx, request.GetDbName(), collectionName) + node.GetMetaCache().RemoveAliasHolders(ctx, request.GetDbName(), collectionName) } } mlog.Info(ctx, "complete to invalidate collection meta cache with collection name", mlog.String("type", request.GetBase().GetMsgType().String())) case commonpb.MsgType_LoadCollection, commonpb.MsgType_ReleaseCollection: // All the request from query use collectionID if request.CollectionID != UniqueID(0) { - aliasName = node.getMetaCache().InvalidateCollectionMeta(ctx, request.GetDbName(), "", collectionID, false) + aliasName = node.GetMetaCache().InvalidateCollectionMeta(ctx, request.GetDbName(), "", collectionID, false) deprecateShardCaches(aliasName...) } mlog.Info(ctx, "complete to invalidate collection meta cache", mlog.String("type", request.GetBase().GetMsgType().String())) @@ -217,20 +217,20 @@ func (node *Proxy) InvalidateCollectionMetaCache(ctx context.Context, request *p mlog.FieldCollectionID(collectionID)) return merr.Status(err), nil } - node.getMetaCache().RemovePartition(ctx, request.GetDbName(), collectionID, collectionName, request.GetPartitionName()) + node.GetMetaCache().RemovePartition(ctx, request.GetDbName(), collectionID, collectionName, request.GetPartitionName()) mlog.Info(ctx, "complete to invalidate collection meta cache", mlog.String("type", request.GetBase().GetMsgType().String())) case commonpb.MsgType_DropDatabase: node.shardMgr.RemoveDatabase(request.GetDbName()) - node.getMetaCache().RemoveDatabase(ctx, request.GetDbName()) + node.GetMetaCache().RemoveDatabase(ctx, request.GetDbName()) case commonpb.MsgType_AlterDatabase: - node.getMetaCache().RemoveDatabaseInfo(ctx, request.GetDbName()) + node.GetMetaCache().RemoveDatabaseInfo(ctx, request.GetDbName()) case commonpb.MsgType_AlterCollection, commonpb.MsgType_AlterCollectionField: - aliasName = node.getMetaCache().InvalidateCollectionMeta(ctx, request.GetDbName(), collectionName, collectionID, false) + aliasName = node.GetMetaCache().InvalidateCollectionMeta(ctx, request.GetDbName(), collectionName, collectionID, false) deprecateShardCaches(append(aliasName, collectionName)...) mlog.Info(ctx, "complete to invalidate collection meta cache", mlog.String("type", request.GetBase().GetMsgType().String())) default: mlog.Warn(ctx, "receive unexpected msgType of invalidate collection meta cache", mlog.String("msgType", request.GetBase().GetMsgType().String())) - aliasName = node.getMetaCache().InvalidateCollectionMeta(ctx, request.GetDbName(), collectionName, collectionID, false) + aliasName = node.GetMetaCache().InvalidateCollectionMeta(ctx, request.GetDbName(), collectionName, collectionID, false) deprecateShardCaches(append(aliasName, collectionName)...) } } @@ -283,18 +283,18 @@ func (node *Proxy) ClearReadTaskQueue(ctx context.Context, request *internalpb.C ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ClearReadTaskQueue") defer sp.End() - result := node.sched.clearDQLQueue(request.GetTaskType(), request.GetReason()) - resp.ProxyQueuedCleared = result.queuedCleared + result := node.sched.ClearDQLQueue(request.GetTaskType(), request.GetReason()) + resp.ProxyQueuedCleared = result.QueuedCleared resp.Results = append(resp.Results, &internalpb.ClearReadTaskQueueComponentResult{ Status: merr.Success(), Role: typeutil.ProxyRole, NodeID: paramtable.GetNodeID(), - QueuedCleared: result.queuedCleared, + QueuedCleared: result.QueuedCleared, }) mlog.Info(ctx, "cleared proxy read task queue", mlog.String("taskType", request.GetTaskType()), mlog.String("reason", request.GetReason()), - mlog.Int64("queuedCleared", result.queuedCleared)) + mlog.Int64("queuedCleared", result.QueuedCleared)) return resp, nil } @@ -317,7 +317,7 @@ func (node *Proxy) CreateDatabase(ctx context.Context, request *milvuspb.CreateD } mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(cct); err != nil { + if err := node.sched.DdQueue.Enqueue(cct); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) return merr.Status(err), nil @@ -352,7 +352,7 @@ func (node *Proxy) DropDatabase(ctx context.Context, request *milvuspb.DropDatab tr := timerecord.NewTimeRecorder(method) dct := &dropDatabaseTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), DropDatabaseRequest: request, @@ -360,7 +360,7 @@ func (node *Proxy) DropDatabase(ctx context.Context, request *milvuspb.DropDatab } mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(dct); err != nil { + if err := node.sched.DdQueue.Enqueue(dct); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) return merr.Status(err), nil } @@ -403,7 +403,7 @@ func (node *Proxy) ListDatabases(ctx context.Context, request *milvuspb.ListData mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(dct); err != nil { + if err := node.sched.DdQueue.Enqueue(dct); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) resp.Status = merr.Status(err) return resp, nil @@ -445,7 +445,7 @@ func (node *Proxy) AlterDatabase(ctx context.Context, request *milvuspb.AlterDat mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(act); err != nil { + if err := node.sched.DdQueue.Enqueue(act); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -496,7 +496,7 @@ func (node *Proxy) DescribeDatabase(ctx context.Context, request *milvuspb.Descr mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(act); err != nil { + if err := node.sched.DdQueue.Enqueue(act); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) resp.Status = merr.Status(err) @@ -547,7 +547,7 @@ func (node *Proxy) CreateCollection(ctx context.Context, request *milvuspb.Creat mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(cct); err != nil { + if err := node.sched.DdQueue.Enqueue(cct); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -598,7 +598,7 @@ func (node *Proxy) DropCollection(ctx context.Context, request *milvuspb.DropCol tr := timerecord.NewTimeRecorder(method) dct := &dropCollectionTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), DropCollectionRequest: request, @@ -608,7 +608,7 @@ func (node *Proxy) DropCollection(ctx context.Context, request *milvuspb.DropCol mlog.Info(context.TODO(), "DropCollection received") - if err := node.sched.ddQueue.Enqueue(dct); err != nil { + if err := node.sched.DdQueue.Enqueue(dct); err != nil { mlog.Warn(context.TODO(), "DropCollection failed to enqueue", mlog.Err(err)) @@ -657,7 +657,7 @@ func (node *Proxy) TruncateCollection(ctx context.Context, request *milvuspb.Tru tr := timerecord.NewTimeRecorder(method) dct := &truncateCollectionTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), TruncateCollectionRequest: request, @@ -667,7 +667,7 @@ func (node *Proxy) TruncateCollection(ctx context.Context, request *milvuspb.Tru mlog.Info(context.TODO(), "TruncateCollection received") - if err := node.sched.ddQueue.Enqueue(dct); err != nil { + if err := node.sched.DdQueue.Enqueue(dct); err != nil { mlog.Warn(context.TODO(), "TruncateCollection failed to enqueue", mlog.Err(err)) @@ -724,14 +724,14 @@ func (node *Proxy) HasCollection(ctx context.Context, request *milvuspb.HasColle mlog.Debug(context.TODO(), "HasCollection received") hct := &hasCollectionTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), HasCollectionRequest: request, mixCoord: node.mixCoord, } - if err := node.sched.ddQueue.Enqueue(hct); err != nil { + if err := node.sched.DdQueue.Enqueue(hct); err != nil { mlog.Warn(context.TODO(), "HasCollection failed to enqueue", mlog.Err(err)) @@ -781,7 +781,7 @@ func (node *Proxy) LoadCollection(ctx context.Context, request *milvuspb.LoadCol tr := timerecord.NewTimeRecorder(method) lct := &loadCollectionTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), LoadCollectionRequest: request, @@ -790,7 +790,7 @@ func (node *Proxy) LoadCollection(ctx context.Context, request *milvuspb.LoadCol mlog.Info(context.TODO(), "LoadCollection received") - if err := node.sched.ddQueue.Enqueue(lct); err != nil { + if err := node.sched.DdQueue.Enqueue(lct); err != nil { mlog.Warn(context.TODO(), "LoadCollection failed to enqueue", mlog.Err(err)) @@ -834,7 +834,7 @@ func (node *Proxy) ReleaseCollection(ctx context.Context, request *milvuspb.Rele method := "ReleaseCollection" tr := timerecord.NewTimeRecorder(method) rct := &releaseCollectionTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), ReleaseCollectionRequest: request, @@ -843,7 +843,7 @@ func (node *Proxy) ReleaseCollection(ctx context.Context, request *milvuspb.Rele mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(rct); err != nil { + if err := node.sched.DdQueue.Enqueue(rct); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -948,7 +948,7 @@ func (node *Proxy) AddCollectionField(ctx context.Context, request *milvuspb.Add mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(task); err != nil { + if err := node.sched.DdQueue.Enqueue(task); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1012,7 +1012,7 @@ func (node *Proxy) AddCollectionStructField(ctx context.Context, request *milvus mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(task); err != nil { + if err := node.sched.DdQueue.Enqueue(task); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1081,7 +1081,7 @@ func (node *Proxy) AlterCollectionSchema(ctx context.Context, request *milvuspb. mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(task); err != nil { + if err := node.sched.DdQueue.Enqueue(task); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1129,7 +1129,7 @@ func (node *Proxy) GetStatistics(ctx context.Context, request *milvuspb.GetStati method := "GetStatistics" tr := timerecord.NewTimeRecorder(method) g := &getStatisticsTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, request: request, Condition: NewTaskCondition(ctx), ctx: ctx, @@ -1142,7 +1142,7 @@ func (node *Proxy) GetStatistics(ctx context.Context, request *milvuspb.GetStati rpcReceived(method), mlog.Strings("partitions", request.PartitionNames)) - if err := node.sched.ddQueue.Enqueue(g); err != nil { + if err := node.sched.DdQueue.Enqueue(g); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err), @@ -1194,7 +1194,7 @@ func (node *Proxy) GetCollectionStatistics(ctx context.Context, request *milvusp method := "GetCollectionStatistics" tr := timerecord.NewTimeRecorder(method) g := &getCollectionStatisticsTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), GetCollectionStatisticsRequest: request, @@ -1203,7 +1203,7 @@ func (node *Proxy) GetCollectionStatistics(ctx context.Context, request *milvusp mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(g); err != nil { + if err := node.sched.DdQueue.Enqueue(g); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1252,7 +1252,7 @@ func (node *Proxy) ShowCollections(ctx context.Context, request *milvuspb.ShowCo tr := timerecord.NewTimeRecorder(method) sct := &showCollectionsTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), ShowCollectionsRequest: request, @@ -1262,7 +1262,7 @@ func (node *Proxy) ShowCollections(ctx context.Context, request *milvuspb.ShowCo mlog.Debug(context.TODO(), "ShowCollections received", mlog.Any("CollectionNames", request.CollectionNames)) - err := node.sched.ddQueue.Enqueue(sct) + err := node.sched.DdQueue.Enqueue(sct) if err != nil { mlog.Warn(context.TODO(), "ShowCollections failed to enqueue", mlog.Err(err), @@ -1307,7 +1307,7 @@ func (node *Proxy) AlterCollection(ctx context.Context, request *milvuspb.AlterC tr := timerecord.NewTimeRecorder(method) act := &alterCollectionTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), AlterCollectionRequest: request, @@ -1317,7 +1317,7 @@ func (node *Proxy) AlterCollection(ctx context.Context, request *milvuspb.AlterC mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(act); err != nil { + if err := node.sched.DdQueue.Enqueue(act); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1373,7 +1373,7 @@ func (node *Proxy) AlterCollectionFunction(ctx context.Context, request *milvusp method := "AlterCollectionFunction" tr := timerecord.NewTimeRecorder(method) task := &alterCollectionFunctionTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), AlterCollectionFunctionRequest: request, @@ -1381,7 +1381,7 @@ func (node *Proxy) AlterCollectionFunction(ctx context.Context, request *milvusp } mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(task); err != nil { + if err := node.sched.DdQueue.Enqueue(task); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1436,12 +1436,12 @@ func (node *Proxy) AlterCollectionField(ctx context.Context, request *milvuspb.A tr := timerecord.NewTimeRecorder(method) // Check for external collection - alter field is not supported - if err := checkExternalCollectionBlockedForWrite(ctx, node.getMetaCache(), request.GetDbName(), request.GetCollectionName(), "alter field"); err != nil { + if err := checkExternalCollectionBlockedForWrite(ctx, node.GetMetaCache(), request.GetDbName(), request.GetCollectionName(), "alter field"); err != nil { return merr.Status(err), nil } act := &alterCollectionFieldTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), AlterCollectionFieldRequest: request, @@ -1450,7 +1450,7 @@ func (node *Proxy) AlterCollectionField(ctx context.Context, request *milvuspb.A mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(act); err != nil { + if err := node.sched.DdQueue.Enqueue(act); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1490,7 +1490,7 @@ func (node *Proxy) CreatePartition(ctx context.Context, request *milvuspb.Create } // Check for external collection - create partition is not supported - if err := checkExternalCollectionBlockedForWrite(ctx, node.getMetaCache(), request.GetDbName(), request.GetCollectionName(), "create partition"); err != nil { + if err := checkExternalCollectionBlockedForWrite(ctx, node.GetMetaCache(), request.GetDbName(), request.GetCollectionName(), "create partition"); err != nil { return merr.Status(err), nil } @@ -1500,7 +1500,7 @@ func (node *Proxy) CreatePartition(ctx context.Context, request *milvuspb.Create tr := timerecord.NewTimeRecorder(method) cpt := &createPartitionTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), CreatePartitionRequest: request, @@ -1510,7 +1510,7 @@ func (node *Proxy) CreatePartition(ctx context.Context, request *milvuspb.Create mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(cpt); err != nil { + if err := node.sched.DdQueue.Enqueue(cpt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1549,7 +1549,7 @@ func (node *Proxy) DropPartition(ctx context.Context, request *milvuspb.DropPart } // Check for external collection - drop partition is not supported - if err := checkExternalCollectionBlockedForWrite(ctx, node.getMetaCache(), request.GetDbName(), request.GetCollectionName(), "drop partition"); err != nil { + if err := checkExternalCollectionBlockedForWrite(ctx, node.GetMetaCache(), request.GetDbName(), request.GetCollectionName(), "drop partition"); err != nil { return merr.Status(err), nil } @@ -1559,7 +1559,7 @@ func (node *Proxy) DropPartition(ctx context.Context, request *milvuspb.DropPart tr := timerecord.NewTimeRecorder(method) dpt := &dropPartitionTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), DropPartitionRequest: request, @@ -1569,7 +1569,7 @@ func (node *Proxy) DropPartition(ctx context.Context, request *milvuspb.DropPart mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(dpt); err != nil { + if err := node.sched.DdQueue.Enqueue(dpt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1624,7 +1624,7 @@ func (node *Proxy) HasPartition(ctx context.Context, request *milvuspb.HasPartit mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(hpt); err != nil { + if err := node.sched.DdQueue.Enqueue(hpt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1673,7 +1673,7 @@ func (node *Proxy) LoadPartitions(ctx context.Context, request *milvuspb.LoadPar method := "LoadPartitions" tr := timerecord.NewTimeRecorder(method) lpt := &loadPartitionsTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), LoadPartitionsRequest: request, @@ -1682,7 +1682,7 @@ func (node *Proxy) LoadPartitions(ctx context.Context, request *milvuspb.LoadPar mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(lpt); err != nil { + if err := node.sched.DdQueue.Enqueue(lpt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1724,7 +1724,7 @@ func (node *Proxy) ReleasePartitions(ctx context.Context, request *milvuspb.Rele defer sp.End() rpt := &releasePartitionsTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), ReleasePartitionsRequest: request, @@ -1736,7 +1736,7 @@ func (node *Proxy) ReleasePartitions(ctx context.Context, request *milvuspb.Rele mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(rpt); err != nil { + if err := node.sched.DdQueue.Enqueue(rpt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1782,7 +1782,7 @@ func (node *Proxy) GetPartitionStatistics(ctx context.Context, request *milvuspb tr := timerecord.NewTimeRecorder(method) g := &getPartitionStatisticsTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), GetPartitionStatisticsRequest: request, @@ -1791,7 +1791,7 @@ func (node *Proxy) GetPartitionStatistics(ctx context.Context, request *milvuspb mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(g); err != nil { + if err := node.sched.DdQueue.Enqueue(g); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -1839,7 +1839,7 @@ func (node *Proxy) ShowPartitions(ctx context.Context, request *milvuspb.ShowPar defer sp.End() spt := &showPartitionsTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), ShowPartitionsRequest: request, @@ -1854,7 +1854,7 @@ func (node *Proxy) ShowPartitions(ctx context.Context, request *milvuspb.ShowPar rpcReceived(method), mlog.Any("request", request)) - if err := node.sched.ddQueue.Enqueue(spt); err != nil { + if err := node.sched.DdQueue.Enqueue(spt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err), @@ -1929,7 +1929,7 @@ func (node *Proxy) GetLoadingProgress(ctx context.Context, request *milvuspb.Get if err := validateCollectionName(request.CollectionName); err != nil { return getErrResponse(err), nil } - collectionID, err := node.getMetaCache().GetCollectionID(ctx, request.GetDbName(), request.CollectionName) + collectionID, err := node.GetMetaCache().GetCollectionID(ctx, request.GetDbName(), request.CollectionName) if err != nil { return getErrResponse(err), nil } @@ -1955,7 +1955,7 @@ func (node *Proxy) GetLoadingProgress(ctx context.Context, request *milvuspb.Get return getErrResponse(err), nil } } else { - if loadProgress, refreshProgress, err = getPartitionProgress(ctx, node.getMetaCache(), node.mixCoord, request.GetBase(), + if loadProgress, refreshProgress, err = getPartitionProgress(ctx, node.GetMetaCache(), node.mixCoord, request.GetBase(), request.GetPartitionNames(), request.GetCollectionName(), collectionID, request.GetDbName()); err != nil { return getErrResponse(err), nil } @@ -2013,7 +2013,7 @@ func (node *Proxy) GetLoadState(ctx context.Context, request *milvuspb.GetLoadSt metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds())) }() - collectionID, err := node.getMetaCache().GetCollectionID(ctx, request.GetDbName(), request.CollectionName) + collectionID, err := node.GetMetaCache().GetCollectionID(ctx, request.GetDbName(), request.CollectionName) if err != nil { mlog.Warn(context.TODO(), "failed to get collection id", mlog.String("dbName", request.GetDbName()), @@ -2047,7 +2047,7 @@ func (node *Proxy) GetLoadState(ctx context.Context, request *milvuspb.GetLoadSt }, nil } } else { - if progress, _, err = getPartitionProgress(ctx, node.getMetaCache(), node.mixCoord, request.GetBase(), + if progress, _, err = getPartitionProgress(ctx, node.GetMetaCache(), node.mixCoord, request.GetBase(), request.GetPartitionNames(), request.GetCollectionName(), collectionID, request.GetDbName()); err != nil { if errors.IsAny(err, merr.ErrCollectionNotLoaded, @@ -2078,7 +2078,7 @@ func (node *Proxy) CreateIndex(ctx context.Context, request *milvuspb.CreateInde defer sp.End() cit := &createIndexTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), req: request, @@ -2090,7 +2090,7 @@ func (node *Proxy) CreateIndex(ctx context.Context, request *milvuspb.CreateInde mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(cit); err != nil { + if err := node.sched.DdQueue.Enqueue(cit); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -2131,7 +2131,7 @@ func (node *Proxy) AlterIndex(ctx context.Context, request *milvuspb.AlterIndexR defer sp.End() task := &alterIndexTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), req: request, @@ -2143,7 +2143,7 @@ func (node *Proxy) AlterIndex(ctx context.Context, request *milvuspb.AlterIndexR mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(task); err != nil { + if err := node.sched.DdQueue.Enqueue(task); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -2187,7 +2187,7 @@ func (node *Proxy) DescribeIndex(ctx context.Context, request *milvuspb.Describe defer sp.End() dit := &describeIndexTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), DescribeIndexRequest: request, @@ -2200,7 +2200,7 @@ func (node *Proxy) DescribeIndex(ctx context.Context, request *milvuspb.Describe mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(dit); err != nil { + if err := node.sched.DdQueue.Enqueue(dit); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -2248,7 +2248,7 @@ func (node *Proxy) GetIndexStatistics(ctx context.Context, request *milvuspb.Get defer sp.End() dit := &getIndexStatisticsTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), GetIndexStatisticsRequest: request, @@ -2261,7 +2261,7 @@ func (node *Proxy) GetIndexStatistics(ctx context.Context, request *milvuspb.Get mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(dit); err != nil { + if err := node.sched.DdQueue.Enqueue(dit); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -2303,7 +2303,7 @@ func (node *Proxy) DropIndex(ctx context.Context, request *milvuspb.DropIndexReq defer sp.End() dit := &dropIndexTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), DropIndexRequest: request, @@ -2315,7 +2315,7 @@ func (node *Proxy) DropIndex(ctx context.Context, request *milvuspb.DropIndexReq mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(dit); err != nil { + if err := node.sched.DdQueue.Enqueue(dit); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -2361,7 +2361,7 @@ func (node *Proxy) GetIndexBuildProgress(ctx context.Context, request *milvuspb. defer sp.End() gibpt := &getIndexBuildProgressTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), GetIndexBuildProgressRequest: request, @@ -2373,7 +2373,7 @@ func (node *Proxy) GetIndexBuildProgress(ctx context.Context, request *milvuspb. mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(gibpt); err != nil { + if err := node.sched.DdQueue.Enqueue(gibpt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -2422,7 +2422,7 @@ func (node *Proxy) GetIndexState(ctx context.Context, request *milvuspb.GetIndex defer sp.End() dipt := &getIndexStateTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), GetIndexStateRequest: request, @@ -2434,7 +2434,7 @@ func (node *Proxy) GetIndexState(ctx context.Context, request *milvuspb.GetIndex mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(dipt); err != nil { + if err := node.sched.DdQueue.Enqueue(dipt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -2482,7 +2482,7 @@ func (node *Proxy) Insert(ctx context.Context, request *milvuspb.InsertRequest) } // Check for external collection - insert is not supported - if err := checkExternalCollectionBlockedForWrite(ctx, node.getMetaCache(), request.GetDbName(), request.GetCollectionName(), "insert"); err != nil { + if err := checkExternalCollectionBlockedForWrite(ctx, node.GetMetaCache(), request.GetDbName(), request.GetCollectionName(), "insert"); err != nil { return &milvuspb.MutationResult{ Status: merr.Status(err), }, nil @@ -2498,7 +2498,7 @@ func (node *Proxy) Insert(ctx context.Context, request *milvuspb.InsertRequest) it := &insertTask{ baseTask: baseTask{ - metaCache: node.getMetaCache(), + MetaCache: node.GetMetaCache(), }, ctx: ctx, Condition: NewTaskCondition(ctx), @@ -2540,7 +2540,7 @@ func (node *Proxy) Insert(ctx context.Context, request *milvuspb.InsertRequest) mlog.Debug(context.TODO(), "Enqueue insert request in Proxy") - if err := node.sched.dmQueue.Enqueue(it); err != nil { + if err := node.sched.DmQueue.Enqueue(it); err != nil { mlog.Warn(context.TODO(), "Failed to enqueue insert task: "+err.Error()) return constructFailedResponse(err), nil } @@ -2622,7 +2622,7 @@ func (node *Proxy) Delete(ctx context.Context, request *milvuspb.DeleteRequest) } // Check for external collection - delete is not supported - if err := checkExternalCollectionBlockedForWrite(ctx, node.getMetaCache(), request.GetDbName(), request.GetCollectionName(), "delete"); err != nil { + if err := checkExternalCollectionBlockedForWrite(ctx, node.GetMetaCache(), request.GetDbName(), request.GetCollectionName(), "delete"); err != nil { return &milvuspb.MutationResult{ Status: merr.Status(err), }, nil @@ -2637,11 +2637,11 @@ func (node *Proxy) Delete(ctx context.Context, request *milvuspb.DeleteRequest) dr := &deleteRunner{ req: request, - metaCache: node.getMetaCache(), + metaCache: node.GetMetaCache(), idAllocator: node.rowIDAllocator, tsoAllocatorIns: node.tsoAllocator, chMgr: node.chMgr, - queue: node.sched.dmQueue, + queue: node.sched.DmQueue, lb: node.lbPolicy, limiter: limiter, } @@ -2719,7 +2719,7 @@ func (node *Proxy) Upsert(ctx context.Context, request *milvuspb.UpsertRequest) } // Check for external collection - upsert is not supported - if err := checkExternalCollectionBlockedForWrite(ctx, node.getMetaCache(), request.GetDbName(), request.GetCollectionName(), "upsert"); err != nil { + if err := checkExternalCollectionBlockedForWrite(ctx, node.GetMetaCache(), request.GetDbName(), request.GetCollectionName(), "upsert"); err != nil { return &milvuspb.MutationResult{ Status: merr.Status(err), }, nil @@ -2741,7 +2741,7 @@ func (node *Proxy) Upsert(ctx context.Context, request *milvuspb.UpsertRequest) it := &upsertTask{ baseTask: baseTask{ - metaCache: node.getMetaCache(), + MetaCache: node.GetMetaCache(), }, baseMsg: msgstream.BaseMsg{ HashValues: request.HashKeys, @@ -2766,7 +2766,7 @@ func (node *Proxy) Upsert(ctx context.Context, request *milvuspb.UpsertRequest) mlog.Int("len(FieldsData)", len(request.FieldsData)), mlog.Int("len(HashKeys)", len(request.HashKeys))) - if err := node.sched.dmQueue.Enqueue(it); err != nil { + if err := node.sched.DmQueue.Enqueue(it); err != nil { mlog.Info(context.TODO(), "Failed to enqueue upsert task", mlog.Err(err)) return &milvuspb.MutationResult{ @@ -3011,7 +3011,7 @@ func (node *Proxy) search(ctx context.Context, request *milvuspb.SearchRequest, qt := &searchTask{ baseTask: baseTask{ - metaCache: node.getMetaCache(), + MetaCache: node.GetMetaCache(), }, ctx: ctx, Condition: NewTaskCondition(ctx), @@ -3066,7 +3066,7 @@ func (node *Proxy) search(ctx context.Context, request *milvuspb.SearchRequest, mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.dqQueue.Enqueue(qt); err != nil { + if err := node.sched.DqQueue.Enqueue(qt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err), @@ -3246,7 +3246,7 @@ func (node *Proxy) hybridSearch(ctx context.Context, request *milvuspb.HybridSea newSearchReq := convertHybridSearchToSearch(request) qt := &searchTask{ baseTask: baseTask{ - metaCache: node.getMetaCache(), + MetaCache: node.GetMetaCache(), }, ctx: ctx, Condition: NewTaskCondition(ctx), @@ -3303,7 +3303,7 @@ func (node *Proxy) hybridSearch(ctx context.Context, request *milvuspb.HybridSea mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.dqQueue.Enqueue(qt); err != nil { + if err := node.sched.DqQueue.Enqueue(qt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err), @@ -3465,7 +3465,7 @@ func (node *Proxy) handleIfSearchByPK(ctx context.Context, request *milvuspb.Sea } // Get collection schema for validation and plan building - collectionInfo, err := node.getMetaCache().GetCollectionInfo(ctx, + collectionInfo, err := node.GetMetaCache().GetCollectionInfo(ctx, request.GetDbName(), request.GetCollectionName(), 0) if err != nil { return nil, err @@ -3548,7 +3548,7 @@ func (node *Proxy) handleIfSearchByPK(ctx context.Context, request *milvuspb.Sea // Create queryTask to execute the retrieval qt := &queryTask{ baseTask: baseTask{ - metaCache: node.getMetaCache(), + MetaCache: node.GetMetaCache(), }, ctx: ctx, Condition: NewTaskCondition(ctx), @@ -3715,7 +3715,7 @@ func (node *Proxy) Flush(ctx context.Context, request *milvuspb.FlushRequest) (* // Check for external collection - flush is not supported for _, collName := range request.GetCollectionNames() { - if err := checkExternalCollectionBlockedForWrite(ctx, node.getMetaCache(), request.GetDbName(), collName, "flush"); err != nil { + if err := checkExternalCollectionBlockedForWrite(ctx, node.GetMetaCache(), request.GetDbName(), collName, "flush"); err != nil { resp.Status = merr.Status(err) return resp, nil } @@ -3725,7 +3725,7 @@ func (node *Proxy) Flush(ctx context.Context, request *milvuspb.FlushRequest) (* defer sp.End() ft := &flushTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), FlushRequest: request, @@ -3737,7 +3737,7 @@ func (node *Proxy) Flush(ctx context.Context, request *milvuspb.FlushRequest) (* tr := timerecord.NewTimeRecorder(method) mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.dcQueue.Enqueue(ft); err != nil { + if err := node.sched.DcQueue.Enqueue(ft); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -3827,7 +3827,7 @@ func (node *Proxy) query(ctx context.Context, qt *queryTask, sp trace.Span) (*mi } }() - if err := node.sched.dqQueue.Enqueue(qt); err != nil { + if err := node.sched.DqQueue.Enqueue(qt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err), @@ -3881,7 +3881,7 @@ func (node *Proxy) query(ctx context.Context, qt *queryTask, sp trace.Span) (*mi func (node *Proxy) Query(ctx context.Context, request *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) { qt := &queryTask{ baseTask: baseTask{ - metaCache: node.getMetaCache(), + MetaCache: node.GetMetaCache(), }, ctx: ctx, Condition: NewTaskCondition(ctx), @@ -3985,7 +3985,7 @@ func (node *Proxy) CreateAlias(ctx context.Context, request *milvuspb.CreateAlia defer sp.End() cat := &CreateAliasTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), CreateAliasRequest: request, @@ -3997,7 +3997,7 @@ func (node *Proxy) CreateAlias(ctx context.Context, request *milvuspb.CreateAlia mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(cat); err != nil { + if err := node.sched.DdQueue.Enqueue(cat); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -4053,7 +4053,7 @@ func (node *Proxy) DescribeAlias(ctx context.Context, request *milvuspb.Describe mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(dat); err != nil { + if err := node.sched.DdQueue.Enqueue(dat); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -4096,7 +4096,7 @@ func (node *Proxy) ListAliases(ctx context.Context, request *milvuspb.ListAliase defer sp.End() lat := &ListAliasesTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), nodeID: node.session.ServerID, @@ -4109,7 +4109,7 @@ func (node *Proxy) ListAliases(ctx context.Context, request *milvuspb.ListAliase mlog.Debug(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(lat); err != nil { + if err := node.sched.DdQueue.Enqueue(lat); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -4161,7 +4161,7 @@ func (node *Proxy) DropAlias(ctx context.Context, request *milvuspb.DropAliasReq mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(dat); err != nil { + if err := node.sched.DdQueue.Enqueue(dat); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -4203,7 +4203,7 @@ func (node *Proxy) AlterAlias(ctx context.Context, request *milvuspb.AlterAliasR defer sp.End() aat := &AlterAliasTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), AlterAliasRequest: request, @@ -4215,7 +4215,7 @@ func (node *Proxy) AlterAlias(ctx context.Context, request *milvuspb.AlterAliasR mlog.Info(ctx, rpcReceived(method)) - if err := node.sched.ddQueue.Enqueue(aat); err != nil { + if err := node.sched.DdQueue.Enqueue(aat); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -4281,7 +4281,7 @@ func (node *Proxy) FlushAll(ctx context.Context, request *milvuspb.FlushAllReque logger.Debug(ctx, rpcReceived(method)) - if err := node.sched.dcQueue.Enqueue(ft); err != nil { + if err := node.sched.DcQueue.Enqueue(ft); err != nil { logger.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) resp.Status = merr.Status(err) return resp, nil @@ -4332,7 +4332,7 @@ func (node *Proxy) GetPersistentSegmentInfo(ctx context.Context, req *milvuspb.G tr := timerecord.NewTimeRecorder(method) // list segments - collectionID, err := node.getMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) + collectionID, err := node.GetMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) if err != nil { resp.Status = merr.Status(err) return resp, nil @@ -4491,7 +4491,7 @@ func (node *Proxy) GetQuerySegmentInfo(ctx context.Context, req *milvuspb.GetQue method := "GetQuerySegmentInfo" tr := timerecord.NewTimeRecorder(method) - collID, err := node.getMetaCache().GetCollectionID(ctx, req.GetDbName(), req.CollectionName) + collID, err := node.GetMetaCache().GetCollectionID(ctx, req.GetDbName(), req.CollectionName) if err != nil { resp.Status = merr.Status(err) return resp, nil @@ -4722,7 +4722,7 @@ func (node *Proxy) LoadBalance(ctx context.Context, req *milvuspb.LoadBalanceReq status := merr.Success() - collectionID, err := node.getMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) + collectionID, err := node.GetMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) if err != nil { mlog.Warn(context.TODO(), "failed to get collection id", mlog.String("collectionName", req.GetCollectionName()), @@ -4781,7 +4781,7 @@ func (node *Proxy) GetReplicas(ctx context.Context, req *milvuspb.GetReplicasReq if req.GetCollectionName() != "" { var err error - req.CollectionID, err = node.getMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) + req.CollectionID, err = node.GetMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) if err != nil { resp.Status = merr.Status(err) return resp, nil @@ -4833,7 +4833,7 @@ func (node *Proxy) ManualCompaction(ctx context.Context, req *milvuspb.ManualCom // before v2.4.18, manual compact request only pass collectionID, should correct sdk's behavior to pass collectionName if req.GetCollectionName() != "" { var err error - req.CollectionID, err = node.getMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) + req.CollectionID, err = node.GetMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) if err != nil { resp.Status = merr.Status(err) return resp, nil @@ -4843,7 +4843,7 @@ func (node *Proxy) ManualCompaction(ctx context.Context, req *milvuspb.ManualCom // Check for external collection - manual compaction is not supported // Only check if we have collection identifier (collectionID > 0 or collectionName not empty) if req.GetCollectionID() > 0 || req.GetCollectionName() != "" { - collInfo, err := node.getMetaCache().GetCollectionInfo(ctx, req.GetDbName(), req.GetCollectionName(), req.GetCollectionID()) + collInfo, err := node.GetMetaCache().GetCollectionInfo(ctx, req.GetDbName(), req.GetCollectionName(), req.GetCollectionID()) if err != nil { resp.Status = merr.Status(err) return resp, nil @@ -4913,7 +4913,7 @@ func (node *Proxy) GetFlushState(ctx context.Context, req *milvuspb.GetFlushStat failResp.Status = merr.Status(err) return failResp, nil } - collectionID, err := node.getMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) + collectionID, err := node.GetMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) if err != nil { failResp.Status = merr.Status(err) return failResp, nil @@ -5612,7 +5612,7 @@ func (node *Proxy) OperatePrivilegeV2(ctx context.Context, req *milvuspb.Operate // Resolve alias to actual collection name so grants are stored under the real name if Params.ProxyCfg.ResolveAliasForPrivilege.GetAsBool() && req.CollectionName != util.AnyWord && req.CollectionName != "" { - resolved, resolveErr := node.getMetaCache().ResolveCollectionAlias(ctx, req.DbName, req.CollectionName) + resolved, resolveErr := node.GetMetaCache().ResolveCollectionAlias(ctx, req.DbName, req.CollectionName) if resolveErr != nil { mlog.Warn(context.TODO(), "failed to resolve collection alias for privilege operation", mlog.String("collectionName", req.CollectionName), mlog.String("dbName", req.DbName), mlog.Err(resolveErr)) @@ -5684,7 +5684,7 @@ func (node *Proxy) OperatePrivilege(ctx context.Context, req *milvuspb.OperatePr if Params.ProxyCfg.ResolveAliasForPrivilege.GetAsBool() && req.Entity.Object != nil && req.Entity.Object.Name == commonpb.ObjectType_Collection.String() && req.Entity.ObjectName != util.AnyWord && req.Entity.ObjectName != "" { - resolved, resolveErr := node.getMetaCache().ResolveCollectionAlias(ctx, req.Entity.DbName, req.Entity.ObjectName) + resolved, resolveErr := node.GetMetaCache().ResolveCollectionAlias(ctx, req.Entity.DbName, req.Entity.ObjectName) if resolveErr != nil { mlog.Warn(context.TODO(), "failed to resolve collection alias for privilege operation", mlog.String("objectName", req.Entity.ObjectName), mlog.String("dbName", req.Entity.DbName), mlog.Err(resolveErr)) @@ -5767,7 +5767,7 @@ func (node *Proxy) SelectGrant(ctx context.Context, req *milvuspb.SelectGrantReq if Params.ProxyCfg.ResolveAliasForPrivilege.GetAsBool() && req.Entity != nil && req.Entity.Object != nil && req.Entity.Object.Name == commonpb.ObjectType_Collection.String() && req.Entity.ObjectName != util.AnyWord && req.Entity.ObjectName != "" { - resolved, resolveErr := node.getMetaCache().ResolveCollectionAlias(ctx, req.Entity.DbName, req.Entity.ObjectName) + resolved, resolveErr := node.GetMetaCache().ResolveCollectionAlias(ctx, req.Entity.DbName, req.Entity.ObjectName) if resolveErr != nil { mlog.Warn(context.TODO(), "failed to resolve collection alias for select grant", mlog.String("objectName", req.Entity.ObjectName), mlog.String("dbName", req.Entity.DbName), mlog.Err(resolveErr)) @@ -5985,7 +5985,7 @@ func (node *Proxy) CreateResourceGroup(ctx context.Context, request *milvuspb.Cr mlog.Info(context.TODO(), "CreateResourceGroup received") - if err := node.sched.ddQueue.Enqueue(t); err != nil { + if err := node.sched.DdQueue.Enqueue(t); err != nil { mlog.Warn(context.TODO(), "CreateResourceGroup failed to enqueue", mlog.Err(err)) return getErrResponse(err, method, "", ""), nil @@ -6038,7 +6038,7 @@ func (node *Proxy) UpdateResourceGroups(ctx context.Context, request *milvuspb.U mlog.Info(context.TODO(), "UpdateResourceGroups received") - if err := node.sched.ddQueue.Enqueue(t); err != nil { + if err := node.sched.DdQueue.Enqueue(t); err != nil { mlog.Warn(context.TODO(), "UpdateResourceGroups failed to enqueue", mlog.Err(err)) return getErrResponse(err, method, "", ""), nil @@ -6086,7 +6086,7 @@ func (node *Proxy) DropResourceGroup(ctx context.Context, request *milvuspb.Drop mlog.Info(context.TODO(), "DropResourceGroup received") - if err := node.sched.ddQueue.Enqueue(t); err != nil { + if err := node.sched.DdQueue.Enqueue(t); err != nil { mlog.Warn(context.TODO(), "DropResourceGroup failed to enqueue", mlog.Err(err)) @@ -6145,7 +6145,7 @@ func (node *Proxy) TransferNode(ctx context.Context, request *milvuspb.TransferN mlog.Info(context.TODO(), "TransferNode received") - if err := node.sched.ddQueue.Enqueue(t); err != nil { + if err := node.sched.DdQueue.Enqueue(t); err != nil { mlog.Warn(context.TODO(), "TransferNode failed to enqueue", mlog.Err(err)) @@ -6196,7 +6196,7 @@ func (node *Proxy) TransferReplica(ctx context.Context, request *milvuspb.Transf defer sp.End() tr := timerecord.NewTimeRecorder(method) t := &TransferReplicaTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), TransferReplicaRequest: request, @@ -6205,7 +6205,7 @@ func (node *Proxy) TransferReplica(ctx context.Context, request *milvuspb.Transf mlog.Info(context.TODO(), "TransferReplica received") - if err := node.sched.ddQueue.Enqueue(t); err != nil { + if err := node.sched.DdQueue.Enqueue(t); err != nil { mlog.Warn(context.TODO(), "TransferReplica failed to enqueue", mlog.Err(err)) @@ -6252,7 +6252,7 @@ func (node *Proxy) ListResourceGroups(ctx context.Context, request *milvuspb.Lis mlog.Debug(context.TODO(), "ListResourceGroups received") - if err := node.sched.ddQueue.Enqueue(t); err != nil { + if err := node.sched.DdQueue.Enqueue(t); err != nil { mlog.Warn(context.TODO(), "ListResourceGroups failed to enqueue", mlog.Err(err)) @@ -6301,7 +6301,7 @@ func (node *Proxy) DescribeResourceGroup(ctx context.Context, request *milvuspb. defer sp.End() tr := timerecord.NewTimeRecorder(method) t := &DescribeResourceGroupTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), DescribeResourceGroupRequest: request, @@ -6310,7 +6310,7 @@ func (node *Proxy) DescribeResourceGroup(ctx context.Context, request *milvuspb. mlog.Debug(context.TODO(), "DescribeResourceGroup received") - if err := node.sched.ddQueue.Enqueue(t); err != nil { + if err := node.sched.DdQueue.Enqueue(t); err != nil { mlog.Warn(context.TODO(), "DescribeResourceGroup failed to enqueue", mlog.Err(err)) @@ -6464,7 +6464,7 @@ func (node *Proxy) ImportV2(ctx context.Context, req *internalpb.ImportRequest) } // Check for external collection - import is not supported - if err := checkExternalCollectionBlockedForWrite(ctx, node.getMetaCache(), req.GetDbName(), req.GetCollectionName(), "import"); err != nil { + if err := checkExternalCollectionBlockedForWrite(ctx, node.GetMetaCache(), req.GetDbName(), req.GetCollectionName(), "import"); err != nil { return &internalpb.ImportResponse{Status: merr.Status(err)}, nil } @@ -6478,7 +6478,7 @@ func (node *Proxy) ImportV2(ctx context.Context, req *internalpb.ImportRequest) nodeID := paramtable.GetStringNodeID() it := &importTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), req: req, @@ -6487,7 +6487,7 @@ func (node *Proxy) ImportV2(ctx context.Context, req *internalpb.ImportRequest) resp: resp, } - if err := node.sched.dmQueue.Enqueue(it); err != nil { + if err := node.sched.DmQueue.Enqueue(it); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) @@ -6554,7 +6554,7 @@ func (node *Proxy) ListImports(ctx context.Context, req *internalpb.ListImportsR collectionID UniqueID ) if req.GetCollectionName() != "" { - collectionID, err = node.getMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) + collectionID, err = node.GetMetaCache().GetCollectionID(ctx, req.GetDbName(), req.GetCollectionName()) if err != nil { resp.Status = merr.Status(err) return resp, nil @@ -6807,14 +6807,14 @@ func (node *Proxy) RunAnalyzer(ctx context.Context, req *milvuspb.RunAnalyzerReq method := "RunAnalyzer" task := &RunAnalyzerTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, lb: node.lbPolicy, Condition: NewTaskCondition(ctx), RunAnalyzerRequest: req, } - if err := node.sched.dqQueue.Enqueue(task); err != nil { + if err := node.sched.DqQueue.Enqueue(task); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err), @@ -7471,14 +7471,14 @@ func (node *Proxy) BatchUpdateManifest(ctx context.Context, req *milvuspb.BatchU nodeID := fmt.Sprint(paramtable.GetNodeID()) bt := &batchUpdateManifestTask{ - baseTask: baseTask{metaCache: node.getMetaCache()}, + baseTask: baseTask{MetaCache: node.GetMetaCache()}, ctx: ctx, Condition: NewTaskCondition(ctx), req: req, mixCoord: node.mixCoord, } - if err := node.sched.ddQueue.Enqueue(bt); err != nil { + if err := node.sched.DdQueue.Enqueue(bt); err != nil { mlog.Warn(ctx, rpcFailedToEnqueue(method), mlog.Err(err)) return merr.Status(err), nil } @@ -7550,7 +7550,7 @@ func (node *Proxy) RefreshExternalCollection(ctx context.Context, req *milvuspb. } // Get collection info from cache (includes schema for validation) - collectionInfo, err := node.getMetaCache().GetCollectionInfo(ctx, req.GetDbName(), req.GetCollectionName(), 0) + collectionInfo, err := node.GetMetaCache().GetCollectionInfo(ctx, req.GetDbName(), req.GetCollectionName(), 0) if err != nil { mlog.Warn(context.TODO(), "failed to get collection info", mlog.Err(err)) return &milvuspb.RefreshExternalCollectionResponse{ @@ -7684,7 +7684,7 @@ func (node *Proxy) ListRefreshExternalCollectionJobs(ctx context.Context, req *m } // Get collection info from cache and validate it's an external collection - collectionInfo, err := node.getMetaCache().GetCollectionInfo(ctx, req.GetDbName(), req.GetCollectionName(), 0) + collectionInfo, err := node.GetMetaCache().GetCollectionInfo(ctx, req.GetDbName(), req.GetCollectionName(), 0) if err != nil { mlog.Warn(context.TODO(), "failed to get collection info", mlog.Err(err)) return &milvuspb.ListRefreshExternalCollectionJobsResponse{ diff --git a/internal/proxy/impl_test.go b/internal/proxy/impl_test.go index 5e1488009c..f6567c929c 100644 --- a/internal/proxy/impl_test.go +++ b/internal/proxy/impl_test.go @@ -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 = ×tampAllocator{ 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() diff --git a/internal/proxy/interface_def.go b/internal/proxy/interface_def.go index baac2a1e19..9494c79fdf 100644 --- a/internal/proxy/interface_def.go +++ b/internal/proxy/interface_def.go @@ -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 diff --git a/internal/proxy/membership_filter_plan_size_test.go b/internal/proxy/membership_filter_plan_size_test.go index 7b43a91b7c..bd3a6f6ea6 100644 --- a/internal/proxy/membership_filter_plan_size_test.go +++ b/internal/proxy/membership_filter_plan_size_test.go @@ -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{ diff --git a/internal/proxy/metrics_info.go b/internal/proxy/metrics_info.go index c6ad6cbcbd..dcbf8261e4 100644 --- a/internal/proxy/metrics_info.go +++ b/internal/proxy/metrics_info.go @@ -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 } diff --git a/internal/proxy/metrics_info_test.go b/internal/proxy/metrics_info_test.go index 7467aebfe9..7f02b6ef5b 100644 --- a/internal/proxy/metrics_info_test.go +++ b/internal/proxy/metrics_info_test.go @@ -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) { diff --git a/internal/proxy/mock_test.go b/internal/proxy/mock_test.go index 4fd3cf7b43..eea3bf056c 100644 --- a/internal/proxy/mock_test.go +++ b/internal/proxy/mock_test.go @@ -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 diff --git a/internal/proxy/privilege_interceptor.go b/internal/proxy/privilege_interceptor.go index 89615f6c90..2a5721b547 100644 --- a/internal/proxy/privilege_interceptor.go +++ b/internal/proxy/privilege_interceptor.go @@ -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) diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 11b69c5117..41b19383cb 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -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() } diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index 9aebecfecb..2f323f52d7 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -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 = ×tampAllocator{ 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()) +} diff --git a/internal/proxy/rate_limit_interceptor.go b/internal/proxy/rate_limit_interceptor.go index b24bec7d6e..507feda335 100644 --- a/internal/proxy/rate_limit_interceptor.go +++ b/internal/proxy/rate_limit_interceptor.go @@ -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) diff --git a/internal/proxy/scheduler/OWNERS b/internal/proxy/scheduler/OWNERS new file mode 100644 index 0000000000..2e385bfde2 --- /dev/null +++ b/internal/proxy/scheduler/OWNERS @@ -0,0 +1,7 @@ +reviewers: + - czs007 + - bigsheeper + - congqixia + +approvers: + - maintainers diff --git a/internal/proxy/scheduler/README.md b/internal/proxy/scheduler/README.md new file mode 100644 index 0000000000..eff29eb32e --- /dev/null +++ b/internal/proxy/scheduler/README.md @@ -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`. diff --git a/internal/proxy/scheduler/mock_cache_test.go b/internal/proxy/scheduler/mock_cache_test.go new file mode 100644 index 0000000000..4343a71927 --- /dev/null +++ b/internal/proxy/scheduler/mock_cache_test.go @@ -0,0 +1,1137 @@ +// Code generated by mockery v2.53.3. DO NOT EDIT. + +package scheduler + +import ( + context "context" + + mock "github.com/stretchr/testify/mock" + + "github.com/milvus-io/milvus/internal/proxy/metacache" +) + +// MockCache is an autogenerated mock type for the Cache type +type MockCache struct { + mock.Mock +} + +type MockCache_Expecter struct { + mock *mock.Mock +} + +func (_m *MockCache) EXPECT() *MockCache_Expecter { + return &MockCache_Expecter{mock: &_m.Mock} +} + +// AllocID provides a mock function with given fields: ctx +func (_m *MockCache) AllocID(ctx context.Context) (int64, error) { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for AllocID") + } + + var r0 int64 + var r1 error + if rf, ok := ret.Get(0).(func(context.Context) (int64, error)); ok { + return rf(ctx) + } + if rf, ok := ret.Get(0).(func(context.Context) int64); ok { + r0 = rf(ctx) + } else { + r0 = ret.Get(0).(int64) + } + + if rf, ok := ret.Get(1).(func(context.Context) error); ok { + r1 = rf(ctx) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_AllocID_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AllocID' +type MockCache_AllocID_Call struct { + *mock.Call +} + +// AllocID is a helper method to define mock.On call +// - ctx context.Context +func (_e *MockCache_Expecter) AllocID(ctx interface{}) *MockCache_AllocID_Call { + return &MockCache_AllocID_Call{Call: _e.mock.On("AllocID", ctx)} +} + +func (_c *MockCache_AllocID_Call) Run(run func(ctx context.Context)) *MockCache_AllocID_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context)) + }) + return _c +} + +func (_c *MockCache_AllocID_Call) Return(_a0 int64, _a1 error) *MockCache_AllocID_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_AllocID_Call) RunAndReturn(run func(context.Context) (int64, error)) *MockCache_AllocID_Call { + _c.Call.Return(run) + return _c +} + +// Close provides a mock function with no fields +func (_m *MockCache) Close() { + _m.Called() +} + +// MockCache_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close' +type MockCache_Close_Call struct { + *mock.Call +} + +// Close is a helper method to define mock.On call +func (_e *MockCache_Expecter) Close() *MockCache_Close_Call { + return &MockCache_Close_Call{Call: _e.mock.On("Close")} +} + +func (_c *MockCache_Close_Call) Run(run func()) *MockCache_Close_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockCache_Close_Call) Return() *MockCache_Close_Call { + _c.Call.Return() + return _c +} + +func (_c *MockCache_Close_Call) RunAndReturn(run func()) *MockCache_Close_Call { + _c.Run(run) + return _c +} + +// GetCollectionID provides a mock function with given fields: ctx, database, collectionName +func (_m *MockCache) GetCollectionID(ctx context.Context, database string, collectionName string) (int64, error) { + ret := _m.Called(ctx, database, collectionName) + + if len(ret) == 0 { + panic("no return value specified for GetCollectionID") + } + + var r0 int64 + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string, string) (int64, error)); ok { + return rf(ctx, database, collectionName) + } + if rf, ok := ret.Get(0).(func(context.Context, string, string) int64); ok { + r0 = rf(ctx, database, collectionName) + } else { + r0 = ret.Get(0).(int64) + } + + if rf, ok := ret.Get(1).(func(context.Context, string, string) error); ok { + r1 = rf(ctx, database, collectionName) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_GetCollectionID_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetCollectionID' +type MockCache_GetCollectionID_Call struct { + *mock.Call +} + +// GetCollectionID is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionName string +func (_e *MockCache_Expecter) GetCollectionID(ctx interface{}, database interface{}, collectionName interface{}) *MockCache_GetCollectionID_Call { + return &MockCache_GetCollectionID_Call{Call: _e.mock.On("GetCollectionID", ctx, database, collectionName)} +} + +func (_c *MockCache_GetCollectionID_Call) Run(run func(ctx context.Context, database string, collectionName string)) *MockCache_GetCollectionID_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string)) + }) + return _c +} + +func (_c *MockCache_GetCollectionID_Call) Return(_a0 int64, _a1 error) *MockCache_GetCollectionID_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_GetCollectionID_Call) RunAndReturn(run func(context.Context, string, string) (int64, error)) *MockCache_GetCollectionID_Call { + _c.Call.Return(run) + return _c +} + +// GetCollectionInfo provides a mock function with given fields: ctx, database, collectionName, collectionID +func (_m *MockCache) GetCollectionInfo(ctx context.Context, database string, collectionName string, collectionID int64) (*metacache.CollectionInfo, error) { + ret := _m.Called(ctx, database, collectionName, collectionID) + + if len(ret) == 0 { + panic("no return value specified for GetCollectionInfo") + } + + var r0 *metacache.CollectionInfo + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) (*metacache.CollectionInfo, error)); ok { + return rf(ctx, database, collectionName, collectionID) + } + if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) *metacache.CollectionInfo); ok { + r0 = rf(ctx, database, collectionName, collectionID) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*metacache.CollectionInfo) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, string, string, int64) error); ok { + r1 = rf(ctx, database, collectionName, collectionID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_GetCollectionInfo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetCollectionInfo' +type MockCache_GetCollectionInfo_Call struct { + *mock.Call +} + +// GetCollectionInfo is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionName string +// - collectionID int64 +func (_e *MockCache_Expecter) GetCollectionInfo(ctx interface{}, database interface{}, collectionName interface{}, collectionID interface{}) *MockCache_GetCollectionInfo_Call { + return &MockCache_GetCollectionInfo_Call{Call: _e.mock.On("GetCollectionInfo", ctx, database, collectionName, collectionID)} +} + +func (_c *MockCache_GetCollectionInfo_Call) Run(run func(ctx context.Context, database string, collectionName string, collectionID int64)) *MockCache_GetCollectionInfo_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string), args[3].(int64)) + }) + return _c +} + +func (_c *MockCache_GetCollectionInfo_Call) Return(_a0 *metacache.CollectionInfo, _a1 error) *MockCache_GetCollectionInfo_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_GetCollectionInfo_Call) RunAndReturn(run func(context.Context, string, string, int64) (*metacache.CollectionInfo, error)) *MockCache_GetCollectionInfo_Call { + _c.Call.Return(run) + return _c +} + +// GetCollectionName provides a mock function with given fields: ctx, database, collectionID +func (_m *MockCache) GetCollectionName(ctx context.Context, database string, collectionID int64) (string, error) { + ret := _m.Called(ctx, database, collectionID) + + if len(ret) == 0 { + panic("no return value specified for GetCollectionName") + } + + var r0 string + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string, int64) (string, error)); ok { + return rf(ctx, database, collectionID) + } + if rf, ok := ret.Get(0).(func(context.Context, string, int64) string); ok { + r0 = rf(ctx, database, collectionID) + } else { + r0 = ret.Get(0).(string) + } + + if rf, ok := ret.Get(1).(func(context.Context, string, int64) error); ok { + r1 = rf(ctx, database, collectionID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_GetCollectionName_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetCollectionName' +type MockCache_GetCollectionName_Call struct { + *mock.Call +} + +// GetCollectionName is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionID int64 +func (_e *MockCache_Expecter) GetCollectionName(ctx interface{}, database interface{}, collectionID interface{}) *MockCache_GetCollectionName_Call { + return &MockCache_GetCollectionName_Call{Call: _e.mock.On("GetCollectionName", ctx, database, collectionID)} +} + +func (_c *MockCache_GetCollectionName_Call) Run(run func(ctx context.Context, database string, collectionID int64)) *MockCache_GetCollectionName_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(int64)) + }) + return _c +} + +func (_c *MockCache_GetCollectionName_Call) Return(_a0 string, _a1 error) *MockCache_GetCollectionName_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_GetCollectionName_Call) RunAndReturn(run func(context.Context, string, int64) (string, error)) *MockCache_GetCollectionName_Call { + _c.Call.Return(run) + return _c +} + +// GetCollectionSchema provides a mock function with given fields: ctx, database, collectionName +func (_m *MockCache) GetCollectionSchema(ctx context.Context, database string, collectionName string) (*metacache.SchemaInfo, error) { + ret := _m.Called(ctx, database, collectionName) + + if len(ret) == 0 { + panic("no return value specified for GetCollectionSchema") + } + + var r0 *metacache.SchemaInfo + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string, string) (*metacache.SchemaInfo, error)); ok { + return rf(ctx, database, collectionName) + } + if rf, ok := ret.Get(0).(func(context.Context, string, string) *metacache.SchemaInfo); ok { + r0 = rf(ctx, database, collectionName) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*metacache.SchemaInfo) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, string, string) error); ok { + r1 = rf(ctx, database, collectionName) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_GetCollectionSchema_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetCollectionSchema' +type MockCache_GetCollectionSchema_Call struct { + *mock.Call +} + +// GetCollectionSchema is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionName string +func (_e *MockCache_Expecter) GetCollectionSchema(ctx interface{}, database interface{}, collectionName interface{}) *MockCache_GetCollectionSchema_Call { + return &MockCache_GetCollectionSchema_Call{Call: _e.mock.On("GetCollectionSchema", ctx, database, collectionName)} +} + +func (_c *MockCache_GetCollectionSchema_Call) Run(run func(ctx context.Context, database string, collectionName string)) *MockCache_GetCollectionSchema_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string)) + }) + return _c +} + +func (_c *MockCache_GetCollectionSchema_Call) Return(_a0 *metacache.SchemaInfo, _a1 error) *MockCache_GetCollectionSchema_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_GetCollectionSchema_Call) RunAndReturn(run func(context.Context, string, string) (*metacache.SchemaInfo, error)) *MockCache_GetCollectionSchema_Call { + _c.Call.Return(run) + return _c +} + +// GetDatabaseInfo provides a mock function with given fields: ctx, database +func (_m *MockCache) GetDatabaseInfo(ctx context.Context, database string) (*metacache.DatabaseInfo, error) { + ret := _m.Called(ctx, database) + + if len(ret) == 0 { + panic("no return value specified for GetDatabaseInfo") + } + + var r0 *metacache.DatabaseInfo + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string) (*metacache.DatabaseInfo, error)); ok { + return rf(ctx, database) + } + if rf, ok := ret.Get(0).(func(context.Context, string) *metacache.DatabaseInfo); ok { + r0 = rf(ctx, database) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*metacache.DatabaseInfo) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, database) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_GetDatabaseInfo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetDatabaseInfo' +type MockCache_GetDatabaseInfo_Call struct { + *mock.Call +} + +// GetDatabaseInfo is a helper method to define mock.On call +// - ctx context.Context +// - database string +func (_e *MockCache_Expecter) GetDatabaseInfo(ctx interface{}, database interface{}) *MockCache_GetDatabaseInfo_Call { + return &MockCache_GetDatabaseInfo_Call{Call: _e.mock.On("GetDatabaseInfo", ctx, database)} +} + +func (_c *MockCache_GetDatabaseInfo_Call) Run(run func(ctx context.Context, database string)) *MockCache_GetDatabaseInfo_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string)) + }) + return _c +} + +func (_c *MockCache_GetDatabaseInfo_Call) Return(_a0 *metacache.DatabaseInfo, _a1 error) *MockCache_GetDatabaseInfo_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_GetDatabaseInfo_Call) RunAndReturn(run func(context.Context, string) (*metacache.DatabaseInfo, error)) *MockCache_GetDatabaseInfo_Call { + _c.Call.Return(run) + return _c +} + +// GetPartitionID provides a mock function with given fields: ctx, database, collectionName, partitionName +func (_m *MockCache) GetPartitionID(ctx context.Context, database string, collectionName string, partitionName string) (int64, error) { + ret := _m.Called(ctx, database, collectionName, partitionName) + + if len(ret) == 0 { + panic("no return value specified for GetPartitionID") + } + + var r0 int64 + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string, string, string) (int64, error)); ok { + return rf(ctx, database, collectionName, partitionName) + } + if rf, ok := ret.Get(0).(func(context.Context, string, string, string) int64); ok { + r0 = rf(ctx, database, collectionName, partitionName) + } else { + r0 = ret.Get(0).(int64) + } + + if rf, ok := ret.Get(1).(func(context.Context, string, string, string) error); ok { + r1 = rf(ctx, database, collectionName, partitionName) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_GetPartitionID_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetPartitionID' +type MockCache_GetPartitionID_Call struct { + *mock.Call +} + +// GetPartitionID is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionName string +// - partitionName string +func (_e *MockCache_Expecter) GetPartitionID(ctx interface{}, database interface{}, collectionName interface{}, partitionName interface{}) *MockCache_GetPartitionID_Call { + return &MockCache_GetPartitionID_Call{Call: _e.mock.On("GetPartitionID", ctx, database, collectionName, partitionName)} +} + +func (_c *MockCache_GetPartitionID_Call) Run(run func(ctx context.Context, database string, collectionName string, partitionName string)) *MockCache_GetPartitionID_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string), args[3].(string)) + }) + return _c +} + +func (_c *MockCache_GetPartitionID_Call) Return(_a0 int64, _a1 error) *MockCache_GetPartitionID_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_GetPartitionID_Call) RunAndReturn(run func(context.Context, string, string, string) (int64, error)) *MockCache_GetPartitionID_Call { + _c.Call.Return(run) + return _c +} + +// GetPartitionInfo provides a mock function with given fields: ctx, database, collectionName, partitionName +func (_m *MockCache) GetPartitionInfo(ctx context.Context, database string, collectionName string, partitionName string) (*metacache.PartitionInfo, error) { + ret := _m.Called(ctx, database, collectionName, partitionName) + + if len(ret) == 0 { + panic("no return value specified for GetPartitionInfo") + } + + var r0 *metacache.PartitionInfo + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string, string, string) (*metacache.PartitionInfo, error)); ok { + return rf(ctx, database, collectionName, partitionName) + } + if rf, ok := ret.Get(0).(func(context.Context, string, string, string) *metacache.PartitionInfo); ok { + r0 = rf(ctx, database, collectionName, partitionName) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*metacache.PartitionInfo) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, string, string, string) error); ok { + r1 = rf(ctx, database, collectionName, partitionName) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_GetPartitionInfo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetPartitionInfo' +type MockCache_GetPartitionInfo_Call struct { + *mock.Call +} + +// GetPartitionInfo is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionName string +// - partitionName string +func (_e *MockCache_Expecter) GetPartitionInfo(ctx interface{}, database interface{}, collectionName interface{}, partitionName interface{}) *MockCache_GetPartitionInfo_Call { + return &MockCache_GetPartitionInfo_Call{Call: _e.mock.On("GetPartitionInfo", ctx, database, collectionName, partitionName)} +} + +func (_c *MockCache_GetPartitionInfo_Call) Run(run func(ctx context.Context, database string, collectionName string, partitionName string)) *MockCache_GetPartitionInfo_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string), args[3].(string)) + }) + return _c +} + +func (_c *MockCache_GetPartitionInfo_Call) Return(_a0 *metacache.PartitionInfo, _a1 error) *MockCache_GetPartitionInfo_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_GetPartitionInfo_Call) RunAndReturn(run func(context.Context, string, string, string) (*metacache.PartitionInfo, error)) *MockCache_GetPartitionInfo_Call { + _c.Call.Return(run) + return _c +} + +// GetPartitionName provides a mock function with given fields: ctx, database, collectionName, partitionID +func (_m *MockCache) GetPartitionName(ctx context.Context, database string, collectionName string, partitionID int64) (string, error) { + ret := _m.Called(ctx, database, collectionName, partitionID) + + if len(ret) == 0 { + panic("no return value specified for GetPartitionName") + } + + var r0 string + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) (string, error)); ok { + return rf(ctx, database, collectionName, partitionID) + } + if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) string); ok { + r0 = rf(ctx, database, collectionName, partitionID) + } else { + r0 = ret.Get(0).(string) + } + + if rf, ok := ret.Get(1).(func(context.Context, string, string, int64) error); ok { + r1 = rf(ctx, database, collectionName, partitionID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_GetPartitionName_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetPartitionName' +type MockCache_GetPartitionName_Call struct { + *mock.Call +} + +// GetPartitionName is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionName string +// - partitionID int64 +func (_e *MockCache_Expecter) GetPartitionName(ctx interface{}, database interface{}, collectionName interface{}, partitionID interface{}) *MockCache_GetPartitionName_Call { + return &MockCache_GetPartitionName_Call{Call: _e.mock.On("GetPartitionName", ctx, database, collectionName, partitionID)} +} + +func (_c *MockCache_GetPartitionName_Call) Run(run func(ctx context.Context, database string, collectionName string, partitionID int64)) *MockCache_GetPartitionName_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string), args[3].(int64)) + }) + return _c +} + +func (_c *MockCache_GetPartitionName_Call) Return(_a0 string, _a1 error) *MockCache_GetPartitionName_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_GetPartitionName_Call) RunAndReturn(run func(context.Context, string, string, int64) (string, error)) *MockCache_GetPartitionName_Call { + _c.Call.Return(run) + return _c +} + +// GetPartitions provides a mock function with given fields: ctx, database, collectionName +func (_m *MockCache) GetPartitions(ctx context.Context, database string, collectionName string) (map[string]int64, error) { + ret := _m.Called(ctx, database, collectionName) + + if len(ret) == 0 { + panic("no return value specified for GetPartitions") + } + + var r0 map[string]int64 + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string, string) (map[string]int64, error)); ok { + return rf(ctx, database, collectionName) + } + if rf, ok := ret.Get(0).(func(context.Context, string, string) map[string]int64); ok { + r0 = rf(ctx, database, collectionName) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(map[string]int64) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, string, string) error); ok { + r1 = rf(ctx, database, collectionName) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_GetPartitions_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetPartitions' +type MockCache_GetPartitions_Call struct { + *mock.Call +} + +// GetPartitions is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionName string +func (_e *MockCache_Expecter) GetPartitions(ctx interface{}, database interface{}, collectionName interface{}) *MockCache_GetPartitions_Call { + return &MockCache_GetPartitions_Call{Call: _e.mock.On("GetPartitions", ctx, database, collectionName)} +} + +func (_c *MockCache_GetPartitions_Call) Run(run func(ctx context.Context, database string, collectionName string)) *MockCache_GetPartitions_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string)) + }) + return _c +} + +func (_c *MockCache_GetPartitions_Call) Return(_a0 map[string]int64, _a1 error) *MockCache_GetPartitions_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_GetPartitions_Call) RunAndReturn(run func(context.Context, string, string) (map[string]int64, error)) *MockCache_GetPartitions_Call { + _c.Call.Return(run) + return _c +} + +// GetPartitionsIndex provides a mock function with given fields: ctx, database, collectionName +func (_m *MockCache) GetPartitionsIndex(ctx context.Context, database string, collectionName string) ([]string, error) { + ret := _m.Called(ctx, database, collectionName) + + if len(ret) == 0 { + panic("no return value specified for GetPartitionsIndex") + } + + var r0 []string + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string, string) ([]string, error)); ok { + return rf(ctx, database, collectionName) + } + if rf, ok := ret.Get(0).(func(context.Context, string, string) []string); ok { + r0 = rf(ctx, database, collectionName) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]string) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, string, string) error); ok { + r1 = rf(ctx, database, collectionName) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_GetPartitionsIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetPartitionsIndex' +type MockCache_GetPartitionsIndex_Call struct { + *mock.Call +} + +// GetPartitionsIndex is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionName string +func (_e *MockCache_Expecter) GetPartitionsIndex(ctx interface{}, database interface{}, collectionName interface{}) *MockCache_GetPartitionsIndex_Call { + return &MockCache_GetPartitionsIndex_Call{Call: _e.mock.On("GetPartitionsIndex", ctx, database, collectionName)} +} + +func (_c *MockCache_GetPartitionsIndex_Call) Run(run func(ctx context.Context, database string, collectionName string)) *MockCache_GetPartitionsIndex_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string)) + }) + return _c +} + +func (_c *MockCache_GetPartitionsIndex_Call) Return(_a0 []string, _a1 error) *MockCache_GetPartitionsIndex_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_GetPartitionsIndex_Call) RunAndReturn(run func(context.Context, string, string) ([]string, error)) *MockCache_GetPartitionsIndex_Call { + _c.Call.Return(run) + return _c +} + +// HasDatabase provides a mock function with given fields: ctx, database +func (_m *MockCache) HasDatabase(ctx context.Context, database string) bool { + ret := _m.Called(ctx, database) + + if len(ret) == 0 { + panic("no return value specified for HasDatabase") + } + + var r0 bool + if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok { + r0 = rf(ctx, database) + } else { + r0 = ret.Get(0).(bool) + } + + return r0 +} + +// MockCache_HasDatabase_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HasDatabase' +type MockCache_HasDatabase_Call struct { + *mock.Call +} + +// HasDatabase is a helper method to define mock.On call +// - ctx context.Context +// - database string +func (_e *MockCache_Expecter) HasDatabase(ctx interface{}, database interface{}) *MockCache_HasDatabase_Call { + return &MockCache_HasDatabase_Call{Call: _e.mock.On("HasDatabase", ctx, database)} +} + +func (_c *MockCache_HasDatabase_Call) Run(run func(ctx context.Context, database string)) *MockCache_HasDatabase_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string)) + }) + return _c +} + +func (_c *MockCache_HasDatabase_Call) Return(_a0 bool) *MockCache_HasDatabase_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockCache_HasDatabase_Call) RunAndReturn(run func(context.Context, string) bool) *MockCache_HasDatabase_Call { + _c.Call.Return(run) + return _c +} + +// InvalidateCollectionMeta provides a mock function with given fields: ctx, database, collectionName, collectionID, removeAlias +func (_m *MockCache) InvalidateCollectionMeta(ctx context.Context, database string, collectionName string, collectionID int64, removeAlias bool) []string { + ret := _m.Called(ctx, database, collectionName, collectionID, removeAlias) + + if len(ret) == 0 { + panic("no return value specified for InvalidateCollectionMeta") + } + + var r0 []string + if rf, ok := ret.Get(0).(func(context.Context, string, string, int64, bool) []string); ok { + r0 = rf(ctx, database, collectionName, collectionID, removeAlias) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]string) + } + } + + return r0 +} + +// MockCache_InvalidateCollectionMeta_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'InvalidateCollectionMeta' +type MockCache_InvalidateCollectionMeta_Call struct { + *mock.Call +} + +// InvalidateCollectionMeta is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionName string +// - collectionID int64 +// - removeAlias bool +func (_e *MockCache_Expecter) InvalidateCollectionMeta(ctx interface{}, database interface{}, collectionName interface{}, collectionID interface{}, removeAlias interface{}) *MockCache_InvalidateCollectionMeta_Call { + return &MockCache_InvalidateCollectionMeta_Call{Call: _e.mock.On("InvalidateCollectionMeta", ctx, database, collectionName, collectionID, removeAlias)} +} + +func (_c *MockCache_InvalidateCollectionMeta_Call) Run(run func(ctx context.Context, database string, collectionName string, collectionID int64, removeAlias bool)) *MockCache_InvalidateCollectionMeta_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string), args[3].(int64), args[4].(bool)) + }) + return _c +} + +func (_c *MockCache_InvalidateCollectionMeta_Call) Return(_a0 []string) *MockCache_InvalidateCollectionMeta_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockCache_InvalidateCollectionMeta_Call) RunAndReturn(run func(context.Context, string, string, int64, bool) []string) *MockCache_InvalidateCollectionMeta_Call { + _c.Call.Return(run) + return _c +} + +// RemoveAlias provides a mock function with given fields: ctx, database, alias +func (_m *MockCache) RemoveAlias(ctx context.Context, database string, alias string) { + _m.Called(ctx, database, alias) +} + +// MockCache_RemoveAlias_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveAlias' +type MockCache_RemoveAlias_Call struct { + *mock.Call +} + +// RemoveAlias is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - alias string +func (_e *MockCache_Expecter) RemoveAlias(ctx interface{}, database interface{}, alias interface{}) *MockCache_RemoveAlias_Call { + return &MockCache_RemoveAlias_Call{Call: _e.mock.On("RemoveAlias", ctx, database, alias)} +} + +func (_c *MockCache_RemoveAlias_Call) Run(run func(ctx context.Context, database string, alias string)) *MockCache_RemoveAlias_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string)) + }) + return _c +} + +func (_c *MockCache_RemoveAlias_Call) Return() *MockCache_RemoveAlias_Call { + _c.Call.Return() + return _c +} + +func (_c *MockCache_RemoveAlias_Call) RunAndReturn(run func(context.Context, string, string)) *MockCache_RemoveAlias_Call { + _c.Run(run) + return _c +} + +// RemoveAliasHolders provides a mock function with given fields: ctx, database, alias +func (_m *MockCache) RemoveAliasHolders(ctx context.Context, database string, alias string) { + _m.Called(ctx, database, alias) +} + +// MockCache_RemoveAliasHolders_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveAliasHolders' +type MockCache_RemoveAliasHolders_Call struct { + *mock.Call +} + +// RemoveAliasHolders is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - alias string +func (_e *MockCache_Expecter) RemoveAliasHolders(ctx interface{}, database interface{}, alias interface{}) *MockCache_RemoveAliasHolders_Call { + return &MockCache_RemoveAliasHolders_Call{Call: _e.mock.On("RemoveAliasHolders", ctx, database, alias)} +} + +func (_c *MockCache_RemoveAliasHolders_Call) Run(run func(ctx context.Context, database string, alias string)) *MockCache_RemoveAliasHolders_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string)) + }) + return _c +} + +func (_c *MockCache_RemoveAliasHolders_Call) Return() *MockCache_RemoveAliasHolders_Call { + _c.Call.Return() + return _c +} + +func (_c *MockCache_RemoveAliasHolders_Call) RunAndReturn(run func(context.Context, string, string)) *MockCache_RemoveAliasHolders_Call { + _c.Run(run) + return _c +} + +// RemoveCollection provides a mock function with given fields: ctx, database, collectionName +func (_m *MockCache) RemoveCollection(ctx context.Context, database string, collectionName string) { + _m.Called(ctx, database, collectionName) +} + +// MockCache_RemoveCollection_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveCollection' +type MockCache_RemoveCollection_Call struct { + *mock.Call +} + +// RemoveCollection is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionName string +func (_e *MockCache_Expecter) RemoveCollection(ctx interface{}, database interface{}, collectionName interface{}) *MockCache_RemoveCollection_Call { + return &MockCache_RemoveCollection_Call{Call: _e.mock.On("RemoveCollection", ctx, database, collectionName)} +} + +func (_c *MockCache_RemoveCollection_Call) Run(run func(ctx context.Context, database string, collectionName string)) *MockCache_RemoveCollection_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string)) + }) + return _c +} + +func (_c *MockCache_RemoveCollection_Call) Return() *MockCache_RemoveCollection_Call { + _c.Call.Return() + return _c +} + +func (_c *MockCache_RemoveCollection_Call) RunAndReturn(run func(context.Context, string, string)) *MockCache_RemoveCollection_Call { + _c.Run(run) + return _c +} + +// RemoveCollectionsByID provides a mock function with given fields: ctx, collectionID +func (_m *MockCache) RemoveCollectionsByID(ctx context.Context, collectionID int64) []string { + ret := _m.Called(ctx, collectionID) + + if len(ret) == 0 { + panic("no return value specified for RemoveCollectionsByID") + } + + var r0 []string + if rf, ok := ret.Get(0).(func(context.Context, int64) []string); ok { + r0 = rf(ctx, collectionID) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]string) + } + } + + return r0 +} + +// MockCache_RemoveCollectionsByID_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveCollectionsByID' +type MockCache_RemoveCollectionsByID_Call struct { + *mock.Call +} + +// RemoveCollectionsByID is a helper method to define mock.On call +// - ctx context.Context +// - collectionID int64 +func (_e *MockCache_Expecter) RemoveCollectionsByID(ctx interface{}, collectionID interface{}) *MockCache_RemoveCollectionsByID_Call { + return &MockCache_RemoveCollectionsByID_Call{Call: _e.mock.On("RemoveCollectionsByID", ctx, collectionID)} +} + +func (_c *MockCache_RemoveCollectionsByID_Call) Run(run func(ctx context.Context, collectionID int64)) *MockCache_RemoveCollectionsByID_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(int64)) + }) + return _c +} + +func (_c *MockCache_RemoveCollectionsByID_Call) Return(_a0 []string) *MockCache_RemoveCollectionsByID_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockCache_RemoveCollectionsByID_Call) RunAndReturn(run func(context.Context, int64) []string) *MockCache_RemoveCollectionsByID_Call { + _c.Call.Return(run) + return _c +} + +// RemoveDatabase provides a mock function with given fields: ctx, database +func (_m *MockCache) RemoveDatabase(ctx context.Context, database string) { + _m.Called(ctx, database) +} + +// MockCache_RemoveDatabase_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveDatabase' +type MockCache_RemoveDatabase_Call struct { + *mock.Call +} + +// RemoveDatabase is a helper method to define mock.On call +// - ctx context.Context +// - database string +func (_e *MockCache_Expecter) RemoveDatabase(ctx interface{}, database interface{}) *MockCache_RemoveDatabase_Call { + return &MockCache_RemoveDatabase_Call{Call: _e.mock.On("RemoveDatabase", ctx, database)} +} + +func (_c *MockCache_RemoveDatabase_Call) Run(run func(ctx context.Context, database string)) *MockCache_RemoveDatabase_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string)) + }) + return _c +} + +func (_c *MockCache_RemoveDatabase_Call) Return() *MockCache_RemoveDatabase_Call { + _c.Call.Return() + return _c +} + +func (_c *MockCache_RemoveDatabase_Call) RunAndReturn(run func(context.Context, string)) *MockCache_RemoveDatabase_Call { + _c.Run(run) + return _c +} + +// RemoveDatabaseInfo provides a mock function with given fields: ctx, database +func (_m *MockCache) RemoveDatabaseInfo(ctx context.Context, database string) { + _m.Called(ctx, database) +} + +// MockCache_RemoveDatabaseInfo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveDatabaseInfo' +type MockCache_RemoveDatabaseInfo_Call struct { + *mock.Call +} + +// RemoveDatabaseInfo is a helper method to define mock.On call +// - ctx context.Context +// - database string +func (_e *MockCache_Expecter) RemoveDatabaseInfo(ctx interface{}, database interface{}) *MockCache_RemoveDatabaseInfo_Call { + return &MockCache_RemoveDatabaseInfo_Call{Call: _e.mock.On("RemoveDatabaseInfo", ctx, database)} +} + +func (_c *MockCache_RemoveDatabaseInfo_Call) Run(run func(ctx context.Context, database string)) *MockCache_RemoveDatabaseInfo_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string)) + }) + return _c +} + +func (_c *MockCache_RemoveDatabaseInfo_Call) Return() *MockCache_RemoveDatabaseInfo_Call { + _c.Call.Return() + return _c +} + +func (_c *MockCache_RemoveDatabaseInfo_Call) RunAndReturn(run func(context.Context, string)) *MockCache_RemoveDatabaseInfo_Call { + _c.Run(run) + return _c +} + +// RemovePartition provides a mock function with given fields: ctx, database, collectionID, collectionName, partitionName +func (_m *MockCache) RemovePartition(ctx context.Context, database string, collectionID int64, collectionName string, partitionName string) { + _m.Called(ctx, database, collectionID, collectionName, partitionName) +} + +// MockCache_RemovePartition_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemovePartition' +type MockCache_RemovePartition_Call struct { + *mock.Call +} + +// RemovePartition is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - collectionID int64 +// - collectionName string +// - partitionName string +func (_e *MockCache_Expecter) RemovePartition(ctx interface{}, database interface{}, collectionID interface{}, collectionName interface{}, partitionName interface{}) *MockCache_RemovePartition_Call { + return &MockCache_RemovePartition_Call{Call: _e.mock.On("RemovePartition", ctx, database, collectionID, collectionName, partitionName)} +} + +func (_c *MockCache_RemovePartition_Call) Run(run func(ctx context.Context, database string, collectionID int64, collectionName string, partitionName string)) *MockCache_RemovePartition_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(int64), args[3].(string), args[4].(string)) + }) + return _c +} + +func (_c *MockCache_RemovePartition_Call) Return() *MockCache_RemovePartition_Call { + _c.Call.Return() + return _c +} + +func (_c *MockCache_RemovePartition_Call) RunAndReturn(run func(context.Context, string, int64, string, string)) *MockCache_RemovePartition_Call { + _c.Run(run) + return _c +} + +// ResolveCollectionAlias provides a mock function with given fields: ctx, database, nameOrAlias +func (_m *MockCache) ResolveCollectionAlias(ctx context.Context, database string, nameOrAlias string) (string, error) { + ret := _m.Called(ctx, database, nameOrAlias) + + if len(ret) == 0 { + panic("no return value specified for ResolveCollectionAlias") + } + + var r0 string + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string, string) (string, error)); ok { + return rf(ctx, database, nameOrAlias) + } + if rf, ok := ret.Get(0).(func(context.Context, string, string) string); ok { + r0 = rf(ctx, database, nameOrAlias) + } else { + r0 = ret.Get(0).(string) + } + + if rf, ok := ret.Get(1).(func(context.Context, string, string) error); ok { + r1 = rf(ctx, database, nameOrAlias) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockCache_ResolveCollectionAlias_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ResolveCollectionAlias' +type MockCache_ResolveCollectionAlias_Call struct { + *mock.Call +} + +// ResolveCollectionAlias is a helper method to define mock.On call +// - ctx context.Context +// - database string +// - nameOrAlias string +func (_e *MockCache_Expecter) ResolveCollectionAlias(ctx interface{}, database interface{}, nameOrAlias interface{}) *MockCache_ResolveCollectionAlias_Call { + return &MockCache_ResolveCollectionAlias_Call{Call: _e.mock.On("ResolveCollectionAlias", ctx, database, nameOrAlias)} +} + +func (_c *MockCache_ResolveCollectionAlias_Call) Run(run func(ctx context.Context, database string, nameOrAlias string)) *MockCache_ResolveCollectionAlias_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(string)) + }) + return _c +} + +func (_c *MockCache_ResolveCollectionAlias_Call) Return(_a0 string, _a1 error) *MockCache_ResolveCollectionAlias_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockCache_ResolveCollectionAlias_Call) RunAndReturn(run func(context.Context, string, string) (string, error)) *MockCache_ResolveCollectionAlias_Call { + _c.Call.Return(run) + return _c +} + +// NewMockCache creates a new instance of MockCache. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewMockCache(t interface { + mock.TestingT + Cleanup(func()) +}) *MockCache { + mock := &MockCache{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/proxy/scheduler/mock_task_test.go b/internal/proxy/scheduler/mock_task_test.go new file mode 100644 index 0000000000..43b212289f --- /dev/null +++ b/internal/proxy/scheduler/mock_task_test.go @@ -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()) +} diff --git a/internal/proxy/task_scheduler.go b/internal/proxy/scheduler/task_scheduler.go similarity index 66% rename from internal/proxy/task_scheduler.go rename to internal/proxy/scheduler/task_scheduler.go index df5aadf0e6..3311d7576e 100644 --- a/internal/proxy/task_scheduler.go +++ b/internal/proxy/scheduler/task_scheduler.go @@ -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, diff --git a/internal/proxy/task_scheduler_test.go b/internal/proxy/scheduler/task_scheduler_test.go similarity index 76% rename from internal/proxy/task_scheduler_test.go rename to internal/proxy/scheduler/task_scheduler_test.go index 93c9341819..1c5bc151d7 100644 --- a/internal/proxy/task_scheduler_test.go +++ b/internal/proxy/scheduler/task_scheduler_test.go @@ -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() { diff --git a/internal/proxy/search_pipeline.go b/internal/proxy/search_pipeline.go index 879ace22cd..cbb33451b6 100644 --- a/internal/proxy/search_pipeline.go +++ b/internal/proxy/search_pipeline.go @@ -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), diff --git a/internal/proxy/search_pipeline_test.go b/internal/proxy/search_pipeline_test.go index 760dcf3a22..838b06c2b9 100644 --- a/internal/proxy/search_pipeline_test.go +++ b/internal/proxy/search_pipeline_test.go @@ -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 = ×tampAllocator{ 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 = ×tampAllocator{ tso: newMockTimestampAllocatorInterface(), } - sched, err := newTaskScheduler(ctx, proxy.tsoAllocator) + sched, err := scheduler.NewTaskScheduler(ctx, proxy.tsoAllocator) s.Require().NoError(err) err = sched.Start() diff --git a/internal/proxy/service_provider.go b/internal/proxy/service_provider.go index 1425a6c88f..ed84ccd5da 100644 --- a/internal/proxy/service_provider.go +++ b/internal/proxy/service_provider.go @@ -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)) diff --git a/internal/proxy/snapshot_impl.go b/internal/proxy/snapshot_impl.go index 505eaa881a..1a527a55ea 100644 --- a/internal/proxy/snapshot_impl.go +++ b/internal/proxy/snapshot_impl.go @@ -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)) diff --git a/internal/proxy/snapshot_impl_test.go b/internal/proxy/snapshot_impl_test.go index 7e78dca5f7..db9f836355 100644 --- a/internal/proxy/snapshot_impl_test.go +++ b/internal/proxy/snapshot_impl_test.go @@ -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() } diff --git a/internal/proxy/task.go b/internal/proxy/task.go index 7b1e4a2774..e853589365 100644 --- a/internal/proxy/task.go +++ b/internal/proxy/task.go @@ -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 } diff --git a/internal/proxy/task_alias.go b/internal/proxy/task_alias.go index cf202dd077..228998b08b 100644 --- a/internal/proxy/task_alias.go +++ b/internal/proxy/task_alias.go @@ -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 } } diff --git a/internal/proxy/task_batch_update_manifest.go b/internal/proxy/task_batch_update_manifest.go index 0d49b87c4e..17a06cea33 100644 --- a/internal/proxy/task_batch_update_manifest.go +++ b/internal/proxy/task_batch_update_manifest.go @@ -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 } diff --git a/internal/proxy/task_database.go b/internal/proxy/task_database.go index a0734c0286..42b544a54d 100644 --- a/internal/proxy/task_database.go +++ b/internal/proxy/task_database.go @@ -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 } diff --git a/internal/proxy/task_database_test.go b/internal/proxy/task_database_test.go index ef9e00f75a..39268bc5e0 100644 --- a/internal/proxy/task_database_test.go +++ b/internal/proxy/task_database_test.go @@ -80,7 +80,7 @@ func TestDropDatabaseTask(t *testing.T) { ).Maybe() task := &dropDatabaseTask{ baseTask: baseTask{ - metaCache: cache, + MetaCache: cache, }, Condition: NewTaskCondition(ctx), DropDatabaseRequest: &milvuspb.DropDatabaseRequest{ diff --git a/internal/proxy/task_delete.go b/internal/proxy/task_delete.go index 287cadc9dd..9931ab5c15 100644 --- a/internal/proxy/task_delete.go +++ b/internal/proxy/task_delete.go @@ -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, diff --git a/internal/proxy/task_delete_streaming.go b/internal/proxy/task_delete_streaming.go index 64e1bf0afd..5efdffd169 100644 --- a/internal/proxy/task_delete_streaming.go +++ b/internal/proxy/task_delete_streaming.go @@ -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 diff --git a/internal/proxy/task_delete_test.go b/internal/proxy/task_delete_test.go index 7e2ed456e5..86732870b1 100644 --- a/internal/proxy/task_delete_test.go +++ b/internal/proxy/task_delete_test.go @@ -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, diff --git a/internal/proxy/task_flush_all_test.go b/internal/proxy/task_flush_all_test.go index 6759d194c3..88f22f1960 100644 --- a/internal/proxy/task_flush_all_test.go +++ b/internal/proxy/task_flush_all_test.go @@ -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) diff --git a/internal/proxy/task_flush_streaming.go b/internal/proxy/task_flush_streaming.go index 6110e6caf8..e90922924b 100644 --- a/internal/proxy/task_flush_streaming.go +++ b/internal/proxy/task_flush_streaming.go @@ -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 } diff --git a/internal/proxy/task_import.go b/internal/proxy/task_import.go index d741d26632..f8c12df62b 100644 --- a/internal/proxy/task_import.go +++ b/internal/proxy/task_import.go @@ -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 diff --git a/internal/proxy/task_import_test.go b/internal/proxy/task_import_test.go index 9dfa1adb10..91a24e8dca 100644 --- a/internal/proxy/task_import_test.go +++ b/internal/proxy/task_import_test.go @@ -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) diff --git a/internal/proxy/task_index.go b/internal/proxy/task_index.go index 0bb1a6279e..11e5805b52 100644 --- a/internal/proxy/task_index.go +++ b/internal/proxy/task_index.go @@ -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 } diff --git a/internal/proxy/task_index_test.go b/internal/proxy/task_index_test.go index 4d68a64470..5de3f0e7bc 100644 --- a/internal/proxy/task_index_test.go +++ b/internal/proxy/task_index_test.go @@ -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{ diff --git a/internal/proxy/task_insert.go b/internal/proxy/task_insert.go index 14a62a7703..5d438922ec 100644 --- a/internal/proxy/task_insert.go +++ b/internal/proxy/task_insert.go @@ -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 diff --git a/internal/proxy/task_insert_streaming.go b/internal/proxy/task_insert_streaming.go index e66502b7ee..a9445e6090 100644 --- a/internal/proxy/task_insert_streaming.go +++ b/internal/proxy/task_insert_streaming.go @@ -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)) diff --git a/internal/proxy/task_insert_test.go b/internal/proxy/task_insert_test.go index 0ed81dbfa0..6bd3050a52 100644 --- a/internal/proxy/task_insert_test.go +++ b/internal/proxy/task_insert_test.go @@ -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) diff --git a/internal/proxy/task_query.go b/internal/proxy/task_query.go index 4d2257c711..348cbb714a 100644 --- a/internal/proxy/task_query.go +++ b/internal/proxy/task_query.go @@ -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), diff --git a/internal/proxy/task_query_namespace_test.go b/internal/proxy/task_query_namespace_test.go index ba2bed964e..1aa8348f80 100644 --- a/internal/proxy/task_query_namespace_test.go +++ b/internal/proxy/task_query_namespace_test.go @@ -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(), diff --git a/internal/proxy/task_query_test.go b/internal/proxy/task_query_test.go index 9cf803f7b5..dea08f8b8e 100644 --- a/internal/proxy/task_query_test.go +++ b/internal/proxy/task_query_test.go @@ -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, diff --git a/internal/proxy/task_search.go b/internal/proxy/task_search.go index d0ae169b9d..29e4a75f04 100644 --- a/internal/proxy/task_search.go +++ b/internal/proxy/task_search.go @@ -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 { diff --git a/internal/proxy/task_search_namespace_test.go b/internal/proxy/task_search_namespace_test.go index 09b13c66f2..fc5fae5b94 100644 --- a/internal/proxy/task_search_namespace_test.go +++ b/internal/proxy/task_search_namespace_test.go @@ -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(), diff --git a/internal/proxy/task_search_test.go b/internal/proxy/task_search_test.go index a65baeed56..0b95556766 100644 --- a/internal/proxy/task_search_test.go +++ b/internal/proxy/task_search_test.go @@ -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 = ×tampAllocator{ 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{}, diff --git a/internal/proxy/task_snapshot.go b/internal/proxy/task_snapshot.go index a27bea289b..26de9e8355 100644 --- a/internal/proxy/task_snapshot.go +++ b/internal/proxy/task_snapshot.go @@ -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 } diff --git a/internal/proxy/task_snapshot_test.go b/internal/proxy/task_snapshot_test.go index 40e9a4e19c..240a8f986e 100644 --- a/internal/proxy/task_snapshot_test.go +++ b/internal/proxy/task_snapshot_test.go @@ -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", diff --git a/internal/proxy/task_statistic.go b/internal/proxy/task_statistic.go index 3c9b60d799..7f4e66e9a1 100644 --- a/internal/proxy/task_statistic.go +++ b/internal/proxy/task_statistic.go @@ -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 } diff --git a/internal/proxy/task_statistic_test.go b/internal/proxy/task_statistic_test.go index 93048a7c96..fa122ee102 100644 --- a/internal/proxy/task_statistic_test.go +++ b/internal/proxy/task_statistic_test.go @@ -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, diff --git a/internal/proxy/task_test.go b/internal/proxy/task_test.go index 5085d6d496..b90a4f1bf3 100644 --- a/internal/proxy/task_test.go +++ b/internal/proxy/task_test.go @@ -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, diff --git a/internal/proxy/task_upsert.go b/internal/proxy/task_upsert.go index 357c097ee8..72c619e3a8 100644 --- a/internal/proxy/task_upsert.go +++ b/internal/proxy/task_upsert.go @@ -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 diff --git a/internal/proxy/task_upsert_streaming.go b/internal/proxy/task_upsert_streaming.go index cc338f30ef..ff7857ad14 100644 --- a/internal/proxy/task_upsert_streaming.go +++ b/internal/proxy/task_upsert_streaming.go @@ -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)) diff --git a/internal/proxy/task_upsert_test.go b/internal/proxy/task_upsert_test.go index bd34d252aa..3930fd767a 100644 --- a/internal/proxy/task_upsert_test.go +++ b/internal/proxy/task_upsert_test.go @@ -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{ diff --git a/internal/proxy/taskmodel/OWNERS b/internal/proxy/taskmodel/OWNERS new file mode 100644 index 0000000000..33d8e737c1 --- /dev/null +++ b/internal/proxy/taskmodel/OWNERS @@ -0,0 +1,8 @@ +reviewers: + - weiliu1031 + - SpadeA-Tang + - czs007 + - congqixia + +approvers: + - maintainers diff --git a/internal/proxy/taskmodel/README.md b/internal/proxy/taskmodel/README.md new file mode 100644 index 0000000000..7666c447f8 --- /dev/null +++ b/internal/proxy/taskmodel/README.md @@ -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()`. diff --git a/internal/proxy/condition.go b/internal/proxy/taskmodel/condition.go similarity index 99% rename from internal/proxy/condition.go rename to internal/proxy/taskmodel/condition.go index ce0a00cb2b..868fb9905e 100644 --- a/internal/proxy/condition.go +++ b/internal/proxy/taskmodel/condition.go @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package proxy +package taskmodel import ( "context" diff --git a/internal/proxy/taskmodel/task.go b/internal/proxy/taskmodel/task.go new file mode 100644 index 0000000000..6d4d226551 --- /dev/null +++ b/internal/proxy/taskmodel/task.go @@ -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 diff --git a/internal/proxy/taskmodel/types.go b/internal/proxy/taskmodel/types.go new file mode 100644 index 0000000000..9d35f2e204 --- /dev/null +++ b/internal/proxy/taskmodel/types.go @@ -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) +} diff --git a/internal/proxy/type_def.go b/internal/proxy/taskmodel_alias.go similarity index 53% rename from internal/proxy/type_def.go rename to internal/proxy/taskmodel_alias.go index cee5b5adbd..e8b7d32c6a 100644 --- a/internal/proxy/type_def.go +++ b/internal/proxy/taskmodel_alias.go @@ -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) }