impl.go 155.0 KB
Newer Older
1 2 3 4 5 6
// 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
7 8
// with the License. You may obtain a copy of the License at
//
9
//     http://www.apache.org/licenses/LICENSE-2.0
10
//
11 12 13 14 15
// 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.
16

C
Cai Yudong 已提交
17
package proxy
18 19 20

import (
	"context"
21
	"fmt"
C
cai.zhang 已提交
22
	"os"
23
	"strconv"
24 25
	"sync"

26
	"github.com/cockroachdb/errors"
J
jaime 已提交
27
	"github.com/golang/protobuf/proto"
28
	"github.com/samber/lo"
E
Enwei Jiao 已提交
29
	"go.opentelemetry.io/otel"
J
jaime 已提交
30
	"go.uber.org/zap"
31
	"golang.org/x/sync/errgroup"
32

S
SimFG 已提交
33 34
	"github.com/milvus-io/milvus-proto/go-api/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/milvuspb"
35
	"github.com/milvus-io/milvus-proto/go-api/msgpb"
S
smellthemoon 已提交
36
	"github.com/milvus-io/milvus-proto/go-api/schemapb"
37
	"github.com/milvus-io/milvus/internal/common"
X
Xiangyu Wang 已提交
38
	"github.com/milvus-io/milvus/internal/log"
39
	"github.com/milvus-io/milvus/internal/metrics"
J
jaime 已提交
40
	"github.com/milvus-io/milvus/internal/mq/msgstream"
X
Xiangyu Wang 已提交
41 42 43 44
	"github.com/milvus-io/milvus/internal/proto/datapb"
	"github.com/milvus-io/milvus/internal/proto/internalpb"
	"github.com/milvus-io/milvus/internal/proto/proxypb"
	"github.com/milvus-io/milvus/internal/proto/querypb"
45
	"github.com/milvus-io/milvus/internal/util"
46
	"github.com/milvus-io/milvus/internal/util/commonpbutil"
47
	"github.com/milvus-io/milvus/internal/util/crypto"
48
	"github.com/milvus-io/milvus/internal/util/errorutil"
49
	"github.com/milvus-io/milvus/internal/util/importutil"
50
	"github.com/milvus-io/milvus/internal/util/logutil"
S
smellthemoon 已提交
51
	"github.com/milvus-io/milvus/internal/util/merr"
52
	"github.com/milvus-io/milvus/internal/util/metricsinfo"
E
Enwei Jiao 已提交
53
	"github.com/milvus-io/milvus/internal/util/paramtable"
54
	"github.com/milvus-io/milvus/internal/util/ratelimitutil"
55
	"github.com/milvus-io/milvus/internal/util/timerecord"
X
Xiangyu Wang 已提交
56
	"github.com/milvus-io/milvus/internal/util/typeutil"
57 58
)

59 60
const moduleName = "Proxy"

61
// UpdateStateCode updates the state code of Proxy.
62
func (node *Proxy) UpdateStateCode(code commonpb.StateCode) {
63
	node.stateCode.Store(code)
Z
zhenshan.cao 已提交
64 65
}

66
// GetComponentStates get state of Proxy.
67 68
func (node *Proxy) GetComponentStates(ctx context.Context) (*milvuspb.ComponentStates, error) {
	stats := &milvuspb.ComponentStates{
G
godchen 已提交
69 70 71 72
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
		},
	}
73
	code, ok := node.stateCode.Load().(commonpb.StateCode)
G
godchen 已提交
74 75 76 77 78 79
	if !ok {
		errMsg := "unexpected error in type assertion"
		stats.Status = &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    errMsg,
		}
G
godchen 已提交
80
		return stats, nil
G
godchen 已提交
81
	}
82 83 84 85
	nodeID := common.NotRegisteredID
	if node.session != nil && node.session.Registered() {
		nodeID = node.session.ServerID
	}
86
	info := &milvuspb.ComponentInfo{
87 88
		// NodeID:    Params.ProxyID, // will race with Proxy.Register()
		NodeID:    nodeID,
C
Cai Yudong 已提交
89
		Role:      typeutil.ProxyRole,
G
godchen 已提交
90 91 92 93 94 95
		StateCode: code,
	}
	stats.State = info
	return stats, nil
}

C
cxytz01 已提交
96
// GetStatisticsChannel gets statistics channel of Proxy.
C
Cai Yudong 已提交
97
func (node *Proxy) GetStatisticsChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
G
godchen 已提交
98 99 100 101 102 103 104 105 106
	return &milvuspb.StringResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
			Reason:    "",
		},
		Value: "",
	}, nil
}

107
// InvalidateCollectionMetaCache invalidate the meta cache of specific collection.
C
Cai Yudong 已提交
108
func (node *Proxy) InvalidateCollectionMetaCache(ctx context.Context, request *proxypb.InvalidateCollMetaCacheRequest) (*commonpb.Status, error) {
109 110 111
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
112
	ctx = logutil.WithModule(ctx, moduleName)
E
Enwei Jiao 已提交
113 114 115

	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-InvalidateCollectionMetaCache")
	defer sp.End()
116
	log := log.Ctx(ctx).With(
117
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
118
		zap.String("db", request.DbName),
119 120
		zap.String("collectionName", request.CollectionName),
		zap.Int64("collectionID", request.CollectionID))
D
dragondriver 已提交
121

122 123
	log.Info("received request to invalidate collection meta cache")

124
	collectionName := request.CollectionName
125
	collectionID := request.CollectionID
X
Xiaofan 已提交
126 127

	var aliasName []string
N
neza2017 已提交
128
	if globalMetaCache != nil {
129 130 131 132
		if collectionName != "" {
			globalMetaCache.RemoveCollection(ctx, collectionName) // no need to return error, though collection may be not cached
		}
		if request.CollectionID != UniqueID(0) {
X
Xiaofan 已提交
133
			aliasName = globalMetaCache.RemoveCollectionsByID(ctx, collectionID)
134
		}
N
neza2017 已提交
135
	}
136 137
	if request.GetBase().GetMsgType() == commonpb.MsgType_DropCollection {
		// no need to handle error, since this Proxy may not create dml stream for the collection.
138 139
		node.chMgr.removeDMLStream(request.GetCollectionID())
		// clean up collection level metrics
E
Enwei Jiao 已提交
140
		metrics.CleanupCollectionMetrics(paramtable.GetNodeID(), collectionName)
X
Xiaofan 已提交
141
		for _, alias := range aliasName {
E
Enwei Jiao 已提交
142
			metrics.CleanupCollectionMetrics(paramtable.GetNodeID(), alias)
X
Xiaofan 已提交
143
		}
144
	}
145
	log.Info("complete to invalidate collection meta cache")
D
dragondriver 已提交
146

147
	return &commonpb.Status{
148
		ErrorCode: commonpb.ErrorCode_Success,
149 150
		Reason:    "",
	}, nil
151 152
}

153
// CreateCollection create a collection by the schema.
154
// TODO(dragondriver): add more detailed ut for ConsistencyLevel, should we support multiple consistency level in Proxy?
C
Cai Yudong 已提交
155
func (node *Proxy) CreateCollection(ctx context.Context, request *milvuspb.CreateCollectionRequest) (*commonpb.Status, error) {
156 157 158
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
159

E
Enwei Jiao 已提交
160 161
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CreateCollection")
	defer sp.End()
162 163 164
	method := "CreateCollection"
	tr := timerecord.NewTimeRecorder(method)

E
Enwei Jiao 已提交
165
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
166

167
	cct := &createCollectionTask{
S
sunby 已提交
168
		ctx:                     ctx,
169 170
		Condition:               NewTaskCondition(ctx),
		CreateCollectionRequest: request,
171
		rootCoord:               node.rootCoord,
172 173
	}

174 175 176
	// avoid data race
	lenOfSchema := len(request.Schema)

177
	log := log.Ctx(ctx).With(
178
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
179 180
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
181
		zap.Int("len(schema)", lenOfSchema),
182 183
		zap.Int32("shards_num", request.ShardsNum),
		zap.String("consistency_level", request.ConsistencyLevel.String()))
184

185 186
	log.Debug(rpcReceived(method))

187 188 189
	if err := node.sched.ddQueue.Enqueue(cct); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
190
			zap.Error(err))
191

E
Enwei Jiao 已提交
192
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
193
		return &commonpb.Status{
194
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
195 196 197 198
			Reason:    err.Error(),
		}, nil
	}

199 200
	log.Debug(
		rpcEnqueued(method),
201 202
		zap.Uint64("BeginTs", cct.BeginTs()),
		zap.Uint64("EndTs", cct.EndTs()),
203
		zap.Uint64("timestamp", request.Base.Timestamp))
204

205 206 207
	if err := cct.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
208
			zap.Error(err),
209
			zap.Uint64("BeginTs", cct.BeginTs()),
210
			zap.Uint64("EndTs", cct.EndTs()))
D
dragondriver 已提交
211

E
Enwei Jiao 已提交
212
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
213
		return &commonpb.Status{
214
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
215 216 217 218
			Reason:    err.Error(),
		}, nil
	}

219 220
	log.Debug(
		rpcDone(method),
221
		zap.Uint64("BeginTs", cct.BeginTs()),
222
		zap.Uint64("EndTs", cct.EndTs()))
223

E
Enwei Jiao 已提交
224 225
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
226 227 228
	return cct.result, nil
}

229
// DropCollection drop a collection.
C
Cai Yudong 已提交
230
func (node *Proxy) DropCollection(ctx context.Context, request *milvuspb.DropCollectionRequest) (*commonpb.Status, error) {
231 232 233
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
234

E
Enwei Jiao 已提交
235 236
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DropCollection")
	defer sp.End()
237 238
	method := "DropCollection"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
239
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
240

241
	dct := &dropCollectionTask{
S
sunby 已提交
242
		ctx:                   ctx,
243 244
		Condition:             NewTaskCondition(ctx),
		DropCollectionRequest: request,
245
		rootCoord:             node.rootCoord,
246
		chMgr:                 node.chMgr,
S
sunby 已提交
247
		chTicker:              node.chTicker,
248 249
	}

250
	log := log.Ctx(ctx).With(
251
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
252 253
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
254

255 256
	log.Debug("DropCollection received")

257 258
	if err := node.sched.ddQueue.Enqueue(dct); err != nil {
		log.Warn("DropCollection failed to enqueue",
259
			zap.Error(err))
260

E
Enwei Jiao 已提交
261
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
262
		return &commonpb.Status{
263
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
264 265 266 267
			Reason:    err.Error(),
		}, nil
	}

268 269
	log.Debug("DropCollection enqueued",
		zap.Uint64("BeginTs", dct.BeginTs()),
270
		zap.Uint64("EndTs", dct.EndTs()))
271 272 273

	if err := dct.WaitToFinish(); err != nil {
		log.Warn("DropCollection failed to WaitToFinish",
D
dragondriver 已提交
274
			zap.Error(err),
275
			zap.Uint64("BeginTs", dct.BeginTs()),
276
			zap.Uint64("EndTs", dct.EndTs()))
D
dragondriver 已提交
277

E
Enwei Jiao 已提交
278
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
279
		return &commonpb.Status{
280
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
281 282 283 284
			Reason:    err.Error(),
		}, nil
	}

285 286
	log.Debug("DropCollection done",
		zap.Uint64("BeginTs", dct.BeginTs()),
287
		zap.Uint64("EndTs", dct.EndTs()))
288

E
Enwei Jiao 已提交
289 290
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
291 292 293
	return dct.result, nil
}

294
// HasCollection check if the specific collection exists in Milvus.
C
Cai Yudong 已提交
295
func (node *Proxy) HasCollection(ctx context.Context, request *milvuspb.HasCollectionRequest) (*milvuspb.BoolResponse, error) {
296 297 298 299 300
	if !node.checkHealthy() {
		return &milvuspb.BoolResponse{
			Status: unhealthyStatus(),
		}, nil
	}
301

E
Enwei Jiao 已提交
302 303
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-HasCollection")
	defer sp.End()
304 305
	method := "HasCollection"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
306
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
307
		metrics.TotalLabel).Inc()
308

309
	log := log.Ctx(ctx).With(
310
		zap.String("role", typeutil.ProxyRole),
311 312 313
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

314 315
	log.Debug("HasCollection received")

316
	hct := &hasCollectionTask{
S
sunby 已提交
317
		ctx:                  ctx,
318 319
		Condition:            NewTaskCondition(ctx),
		HasCollectionRequest: request,
320
		rootCoord:            node.rootCoord,
321 322
	}

323 324
	if err := node.sched.ddQueue.Enqueue(hct); err != nil {
		log.Warn("HasCollection failed to enqueue",
325
			zap.Error(err))
326

E
Enwei Jiao 已提交
327
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
328
			metrics.AbandonLabel).Inc()
329 330
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
331
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
332 333 334 335 336
				Reason:    err.Error(),
			},
		}, nil
	}

337 338
	log.Debug("HasCollection enqueued",
		zap.Uint64("BeginTS", hct.BeginTs()),
339
		zap.Uint64("EndTS", hct.EndTs()))
340 341 342

	if err := hct.WaitToFinish(); err != nil {
		log.Warn("HasCollection failed to WaitToFinish",
D
dragondriver 已提交
343
			zap.Error(err),
344
			zap.Uint64("BeginTS", hct.BeginTs()),
345
			zap.Uint64("EndTS", hct.EndTs()))
D
dragondriver 已提交
346

E
Enwei Jiao 已提交
347
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
348
			metrics.FailLabel).Inc()
349 350
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
351
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
352 353 354 355 356
				Reason:    err.Error(),
			},
		}, nil
	}

357 358
	log.Debug("HasCollection done",
		zap.Uint64("BeginTS", hct.BeginTs()),
359
		zap.Uint64("EndTS", hct.EndTs()))
360

E
Enwei Jiao 已提交
361
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
362
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
363
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
364 365 366
	return hct.result, nil
}

367
// LoadCollection load a collection into query nodes.
C
Cai Yudong 已提交
368
func (node *Proxy) LoadCollection(ctx context.Context, request *milvuspb.LoadCollectionRequest) (*commonpb.Status, error) {
369 370 371
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
372

E
Enwei Jiao 已提交
373 374
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-LoadCollection")
	defer sp.End()
375 376
	method := "LoadCollection"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
377
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
378
		metrics.TotalLabel).Inc()
379
	lct := &loadCollectionTask{
S
sunby 已提交
380
		ctx:                   ctx,
381 382
		Condition:             NewTaskCondition(ctx),
		LoadCollectionRequest: request,
383
		queryCoord:            node.queryCoord,
384
		datacoord:             node.dataCoord,
385 386
	}

387
	log := log.Ctx(ctx).With(
388
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
389
		zap.String("db", request.DbName),
390 391
		zap.String("collection", request.CollectionName),
		zap.Bool("refreshMode", request.Refresh))
392

393 394
	log.Debug("LoadCollection received")

395 396
	if err := node.sched.ddQueue.Enqueue(lct); err != nil {
		log.Warn("LoadCollection failed to enqueue",
397
			zap.Error(err))
398

E
Enwei Jiao 已提交
399
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
400
			metrics.AbandonLabel).Inc()
401
		return &commonpb.Status{
402
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
403 404 405
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
406

407 408
	log.Debug("LoadCollection enqueued",
		zap.Uint64("BeginTS", lct.BeginTs()),
409
		zap.Uint64("EndTS", lct.EndTs()))
410 411 412

	if err := lct.WaitToFinish(); err != nil {
		log.Warn("LoadCollection failed to WaitToFinish",
D
dragondriver 已提交
413
			zap.Error(err),
414
			zap.Uint64("BeginTS", lct.BeginTs()),
415
			zap.Uint64("EndTS", lct.EndTs()))
E
Enwei Jiao 已提交
416
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
417
			metrics.FailLabel).Inc()
418
		return &commonpb.Status{
419
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
420 421 422 423
			Reason:    err.Error(),
		}, nil
	}

424 425
	log.Debug("LoadCollection done",
		zap.Uint64("BeginTS", lct.BeginTs()),
426
		zap.Uint64("EndTS", lct.EndTs()))
427

E
Enwei Jiao 已提交
428
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
429
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
430
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
431
	return lct.result, nil
432 433
}

434
// ReleaseCollection remove the loaded collection from query nodes.
C
Cai Yudong 已提交
435
func (node *Proxy) ReleaseCollection(ctx context.Context, request *milvuspb.ReleaseCollectionRequest) (*commonpb.Status, error) {
436 437 438
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
439

E
Enwei Jiao 已提交
440 441
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ReleaseCollection")
	defer sp.End()
442 443
	method := "ReleaseCollection"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
444
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
445
		metrics.TotalLabel).Inc()
446
	rct := &releaseCollectionTask{
S
sunby 已提交
447
		ctx:                      ctx,
448 449
		Condition:                NewTaskCondition(ctx),
		ReleaseCollectionRequest: request,
450
		queryCoord:               node.queryCoord,
451
		chMgr:                    node.chMgr,
452 453
	}

454
	log := log.Ctx(ctx).With(
455
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
456 457
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
458

459 460
	log.Debug(rpcReceived(method))

461
	if err := node.sched.ddQueue.Enqueue(rct); err != nil {
462 463
		log.Warn(
			rpcFailedToEnqueue(method),
464
			zap.Error(err))
465

E
Enwei Jiao 已提交
466
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
467
			metrics.AbandonLabel).Inc()
468
		return &commonpb.Status{
469
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
470 471 472 473
			Reason:    err.Error(),
		}, nil
	}

474 475
	log.Debug(
		rpcEnqueued(method),
476
		zap.Uint64("BeginTS", rct.BeginTs()),
477
		zap.Uint64("EndTS", rct.EndTs()))
478 479

	if err := rct.WaitToFinish(); err != nil {
480 481
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
482
			zap.Error(err),
483
			zap.Uint64("BeginTS", rct.BeginTs()),
484
			zap.Uint64("EndTS", rct.EndTs()))
D
dragondriver 已提交
485

E
Enwei Jiao 已提交
486
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
487
			metrics.FailLabel).Inc()
488
		return &commonpb.Status{
489
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
490 491 492 493
			Reason:    err.Error(),
		}, nil
	}

494 495
	log.Debug(
		rpcDone(method),
496
		zap.Uint64("BeginTS", rct.BeginTs()),
497
		zap.Uint64("EndTS", rct.EndTs()))
498

E
Enwei Jiao 已提交
499
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
500
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
501
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
502
	return rct.result, nil
503 504
}

505
// DescribeCollection get the meta information of specific collection, such as schema, created timestamp and etc.
C
Cai Yudong 已提交
506
func (node *Proxy) DescribeCollection(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
507 508 509 510 511
	if !node.checkHealthy() {
		return &milvuspb.DescribeCollectionResponse{
			Status: unhealthyStatus(),
		}, nil
	}
512

E
Enwei Jiao 已提交
513 514
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DescribeCollection")
	defer sp.End()
515 516
	method := "DescribeCollection"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
517
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
518
		metrics.TotalLabel).Inc()
519

520
	dct := &describeCollectionTask{
S
sunby 已提交
521
		ctx:                       ctx,
522 523
		Condition:                 NewTaskCondition(ctx),
		DescribeCollectionRequest: request,
524
		rootCoord:                 node.rootCoord,
525 526
	}

527
	log := log.Ctx(ctx).With(
528
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
529 530
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
531

532 533
	log.Debug("DescribeCollection received")

534 535
	if err := node.sched.ddQueue.Enqueue(dct); err != nil {
		log.Warn("DescribeCollection failed to enqueue",
536
			zap.Error(err))
537

E
Enwei Jiao 已提交
538
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
539
			metrics.AbandonLabel).Inc()
540 541
		return &milvuspb.DescribeCollectionResponse{
			Status: &commonpb.Status{
542
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
543 544 545 546 547
				Reason:    err.Error(),
			},
		}, nil
	}

548 549
	log.Debug("DescribeCollection enqueued",
		zap.Uint64("BeginTS", dct.BeginTs()),
550
		zap.Uint64("EndTS", dct.EndTs()))
551 552 553

	if err := dct.WaitToFinish(); err != nil {
		log.Warn("DescribeCollection failed to WaitToFinish",
D
dragondriver 已提交
554
			zap.Error(err),
555
			zap.Uint64("BeginTS", dct.BeginTs()),
556
			zap.Uint64("EndTS", dct.EndTs()))
D
dragondriver 已提交
557

E
Enwei Jiao 已提交
558
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
559
			metrics.FailLabel).Inc()
560

561 562
		return &milvuspb.DescribeCollectionResponse{
			Status: &commonpb.Status{
563
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
564 565 566 567 568
				Reason:    err.Error(),
			},
		}, nil
	}

569 570
	log.Debug("DescribeCollection done",
		zap.Uint64("BeginTS", dct.BeginTs()),
571
		zap.Uint64("EndTS", dct.EndTs()))
572

E
Enwei Jiao 已提交
573
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
574
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
575
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
576 577 578
	return dct.result, nil
}

579 580 581 582 583 584 585 586 587
// GetStatistics get the statistics, such as `num_rows`.
// WARNING: It is an experimental API
func (node *Proxy) GetStatistics(ctx context.Context, request *milvuspb.GetStatisticsRequest) (*milvuspb.GetStatisticsResponse, error) {
	if !node.checkHealthy() {
		return &milvuspb.GetStatisticsResponse{
			Status: unhealthyStatus(),
		}, nil
	}

E
Enwei Jiao 已提交
588 589
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetStatistics")
	defer sp.End()
590 591
	method := "GetStatistics"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
592
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
593
		metrics.TotalLabel).Inc()
594 595 596 597 598 599 600 601 602 603
	g := &getStatisticsTask{
		request:   request,
		Condition: NewTaskCondition(ctx),
		ctx:       ctx,
		tr:        tr,
		dc:        node.dataCoord,
		qc:        node.queryCoord,
		shardMgr:  node.shardMgr,
	}

604
	log := log.Ctx(ctx).With(
605 606
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
607 608 609 610
		zap.String("collection", request.CollectionName))

	log.Debug(
		rpcReceived(method),
611 612 613 614 615 616 617 618
		zap.Strings("partitions", request.PartitionNames))

	if err := node.sched.ddQueue.Enqueue(g); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.Strings("partitions", request.PartitionNames))

E
Enwei Jiao 已提交
619
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643
			metrics.AbandonLabel).Inc()

		return &milvuspb.GetStatisticsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", g.BeginTs()),
		zap.Uint64("EndTS", g.EndTs()),
		zap.Strings("partitions", request.PartitionNames))

	if err := g.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
			zap.Error(err),
			zap.Uint64("BeginTS", g.BeginTs()),
			zap.Uint64("EndTS", g.EndTs()),
			zap.Strings("partitions", request.PartitionNames))

E
Enwei Jiao 已提交
644
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
645 646 647 648 649 650 651 652 653 654 655 656 657
			metrics.FailLabel).Inc()

		return &milvuspb.GetStatisticsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", g.BeginTs()),
658
		zap.Uint64("EndTS", g.EndTs()))
659

E
Enwei Jiao 已提交
660
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
661
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
662
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
663 664 665
	return g.result, nil
}

666
// GetCollectionStatistics get the collection statistics, such as `num_rows`.
C
Cai Yudong 已提交
667
func (node *Proxy) GetCollectionStatistics(ctx context.Context, request *milvuspb.GetCollectionStatisticsRequest) (*milvuspb.GetCollectionStatisticsResponse, error) {
668 669 670 671 672
	if !node.checkHealthy() {
		return &milvuspb.GetCollectionStatisticsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
673

E
Enwei Jiao 已提交
674 675
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetCollectionStatistics")
	defer sp.End()
676 677
	method := "GetCollectionStatistics"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
678
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
679
		metrics.TotalLabel).Inc()
680
	g := &getCollectionStatisticsTask{
G
godchen 已提交
681 682 683
		ctx:                            ctx,
		Condition:                      NewTaskCondition(ctx),
		GetCollectionStatisticsRequest: request,
684
		dataCoord:                      node.dataCoord,
685 686
	}

687
	log := log.Ctx(ctx).With(
688
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
689 690
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
691

692 693
	log.Debug(rpcReceived(method))

694
	if err := node.sched.ddQueue.Enqueue(g); err != nil {
695 696
		log.Warn(
			rpcFailedToEnqueue(method),
697
			zap.Error(err))
698

E
Enwei Jiao 已提交
699
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
700
			metrics.AbandonLabel).Inc()
701

G
godchen 已提交
702
		return &milvuspb.GetCollectionStatisticsResponse{
703
			Status: &commonpb.Status{
704
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
705 706 707 708 709
				Reason:    err.Error(),
			},
		}, nil
	}

710 711
	log.Debug(
		rpcEnqueued(method),
712
		zap.Uint64("BeginTS", g.BeginTs()),
713
		zap.Uint64("EndTS", g.EndTs()))
714 715

	if err := g.WaitToFinish(); err != nil {
716 717
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
718
			zap.Error(err),
719
			zap.Uint64("BeginTS", g.BeginTs()),
720
			zap.Uint64("EndTS", g.EndTs()))
D
dragondriver 已提交
721

E
Enwei Jiao 已提交
722
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
723
			metrics.FailLabel).Inc()
724

G
godchen 已提交
725
		return &milvuspb.GetCollectionStatisticsResponse{
726
			Status: &commonpb.Status{
727
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
728 729 730 731 732
				Reason:    err.Error(),
			},
		}, nil
	}

733 734
	log.Debug(
		rpcDone(method),
735
		zap.Uint64("BeginTS", g.BeginTs()),
736
		zap.Uint64("EndTS", g.EndTs()))
737

E
Enwei Jiao 已提交
738
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
739
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
740
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
741
	return g.result, nil
742 743
}

744
// ShowCollections list all collections in Milvus.
C
Cai Yudong 已提交
745
func (node *Proxy) ShowCollections(ctx context.Context, request *milvuspb.ShowCollectionsRequest) (*milvuspb.ShowCollectionsResponse, error) {
746 747 748 749 750
	if !node.checkHealthy() {
		return &milvuspb.ShowCollectionsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
E
Enwei Jiao 已提交
751 752
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ShowCollections")
	defer sp.End()
753 754
	method := "ShowCollections"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
755
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
756

757
	sct := &showCollectionsTask{
G
godchen 已提交
758 759 760
		ctx:                    ctx,
		Condition:              NewTaskCondition(ctx),
		ShowCollectionsRequest: request,
761
		queryCoord:             node.queryCoord,
762
		rootCoord:              node.rootCoord,
763 764
	}

765
	log := log.Ctx(ctx).With(
766
		zap.String("role", typeutil.ProxyRole),
767 768
		zap.String("DbName", request.DbName),
		zap.Uint64("TimeStamp", request.TimeStamp),
769 770 771 772
		zap.String("ShowType", request.Type.String()))

	log.Debug("ShowCollections received",
		zap.Any("CollectionNames", request.CollectionNames))
773

774
	err := node.sched.ddQueue.Enqueue(sct)
775
	if err != nil {
776 777
		log.Warn("ShowCollections failed to enqueue",
			zap.Error(err),
778
			zap.Any("CollectionNames", request.CollectionNames))
779

E
Enwei Jiao 已提交
780
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
G
godchen 已提交
781
		return &milvuspb.ShowCollectionsResponse{
782
			Status: &commonpb.Status{
783
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
784 785 786 787 788
				Reason:    err.Error(),
			},
		}, nil
	}

789
	log.Debug("ShowCollections enqueued",
790
		zap.Any("CollectionNames", request.CollectionNames))
D
dragondriver 已提交
791

792 793
	err = sct.WaitToFinish()
	if err != nil {
794 795
		log.Warn("ShowCollections failed to WaitToFinish",
			zap.Error(err),
796
			zap.Any("CollectionNames", request.CollectionNames))
797

E
Enwei Jiao 已提交
798
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
799

G
godchen 已提交
800
		return &milvuspb.ShowCollectionsResponse{
801
			Status: &commonpb.Status{
802
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
803 804 805 806 807
				Reason:    err.Error(),
			},
		}, nil
	}

808
	log.Debug("ShowCollections Done",
809 810
		zap.Int("len(CollectionNames)", len(request.CollectionNames)),
		zap.Int("num_collections", len(sct.result.CollectionNames)))
811

E
Enwei Jiao 已提交
812 813
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
814 815 816
	return sct.result, nil
}

J
jaime 已提交
817 818 819 820 821
func (node *Proxy) AlterCollection(ctx context.Context, request *milvuspb.AlterCollectionRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

E
Enwei Jiao 已提交
822 823
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-AlterCollection")
	defer sp.End()
J
jaime 已提交
824 825 826
	method := "AlterCollection"
	tr := timerecord.NewTimeRecorder(method)

E
Enwei Jiao 已提交
827
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
J
jaime 已提交
828 829 830 831 832 833 834 835

	act := &alterCollectionTask{
		ctx:                    ctx,
		Condition:              NewTaskCondition(ctx),
		AlterCollectionRequest: request,
		rootCoord:              node.rootCoord,
	}

836
	log := log.Ctx(ctx).With(
J
jaime 已提交
837 838 839 840
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

841 842 843
	log.Debug(
		rpcReceived(method))

J
jaime 已提交
844 845 846
	if err := node.sched.ddQueue.Enqueue(act); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
847
			zap.Error(err))
J
jaime 已提交
848

E
Enwei Jiao 已提交
849
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
J
jaime 已提交
850 851 852 853 854 855 856 857 858 859
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", act.BeginTs()),
		zap.Uint64("EndTs", act.EndTs()),
860
		zap.Uint64("timestamp", request.Base.Timestamp))
J
jaime 已提交
861 862 863 864 865 866

	if err := act.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
			zap.Error(err),
			zap.Uint64("BeginTs", act.BeginTs()),
867
			zap.Uint64("EndTs", act.EndTs()))
J
jaime 已提交
868

E
Enwei Jiao 已提交
869
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
J
jaime 已提交
870 871 872 873 874 875 876 877 878
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", act.BeginTs()),
879
		zap.Uint64("EndTs", act.EndTs()))
J
jaime 已提交
880

E
Enwei Jiao 已提交
881 882
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
J
jaime 已提交
883 884 885
	return act.result, nil
}

886
// CreatePartition create a partition in specific collection.
C
Cai Yudong 已提交
887
func (node *Proxy) CreatePartition(ctx context.Context, request *milvuspb.CreatePartitionRequest) (*commonpb.Status, error) {
888 889 890
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
891

E
Enwei Jiao 已提交
892 893
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CreatePartition")
	defer sp.End()
894 895
	method := "CreatePartition"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
896
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
897

898
	cpt := &createPartitionTask{
S
sunby 已提交
899
		ctx:                    ctx,
900 901
		Condition:              NewTaskCondition(ctx),
		CreatePartitionRequest: request,
902
		rootCoord:              node.rootCoord,
903 904 905
		result:                 nil,
	}

906
	log := log.Ctx(ctx).With(
907
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
908 909 910
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
911

912 913
	log.Debug(rpcReceived("CreatePartition"))

914 915 916
	if err := node.sched.ddQueue.Enqueue(cpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue("CreatePartition"),
917
			zap.Error(err))
918

E
Enwei Jiao 已提交
919
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
920

921
		return &commonpb.Status{
922
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
923 924 925
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
926

927 928 929
	log.Debug(
		rpcEnqueued("CreatePartition"),
		zap.Uint64("BeginTS", cpt.BeginTs()),
930
		zap.Uint64("EndTS", cpt.EndTs()))
931 932 933 934

	if err := cpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish("CreatePartition"),
D
dragondriver 已提交
935
			zap.Error(err),
936
			zap.Uint64("BeginTS", cpt.BeginTs()),
937
			zap.Uint64("EndTS", cpt.EndTs()))
D
dragondriver 已提交
938

E
Enwei Jiao 已提交
939
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
940

941
		return &commonpb.Status{
942
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
943 944 945
			Reason:    err.Error(),
		}, nil
	}
946 947 948 949

	log.Debug(
		rpcDone("CreatePartition"),
		zap.Uint64("BeginTS", cpt.BeginTs()),
950
		zap.Uint64("EndTS", cpt.EndTs()))
951

E
Enwei Jiao 已提交
952 953
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
954 955 956
	return cpt.result, nil
}

957
// DropPartition drop a partition in specific collection.
C
Cai Yudong 已提交
958
func (node *Proxy) DropPartition(ctx context.Context, request *milvuspb.DropPartitionRequest) (*commonpb.Status, error) {
959 960 961
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
962

E
Enwei Jiao 已提交
963 964
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DropPartition")
	defer sp.End()
965 966
	method := "DropPartition"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
967
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
968

969
	dpt := &dropPartitionTask{
S
sunby 已提交
970
		ctx:                  ctx,
971 972
		Condition:            NewTaskCondition(ctx),
		DropPartitionRequest: request,
973
		rootCoord:            node.rootCoord,
C
cai.zhang 已提交
974
		queryCoord:           node.queryCoord,
975 976 977
		result:               nil,
	}

978
	log := log.Ctx(ctx).With(
979
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
980 981 982
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
983

984 985
	log.Debug(rpcReceived(method))

986 987 988
	if err := node.sched.ddQueue.Enqueue(dpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
989
			zap.Error(err))
990

E
Enwei Jiao 已提交
991
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
992

993
		return &commonpb.Status{
994
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
995 996 997
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
998

999 1000 1001
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", dpt.BeginTs()),
1002
		zap.Uint64("EndTS", dpt.EndTs()))
1003 1004 1005 1006

	if err := dpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1007
			zap.Error(err),
1008
			zap.Uint64("BeginTS", dpt.BeginTs()),
1009
			zap.Uint64("EndTS", dpt.EndTs()))
D
dragondriver 已提交
1010

E
Enwei Jiao 已提交
1011
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
1012

1013
		return &commonpb.Status{
1014
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1015 1016 1017
			Reason:    err.Error(),
		}, nil
	}
1018 1019 1020 1021

	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", dpt.BeginTs()),
1022
		zap.Uint64("EndTS", dpt.EndTs()))
1023

E
Enwei Jiao 已提交
1024 1025
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1026 1027 1028
	return dpt.result, nil
}

1029
// HasPartition check if partition exist.
C
Cai Yudong 已提交
1030
func (node *Proxy) HasPartition(ctx context.Context, request *milvuspb.HasPartitionRequest) (*milvuspb.BoolResponse, error) {
1031 1032 1033 1034 1035
	if !node.checkHealthy() {
		return &milvuspb.BoolResponse{
			Status: unhealthyStatus(),
		}, nil
	}
D
dragondriver 已提交
1036

E
Enwei Jiao 已提交
1037 1038
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-HasPartition")
	defer sp.End()
1039 1040 1041
	method := "HasPartition"
	tr := timerecord.NewTimeRecorder(method)
	//TODO: use collectionID instead of collectionName
E
Enwei Jiao 已提交
1042
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1043
		metrics.TotalLabel).Inc()
D
dragondriver 已提交
1044

1045
	hpt := &hasPartitionTask{
S
sunby 已提交
1046
		ctx:                 ctx,
1047 1048
		Condition:           NewTaskCondition(ctx),
		HasPartitionRequest: request,
1049
		rootCoord:           node.rootCoord,
1050 1051 1052
		result:              nil,
	}

1053
	log := log.Ctx(ctx).With(
1054
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1055 1056 1057
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
D
dragondriver 已提交
1058

1059 1060
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
1061 1062 1063
	if err := node.sched.ddQueue.Enqueue(hpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1064
			zap.Error(err))
D
dragondriver 已提交
1065

E
Enwei Jiao 已提交
1066
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1067
			metrics.AbandonLabel).Inc()
1068

1069 1070
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
1071
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1072 1073 1074 1075 1076
				Reason:    err.Error(),
			},
			Value: false,
		}, nil
	}
D
dragondriver 已提交
1077

D
dragondriver 已提交
1078 1079 1080
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", hpt.BeginTs()),
1081
		zap.Uint64("EndTS", hpt.EndTs()))
D
dragondriver 已提交
1082 1083 1084 1085

	if err := hpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1086
			zap.Error(err),
D
dragondriver 已提交
1087
			zap.Uint64("BeginTS", hpt.BeginTs()),
1088
			zap.Uint64("EndTS", hpt.EndTs()))
D
dragondriver 已提交
1089

E
Enwei Jiao 已提交
1090
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1091
			metrics.FailLabel).Inc()
1092

1093 1094
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
1095
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1096 1097 1098 1099 1100
				Reason:    err.Error(),
			},
			Value: false,
		}, nil
	}
D
dragondriver 已提交
1101 1102 1103 1104

	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", hpt.BeginTs()),
1105
		zap.Uint64("EndTS", hpt.EndTs()))
D
dragondriver 已提交
1106

E
Enwei Jiao 已提交
1107
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1108
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1109
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1110 1111 1112
	return hpt.result, nil
}

1113
// LoadPartitions load specific partitions into query nodes.
C
Cai Yudong 已提交
1114
func (node *Proxy) LoadPartitions(ctx context.Context, request *milvuspb.LoadPartitionsRequest) (*commonpb.Status, error) {
1115 1116 1117
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
1118

E
Enwei Jiao 已提交
1119 1120
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-LoadPartitions")
	defer sp.End()
1121 1122
	method := "LoadPartitions"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1123
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1124
		metrics.TotalLabel).Inc()
1125
	lpt := &loadPartitionsTask{
G
godchen 已提交
1126 1127 1128
		ctx:                   ctx,
		Condition:             NewTaskCondition(ctx),
		LoadPartitionsRequest: request,
1129
		queryCoord:            node.queryCoord,
1130
		datacoord:             node.dataCoord,
1131 1132
	}

1133
	log := log.Ctx(ctx).With(
1134
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1135 1136
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
1137 1138
		zap.Strings("partitions", request.PartitionNames),
		zap.Bool("refreshMode", request.Refresh))
1139

1140 1141
	log.Debug(rpcReceived(method))

1142 1143 1144
	if err := node.sched.ddQueue.Enqueue(lpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1145
			zap.Error(err))
1146

E
Enwei Jiao 已提交
1147
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1148
			metrics.AbandonLabel).Inc()
1149

1150
		return &commonpb.Status{
1151
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1152 1153 1154 1155
			Reason:    err.Error(),
		}, nil
	}

1156 1157 1158
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", lpt.BeginTs()),
1159
		zap.Uint64("EndTS", lpt.EndTs()))
1160 1161 1162 1163

	if err := lpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1164
			zap.Error(err),
1165
			zap.Uint64("BeginTS", lpt.BeginTs()),
1166
			zap.Uint64("EndTS", lpt.EndTs()))
D
dragondriver 已提交
1167

E
Enwei Jiao 已提交
1168
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1169
			metrics.FailLabel).Inc()
1170

1171
		return &commonpb.Status{
1172
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1173 1174 1175 1176
			Reason:    err.Error(),
		}, nil
	}

1177 1178 1179
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", lpt.BeginTs()),
1180
		zap.Uint64("EndTS", lpt.EndTs()))
1181

E
Enwei Jiao 已提交
1182
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1183
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1184
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1185
	return lpt.result, nil
1186 1187
}

1188
// ReleasePartitions release specific partitions from query nodes.
C
Cai Yudong 已提交
1189
func (node *Proxy) ReleasePartitions(ctx context.Context, request *milvuspb.ReleasePartitionsRequest) (*commonpb.Status, error) {
1190 1191 1192
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
1193

E
Enwei Jiao 已提交
1194 1195
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ReleasePartitions")
	defer sp.End()
1196

1197
	rpt := &releasePartitionsTask{
G
godchen 已提交
1198 1199 1200
		ctx:                      ctx,
		Condition:                NewTaskCondition(ctx),
		ReleasePartitionsRequest: request,
1201
		queryCoord:               node.queryCoord,
1202 1203
	}

1204
	method := "ReleasePartitions"
1205
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1206
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1207
		metrics.TotalLabel).Inc()
1208 1209

	log := log.Ctx(ctx).With(
1210
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1211 1212 1213
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
1214

1215 1216
	log.Debug(rpcReceived(method))

1217 1218 1219
	if err := node.sched.ddQueue.Enqueue(rpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1220
			zap.Error(err))
1221

E
Enwei Jiao 已提交
1222
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1223
			metrics.AbandonLabel).Inc()
1224

1225
		return &commonpb.Status{
1226
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1227 1228 1229 1230
			Reason:    err.Error(),
		}, nil
	}

1231 1232 1233
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", rpt.BeginTs()),
1234
		zap.Uint64("EndTS", rpt.EndTs()))
1235 1236 1237 1238

	if err := rpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1239
			zap.Error(err),
1240
			zap.Uint64("BeginTS", rpt.BeginTs()),
1241
			zap.Uint64("EndTS", rpt.EndTs()))
D
dragondriver 已提交
1242

E
Enwei Jiao 已提交
1243
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1244
			metrics.FailLabel).Inc()
1245

1246
		return &commonpb.Status{
1247
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1248 1249 1250 1251
			Reason:    err.Error(),
		}, nil
	}

1252 1253 1254
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", rpt.BeginTs()),
1255
		zap.Uint64("EndTS", rpt.EndTs()))
1256

E
Enwei Jiao 已提交
1257
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1258
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1259
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1260
	return rpt.result, nil
1261 1262
}

1263
// GetPartitionStatistics get the statistics of partition, such as num_rows.
C
Cai Yudong 已提交
1264
func (node *Proxy) GetPartitionStatistics(ctx context.Context, request *milvuspb.GetPartitionStatisticsRequest) (*milvuspb.GetPartitionStatisticsResponse, error) {
1265 1266 1267 1268 1269
	if !node.checkHealthy() {
		return &milvuspb.GetPartitionStatisticsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1270

E
Enwei Jiao 已提交
1271 1272
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetPartitionStatistics")
	defer sp.End()
1273 1274
	method := "GetPartitionStatistics"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1275
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1276
		metrics.TotalLabel).Inc()
1277

1278
	g := &getPartitionStatisticsTask{
1279 1280 1281
		ctx:                           ctx,
		Condition:                     NewTaskCondition(ctx),
		GetPartitionStatisticsRequest: request,
1282
		dataCoord:                     node.dataCoord,
1283 1284
	}

1285
	log := log.Ctx(ctx).With(
1286
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1287 1288 1289
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
1290

1291 1292
	log.Debug(rpcReceived(method))

1293 1294 1295
	if err := node.sched.ddQueue.Enqueue(g); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1296
			zap.Error(err))
1297

E
Enwei Jiao 已提交
1298
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1299
			metrics.AbandonLabel).Inc()
1300

1301 1302 1303 1304 1305 1306 1307 1308
		return &milvuspb.GetPartitionStatisticsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

1309 1310 1311
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", g.BeginTs()),
1312
		zap.Uint64("EndTS", g.EndTs()))
1313 1314 1315 1316

	if err := g.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
1317
			zap.Error(err),
1318
			zap.Uint64("BeginTS", g.BeginTs()),
1319
			zap.Uint64("EndTS", g.EndTs()))
1320

E
Enwei Jiao 已提交
1321
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1322
			metrics.FailLabel).Inc()
1323

1324 1325 1326 1327 1328 1329 1330 1331
		return &milvuspb.GetPartitionStatisticsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

1332 1333 1334
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", g.BeginTs()),
1335
		zap.Uint64("EndTS", g.EndTs()))
1336

E
Enwei Jiao 已提交
1337
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1338
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1339
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1340
	return g.result, nil
1341 1342
}

1343
// ShowPartitions list all partitions in the specific collection.
C
Cai Yudong 已提交
1344
func (node *Proxy) ShowPartitions(ctx context.Context, request *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
1345 1346 1347 1348 1349
	if !node.checkHealthy() {
		return &milvuspb.ShowPartitionsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1350

E
Enwei Jiao 已提交
1351 1352
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ShowPartitions")
	defer sp.End()
1353

1354
	spt := &showPartitionsTask{
G
godchen 已提交
1355 1356 1357
		ctx:                   ctx,
		Condition:             NewTaskCondition(ctx),
		ShowPartitionsRequest: request,
1358
		rootCoord:             node.rootCoord,
1359
		queryCoord:            node.queryCoord,
G
godchen 已提交
1360
		result:                nil,
1361 1362
	}

1363
	method := "ShowPartitions"
1364 1365
	tr := timerecord.NewTimeRecorder(method)
	//TODO: use collectionID instead of collectionName
E
Enwei Jiao 已提交
1366
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1367
		metrics.TotalLabel).Inc()
1368

1369 1370
	log := log.Ctx(ctx).With(zap.String("role", typeutil.ProxyRole))

1371 1372
	log.Debug(
		rpcReceived(method),
G
godchen 已提交
1373
		zap.Any("request", request))
1374 1375 1376 1377 1378 1379 1380

	if err := node.sched.ddQueue.Enqueue(spt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.Any("request", request))

E
Enwei Jiao 已提交
1381
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1382
			metrics.AbandonLabel).Inc()
1383

G
godchen 已提交
1384
		return &milvuspb.ShowPartitionsResponse{
1385
			Status: &commonpb.Status{
1386
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1387 1388 1389 1390 1391
				Reason:    err.Error(),
			},
		}, nil
	}

1392 1393 1394 1395
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", spt.BeginTs()),
		zap.Uint64("EndTS", spt.EndTs()),
1396 1397
		zap.String("db", spt.ShowPartitionsRequest.DbName),
		zap.String("collection", spt.ShowPartitionsRequest.CollectionName),
1398 1399 1400 1401 1402
		zap.Any("partitions", spt.ShowPartitionsRequest.PartitionNames))

	if err := spt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1403
			zap.Error(err),
1404 1405 1406 1407 1408
			zap.Uint64("BeginTS", spt.BeginTs()),
			zap.Uint64("EndTS", spt.EndTs()),
			zap.String("db", spt.ShowPartitionsRequest.DbName),
			zap.String("collection", spt.ShowPartitionsRequest.CollectionName),
			zap.Any("partitions", spt.ShowPartitionsRequest.PartitionNames))
D
dragondriver 已提交
1409

E
Enwei Jiao 已提交
1410
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1411
			metrics.FailLabel).Inc()
1412

G
godchen 已提交
1413
		return &milvuspb.ShowPartitionsResponse{
1414
			Status: &commonpb.Status{
1415
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1416 1417 1418 1419
				Reason:    err.Error(),
			},
		}, nil
	}
1420 1421 1422 1423 1424 1425 1426 1427 1428

	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", spt.BeginTs()),
		zap.Uint64("EndTS", spt.EndTs()),
		zap.String("db", spt.ShowPartitionsRequest.DbName),
		zap.String("collection", spt.ShowPartitionsRequest.CollectionName),
		zap.Any("partitions", spt.ShowPartitionsRequest.PartitionNames))

E
Enwei Jiao 已提交
1429
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1430
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1431
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1432 1433 1434
	return spt.result, nil
}

S
SimFG 已提交
1435 1436 1437 1438 1439 1440
func (node *Proxy) GetLoadingProgress(ctx context.Context, request *milvuspb.GetLoadingProgressRequest) (*milvuspb.GetLoadingProgressResponse, error) {
	if !node.checkHealthy() {
		return &milvuspb.GetLoadingProgressResponse{Status: unhealthyStatus()}, nil
	}
	method := "GetLoadingProgress"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1441 1442
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetLoadingProgress")
	defer sp.End()
E
Enwei Jiao 已提交
1443
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
1444 1445 1446
	log := log.Ctx(ctx)

	log.Debug(
S
SimFG 已提交
1447 1448 1449 1450
		rpcReceived(method),
		zap.Any("request", request))

	getErrResponse := func(err error) *milvuspb.GetLoadingProgressResponse {
J
Jiquan Long 已提交
1451
		log.Warn("fail to get loading progress",
1452
			zap.String("collection_name", request.CollectionName),
S
SimFG 已提交
1453 1454
			zap.Strings("partition_name", request.PartitionNames),
			zap.Error(err))
E
Enwei Jiao 已提交
1455
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
1456 1457 1458 1459 1460
		if errors.Is(err, ErrInsufficientMemory) {
			return &milvuspb.GetLoadingProgressResponse{
				Status: InSufficientMemoryStatus(request.GetCollectionName()),
			}
		}
S
SimFG 已提交
1461
		return &milvuspb.GetLoadingProgressResponse{
S
smellthemoon 已提交
1462
			Status: merr.Status(err),
S
SimFG 已提交
1463 1464 1465 1466 1467 1468 1469 1470 1471
		}
	}
	if err := validateCollectionName(request.CollectionName); err != nil {
		return getErrResponse(err), nil
	}
	collectionID, err := globalMetaCache.GetCollectionID(ctx, request.CollectionName)
	if err != nil {
		return getErrResponse(err), nil
	}
S
SimFG 已提交
1472

1473 1474 1475
	msgBase := commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_SystemInfo),
		commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
1476
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
1477
	)
S
SimFG 已提交
1478 1479 1480 1481 1482 1483 1484 1485 1486 1487
	if request.Base == nil {
		request.Base = msgBase
	} else {
		request.Base.MsgID = msgBase.MsgID
		request.Base.Timestamp = msgBase.Timestamp
		request.Base.SourceID = msgBase.SourceID
	}

	var progress int64
	if len(request.GetPartitionNames()) == 0 {
S
SimFG 已提交
1488
		if progress, err = getCollectionProgress(ctx, node.queryCoord, request.GetBase(), collectionID); err != nil {
S
SimFG 已提交
1489 1490 1491
			return getErrResponse(err), nil
		}
	} else {
S
SimFG 已提交
1492 1493
		if progress, err = getPartitionProgress(ctx, node.queryCoord, request.GetBase(),
			request.GetPartitionNames(), request.GetCollectionName(), collectionID); err != nil {
S
SimFG 已提交
1494 1495 1496 1497
			return getErrResponse(err), nil
		}
	}

1498
	log.Debug(
S
SimFG 已提交
1499 1500
		rpcDone(method),
		zap.Any("request", request))
E
Enwei Jiao 已提交
1501 1502
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
S
SimFG 已提交
1503 1504 1505 1506 1507 1508 1509 1510
	return &milvuspb.GetLoadingProgressResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
		},
		Progress: progress,
	}, nil
}

1511
func (node *Proxy) GetLoadState(ctx context.Context, request *milvuspb.GetLoadStateRequest) (*milvuspb.GetLoadStateResponse, error) {
S
SimFG 已提交
1512 1513 1514 1515 1516
	if !node.checkHealthy() {
		return &milvuspb.GetLoadStateResponse{Status: unhealthyStatus()}, nil
	}
	method := "GetLoadState"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1517 1518
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetLoadState")
	defer sp.End()
S
SimFG 已提交
1519 1520 1521 1522 1523 1524 1525 1526 1527 1528 1529 1530 1531 1532
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
	log := log.Ctx(ctx)

	log.Debug(
		rpcReceived(method),
		zap.Any("request", request))

	getErrResponse := func(err error) *milvuspb.GetLoadStateResponse {
		log.Warn("fail to get load state",
			zap.String("collection_name", request.CollectionName),
			zap.Strings("partition_name", request.PartitionNames),
			zap.Error(err))
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
		return &milvuspb.GetLoadStateResponse{
S
smellthemoon 已提交
1533
			Status: merr.Status(err),
S
SimFG 已提交
1534 1535 1536 1537 1538 1539 1540
		}
	}

	if err := validateCollectionName(request.CollectionName); err != nil {
		return getErrResponse(err), nil
	}

1541 1542
	// TODO(longjiquan): https://github.com/milvus-io/milvus/issues/21485, Remove `GetComponentStates` after error code
	// 	is ready to distinguish case whether the querycoord is not healthy or the collection is not even loaded.
S
SimFG 已提交
1543 1544 1545 1546 1547 1548 1549 1550 1551 1552 1553 1554 1555 1556 1557 1558 1559 1560 1561 1562 1563 1564 1565 1566 1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583
	if statesResp, err := node.queryCoord.GetComponentStates(ctx); err != nil {
		return getErrResponse(err), nil
	} else if statesResp.State == nil || statesResp.State.StateCode != commonpb.StateCode_Healthy {
		return getErrResponse(fmt.Errorf("the querycoord server isn't healthy, state: %v", statesResp.State)), nil
	}

	successResponse := &milvuspb.GetLoadStateResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
		},
	}
	defer func() {
		log.Debug(
			rpcDone(method),
			zap.Any("request", request))
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
		metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
	}()

	collectionID, err := globalMetaCache.GetCollectionID(ctx, request.CollectionName)
	if err != nil {
		successResponse.State = commonpb.LoadState_LoadStateNotExist
		return successResponse, nil
	}

	msgBase := commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_SystemInfo),
		commonpbutil.WithMsgID(0),
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
	)
	if request.Base == nil {
		request.Base = msgBase
	} else {
		request.Base.MsgID = msgBase.MsgID
		request.Base.Timestamp = msgBase.Timestamp
		request.Base.SourceID = msgBase.SourceID
	}

	var progress int64
	if len(request.GetPartitionNames()) == 0 {
		if progress, err = getCollectionProgress(ctx, node.queryCoord, request.GetBase(), collectionID); err != nil {
1584 1585 1586 1587 1588
			if errors.Is(err, ErrInsufficientMemory) {
				return &milvuspb.GetLoadStateResponse{
					Status: InSufficientMemoryStatus(request.GetCollectionName()),
				}, nil
			}
S
SimFG 已提交
1589 1590 1591 1592 1593 1594
			successResponse.State = commonpb.LoadState_LoadStateNotLoad
			return successResponse, nil
		}
	} else {
		if progress, err = getPartitionProgress(ctx, node.queryCoord, request.GetBase(),
			request.GetPartitionNames(), request.GetCollectionName(), collectionID); err != nil {
1595 1596 1597 1598 1599
			if errors.Is(err, ErrInsufficientMemory) {
				return &milvuspb.GetLoadStateResponse{
					Status: InSufficientMemoryStatus(request.GetCollectionName()),
				}, nil
			}
S
SimFG 已提交
1600 1601 1602 1603 1604 1605 1606 1607 1608 1609
			successResponse.State = commonpb.LoadState_LoadStateNotLoad
			return successResponse, nil
		}
	}
	if progress >= 100 {
		successResponse.State = commonpb.LoadState_LoadStateLoaded
	} else {
		successResponse.State = commonpb.LoadState_LoadStateLoading
	}
	return successResponse, nil
1610 1611
}

1612
// CreateIndex create index for collection.
C
Cai Yudong 已提交
1613
func (node *Proxy) CreateIndex(ctx context.Context, request *milvuspb.CreateIndexRequest) (*commonpb.Status, error) {
1614 1615 1616
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
1617

E
Enwei Jiao 已提交
1618 1619
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CreateIndex")
	defer sp.End()
D
dragondriver 已提交
1620

1621
	cit := &createIndexTask{
Z
zhenshan.cao 已提交
1622 1623 1624 1625
		ctx:        ctx,
		Condition:  NewTaskCondition(ctx),
		req:        request,
		rootCoord:  node.rootCoord,
1626
		datacoord:  node.dataCoord,
1627
		queryCoord: node.queryCoord,
1628 1629
	}

D
dragondriver 已提交
1630
	method := "CreateIndex"
1631
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1632
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1633
		metrics.TotalLabel).Inc()
1634 1635

	log := log.Ctx(ctx).With(
1636
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1637 1638 1639 1640
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.Any("extra_params", request.ExtraParams))
D
dragondriver 已提交
1641

1642 1643
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
1644 1645 1646
	if err := node.sched.ddQueue.Enqueue(cit); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1647
			zap.Error(err))
D
dragondriver 已提交
1648

E
Enwei Jiao 已提交
1649
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1650
			metrics.AbandonLabel).Inc()
1651

1652
		return &commonpb.Status{
1653
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1654 1655 1656 1657
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
1658 1659 1660
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", cit.BeginTs()),
1661
		zap.Uint64("EndTs", cit.EndTs()))
D
dragondriver 已提交
1662 1663 1664 1665

	if err := cit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1666
			zap.Error(err),
D
dragondriver 已提交
1667
			zap.Uint64("BeginTs", cit.BeginTs()),
1668
			zap.Uint64("EndTs", cit.EndTs()))
D
dragondriver 已提交
1669

E
Enwei Jiao 已提交
1670
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1671
			metrics.FailLabel).Inc()
1672

1673
		return &commonpb.Status{
1674
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1675 1676 1677 1678
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
1679 1680 1681
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", cit.BeginTs()),
1682
		zap.Uint64("EndTs", cit.EndTs()))
D
dragondriver 已提交
1683

E
Enwei Jiao 已提交
1684
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1685
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1686
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1687 1688 1689
	return cit.result, nil
}

1690
// DescribeIndex get the meta information of index, such as index state, index id and etc.
C
Cai Yudong 已提交
1691
func (node *Proxy) DescribeIndex(ctx context.Context, request *milvuspb.DescribeIndexRequest) (*milvuspb.DescribeIndexResponse, error) {
1692 1693 1694 1695 1696
	if !node.checkHealthy() {
		return &milvuspb.DescribeIndexResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1697

E
Enwei Jiao 已提交
1698 1699
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DescribeIndex")
	defer sp.End()
1700

1701
	dit := &describeIndexTask{
S
sunby 已提交
1702
		ctx:                  ctx,
1703 1704
		Condition:            NewTaskCondition(ctx),
		DescribeIndexRequest: request,
1705
		datacoord:            node.dataCoord,
1706 1707
	}

1708 1709
	method := "DescribeIndex"
	// avoid data race
1710
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1711
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1712
		metrics.TotalLabel).Inc()
1713 1714

	log := log.Ctx(ctx).With(
1715
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1716 1717 1718
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
1719 1720 1721
		zap.String("index name", request.IndexName))

	log.Debug(rpcReceived(method))
1722 1723 1724 1725

	if err := node.sched.ddQueue.Enqueue(dit); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1726
			zap.Error(err))
1727

E
Enwei Jiao 已提交
1728
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1729
			metrics.AbandonLabel).Inc()
1730

1731 1732
		return &milvuspb.DescribeIndexResponse{
			Status: &commonpb.Status{
1733
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1734 1735 1736 1737 1738
				Reason:    err.Error(),
			},
		}, nil
	}

1739 1740 1741
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", dit.BeginTs()),
1742
		zap.Uint64("EndTs", dit.EndTs()))
1743 1744 1745 1746

	if err := dit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1747
			zap.Error(err),
1748
			zap.Uint64("BeginTs", dit.BeginTs()),
1749
			zap.Uint64("EndTs", dit.EndTs()))
D
dragondriver 已提交
1750

Z
zhenshan.cao 已提交
1751 1752 1753 1754
		errCode := commonpb.ErrorCode_UnexpectedError
		if dit.result != nil {
			errCode = dit.result.Status.GetErrorCode()
		}
E
Enwei Jiao 已提交
1755
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1756
			metrics.FailLabel).Inc()
1757

1758 1759
		return &milvuspb.DescribeIndexResponse{
			Status: &commonpb.Status{
Z
zhenshan.cao 已提交
1760
				ErrorCode: errCode,
1761 1762 1763 1764 1765
				Reason:    err.Error(),
			},
		}, nil
	}

1766 1767 1768
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", dit.BeginTs()),
1769
		zap.Uint64("EndTs", dit.EndTs()))
1770

E
Enwei Jiao 已提交
1771
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1772
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1773
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1774 1775 1776
	return dit.result, nil
}

1777
// DropIndex drop the index of collection.
C
Cai Yudong 已提交
1778
func (node *Proxy) DropIndex(ctx context.Context, request *milvuspb.DropIndexRequest) (*commonpb.Status, error) {
1779 1780 1781
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
1782

E
Enwei Jiao 已提交
1783 1784
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DropIndex")
	defer sp.End()
D
dragondriver 已提交
1785

1786
	dit := &dropIndexTask{
S
sunby 已提交
1787
		ctx:              ctx,
B
BossZou 已提交
1788 1789
		Condition:        NewTaskCondition(ctx),
		DropIndexRequest: request,
1790
		dataCoord:        node.dataCoord,
1791
		queryCoord:       node.queryCoord,
B
BossZou 已提交
1792
	}
G
godchen 已提交
1793

D
dragondriver 已提交
1794
	method := "DropIndex"
1795
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1796
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1797
		metrics.TotalLabel).Inc()
D
dragondriver 已提交
1798

1799
	log := log.Ctx(ctx).With(
1800
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1801 1802 1803 1804 1805
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))

1806 1807
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
1808 1809 1810
	if err := node.sched.ddQueue.Enqueue(dit); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1811
			zap.Error(err))
E
Enwei Jiao 已提交
1812
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1813
			metrics.AbandonLabel).Inc()
D
dragondriver 已提交
1814

B
BossZou 已提交
1815
		return &commonpb.Status{
1816
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
B
BossZou 已提交
1817 1818 1819
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
1820

D
dragondriver 已提交
1821 1822 1823
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", dit.BeginTs()),
1824
		zap.Uint64("EndTs", dit.EndTs()))
D
dragondriver 已提交
1825 1826 1827 1828

	if err := dit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1829
			zap.Error(err),
D
dragondriver 已提交
1830
			zap.Uint64("BeginTs", dit.BeginTs()),
1831
			zap.Uint64("EndTs", dit.EndTs()))
D
dragondriver 已提交
1832

E
Enwei Jiao 已提交
1833
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1834
			metrics.FailLabel).Inc()
1835

B
BossZou 已提交
1836
		return &commonpb.Status{
1837
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
B
BossZou 已提交
1838 1839 1840
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
1841 1842 1843 1844

	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", dit.BeginTs()),
1845
		zap.Uint64("EndTs", dit.EndTs()))
D
dragondriver 已提交
1846

E
Enwei Jiao 已提交
1847
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1848
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1849
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
B
BossZou 已提交
1850 1851 1852
	return dit.result, nil
}

1853 1854
// GetIndexBuildProgress gets index build progress with filed_name and index_name.
// IndexRows is the num of indexed rows. And TotalRows is the total number of segment rows.
1855
// Deprecated: use DescribeIndex instead
C
Cai Yudong 已提交
1856
func (node *Proxy) GetIndexBuildProgress(ctx context.Context, request *milvuspb.GetIndexBuildProgressRequest) (*milvuspb.GetIndexBuildProgressResponse, error) {
1857 1858 1859 1860 1861
	if !node.checkHealthy() {
		return &milvuspb.GetIndexBuildProgressResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1862

E
Enwei Jiao 已提交
1863 1864
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetIndexBuildProgress")
	defer sp.End()
1865

1866
	gibpt := &getIndexBuildProgressTask{
1867 1868 1869
		ctx:                          ctx,
		Condition:                    NewTaskCondition(ctx),
		GetIndexBuildProgressRequest: request,
1870
		rootCoord:                    node.rootCoord,
1871
		dataCoord:                    node.dataCoord,
1872 1873
	}

1874
	method := "GetIndexBuildProgress"
1875
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1876
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1877
		metrics.TotalLabel).Inc()
1878 1879

	log := log.Ctx(ctx).With(
1880
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1881 1882 1883 1884
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
1885

1886 1887
	log.Debug(rpcReceived(method))

1888 1889 1890
	if err := node.sched.ddQueue.Enqueue(gibpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1891
			zap.Error(err))
E
Enwei Jiao 已提交
1892
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1893
			metrics.AbandonLabel).Inc()
1894

1895 1896 1897 1898 1899 1900 1901 1902
		return &milvuspb.GetIndexBuildProgressResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

1903 1904 1905
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", gibpt.BeginTs()),
1906
		zap.Uint64("EndTs", gibpt.EndTs()))
1907 1908 1909 1910

	if err := gibpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
1911
			zap.Error(err),
1912
			zap.Uint64("BeginTs", gibpt.BeginTs()),
1913
			zap.Uint64("EndTs", gibpt.EndTs()))
E
Enwei Jiao 已提交
1914
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1915
			metrics.FailLabel).Inc()
1916 1917 1918 1919 1920 1921 1922 1923

		return &milvuspb.GetIndexBuildProgressResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}
1924 1925 1926 1927

	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", gibpt.BeginTs()),
1928
		zap.Uint64("EndTs", gibpt.EndTs()))
1929

E
Enwei Jiao 已提交
1930
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1931
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1932
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1933
	return gibpt.result, nil
1934 1935
}

1936
// GetIndexState get the build-state of index.
1937
// Deprecated: use DescribeIndex instead
C
Cai Yudong 已提交
1938
func (node *Proxy) GetIndexState(ctx context.Context, request *milvuspb.GetIndexStateRequest) (*milvuspb.GetIndexStateResponse, error) {
1939 1940 1941 1942 1943
	if !node.checkHealthy() {
		return &milvuspb.GetIndexStateResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1944

E
Enwei Jiao 已提交
1945 1946
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Insert")
	defer sp.End()
1947

1948
	dipt := &getIndexStateTask{
G
godchen 已提交
1949 1950 1951
		ctx:                  ctx,
		Condition:            NewTaskCondition(ctx),
		GetIndexStateRequest: request,
1952
		dataCoord:            node.dataCoord,
1953
		rootCoord:            node.rootCoord,
1954 1955
	}

1956
	method := "GetIndexState"
1957
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1958
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1959
		metrics.TotalLabel).Inc()
1960 1961

	log := log.Ctx(ctx).With(
1962
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1963 1964 1965 1966
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
1967

1968 1969
	log.Debug(rpcReceived(method))

1970 1971 1972
	if err := node.sched.ddQueue.Enqueue(dipt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1973
			zap.Error(err))
1974

E
Enwei Jiao 已提交
1975
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1976
			metrics.AbandonLabel).Inc()
1977

G
godchen 已提交
1978
		return &milvuspb.GetIndexStateResponse{
1979
			Status: &commonpb.Status{
1980
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1981 1982 1983 1984 1985
				Reason:    err.Error(),
			},
		}, nil
	}

1986 1987 1988
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", dipt.BeginTs()),
1989
		zap.Uint64("EndTs", dipt.EndTs()))
1990 1991 1992 1993

	if err := dipt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1994
			zap.Error(err),
1995
			zap.Uint64("BeginTs", dipt.BeginTs()),
1996
			zap.Uint64("EndTs", dipt.EndTs()))
E
Enwei Jiao 已提交
1997
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1998
			metrics.FailLabel).Inc()
1999

G
godchen 已提交
2000
		return &milvuspb.GetIndexStateResponse{
2001
			Status: &commonpb.Status{
2002
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2003 2004 2005 2006 2007
				Reason:    err.Error(),
			},
		}, nil
	}

2008 2009 2010
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", dipt.BeginTs()),
2011
		zap.Uint64("EndTs", dipt.EndTs()))
2012

E
Enwei Jiao 已提交
2013
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2014
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
2015
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
2016 2017 2018
	return dipt.result, nil
}

2019
// Insert insert records into collection.
C
Cai Yudong 已提交
2020
func (node *Proxy) Insert(ctx context.Context, request *milvuspb.InsertRequest) (*milvuspb.MutationResult, error) {
E
Enwei Jiao 已提交
2021 2022
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Insert")
	defer sp.End()
X
Xiangyu Wang 已提交
2023

2024 2025 2026 2027 2028
	if !node.checkHealthy() {
		return &milvuspb.MutationResult{
			Status: unhealthyStatus(),
		}, nil
	}
2029 2030
	method := "Insert"
	tr := timerecord.NewTimeRecorder(method)
S
smellthemoon 已提交
2031
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.InsertLabel).Add(float64(proto.Size(request)))
E
Enwei Jiao 已提交
2032
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
S
smellthemoon 已提交
2033

2034
	it := &insertTask{
2035 2036
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),
X
xige-16 已提交
2037
		// req:       request,
2038
		insertMsg: &msgstream.InsertMsg{
2039 2040 2041
			BaseMsg: msgstream.BaseMsg{
				HashValues: request.HashKeys,
			},
2042
			InsertRequest: msgpb.InsertRequest{
2043 2044 2045
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Insert),
					commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
2046
					commonpbutil.WithSourceID(paramtable.GetNodeID()),
2047
				),
2048 2049
				CollectionName: request.CollectionName,
				PartitionName:  request.PartitionName,
X
xige-16 已提交
2050 2051
				FieldsData:     request.FieldsData,
				NumRows:        uint64(request.NumRows),
2052
				Version:        msgpb.InsertDataVersion_ColumnBased,
2053
				// RowData: transfer column based request to this
2054 2055
			},
		},
2056
		idAllocator:   node.rowIDAllocator,
2057 2058 2059
		segIDAssigner: node.segAssigner,
		chMgr:         node.chMgr,
		chTicker:      node.chTicker,
2060
	}
2061

2062 2063
	if len(it.insertMsg.PartitionName) <= 0 {
		it.insertMsg.PartitionName = Params.CommonCfg.DefaultPartitionName.GetValue()
2064 2065
	}

X
Xiangyu Wang 已提交
2066
	constructFailedResponse := func(err error) *milvuspb.MutationResult {
X
xige-16 已提交
2067
		numRows := request.NumRows
2068 2069 2070 2071
		errIndex := make([]uint32, numRows)
		for i := uint32(0); i < numRows; i++ {
			errIndex[i] = i
		}
2072

X
Xiangyu Wang 已提交
2073 2074 2075 2076 2077 2078 2079
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
			ErrIndex: errIndex,
		}
2080 2081
	}

X
Xiangyu Wang 已提交
2082
	log.Debug("Enqueue insert request in Proxy",
2083
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2084 2085 2086 2087 2088
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
		zap.Int("len(FieldsData)", len(request.FieldsData)),
		zap.Int("len(HashKeys)", len(request.HashKeys)),
2089
		zap.Uint32("NumRows", request.NumRows))
D
dragondriver 已提交
2090

X
Xiangyu Wang 已提交
2091
	if err := node.sched.dmQueue.Enqueue(it); err != nil {
J
Jiquan Long 已提交
2092
		log.Warn("Failed to enqueue insert task: " + err.Error())
E
Enwei Jiao 已提交
2093
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2094
			metrics.AbandonLabel).Inc()
X
Xiangyu Wang 已提交
2095
		return constructFailedResponse(err), nil
2096
	}
D
dragondriver 已提交
2097

X
Xiangyu Wang 已提交
2098
	log.Debug("Detail of insert request in Proxy",
2099
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2100 2101 2102 2103 2104
		zap.Uint64("BeginTS", it.BeginTs()),
		zap.Uint64("EndTS", it.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
2105
		zap.Uint32("NumRows", request.NumRows))
X
Xiangyu Wang 已提交
2106 2107

	if err := it.WaitToFinish(); err != nil {
2108
		log.Warn("Failed to execute insert task in task scheduler: " + err.Error())
E
Enwei Jiao 已提交
2109
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2110
			metrics.FailLabel).Inc()
X
Xiangyu Wang 已提交
2111 2112 2113 2114 2115
		return constructFailedResponse(err), nil
	}

	if it.result.Status.ErrorCode != commonpb.ErrorCode_Success {
		setErrorIndex := func() {
X
xige-16 已提交
2116
			numRows := request.NumRows
X
Xiangyu Wang 已提交
2117 2118 2119 2120 2121 2122 2123 2124 2125 2126 2127
			errIndex := make([]uint32, numRows)
			for i := uint32(0); i < numRows; i++ {
				errIndex[i] = i
			}
			it.result.ErrIndex = errIndex
		}

		setErrorIndex()
	}

	// InsertCnt always equals to the number of entities in the request
X
xige-16 已提交
2128
	it.result.InsertCnt = int64(request.NumRows)
D
dragondriver 已提交
2129

S
smellthemoon 已提交
2130 2131 2132
	receiveSize := proto.Size(it.insertMsg)
	rateCol.Add(internalpb.RateType_DMLInsert.String(), float64(receiveSize))

E
Enwei Jiao 已提交
2133
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2134
		metrics.SuccessLabel).Inc()
2135
	successCnt := it.result.InsertCnt - int64(len(it.result.ErrIndex))
E
Enwei Jiao 已提交
2136 2137 2138
	metrics.ProxyInsertVectors.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Add(float64(successCnt))
	metrics.ProxyMutationLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.InsertLabel).Observe(float64(tr.ElapseSpan().Milliseconds()))
	metrics.ProxyCollectionMutationLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.InsertLabel, request.CollectionName).Observe(float64(tr.ElapseSpan().Milliseconds()))
2139 2140 2141
	return it.result, nil
}

2142
// Delete delete records from collection, then these records cannot be searched.
G
groot 已提交
2143
func (node *Proxy) Delete(ctx context.Context, request *milvuspb.DeleteRequest) (*milvuspb.MutationResult, error) {
E
Enwei Jiao 已提交
2144 2145
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Delete")
	defer sp.End()
2146 2147 2148
	log := log.Ctx(ctx)
	log.Debug("Start processing delete request in Proxy")
	defer log.Debug("Finish processing delete request in Proxy")
2149

S
smellthemoon 已提交
2150
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.DeleteLabel).Add(float64(proto.Size(request)))
2151

G
groot 已提交
2152 2153 2154 2155 2156 2157
	if !node.checkHealthy() {
		return &milvuspb.MutationResult{
			Status: unhealthyStatus(),
		}, nil
	}

2158 2159 2160
	method := "Delete"
	tr := timerecord.NewTimeRecorder(method)

E
Enwei Jiao 已提交
2161
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2162
		metrics.TotalLabel).Inc()
2163
	dt := &deleteTask{
X
xige-16 已提交
2164 2165 2166
		ctx:        ctx,
		Condition:  NewTaskCondition(ctx),
		deleteExpr: request.Expr,
2167
		deleteMsg: &BaseDeleteTask{
G
godchen 已提交
2168 2169 2170
			BaseMsg: msgstream.BaseMsg{
				HashValues: request.HashKeys,
			},
2171
			DeleteRequest: msgpb.DeleteRequest{
2172 2173 2174 2175
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Delete),
					commonpbutil.WithMsgID(0),
				),
X
xige-16 已提交
2176
				DbName:         request.DbName,
G
godchen 已提交
2177 2178 2179
				CollectionName: request.CollectionName,
				PartitionName:  request.PartitionName,
				// RowData: transfer column based request to this
C
Cai Yudong 已提交
2180 2181
			},
		},
S
smellthemoon 已提交
2182 2183 2184
		idAllocator: node.rowIDAllocator,
		chMgr:       node.chMgr,
		chTicker:    node.chTicker,
G
groot 已提交
2185 2186
	}

2187
	log.Debug("Enqueue delete request in Proxy",
2188
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2189 2190 2191 2192
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
		zap.String("expr", request.Expr))
2193 2194 2195

	// MsgID will be set by Enqueue()
	if err := node.sched.dmQueue.Enqueue(dt); err != nil {
2196
		log.Error("Failed to enqueue delete task: " + err.Error())
E
Enwei Jiao 已提交
2197
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2198
			metrics.AbandonLabel).Inc()
2199

G
groot 已提交
2200 2201 2202 2203 2204 2205 2206 2207
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

2208
	log.Debug("Detail of delete request in Proxy",
2209
		zap.String("role", typeutil.ProxyRole),
2210
		zap.Uint64("timestamp", dt.deleteMsg.Base.Timestamp),
G
groot 已提交
2211 2212 2213
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
2214
		zap.String("expr", request.Expr))
G
groot 已提交
2215

2216
	if err := dt.WaitToFinish(); err != nil {
2217
		log.Error("Failed to execute delete task in task scheduler: " + err.Error())
E
Enwei Jiao 已提交
2218
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2219
			metrics.FailLabel).Inc()
G
groot 已提交
2220 2221 2222 2223 2224 2225 2226 2227
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

S
smellthemoon 已提交
2228 2229 2230
	receiveSize := proto.Size(dt.deleteMsg)
	rateCol.Add(internalpb.RateType_DMLDelete.String(), float64(receiveSize))

E
Enwei Jiao 已提交
2231
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2232
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
2233 2234
	metrics.ProxyMutationLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.DeleteLabel).Observe(float64(tr.ElapseSpan().Milliseconds()))
	metrics.ProxyCollectionMutationLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.DeleteLabel, request.CollectionName).Observe(float64(tr.ElapseSpan().Milliseconds()))
G
groot 已提交
2235 2236 2237
	return dt.result, nil
}

S
smellthemoon 已提交
2238 2239
// Upsert upsert records into collection.
func (node *Proxy) Upsert(ctx context.Context, request *milvuspb.UpsertRequest) (*milvuspb.MutationResult, error) {
E
Enwei Jiao 已提交
2240 2241
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Upsert")
	defer sp.End()
S
smellthemoon 已提交
2242 2243 2244 2245 2246 2247 2248 2249 2250 2251 2252 2253 2254 2255 2256 2257 2258 2259 2260 2261 2262 2263 2264 2265 2266 2267 2268 2269 2270 2271

	log := log.Ctx(ctx).With(
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
		zap.Uint32("NumRows", request.NumRows),
	)
	log.Debug("Start processing upsert request in Proxy")

	if !node.checkHealthy() {
		return &milvuspb.MutationResult{
			Status: unhealthyStatus(),
		}, nil
	}
	method := "Upsert"
	tr := timerecord.NewTimeRecorder(method)

	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.UpsertLabel).Add(float64(proto.Size(request)))
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()

	it := &upsertTask{
		baseMsg: msgstream.BaseMsg{
			HashValues: request.HashKeys,
		},
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),

		req: &milvuspb.UpsertRequest{
			Base: commonpbutil.NewMsgBase(
2272
				commonpbutil.WithMsgType(commonpb.MsgType_Upsert),
S
smellthemoon 已提交
2273 2274 2275 2276 2277 2278 2279 2280 2281 2282 2283 2284 2285 2286 2287 2288 2289 2290 2291 2292 2293 2294 2295 2296 2297 2298 2299 2300 2301 2302 2303 2304 2305 2306 2307 2308 2309 2310 2311 2312 2313 2314 2315 2316 2317 2318 2319 2320 2321 2322 2323 2324 2325 2326 2327 2328 2329 2330 2331 2332 2333 2334 2335 2336 2337 2338 2339 2340 2341
				commonpbutil.WithSourceID(paramtable.GetNodeID()),
			),
			CollectionName: request.CollectionName,
			PartitionName:  request.PartitionName,
			FieldsData:     request.FieldsData,
			NumRows:        request.NumRows,
		},

		result: &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_Success,
			},
			IDs: &schemapb.IDs{
				IdField: nil,
			},
		},

		idAllocator:   node.rowIDAllocator,
		segIDAssigner: node.segAssigner,
		chMgr:         node.chMgr,
		chTicker:      node.chTicker,
	}

	if len(it.req.PartitionName) <= 0 {
		it.req.PartitionName = Params.CommonCfg.DefaultPartitionName.GetValue()
	}

	constructFailedResponse := func(err error, errCode commonpb.ErrorCode) *milvuspb.MutationResult {
		numRows := request.NumRows
		errIndex := make([]uint32, numRows)
		for i := uint32(0); i < numRows; i++ {
			errIndex[i] = i
		}

		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: errCode,
				Reason:    err.Error(),
			},
			ErrIndex: errIndex,
		}
	}

	log.Debug("Enqueue upsert request in Proxy",
		zap.Int("len(FieldsData)", len(request.FieldsData)),
		zap.Int("len(HashKeys)", len(request.HashKeys)))

	if err := node.sched.dmQueue.Enqueue(it); err != nil {
		log.Info("Failed to enqueue upsert task",
			zap.Error(err))
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
			metrics.AbandonLabel).Inc()
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

	log.Debug("Detail of upsert request in Proxy",
		zap.Uint64("BeginTS", it.BeginTs()),
		zap.Uint64("EndTS", it.EndTs()))

	if err := it.WaitToFinish(); err != nil {
		log.Info("Failed to execute insert task in task scheduler",
			zap.Error(err))
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
			metrics.FailLabel).Inc()
S
smellthemoon 已提交
2342 2343 2344 2345 2346
		// Not every error case changes the status internally
		// change status there to handle it
		if it.result.Status.ErrorCode == commonpb.ErrorCode_Success {
			it.result.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		}
S
smellthemoon 已提交
2347 2348 2349 2350 2351 2352 2353 2354 2355 2356 2357 2358 2359 2360 2361 2362 2363 2364 2365 2366 2367 2368 2369 2370 2371 2372 2373 2374 2375 2376
		return constructFailedResponse(err, it.result.Status.ErrorCode), nil
	}

	if it.result.Status.ErrorCode != commonpb.ErrorCode_Success {
		setErrorIndex := func() {
			numRows := request.NumRows
			errIndex := make([]uint32, numRows)
			for i := uint32(0); i < numRows; i++ {
				errIndex[i] = i
			}
			it.result.ErrIndex = errIndex
		}
		setErrorIndex()
	}

	insertReceiveSize := proto.Size(it.upsertMsg.InsertMsg)
	deleteReceiveSize := proto.Size(it.upsertMsg.DeleteMsg)

	rateCol.Add(internalpb.RateType_DMLDelete.String(), float64(deleteReceiveSize))
	rateCol.Add(internalpb.RateType_DMLInsert.String(), float64(insertReceiveSize))

	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.SuccessLabel).Inc()
	metrics.ProxyMutationLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.UpsertLabel).Observe(float64(tr.ElapseSpan().Milliseconds()))
	metrics.ProxyCollectionMutationLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.UpsertLabel, request.CollectionName).Observe(float64(tr.ElapseSpan().Milliseconds()))

	log.Debug("Finish processing upsert request in Proxy")
	return it.result, nil
}

2377
// Search search the most similar records of requests.
C
Cai Yudong 已提交
2378
func (node *Proxy) Search(ctx context.Context, request *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
2379
	receiveSize := proto.Size(request)
E
Enwei Jiao 已提交
2380
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.SearchLabel).Add(float64(receiveSize))
2381 2382 2383

	rateCol.Add(internalpb.RateType_DQLSearch.String(), float64(request.GetNq()))

2384 2385 2386 2387 2388
	if !node.checkHealthy() {
		return &milvuspb.SearchResults{
			Status: unhealthyStatus(),
		}, nil
	}
2389 2390
	method := "Search"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2391
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2392
		metrics.TotalLabel).Inc()
D
dragondriver 已提交
2393

E
Enwei Jiao 已提交
2394 2395
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Search")
	defer sp.End()
D
dragondriver 已提交
2396

2397
	qt := &searchTask{
S
sunby 已提交
2398
		ctx:       ctx,
2399
		Condition: NewTaskCondition(ctx),
G
godchen 已提交
2400
		SearchRequest: &internalpb.SearchRequest{
2401 2402
			Base: commonpbutil.NewMsgBase(
				commonpbutil.WithMsgType(commonpb.MsgType_Search),
E
Enwei Jiao 已提交
2403
				commonpbutil.WithSourceID(paramtable.GetNodeID()),
2404
			),
E
Enwei Jiao 已提交
2405
			ReqID: paramtable.GetNodeID(),
2406
		},
2407 2408 2409 2410
		request:  request,
		qc:       node.queryCoord,
		tr:       timerecord.NewTimeRecorder("search"),
		shardMgr: node.shardMgr,
2411 2412
	}

2413 2414 2415
	travelTs := request.TravelTimestamp
	guaranteeTs := request.GuaranteeTimestamp

2416
	log := log.Ctx(ctx).With(
2417
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2418 2419 2420 2421 2422
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames),
		zap.Any("dsl", request.Dsl),
		zap.Any("len(PlaceholderGroup)", len(request.PlaceholderGroup)),
2423 2424 2425 2426
		zap.Any("OutputFields", request.OutputFields),
		zap.Any("search_params", request.SearchParams),
		zap.Uint64("travel_timestamp", travelTs),
		zap.Uint64("guarantee_timestamp", guaranteeTs))
D
dragondriver 已提交
2427

2428 2429 2430
	log.Debug(
		rpcReceived(method))

2431
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
2432
		log.Warn(
2433
			rpcFailedToEnqueue(method),
2434
			zap.Error(err))
D
dragondriver 已提交
2435

E
Enwei Jiao 已提交
2436
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2437
			metrics.AbandonLabel).Inc()
2438

2439 2440
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
2441
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2442 2443 2444 2445
				Reason:    err.Error(),
			},
		}, nil
	}
Z
Zach 已提交
2446
	tr.CtxRecord(ctx, "search request enqueue")
2447

2448
	log.Debug(
2449
		rpcEnqueued(method),
2450
		zap.Uint64("timestamp", qt.Base.Timestamp))
D
dragondriver 已提交
2451

2452
	if err := qt.WaitToFinish(); err != nil {
2453
		log.Warn(
2454
			rpcFailedToWaitToFinish(method),
2455
			zap.Error(err))
2456

E
Enwei Jiao 已提交
2457
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2458
			metrics.FailLabel).Inc()
2459

2460 2461
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
2462
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2463 2464 2465 2466 2467
				Reason:    err.Error(),
			},
		}, nil
	}

Z
Zach 已提交
2468
	span := tr.CtxRecord(ctx, "wait search result")
E
Enwei Jiao 已提交
2469
	metrics.ProxyWaitForSearchResultLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2470
		metrics.SearchLabel).Observe(float64(span.Milliseconds()))
2471
	tr.CtxRecord(ctx, "wait search result")
2472
	log.Debug(rpcDone(method))
D
dragondriver 已提交
2473

E
Enwei Jiao 已提交
2474
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2475
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
2476
	metrics.ProxySearchVectors.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Add(float64(qt.result.GetResults().GetNumQueries()))
C
cai.zhang 已提交
2477
	searchDur := tr.ElapseSpan().Milliseconds()
E
Enwei Jiao 已提交
2478
	metrics.ProxySQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2479
		metrics.SearchLabel).Observe(float64(searchDur))
E
Enwei Jiao 已提交
2480
	metrics.ProxyCollectionSQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2481
		metrics.SearchLabel, request.CollectionName).Observe(float64(searchDur))
2482 2483
	if qt.result != nil {
		sentSize := proto.Size(qt.result)
E
Enwei Jiao 已提交
2484
		metrics.ProxyReadReqSendBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Add(float64(sentSize))
2485
		rateCol.Add(metricsinfo.ReadResultThroughput, float64(sentSize))
2486
	}
2487 2488 2489
	return qt.result, nil
}

2490
// Flush notify data nodes to persist the data of collection.
2491 2492 2493 2494 2495 2496 2497
func (node *Proxy) Flush(ctx context.Context, request *milvuspb.FlushRequest) (*milvuspb.FlushResponse, error) {
	resp := &milvuspb.FlushResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    "",
		},
	}
2498
	if !node.checkHealthy() {
2499 2500
		resp.Status.Reason = "proxy is not healthy"
		return resp, nil
2501
	}
D
dragondriver 已提交
2502

E
Enwei Jiao 已提交
2503 2504
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Flush")
	defer sp.End()
D
dragondriver 已提交
2505

2506
	ft := &flushTask{
T
ThreadDao 已提交
2507 2508 2509
		ctx:          ctx,
		Condition:    NewTaskCondition(ctx),
		FlushRequest: request,
2510
		dataCoord:    node.dataCoord,
2511 2512
	}

D
dragondriver 已提交
2513
	method := "Flush"
2514
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2515
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2516

2517
	log := log.Ctx(ctx).With(
2518
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2519 2520
		zap.String("db", request.DbName),
		zap.Any("collections", request.CollectionNames))
D
dragondriver 已提交
2521

2522 2523
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2524 2525 2526
	if err := node.sched.ddQueue.Enqueue(ft); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2527
			zap.Error(err))
D
dragondriver 已提交
2528

E
Enwei Jiao 已提交
2529
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
2530

2531 2532
		resp.Status.Reason = err.Error()
		return resp, nil
2533 2534
	}

D
dragondriver 已提交
2535 2536 2537
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", ft.BeginTs()),
2538
		zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2539 2540 2541 2542

	if err := ft.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
2543
			zap.Error(err),
D
dragondriver 已提交
2544
			zap.Uint64("BeginTs", ft.BeginTs()),
2545
			zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2546

E
Enwei Jiao 已提交
2547
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
2548

D
dragondriver 已提交
2549
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
2550 2551
		resp.Status.Reason = err.Error()
		return resp, nil
2552 2553
	}

D
dragondriver 已提交
2554 2555 2556
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", ft.BeginTs()),
2557
		zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2558

E
Enwei Jiao 已提交
2559 2560
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
2561
	return ft.result, nil
2562 2563
}

2564
// Query get the records by primary keys.
C
Cai Yudong 已提交
2565
func (node *Proxy) Query(ctx context.Context, request *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) {
2566
	receiveSize := proto.Size(request)
E
Enwei Jiao 已提交
2567
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.QueryLabel).Add(float64(receiveSize))
2568 2569 2570

	rateCol.Add(internalpb.RateType_DQLQuery.String(), 1)

2571 2572 2573 2574 2575
	if !node.checkHealthy() {
		return &milvuspb.QueryResults{
			Status: unhealthyStatus(),
		}, nil
	}
2576

E
Enwei Jiao 已提交
2577 2578
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Query")
	defer sp.End()
2579
	tr := timerecord.NewTimeRecorder("Query")
D
dragondriver 已提交
2580

2581
	qt := &queryTask{
2582 2583 2584
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),
		RetrieveRequest: &internalpb.RetrieveRequest{
2585 2586
			Base: commonpbutil.NewMsgBase(
				commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
E
Enwei Jiao 已提交
2587
				commonpbutil.WithSourceID(paramtable.GetNodeID()),
2588
			),
E
Enwei Jiao 已提交
2589
			ReqID: paramtable.GetNodeID(),
2590
		},
2591 2592
		request:          request,
		qc:               node.queryCoord,
2593
		queryShardPolicy: mergeRoundRobinPolicy,
2594
		shardMgr:         node.shardMgr,
2595 2596
	}

D
dragondriver 已提交
2597 2598
	method := "Query"

E
Enwei Jiao 已提交
2599
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2600 2601
		metrics.TotalLabel).Inc()

2602
	log := log.Ctx(ctx).With(
2603
		zap.String("role", typeutil.ProxyRole),
2604 2605
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
2606 2607 2608 2609
		zap.Strings("partitions", request.PartitionNames))

	log.Debug(
		rpcReceived(method),
2610 2611 2612 2613
		zap.String("expr", request.Expr),
		zap.Strings("OutputFields", request.OutputFields),
		zap.Uint64("travel_timestamp", request.TravelTimestamp),
		zap.Uint64("guarantee_timestamp", request.GuaranteeTimestamp))
G
godchen 已提交
2614

D
dragondriver 已提交
2615
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
2616
		log.Warn(
D
dragondriver 已提交
2617
			rpcFailedToEnqueue(method),
2618
			zap.Error(err))
D
dragondriver 已提交
2619

E
Enwei Jiao 已提交
2620
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2621
			metrics.AbandonLabel).Inc()
2622

2623 2624 2625 2626 2627 2628
		return &milvuspb.QueryResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
2629
	}
Z
Zach 已提交
2630
	tr.CtxRecord(ctx, "query request enqueue")
2631

2632
	log.Debug(rpcEnqueued(method))
D
dragondriver 已提交
2633 2634

	if err := qt.WaitToFinish(); err != nil {
2635
		log.Warn(
D
dragondriver 已提交
2636
			rpcFailedToWaitToFinish(method),
2637
			zap.Error(err))
2638

E
Enwei Jiao 已提交
2639
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2640
			metrics.FailLabel).Inc()
2641

2642 2643 2644 2645 2646 2647 2648
		return &milvuspb.QueryResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}
Z
Zach 已提交
2649
	span := tr.CtxRecord(ctx, "wait query result")
E
Enwei Jiao 已提交
2650
	metrics.ProxyWaitForSearchResultLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2651
		metrics.QueryLabel).Observe(float64(span.Milliseconds()))
2652

2653
	log.Debug(rpcDone(method))
D
dragondriver 已提交
2654

E
Enwei Jiao 已提交
2655
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2656 2657
		metrics.SuccessLabel).Inc()

E
Enwei Jiao 已提交
2658
	metrics.ProxySQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2659
		metrics.QueryLabel).Observe(float64(tr.ElapseSpan().Milliseconds()))
E
Enwei Jiao 已提交
2660
	metrics.ProxyCollectionSQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2661
		metrics.QueryLabel, request.CollectionName).Observe(float64(tr.ElapseSpan().Milliseconds()))
2662 2663

	ret := &milvuspb.QueryResults{
2664 2665
		Status:     qt.result.Status,
		FieldsData: qt.result.FieldsData,
2666 2667
	}
	sentSize := proto.Size(qt.result)
2668
	rateCol.Add(metricsinfo.ReadResultThroughput, float64(sentSize))
E
Enwei Jiao 已提交
2669
	metrics.ProxyReadReqSendBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Add(float64(sentSize))
2670
	return ret, nil
2671
}
2672

2673
// CreateAlias create alias for collection, then you can search the collection with alias.
Y
Yusup 已提交
2674 2675 2676 2677
func (node *Proxy) CreateAlias(ctx context.Context, request *milvuspb.CreateAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2678

E
Enwei Jiao 已提交
2679 2680
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CreateAlias")
	defer sp.End()
D
dragondriver 已提交
2681

Y
Yusup 已提交
2682 2683 2684 2685 2686 2687 2688
	cat := &CreateAliasTask{
		ctx:                ctx,
		Condition:          NewTaskCondition(ctx),
		CreateAliasRequest: request,
		rootCoord:          node.rootCoord,
	}

D
dragondriver 已提交
2689
	method := "CreateAlias"
2690
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2691
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2692

2693
	log := log.Ctx(ctx).With(
D
dragondriver 已提交
2694 2695 2696 2697 2698
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

2699 2700
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2701 2702 2703
	if err := node.sched.ddQueue.Enqueue(cat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2704
			zap.Error(err))
D
dragondriver 已提交
2705

E
Enwei Jiao 已提交
2706
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
2707

Y
Yusup 已提交
2708 2709 2710 2711 2712 2713
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2714 2715 2716
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", cat.BeginTs()),
2717
		zap.Uint64("EndTs", cat.EndTs()))
D
dragondriver 已提交
2718 2719 2720 2721

	if err := cat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2722
			zap.Error(err),
D
dragondriver 已提交
2723
			zap.Uint64("BeginTs", cat.BeginTs()),
2724
			zap.Uint64("EndTs", cat.EndTs()))
E
Enwei Jiao 已提交
2725
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
Y
Yusup 已提交
2726 2727 2728 2729 2730 2731 2732

		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2733 2734 2735
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", cat.BeginTs()),
2736
		zap.Uint64("EndTs", cat.EndTs()))
D
dragondriver 已提交
2737

E
Enwei Jiao 已提交
2738 2739
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
Y
Yusup 已提交
2740 2741 2742
	return cat.result, nil
}

2743
// DropAlias alter the alias of collection.
Y
Yusup 已提交
2744 2745 2746 2747
func (node *Proxy) DropAlias(ctx context.Context, request *milvuspb.DropAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2748

E
Enwei Jiao 已提交
2749 2750
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DropAlias")
	defer sp.End()
D
dragondriver 已提交
2751

Y
Yusup 已提交
2752 2753 2754 2755 2756 2757 2758
	dat := &DropAliasTask{
		ctx:              ctx,
		Condition:        NewTaskCondition(ctx),
		DropAliasRequest: request,
		rootCoord:        node.rootCoord,
	}

D
dragondriver 已提交
2759
	method := "DropAlias"
2760
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2761
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2762

2763
	log := log.Ctx(ctx).With(
D
dragondriver 已提交
2764 2765 2766 2767
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias))

2768 2769
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2770 2771 2772
	if err := node.sched.ddQueue.Enqueue(dat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2773
			zap.Error(err))
E
Enwei Jiao 已提交
2774
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
D
dragondriver 已提交
2775

Y
Yusup 已提交
2776 2777 2778 2779 2780 2781
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2782 2783 2784
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", dat.BeginTs()),
2785
		zap.Uint64("EndTs", dat.EndTs()))
D
dragondriver 已提交
2786 2787 2788 2789

	if err := dat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2790
			zap.Error(err),
D
dragondriver 已提交
2791
			zap.Uint64("BeginTs", dat.BeginTs()),
2792
			zap.Uint64("EndTs", dat.EndTs()))
Y
Yusup 已提交
2793

E
Enwei Jiao 已提交
2794
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
2795

Y
Yusup 已提交
2796 2797 2798 2799 2800 2801
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2802 2803 2804
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", dat.BeginTs()),
2805
		zap.Uint64("EndTs", dat.EndTs()))
D
dragondriver 已提交
2806

E
Enwei Jiao 已提交
2807 2808
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
Y
Yusup 已提交
2809 2810 2811
	return dat.result, nil
}

2812
// AlterAlias alter alias of collection.
Y
Yusup 已提交
2813 2814 2815 2816
func (node *Proxy) AlterAlias(ctx context.Context, request *milvuspb.AlterAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2817

E
Enwei Jiao 已提交
2818 2819
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-AlterAlias")
	defer sp.End()
D
dragondriver 已提交
2820

Y
Yusup 已提交
2821 2822 2823 2824 2825 2826 2827
	aat := &AlterAliasTask{
		ctx:               ctx,
		Condition:         NewTaskCondition(ctx),
		AlterAliasRequest: request,
		rootCoord:         node.rootCoord,
	}

D
dragondriver 已提交
2828
	method := "AlterAlias"
2829
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2830
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2831

2832
	log := log.Ctx(ctx).With(
D
dragondriver 已提交
2833 2834 2835 2836 2837
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

2838 2839
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2840 2841 2842
	if err := node.sched.ddQueue.Enqueue(aat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2843
			zap.Error(err))
E
Enwei Jiao 已提交
2844
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
D
dragondriver 已提交
2845

Y
Yusup 已提交
2846 2847 2848 2849 2850 2851
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2852 2853 2854
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", aat.BeginTs()),
2855
		zap.Uint64("EndTs", aat.EndTs()))
D
dragondriver 已提交
2856 2857 2858 2859

	if err := aat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2860
			zap.Error(err),
D
dragondriver 已提交
2861
			zap.Uint64("BeginTs", aat.BeginTs()),
2862
			zap.Uint64("EndTs", aat.EndTs()))
Y
Yusup 已提交
2863

E
Enwei Jiao 已提交
2864
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
2865

Y
Yusup 已提交
2866 2867 2868 2869 2870 2871
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2872 2873 2874
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", aat.BeginTs()),
2875
		zap.Uint64("EndTs", aat.EndTs()))
D
dragondriver 已提交
2876

E
Enwei Jiao 已提交
2877 2878
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
Y
Yusup 已提交
2879 2880 2881
	return aat.result, nil
}

2882
// CalcDistance calculates the distances between vectors.
2883
func (node *Proxy) CalcDistance(ctx context.Context, request *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) {
2884 2885 2886 2887 2888
	if !node.checkHealthy() {
		return &milvuspb.CalcDistanceResults{
			Status: unhealthyStatus(),
		}, nil
	}
2889

E
Enwei Jiao 已提交
2890 2891
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CalcDistance")
	defer sp.End()
2892

2893 2894
	query := func(ids *milvuspb.VectorIDs) (*milvuspb.QueryResults, error) {
		outputFields := []string{ids.FieldName}
2895

2896 2897 2898 2899 2900
		queryRequest := &milvuspb.QueryRequest{
			DbName:         "",
			CollectionName: ids.CollectionName,
			PartitionNames: ids.PartitionNames,
			OutputFields:   outputFields,
2901 2902
		}

2903
		qt := &queryTask{
2904 2905 2906
			ctx:       ctx,
			Condition: NewTaskCondition(ctx),
			RetrieveRequest: &internalpb.RetrieveRequest{
2907 2908
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
E
Enwei Jiao 已提交
2909
					commonpbutil.WithSourceID(paramtable.GetNodeID()),
2910
				),
E
Enwei Jiao 已提交
2911
				ReqID: paramtable.GetNodeID(),
2912
			},
2913 2914 2915 2916
			request: queryRequest,
			qc:      node.queryCoord,
			ids:     ids.IdArray,

2917
			queryShardPolicy: mergeRoundRobinPolicy,
2918
			shardMgr:         node.shardMgr,
2919 2920
		}

2921
		log := log.Ctx(ctx).With(
G
groot 已提交
2922 2923
			zap.String("collection", queryRequest.CollectionName),
			zap.Any("partitions", queryRequest.PartitionNames),
2924
			zap.Any("OutputFields", queryRequest.OutputFields))
G
groot 已提交
2925

2926
		err := node.sched.dqQueue.Enqueue(qt)
2927
		if err != nil {
2928 2929
			log.Error("CalcDistance queryTask failed to enqueue",
				zap.Error(err))
2930

2931 2932 2933 2934 2935
			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2936
			}, err
2937
		}
2938

2939
		log.Debug("CalcDistance queryTask enqueued")
2940 2941 2942

		err = qt.WaitToFinish()
		if err != nil {
2943 2944
			log.Error("CalcDistance queryTask failed to WaitToFinish",
				zap.Error(err))
2945 2946 2947 2948 2949 2950

			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2951
			}, err
2952
		}
2953

2954
		log.Debug("CalcDistance queryTask Done")
2955 2956

		return &milvuspb.QueryResults{
2957 2958
			Status:     qt.result.Status,
			FieldsData: qt.result.FieldsData,
2959 2960 2961
		}, nil
	}

G
groot 已提交
2962 2963
	// calcDistanceTask is not a standard task, no need to enqueue
	task := &calcDistanceTask{
E
Enwei Jiao 已提交
2964
		traceID:   sp.SpanContext().TraceID().String(),
G
groot 已提交
2965
		queryFunc: query,
2966 2967
	}

G
groot 已提交
2968
	return task.Execute(ctx, request)
2969 2970
}

2971
// GetDdChannel returns the used channel for dd operations.
C
Cai Yudong 已提交
2972
func (node *Proxy) GetDdChannel(ctx context.Context, request *internalpb.GetDdChannelRequest) (*milvuspb.StringResponse, error) {
2973 2974
	panic("implement me")
}
X
XuanYang-cn 已提交
2975

2976
// GetPersistentSegmentInfo get the information of sealed segment.
C
Cai Yudong 已提交
2977
func (node *Proxy) GetPersistentSegmentInfo(ctx context.Context, req *milvuspb.GetPersistentSegmentInfoRequest) (*milvuspb.GetPersistentSegmentInfoResponse, error) {
E
Enwei Jiao 已提交
2978 2979
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetPersistentSegmentInfo")
	defer sp.End()
2980 2981 2982

	log := log.Ctx(ctx)

D
dragondriver 已提交
2983
	log.Debug("GetPersistentSegmentInfo",
2984
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2985 2986 2987
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
2988
	resp := &milvuspb.GetPersistentSegmentInfoResponse{
X
XuanYang-cn 已提交
2989
		Status: &commonpb.Status{
2990
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
X
XuanYang-cn 已提交
2991 2992
		},
	}
2993 2994 2995 2996
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
2997 2998
	method := "GetPersistentSegmentInfo"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2999
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3000
		metrics.TotalLabel).Inc()
3001 3002 3003

	// list segments
	collectionID, err := globalMetaCache.GetCollectionID(ctx, req.GetCollectionName())
X
XuanYang-cn 已提交
3004
	if err != nil {
E
Enwei Jiao 已提交
3005
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3006 3007 3008 3009 3010 3011 3012 3013 3014 3015 3016
		resp.Status.Reason = fmt.Errorf("getCollectionID failed, err:%w", err).Error()
		return resp, nil
	}

	getSegmentsByStatesResponse, err := node.dataCoord.GetSegmentsByStates(ctx, &datapb.GetSegmentsByStatesRequest{
		CollectionID: collectionID,
		// -1 means list all partition segemnts
		PartitionID: -1,
		States:      []commonpb.SegmentState{commonpb.SegmentState_Flushing, commonpb.SegmentState_Flushed, commonpb.SegmentState_Sealed},
	})
	if err != nil {
E
Enwei Jiao 已提交
3017
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3018
		resp.Status.Reason = fmt.Errorf("getSegmentsOfCollection, err:%w", err).Error()
X
XuanYang-cn 已提交
3019 3020
		return resp, nil
	}
3021 3022

	// get Segment info
3023
	infoResp, err := node.dataCoord.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{
3024 3025 3026
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_SegmentInfo),
			commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
3027
			commonpbutil.WithSourceID(paramtable.GetNodeID()),
3028
		),
3029
		SegmentIDs: getSegmentsByStatesResponse.Segments,
X
XuanYang-cn 已提交
3030 3031
	})
	if err != nil {
E
Enwei Jiao 已提交
3032
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3033
			metrics.FailLabel).Inc()
3034 3035
		log.Warn("GetPersistentSegmentInfo fail",
			zap.Error(err))
3036
		resp.Status.Reason = fmt.Errorf("dataCoord:GetSegmentInfo, err:%w", err).Error()
X
XuanYang-cn 已提交
3037 3038
		return resp, nil
	}
3039 3040 3041
	log.Debug("GetPersistentSegmentInfo",
		zap.Int("len(infos)", len(infoResp.Infos)),
		zap.Any("status", infoResp.Status))
3042
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
E
Enwei Jiao 已提交
3043
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3044
			metrics.FailLabel).Inc()
X
XuanYang-cn 已提交
3045 3046 3047 3048 3049 3050
		resp.Status.Reason = infoResp.Status.Reason
		return resp, nil
	}
	persistentInfos := make([]*milvuspb.PersistentSegmentInfo, len(infoResp.Infos))
	for i, info := range infoResp.Infos {
		persistentInfos[i] = &milvuspb.PersistentSegmentInfo{
S
sunby 已提交
3051
			SegmentID:    info.ID,
X
XuanYang-cn 已提交
3052 3053
			CollectionID: info.CollectionID,
			PartitionID:  info.PartitionID,
S
sunby 已提交
3054
			NumRows:      info.NumOfRows,
X
XuanYang-cn 已提交
3055 3056 3057
			State:        info.State,
		}
	}
E
Enwei Jiao 已提交
3058
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3059
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
3060
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
3061
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
X
XuanYang-cn 已提交
3062 3063 3064 3065
	resp.Infos = persistentInfos
	return resp, nil
}

J
jingkl 已提交
3066
// GetQuerySegmentInfo gets segment information from QueryCoord.
C
Cai Yudong 已提交
3067
func (node *Proxy) GetQuerySegmentInfo(ctx context.Context, req *milvuspb.GetQuerySegmentInfoRequest) (*milvuspb.GetQuerySegmentInfoResponse, error) {
E
Enwei Jiao 已提交
3068 3069
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetQuerySegmentInfo")
	defer sp.End()
3070 3071 3072

	log := log.Ctx(ctx)

D
dragondriver 已提交
3073
	log.Debug("GetQuerySegmentInfo",
3074
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3075 3076 3077
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
3078
	resp := &milvuspb.GetQuerySegmentInfoResponse{
Z
zhenshan.cao 已提交
3079
		Status: &commonpb.Status{
3080
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
Z
zhenshan.cao 已提交
3081 3082
		},
	}
3083 3084 3085 3086
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3087

3088 3089
	method := "GetQuerySegmentInfo"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
3090
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3091 3092
		metrics.TotalLabel).Inc()

3093 3094
	collID, err := globalMetaCache.GetCollectionID(ctx, req.CollectionName)
	if err != nil {
E
Enwei Jiao 已提交
3095
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3096 3097 3098
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3099
	infoResp, err := node.queryCoord.GetSegmentInfo(ctx, &querypb.GetSegmentInfoRequest{
3100 3101 3102
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_SegmentInfo),
			commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
3103
			commonpbutil.WithSourceID(paramtable.GetNodeID()),
3104
		),
3105
		CollectionID: collID,
Z
zhenshan.cao 已提交
3106 3107
	})
	if err != nil {
E
Enwei Jiao 已提交
3108
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3109 3110
		log.Error("Failed to get segment info from QueryCoord",
			zap.Error(err))
Z
zhenshan.cao 已提交
3111 3112 3113
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3114 3115 3116
	log.Debug("GetQuerySegmentInfo",
		zap.Any("infos", infoResp.Infos),
		zap.Any("status", infoResp.Status))
3117
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
E
Enwei Jiao 已提交
3118
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3119 3120
		log.Error("Failed to get segment info from QueryCoord",
			zap.String("errMsg", infoResp.Status.Reason))
Z
zhenshan.cao 已提交
3121 3122 3123 3124 3125 3126 3127 3128 3129 3130 3131 3132 3133
		resp.Status.Reason = infoResp.Status.Reason
		return resp, nil
	}
	queryInfos := make([]*milvuspb.QuerySegmentInfo, len(infoResp.Infos))
	for i, info := range infoResp.Infos {
		queryInfos[i] = &milvuspb.QuerySegmentInfo{
			SegmentID:    info.SegmentID,
			CollectionID: info.CollectionID,
			PartitionID:  info.PartitionID,
			NumRows:      info.NumRows,
			MemSize:      info.MemSize,
			IndexName:    info.IndexName,
			IndexID:      info.IndexID,
X
xige-16 已提交
3134
			State:        info.SegmentState,
3135
			NodeIds:      info.NodeIds,
Z
zhenshan.cao 已提交
3136 3137
		}
	}
3138

E
Enwei Jiao 已提交
3139 3140
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
3141
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
Z
zhenshan.cao 已提交
3142 3143 3144 3145
	resp.Infos = queryInfos
	return resp, nil
}

J
jingkl 已提交
3146
// Dummy handles dummy request
C
Cai Yudong 已提交
3147
func (node *Proxy) Dummy(ctx context.Context, req *milvuspb.DummyRequest) (*milvuspb.DummyResponse, error) {
3148 3149 3150 3151 3152 3153
	failedResponse := &milvuspb.DummyResponse{
		Response: `{"status": "fail"}`,
	}

	// TODO(wxyu): change name RequestType to Request
	drt, err := parseDummyRequestType(req.RequestType)
3154

E
Enwei Jiao 已提交
3155 3156
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Dummy")
	defer sp.End()
3157 3158 3159

	log := log.Ctx(ctx)

3160
	if err != nil {
3161 3162
		log.Warn("Failed to parse dummy request type",
			zap.Error(err))
3163 3164 3165
		return failedResponse, nil
	}

3166 3167
	if drt.RequestType == "query" {
		drr, err := parseDummyQueryRequest(req.RequestType)
3168
		if err != nil {
3169 3170
			log.Warn("Failed to parse dummy query request",
				zap.Error(err))
3171 3172 3173
			return failedResponse, nil
		}

3174
		request := &milvuspb.QueryRequest{
3175 3176 3177
			DbName:         drr.DbName,
			CollectionName: drr.CollectionName,
			PartitionNames: drr.PartitionNames,
3178
			OutputFields:   drr.OutputFields,
X
Xiangyu Wang 已提交
3179 3180
		}

3181
		_, err = node.Query(ctx, request)
3182
		if err != nil {
3183 3184
			log.Warn("Failed to execute dummy query",
				zap.Error(err))
3185 3186
			return failedResponse, err
		}
X
Xiangyu Wang 已提交
3187 3188 3189 3190 3191 3192

		return &milvuspb.DummyResponse{
			Response: `{"status": "success"}`,
		}, nil
	}

3193 3194
	log.Debug("cannot find specify dummy request type")
	return failedResponse, nil
X
Xiangyu Wang 已提交
3195 3196
}

J
jingkl 已提交
3197
// RegisterLink registers a link
C
Cai Yudong 已提交
3198
func (node *Proxy) RegisterLink(ctx context.Context, req *milvuspb.RegisterLinkRequest) (*milvuspb.RegisterLinkResponse, error) {
3199
	code := node.stateCode.Load().(commonpb.StateCode)
3200

E
Enwei Jiao 已提交
3201 3202
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-RegisterLink")
	defer sp.End()
3203 3204

	log := log.Ctx(ctx).With(
3205
		zap.String("role", typeutil.ProxyRole),
C
Cai Yudong 已提交
3206
		zap.Any("state code of proxy", code))
D
dragondriver 已提交
3207

3208 3209
	log.Debug("RegisterLink")

3210
	if code != commonpb.StateCode_Healthy {
3211 3212 3213
		return &milvuspb.RegisterLinkResponse{
			Address: nil,
			Status: &commonpb.Status{
3214
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
C
Cai Yudong 已提交
3215
				Reason:    "proxy not healthy",
3216 3217 3218
			},
		}, nil
	}
E
Enwei Jiao 已提交
3219
	//metrics.ProxyLinkedSDKs.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Inc()
3220 3221 3222
	return &milvuspb.RegisterLinkResponse{
		Address: nil,
		Status: &commonpb.Status{
3223
			ErrorCode: commonpb.ErrorCode_Success,
3224
			Reason:    os.Getenv(metricsinfo.DeployModeEnvKey),
3225 3226 3227
		},
	}, nil
}
3228

3229
// GetMetrics gets the metrics of proxy
3230 3231
// TODO(dragondriver): cache the Metrics and set a retention to the cache
func (node *Proxy) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
E
Enwei Jiao 已提交
3232 3233
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetMetrics")
	defer sp.End()
3234 3235 3236

	log := log.Ctx(ctx)

3237 3238
	log.RatedDebug(60, "Proxy.GetMetrics",
		zap.Int64("nodeID", paramtable.GetNodeID()),
3239 3240 3241
		zap.String("req", req.Request))

	if !node.checkHealthy() {
S
smellthemoon 已提交
3242
		err := merr.WrapErrServiceNotReady(fmt.Sprintf("proxy %d is unhealthy", paramtable.GetNodeID()))
3243
		log.Warn("Proxy.GetMetrics failed",
3244
			zap.Int64("nodeID", paramtable.GetNodeID()),
3245
			zap.String("req", req.Request),
S
smellthemoon 已提交
3246
			zap.Error(err))
3247 3248

		return &milvuspb.GetMetricsResponse{
S
smellthemoon 已提交
3249
			Status:   merr.Status(err),
3250 3251 3252 3253 3254 3255 3256
			Response: "",
		}, nil
	}

	metricType, err := metricsinfo.ParseMetricType(req.Request)
	if err != nil {
		log.Warn("Proxy.GetMetrics failed to parse metric type",
3257
			zap.Int64("nodeID", paramtable.GetNodeID()),
3258 3259 3260 3261 3262 3263 3264 3265 3266 3267 3268 3269
			zap.String("req", req.Request),
			zap.Error(err))

		return &milvuspb.GetMetricsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
			Response: "",
		}, nil
	}

3270 3271 3272
	req.Base = commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_SystemInfo),
		commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
3273
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
3274
	)
3275
	if metricType == metricsinfo.SystemInfoMetrics {
3276 3277 3278
		metrics, err := node.metricsCacheManager.GetSystemInfoMetrics()
		if err != nil {
			metrics, err = getSystemInfoMetrics(ctx, req, node)
3279
		}
3280

3281 3282
		log.RatedDebug(60, "Proxy.GetMetrics",
			zap.Int64("nodeID", paramtable.GetNodeID()),
3283
			zap.String("req", req.Request),
3284
			zap.String("metricType", metricType),
3285 3286 3287
			zap.Any("metrics", metrics), // TODO(dragondriver): necessary? may be very large
			zap.Error(err))

3288 3289
		node.metricsCacheManager.UpdateSystemInfoMetrics(metrics)

G
godchen 已提交
3290
		return metrics, nil
3291 3292
	}

3293 3294
	log.RatedWarn(60, "Proxy.GetMetrics failed, request metric type is not implemented yet",
		zap.Int64("nodeID", paramtable.GetNodeID()),
3295
		zap.String("req", req.Request),
3296
		zap.String("metricType", metricType))
3297 3298 3299 3300 3301 3302 3303 3304 3305 3306

	return &milvuspb.GetMetricsResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    metricsinfo.MsgUnimplementedMetric,
		},
		Response: "",
	}, nil
}

3307 3308 3309
// GetProxyMetrics gets the metrics of proxy, it's an internal interface which is different from GetMetrics interface,
// because it only obtains the metrics of Proxy, not including the topological metrics of Query cluster and Data cluster.
func (node *Proxy) GetProxyMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
E
Enwei Jiao 已提交
3310 3311
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetProxyMetrics")
	defer sp.End()
3312 3313

	log := log.Ctx(ctx).With(
3314
		zap.Int64("nodeID", paramtable.GetNodeID()),
3315 3316
		zap.String("req", req.Request))

3317
	if !node.checkHealthy() {
S
smellthemoon 已提交
3318
		err := merr.WrapErrServiceNotReady(fmt.Sprintf("proxy %d is unhealthy", paramtable.GetNodeID()))
3319
		log.Warn("Proxy.GetProxyMetrics failed",
S
smellthemoon 已提交
3320
			zap.Error(err))
3321 3322

		return &milvuspb.GetMetricsResponse{
S
smellthemoon 已提交
3323
			Status: merr.Status(err),
3324 3325 3326 3327 3328 3329 3330 3331 3332 3333 3334 3335 3336 3337 3338 3339
		}, nil
	}

	metricType, err := metricsinfo.ParseMetricType(req.Request)
	if err != nil {
		log.Warn("Proxy.GetProxyMetrics failed to parse metric type",
			zap.Error(err))

		return &milvuspb.GetMetricsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

3340 3341 3342
	req.Base = commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_SystemInfo),
		commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
3343
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
3344
	)
3345 3346 3347 3348 3349 3350 3351 3352 3353 3354 3355 3356 3357 3358 3359

	if metricType == metricsinfo.SystemInfoMetrics {
		proxyMetrics, err := getProxyMetrics(ctx, req, node)
		if err != nil {
			log.Warn("Proxy.GetProxyMetrics failed to getProxyMetrics",
				zap.Error(err))

			return &milvuspb.GetMetricsResponse{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
		}

3360 3361
		//log.Debug("Proxy.GetProxyMetrics",
		//	zap.String("metricType", metricType))
3362 3363 3364 3365

		return proxyMetrics, nil
	}

J
Jiquan Long 已提交
3366
	log.Warn("Proxy.GetProxyMetrics failed, request metric type is not implemented yet",
3367
		zap.String("metricType", metricType))
3368 3369 3370 3371 3372 3373 3374 3375 3376

	return &milvuspb.GetMetricsResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    metricsinfo.MsgUnimplementedMetric,
		},
	}, nil
}

B
bigsheeper 已提交
3377 3378
// LoadBalance would do a load balancing operation between query nodes
func (node *Proxy) LoadBalance(ctx context.Context, req *milvuspb.LoadBalanceRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
3379 3380
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-LoadBalance")
	defer sp.End()
3381 3382 3383

	log := log.Ctx(ctx)

B
bigsheeper 已提交
3384
	log.Debug("Proxy.LoadBalance",
E
Enwei Jiao 已提交
3385
		zap.Int64("proxy_id", paramtable.GetNodeID()),
B
bigsheeper 已提交
3386 3387 3388 3389 3390 3391 3392 3393 3394
		zap.Any("req", req))

	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

	status := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
	}
3395 3396 3397

	collectionID, err := globalMetaCache.GetCollectionID(ctx, req.GetCollectionName())
	if err != nil {
J
Jiquan Long 已提交
3398
		log.Warn("failed to get collection id",
3399 3400
			zap.String("collection name", req.GetCollectionName()),
			zap.Error(err))
3401 3402 3403
		status.Reason = err.Error()
		return status, nil
	}
B
bigsheeper 已提交
3404
	infoResp, err := node.queryCoord.LoadBalance(ctx, &querypb.LoadBalanceRequest{
3405 3406 3407
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_LoadBalanceSegments),
			commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
3408
			commonpbutil.WithSourceID(paramtable.GetNodeID()),
3409
		),
B
bigsheeper 已提交
3410 3411
		SourceNodeIDs:    []int64{req.SrcNodeID},
		DstNodeIDs:       req.DstNodeIDs,
X
xige-16 已提交
3412
		BalanceReason:    querypb.TriggerCondition_GrpcRequest,
B
bigsheeper 已提交
3413
		SealedSegmentIDs: req.SealedSegmentIDs,
3414
		CollectionID:     collectionID,
B
bigsheeper 已提交
3415 3416
	})
	if err != nil {
J
Jiquan Long 已提交
3417
		log.Warn("Failed to LoadBalance from Query Coordinator",
3418 3419
			zap.Any("req", req),
			zap.Error(err))
B
bigsheeper 已提交
3420 3421 3422 3423
		status.Reason = err.Error()
		return status, nil
	}
	if infoResp.ErrorCode != commonpb.ErrorCode_Success {
J
Jiquan Long 已提交
3424
		log.Warn("Failed to LoadBalance from Query Coordinator",
3425
			zap.String("errMsg", infoResp.Reason))
B
bigsheeper 已提交
3426 3427 3428
		status.Reason = infoResp.Reason
		return status, nil
	}
3429 3430 3431
	log.Debug("LoadBalance Done",
		zap.Any("req", req),
		zap.Any("status", infoResp))
B
bigsheeper 已提交
3432 3433 3434 3435
	status.ErrorCode = commonpb.ErrorCode_Success
	return status, nil
}

3436 3437
// GetReplicas gets replica info
func (node *Proxy) GetReplicas(ctx context.Context, req *milvuspb.GetReplicasRequest) (*milvuspb.GetReplicasResponse, error) {
E
Enwei Jiao 已提交
3438 3439
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetReplicas")
	defer sp.End()
3440 3441 3442 3443 3444 3445

	log := log.Ctx(ctx)

	log.Debug("received get replicas request",
		zap.Int64("collection", req.GetCollectionID()),
		zap.Bool("with shard nodes", req.GetWithShardNodes()))
3446 3447 3448 3449 3450 3451
	resp := &milvuspb.GetReplicasResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

S
smellthemoon 已提交
3452 3453
	req.Base = commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_GetReplicas),
E
Enwei Jiao 已提交
3454
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
S
smellthemoon 已提交
3455
	)
3456

W
wei liu 已提交
3457 3458 3459 3460
	if req.GetCollectionName() != "" {
		req.CollectionID, _ = globalMetaCache.GetCollectionID(ctx, req.GetCollectionName())
	}

3461 3462
	resp, err := node.queryCoord.GetReplicas(ctx, req)
	if err != nil {
3463 3464
		log.Error("Failed to get replicas from Query Coordinator",
			zap.Error(err))
3465 3466 3467 3468
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3469 3470 3471
	log.Debug("received get replicas response",
		zap.Any("resp", resp),
		zap.Error(err))
3472 3473 3474
	return resp, nil
}

3475
// GetCompactionState gets the compaction state of multiple segments
3476
func (node *Proxy) GetCompactionState(ctx context.Context, req *milvuspb.GetCompactionStateRequest) (*milvuspb.GetCompactionStateResponse, error) {
E
Enwei Jiao 已提交
3477 3478
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetCompactionState")
	defer sp.End()
3479 3480 3481 3482 3483

	log := log.Ctx(ctx).With(
		zap.Int64("compactionID", req.GetCompactionID()))

	log.Debug("received GetCompactionState request")
3484 3485 3486 3487 3488 3489 3490
	resp := &milvuspb.GetCompactionStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionState(ctx, req)
3491 3492 3493
	log.Debug("received GetCompactionState response",
		zap.Any("resp", resp),
		zap.Error(err))
3494 3495 3496
	return resp, err
}

3497
// ManualCompaction invokes compaction on specified collection
3498
func (node *Proxy) ManualCompaction(ctx context.Context, req *milvuspb.ManualCompactionRequest) (*milvuspb.ManualCompactionResponse, error) {
E
Enwei Jiao 已提交
3499 3500
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ManualCompaction")
	defer sp.End()
3501 3502 3503 3504 3505

	log := log.Ctx(ctx).With(
		zap.Int64("collectionID", req.GetCollectionID()))

	log.Info("received ManualCompaction request")
3506 3507 3508 3509 3510 3511 3512
	resp := &milvuspb.ManualCompactionResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.ManualCompaction(ctx, req)
3513 3514 3515
	log.Info("received ManualCompaction response",
		zap.Any("resp", resp),
		zap.Error(err))
3516 3517 3518
	return resp, err
}

3519
// GetCompactionStateWithPlans returns the compactions states with the given plan ID
3520
func (node *Proxy) GetCompactionStateWithPlans(ctx context.Context, req *milvuspb.GetCompactionPlansRequest) (*milvuspb.GetCompactionPlansResponse, error) {
E
Enwei Jiao 已提交
3521 3522
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetCompactionStateWithPlans")
	defer sp.End()
3523 3524 3525 3526 3527

	log := log.Ctx(ctx).With(
		zap.Int64("compactionID", req.GetCompactionID()))

	log.Debug("received GetCompactionStateWithPlans request")
3528 3529 3530 3531 3532 3533 3534
	resp := &milvuspb.GetCompactionPlansResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionStateWithPlans(ctx, req)
3535 3536 3537
	log.Debug("received GetCompactionStateWithPlans response",
		zap.Any("resp", resp),
		zap.Error(err))
3538 3539 3540
	return resp, err
}

B
Bingyi Sun 已提交
3541 3542
// GetFlushState gets the flush state of multiple segments
func (node *Proxy) GetFlushState(ctx context.Context, req *milvuspb.GetFlushStateRequest) (*milvuspb.GetFlushStateResponse, error) {
E
Enwei Jiao 已提交
3543 3544
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetFlushState")
	defer sp.End()
3545 3546 3547 3548 3549

	log := log.Ctx(ctx)

	log.Debug("received get flush state request",
		zap.Any("request", req))
3550
	var err error
B
Bingyi Sun 已提交
3551 3552 3553
	resp := &milvuspb.GetFlushStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
J
Jiquan Long 已提交
3554
		log.Warn("unable to get flush state because of closed server")
B
Bingyi Sun 已提交
3555 3556 3557
		return resp, nil
	}

3558
	resp, err = node.dataCoord.GetFlushState(ctx, req)
X
Xiaofan 已提交
3559
	if err != nil {
3560 3561
		log.Warn("failed to get flush state response",
			zap.Error(err))
X
Xiaofan 已提交
3562 3563
		return nil, err
	}
3564 3565
	log.Debug("received get flush state response",
		zap.Any("response", resp))
B
Bingyi Sun 已提交
3566 3567 3568
	return resp, err
}

C
Cai Yudong 已提交
3569 3570
// checkHealthy checks proxy state is Healthy
func (node *Proxy) checkHealthy() bool {
3571 3572
	code := node.stateCode.Load().(commonpb.StateCode)
	return code == commonpb.StateCode_Healthy
3573 3574
}

3575 3576 3577
func (node *Proxy) checkHealthyAndReturnCode() (commonpb.StateCode, bool) {
	code := node.stateCode.Load().(commonpb.StateCode)
	return code, code == commonpb.StateCode_Healthy
3578 3579
}

3580
// unhealthyStatus returns the proxy not healthy status
3581 3582 3583
func unhealthyStatus() *commonpb.Status {
	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
C
Cai Yudong 已提交
3584
		Reason:    "proxy not healthy",
3585 3586
	}
}
G
groot 已提交
3587 3588 3589

// Import data files(json, numpy, etc.) on MinIO/S3 storage, read and parse them into sealed segments
func (node *Proxy) Import(ctx context.Context, req *milvuspb.ImportRequest) (*milvuspb.ImportResponse, error) {
E
Enwei Jiao 已提交
3590 3591
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Import")
	defer sp.End()
3592 3593 3594

	log := log.Ctx(ctx)

3595 3596
	log.Info("received import request",
		zap.String("collection name", req.GetCollectionName()),
G
groot 已提交
3597 3598
		zap.String("partition name", req.GetPartitionName()),
		zap.Strings("files", req.GetFiles()))
3599 3600 3601 3602 3603 3604
	resp := &milvuspb.ImportResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
			Reason:    "",
		},
	}
G
groot 已提交
3605 3606 3607 3608
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3609

3610 3611
	err := importutil.ValidateOptions(req.GetOptions())
	if err != nil {
3612 3613
		log.Error("failed to execute import request",
			zap.Error(err))
3614 3615 3616 3617 3618
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = "request options is not illegal    \n" + err.Error() + "    \nIllegal option format    \n" + importutil.OptionFormat
		return resp, nil
	}

3619 3620
	method := "Import"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
3621
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3622 3623
		metrics.TotalLabel).Inc()

3624
	// Call rootCoord to finish import.
3625 3626
	respFromRC, err := node.rootCoord.Import(ctx, req)
	if err != nil {
E
Enwei Jiao 已提交
3627
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3628 3629
		log.Error("failed to execute bulk insert request",
			zap.Error(err))
3630 3631 3632 3633
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3634

E
Enwei Jiao 已提交
3635 3636
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
3637
	return respFromRC, nil
G
groot 已提交
3638 3639
}

3640
// GetImportState checks import task state from RootCoord.
G
groot 已提交
3641
func (node *Proxy) GetImportState(ctx context.Context, req *milvuspb.GetImportStateRequest) (*milvuspb.GetImportStateResponse, error) {
E
Enwei Jiao 已提交
3642 3643
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetImportState")
	defer sp.End()
3644 3645 3646 3647 3648

	log := log.Ctx(ctx)

	log.Debug("received get import state request",
		zap.Int64("taskID", req.GetTask()))
G
groot 已提交
3649 3650 3651 3652 3653
	resp := &milvuspb.GetImportStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3654 3655
	method := "GetImportState"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
3656
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3657
		metrics.TotalLabel).Inc()
G
groot 已提交
3658 3659

	resp, err := node.rootCoord.GetImportState(ctx, req)
3660
	if err != nil {
E
Enwei Jiao 已提交
3661
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3662 3663
		log.Error("failed to execute get import state",
			zap.Error(err))
3664 3665 3666 3667 3668
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
		return resp, nil
	}

3669 3670 3671
	log.Debug("successfully received get import state response",
		zap.Int64("taskID", req.GetTask()),
		zap.Any("resp", resp), zap.Error(err))
E
Enwei Jiao 已提交
3672 3673
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
3674
	return resp, nil
G
groot 已提交
3675 3676 3677 3678
}

// ListImportTasks get id array of all import tasks from rootcoord
func (node *Proxy) ListImportTasks(ctx context.Context, req *milvuspb.ListImportTasksRequest) (*milvuspb.ListImportTasksResponse, error) {
E
Enwei Jiao 已提交
3679 3680
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ListImportTasks")
	defer sp.End()
3681 3682 3683

	log := log.Ctx(ctx)

J
Jiquan Long 已提交
3684
	log.Debug("received list import tasks request")
G
groot 已提交
3685 3686 3687 3688 3689
	resp := &milvuspb.ListImportTasksResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3690 3691
	method := "ListImportTasks"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
3692
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3693
		metrics.TotalLabel).Inc()
G
groot 已提交
3694
	resp, err := node.rootCoord.ListImportTasks(ctx, req)
3695
	if err != nil {
E
Enwei Jiao 已提交
3696
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3697 3698
		log.Error("failed to execute list import tasks",
			zap.Error(err))
3699 3700
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
X
XuanYang-cn 已提交
3701 3702 3703
		return resp, nil
	}

3704 3705 3706
	log.Debug("successfully received list import tasks response",
		zap.String("collection", req.CollectionName),
		zap.Any("tasks", resp.Tasks))
E
Enwei Jiao 已提交
3707 3708
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
X
XuanYang-cn 已提交
3709 3710 3711
	return resp, err
}

3712 3713
// InvalidateCredentialCache invalidate the credential cache of specified username.
func (node *Proxy) InvalidateCredentialCache(ctx context.Context, request *proxypb.InvalidateCredCacheRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
3714 3715
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-InvalidateCredentialCache")
	defer sp.End()
3716 3717

	log := log.Ctx(ctx).With(
3718 3719
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))
3720 3721

	log.Debug("received request to invalidate credential cache")
3722
	if !node.checkHealthy() {
3723
		return unhealthyStatus(), nil
3724
	}
3725 3726 3727 3728 3729

	username := request.Username
	if globalMetaCache != nil {
		globalMetaCache.RemoveCredential(username) // no need to return error, though credential may be not cached
	}
3730
	log.Debug("complete to invalidate credential cache")
3731 3732 3733 3734 3735 3736 3737 3738 3739

	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
		Reason:    "",
	}, nil
}

// UpdateCredentialCache update the credential cache of specified username.
func (node *Proxy) UpdateCredentialCache(ctx context.Context, request *proxypb.UpdateCredCacheRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
3740 3741
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-UpdateCredentialCache")
	defer sp.End()
3742 3743

	log := log.Ctx(ctx).With(
3744 3745
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))
3746 3747

	log.Debug("received request to update credential cache")
3748
	if !node.checkHealthy() {
3749
		return unhealthyStatus(), nil
3750
	}
3751 3752

	credInfo := &internalpb.CredentialInfo{
3753 3754
		Username:       request.Username,
		Sha256Password: request.Password,
3755 3756 3757 3758
	}
	if globalMetaCache != nil {
		globalMetaCache.UpdateCredential(credInfo) // no need to return error, though credential may be not cached
	}
3759
	log.Debug("complete to update credential cache")
3760 3761 3762 3763 3764 3765 3766 3767

	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
		Reason:    "",
	}, nil
}

func (node *Proxy) CreateCredential(ctx context.Context, req *milvuspb.CreateCredentialRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
3768 3769
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CreateCredential")
	defer sp.End()
3770 3771 3772 3773 3774 3775

	log := log.Ctx(ctx).With(
		zap.String("username", req.Username))

	log.Debug("CreateCredential",
		zap.String("role", typeutil.ProxyRole))
3776
	if !node.checkHealthy() {
3777
		return unhealthyStatus(), nil
3778
	}
3779 3780 3781 3782 3783 3784 3785 3786 3787 3788
	// validate params
	username := req.Username
	if err := ValidateUsername(username); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
		}, nil
	}
	rawPassword, err := crypto.Base64Decode(req.Password)
	if err != nil {
3789 3790
		log.Error("decode password fail",
			zap.Error(err))
3791 3792 3793 3794 3795 3796
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_CreateCredentialFailure,
			Reason:    "decode password fail key:" + req.Username,
		}, nil
	}
	if err = ValidatePassword(rawPassword); err != nil {
3797 3798
		log.Error("illegal password",
			zap.Error(err))
3799 3800 3801 3802 3803 3804 3805
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
		}, nil
	}
	encryptedPassword, err := crypto.PasswordEncrypt(rawPassword)
	if err != nil {
3806 3807
		log.Error("encrypt password fail",
			zap.Error(err))
3808 3809 3810 3811 3812
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_CreateCredentialFailure,
			Reason:    "encrypt password fail key:" + req.Username,
		}, nil
	}
3813

3814 3815 3816
	credInfo := &internalpb.CredentialInfo{
		Username:          req.Username,
		EncryptedPassword: encryptedPassword,
3817
		Sha256Password:    crypto.SHA256(rawPassword, req.Username),
3818 3819 3820
	}
	result, err := node.rootCoord.CreateCredential(ctx, credInfo)
	if err != nil { // for error like conntext timeout etc.
3821 3822
		log.Error("create credential fail",
			zap.Error(err))
3823 3824 3825 3826 3827 3828 3829 3830
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}
	return result, err
}

C
codeman 已提交
3831
func (node *Proxy) UpdateCredential(ctx context.Context, req *milvuspb.UpdateCredentialRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
3832 3833
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-UpdateCredential")
	defer sp.End()
3834 3835 3836 3837 3838 3839

	log := log.Ctx(ctx).With(
		zap.String("username", req.Username))

	log.Debug("UpdateCredential",
		zap.String("role", typeutil.ProxyRole))
3840
	if !node.checkHealthy() {
3841
		return unhealthyStatus(), nil
3842
	}
C
codeman 已提交
3843 3844
	rawOldPassword, err := crypto.Base64Decode(req.OldPassword)
	if err != nil {
3845 3846
		log.Error("decode old password fail",
			zap.Error(err))
C
codeman 已提交
3847 3848 3849 3850 3851 3852
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "decode old password fail when updating:" + req.Username,
		}, nil
	}
	rawNewPassword, err := crypto.Base64Decode(req.NewPassword)
3853
	if err != nil {
3854 3855
		log.Error("decode password fail",
			zap.Error(err))
3856 3857 3858 3859 3860
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "decode password fail when updating:" + req.Username,
		}, nil
	}
C
codeman 已提交
3861 3862
	// valid new password
	if err = ValidatePassword(rawNewPassword); err != nil {
3863 3864
		log.Error("illegal password",
			zap.Error(err))
3865 3866 3867 3868 3869
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
		}, nil
	}
3870 3871

	if !passwordVerify(ctx, req.Username, rawOldPassword, globalMetaCache) {
C
codeman 已提交
3872 3873 3874 3875 3876 3877 3878
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "old password is not correct:" + req.Username,
		}, nil
	}
	// update meta data
	encryptedPassword, err := crypto.PasswordEncrypt(rawNewPassword)
3879
	if err != nil {
3880 3881
		log.Error("encrypt password fail",
			zap.Error(err))
3882 3883 3884 3885 3886
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "encrypt password fail when updating:" + req.Username,
		}, nil
	}
C
codeman 已提交
3887
	updateCredReq := &internalpb.CredentialInfo{
3888
		Username:          req.Username,
3889
		Sha256Password:    crypto.SHA256(rawNewPassword, req.Username),
3890 3891
		EncryptedPassword: encryptedPassword,
	}
C
codeman 已提交
3892
	result, err := node.rootCoord.UpdateCredential(ctx, updateCredReq)
3893
	if err != nil { // for error like conntext timeout etc.
3894 3895
		log.Error("update credential fail",
			zap.Error(err))
3896 3897 3898 3899 3900 3901 3902 3903 3904
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}
	return result, err
}

func (node *Proxy) DeleteCredential(ctx context.Context, req *milvuspb.DeleteCredentialRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
3905 3906
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DeleteCredential")
	defer sp.End()
3907 3908 3909 3910 3911 3912

	log := log.Ctx(ctx).With(
		zap.String("username", req.Username))

	log.Debug("DeleteCredential",
		zap.String("role", typeutil.ProxyRole))
3913
	if !node.checkHealthy() {
3914
		return unhealthyStatus(), nil
3915 3916
	}

3917 3918 3919 3920 3921 3922
	if req.Username == util.UserRoot {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_DeleteCredentialFailure,
			Reason:    "user root cannot be deleted",
		}, nil
	}
3923 3924
	result, err := node.rootCoord.DeleteCredential(ctx, req)
	if err != nil { // for error like conntext timeout etc.
3925 3926
		log.Error("delete credential fail",
			zap.Error(err))
3927 3928 3929 3930 3931 3932 3933 3934 3935
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}
	return result, err
}

func (node *Proxy) ListCredUsers(ctx context.Context, req *milvuspb.ListCredUsersRequest) (*milvuspb.ListCredUsersResponse, error) {
E
Enwei Jiao 已提交
3936 3937
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ListCredUsers")
	defer sp.End()
3938 3939 3940 3941 3942

	log := log.Ctx(ctx).With(
		zap.String("role", typeutil.ProxyRole))

	log.Debug("ListCredUsers")
3943
	if !node.checkHealthy() {
3944
		return &milvuspb.ListCredUsersResponse{Status: unhealthyStatus()}, nil
3945
	}
3946
	rootCoordReq := &milvuspb.ListCredUsersRequest{
3947 3948 3949
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_ListCredUsernames),
		),
3950 3951
	}
	resp, err := node.rootCoord.ListCredUsers(ctx, rootCoordReq)
3952 3953 3954 3955 3956 3957 3958 3959 3960 3961 3962 3963
	if err != nil {
		return &milvuspb.ListCredUsersResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}
	return &milvuspb.ListCredUsersResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
		},
3964
		Usernames: resp.Usernames,
3965 3966
	}, nil
}
3967

3968
func (node *Proxy) CreateRole(ctx context.Context, req *milvuspb.CreateRoleRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
3969 3970
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CreateRole")
	defer sp.End()
3971 3972 3973 3974 3975

	log := log.Ctx(ctx)

	log.Debug("CreateRole",
		zap.Any("req", req))
3976
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
3977
		return errorutil.UnhealthyStatus(code), nil
3978 3979 3980 3981 3982 3983 3984 3985 3986 3987
	}

	var roleName string
	if req.Entity != nil {
		roleName = req.Entity.Name
	}
	if err := ValidateRoleName(roleName); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
3988
		}, nil
3989 3990 3991 3992
	}

	result, err := node.rootCoord.CreateRole(ctx, req)
	if err != nil {
3993 3994
		log.Error("fail to create role",
			zap.Error(err))
3995 3996 3997
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
3998
		}, nil
3999 4000
	}
	return result, nil
4001 4002
}

4003
func (node *Proxy) DropRole(ctx context.Context, req *milvuspb.DropRoleRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
4004 4005
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DropRole")
	defer sp.End()
4006 4007 4008 4009 4010

	log := log.Ctx(ctx)

	log.Debug("DropRole",
		zap.Any("req", req))
4011
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4012
		return errorutil.UnhealthyStatus(code), nil
4013 4014 4015 4016 4017
	}
	if err := ValidateRoleName(req.RoleName); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4018
		}, nil
4019
	}
4020 4021 4022 4023 4024
	if IsDefaultRole(req.RoleName) {
		errMsg := fmt.Sprintf("the role[%s] is a default role, which can't be droped", req.RoleName)
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    errMsg,
4025
		}, nil
4026
	}
4027 4028
	result, err := node.rootCoord.DropRole(ctx, req)
	if err != nil {
4029 4030 4031
		log.Error("fail to drop role",
			zap.String("role_name", req.RoleName),
			zap.Error(err))
4032 4033 4034
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4035
		}, nil
4036 4037
	}
	return result, nil
4038 4039
}

4040
func (node *Proxy) OperateUserRole(ctx context.Context, req *milvuspb.OperateUserRoleRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
4041 4042
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-OperateUserRole")
	defer sp.End()
4043 4044 4045 4046 4047

	log := log.Ctx(ctx)

	log.Debug("OperateUserRole",
		zap.Any("req", req))
4048
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4049
		return errorutil.UnhealthyStatus(code), nil
4050 4051 4052 4053 4054
	}
	if err := ValidateUsername(req.Username); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4055
		}, nil
4056 4057 4058 4059 4060
	}
	if err := ValidateRoleName(req.RoleName); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4061
		}, nil
4062 4063 4064 4065
	}

	result, err := node.rootCoord.OperateUserRole(ctx, req)
	if err != nil {
4066 4067
		logger.Error("fail to operate user role",
			zap.Error(err))
4068 4069 4070
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4071
		}, nil
4072 4073
	}
	return result, nil
4074 4075
}

4076
func (node *Proxy) SelectRole(ctx context.Context, req *milvuspb.SelectRoleRequest) (*milvuspb.SelectRoleResponse, error) {
E
Enwei Jiao 已提交
4077 4078
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-SelectRole")
	defer sp.End()
4079 4080 4081 4082

	log := log.Ctx(ctx)

	log.Debug("SelectRole", zap.Any("req", req))
4083
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4084
		return &milvuspb.SelectRoleResponse{Status: errorutil.UnhealthyStatus(code)}, nil
4085 4086 4087 4088 4089 4090 4091 4092 4093
	}

	if req.Role != nil {
		if err := ValidateRoleName(req.Role.Name); err != nil {
			return &milvuspb.SelectRoleResponse{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_IllegalArgument,
					Reason:    err.Error(),
				},
4094
			}, nil
4095 4096 4097 4098 4099
		}
	}

	result, err := node.rootCoord.SelectRole(ctx, req)
	if err != nil {
4100 4101
		log.Error("fail to select role",
			zap.Error(err))
4102 4103 4104 4105 4106
		return &milvuspb.SelectRoleResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
4107
		}, nil
4108 4109
	}
	return result, nil
4110 4111
}

4112
func (node *Proxy) SelectUser(ctx context.Context, req *milvuspb.SelectUserRequest) (*milvuspb.SelectUserResponse, error) {
E
Enwei Jiao 已提交
4113 4114
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-SelectUser")
	defer sp.End()
4115 4116 4117 4118 4119

	log := log.Ctx(ctx)

	log.Debug("SelectUser",
		zap.Any("req", req))
4120
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4121
		return &milvuspb.SelectUserResponse{Status: errorutil.UnhealthyStatus(code)}, nil
4122 4123 4124 4125 4126 4127 4128 4129 4130
	}

	if req.User != nil {
		if err := ValidateUsername(req.User.Name); err != nil {
			return &milvuspb.SelectUserResponse{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_IllegalArgument,
					Reason:    err.Error(),
				},
4131
			}, nil
4132 4133 4134 4135 4136
		}
	}

	result, err := node.rootCoord.SelectUser(ctx, req)
	if err != nil {
4137 4138
		log.Error("fail to select user",
			zap.Error(err))
4139 4140 4141 4142 4143
		return &milvuspb.SelectUserResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
4144
		}, nil
4145 4146
	}
	return result, nil
4147 4148
}

4149 4150 4151 4152 4153 4154 4155 4156 4157 4158 4159 4160 4161 4162 4163 4164 4165 4166 4167 4168 4169 4170 4171 4172 4173 4174 4175 4176 4177 4178
func (node *Proxy) validPrivilegeParams(req *milvuspb.OperatePrivilegeRequest) error {
	if req.Entity == nil {
		return fmt.Errorf("the entity in the request is nil")
	}
	if req.Entity.Grantor == nil {
		return fmt.Errorf("the grantor entity in the grant entity is nil")
	}
	if req.Entity.Grantor.Privilege == nil {
		return fmt.Errorf("the privilege entity in the grantor entity is nil")
	}
	if err := ValidatePrivilege(req.Entity.Grantor.Privilege.Name); err != nil {
		return err
	}
	if req.Entity.Object == nil {
		return fmt.Errorf("the resource entity in the grant entity is nil")
	}
	if err := ValidateObjectType(req.Entity.Object.Name); err != nil {
		return err
	}
	if err := ValidateObjectName(req.Entity.ObjectName); err != nil {
		return err
	}
	if req.Entity.Role == nil {
		return fmt.Errorf("the object entity in the grant entity is nil")
	}
	if err := ValidateRoleName(req.Entity.Role.Name); err != nil {
		return err
	}

	return nil
4179 4180
}

4181
func (node *Proxy) OperatePrivilege(ctx context.Context, req *milvuspb.OperatePrivilegeRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
4182 4183
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-OperatePrivilege")
	defer sp.End()
4184 4185 4186 4187 4188

	log := log.Ctx(ctx)

	log.Debug("OperatePrivilege",
		zap.Any("req", req))
4189
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4190
		return errorutil.UnhealthyStatus(code), nil
4191 4192 4193 4194 4195
	}
	if err := node.validPrivilegeParams(req); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4196
		}, nil
4197 4198 4199 4200 4201 4202
	}
	curUser, err := GetCurUserFromContext(ctx)
	if err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4203
		}, nil
4204 4205 4206 4207
	}
	req.Entity.Grantor.User = &milvuspb.UserEntity{Name: curUser}
	result, err := node.rootCoord.OperatePrivilege(ctx, req)
	if err != nil {
4208 4209
		log.Error("fail to operate privilege",
			zap.Error(err))
4210 4211 4212
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4213
		}, nil
4214 4215
	}
	return result, nil
4216 4217
}

4218 4219 4220 4221 4222 4223 4224 4225 4226 4227 4228 4229 4230 4231 4232 4233 4234 4235 4236 4237 4238 4239 4240 4241 4242 4243 4244
func (node *Proxy) validGrantParams(req *milvuspb.SelectGrantRequest) error {
	if req.Entity == nil {
		return fmt.Errorf("the grant entity in the request is nil")
	}

	if req.Entity.Object != nil {
		if err := ValidateObjectType(req.Entity.Object.Name); err != nil {
			return err
		}

		if err := ValidateObjectName(req.Entity.ObjectName); err != nil {
			return err
		}
	}

	if req.Entity.Role == nil {
		return fmt.Errorf("the role entity in the grant entity is nil")
	}

	if err := ValidateRoleName(req.Entity.Role.Name); err != nil {
		return err
	}

	return nil
}

func (node *Proxy) SelectGrant(ctx context.Context, req *milvuspb.SelectGrantRequest) (*milvuspb.SelectGrantResponse, error) {
E
Enwei Jiao 已提交
4245 4246
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-SelectGrant")
	defer sp.End()
4247 4248 4249 4250 4251

	log := log.Ctx(ctx)

	log.Debug("SelectGrant",
		zap.Any("req", req))
4252
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4253
		return &milvuspb.SelectGrantResponse{Status: errorutil.UnhealthyStatus(code)}, nil
4254 4255 4256 4257 4258 4259 4260 4261
	}

	if err := node.validGrantParams(req); err != nil {
		return &milvuspb.SelectGrantResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_IllegalArgument,
				Reason:    err.Error(),
			},
4262
		}, nil
4263 4264 4265 4266
	}

	result, err := node.rootCoord.SelectGrant(ctx, req)
	if err != nil {
4267 4268
		log.Error("fail to select grant",
			zap.Error(err))
4269 4270 4271 4272 4273
		return &milvuspb.SelectGrantResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
4274
		}, nil
4275 4276 4277 4278 4279
	}
	return result, nil
}

func (node *Proxy) RefreshPolicyInfoCache(ctx context.Context, req *proxypb.RefreshPolicyInfoCacheRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
4280 4281
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-RefreshPolicyInfoCache")
	defer sp.End()
4282 4283 4284 4285 4286

	log := log.Ctx(ctx)

	log.Debug("RefreshPrivilegeInfoCache",
		zap.Any("req", req))
4287 4288 4289 4290 4291 4292 4293 4294 4295 4296
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
		return errorutil.UnhealthyStatus(code), errorutil.UnhealthyError()
	}

	if globalMetaCache != nil {
		err := globalMetaCache.RefreshPolicyInfo(typeutil.CacheOp{
			OpType: typeutil.CacheOpType(req.OpType),
			OpKey:  req.OpKey,
		})
		if err != nil {
4297 4298
			log.Error("fail to refresh policy info",
				zap.Error(err))
4299 4300 4301 4302 4303 4304
			return &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_RefreshPolicyInfoCacheFailure,
				Reason:    err.Error(),
			}, err
		}
	}
4305
	log.Debug("RefreshPrivilegeInfoCache success")
4306 4307 4308 4309

	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
	}, nil
4310
}
4311 4312 4313 4314 4315 4316 4317 4318 4319 4320 4321 4322 4323 4324 4325 4326 4327

// SetRates limits the rates of requests.
func (node *Proxy) SetRates(ctx context.Context, request *proxypb.SetRatesRequest) (*commonpb.Status, error) {
	resp := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
	}
	if !node.checkHealthy() {
		resp = unhealthyStatus()
		return resp, nil
	}

	err := node.multiRateLimiter.globalRateLimiter.setRates(request.GetRates())
	// TODO: set multiple rate limiter rates
	if err != nil {
		resp.Reason = err.Error()
		return resp, nil
	}
4328
	node.multiRateLimiter.SetQuotaStates(request.GetStates(), request.GetCodes())
4329 4330 4331 4332 4333 4334 4335
	rateStrs := lo.FilterMap(request.GetRates(), func(r *internalpb.Rate, _ int) (string, bool) {
		return fmt.Sprintf("rateType:%s, rate:%s", r.GetRt().String(), ratelimitutil.Limit(r.GetR()).String()),
			ratelimitutil.Limit(r.GetR()) != ratelimitutil.Inf
	})
	if len(rateStrs) > 0 {
		log.RatedInfo(30, "current rates in proxy", zap.Int64("proxyNodeID", paramtable.GetNodeID()), zap.Strings("rates", rateStrs))
	}
4336 4337
	if len(request.GetStates()) != 0 {
		for i := range request.GetStates() {
4338
			log.Warn("Proxy set quota states", zap.String("state", request.GetStates()[i].String()), zap.String("reason", request.GetCodes()[i].String()))
4339 4340
		}
	}
4341 4342 4343
	resp.ErrorCode = commonpb.ErrorCode_Success
	return resp, nil
}
4344 4345 4346 4347

func (node *Proxy) CheckHealth(ctx context.Context, request *milvuspb.CheckHealthRequest) (*milvuspb.CheckHealthResponse, error) {
	if !node.checkHealthy() {
		reason := errorutil.UnHealthReason("proxy", node.session.ServerID, "proxy is unhealthy")
4348 4349 4350 4351
		return &milvuspb.CheckHealthResponse{
			Status:    unhealthyStatus(),
			IsHealthy: false,
			Reasons:   []string{reason}}, nil
4352 4353 4354 4355 4356 4357 4358 4359 4360 4361
	}

	group, ctx := errgroup.WithContext(ctx)
	errReasons := make([]string, 0)

	mu := &sync.Mutex{}
	fn := func(role string, resp *milvuspb.CheckHealthResponse, err error) error {
		mu.Lock()
		defer mu.Unlock()

E
Enwei Jiao 已提交
4362 4363
		ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-RefreshPolicyInfoCache")
		defer sp.End()
4364 4365 4366

		log := log.Ctx(ctx).With(zap.String("role", role))

4367
		if err != nil {
4368 4369
			log.Warn("check health fail",
				zap.Error(err))
4370 4371 4372 4373 4374
			errReasons = append(errReasons, fmt.Sprintf("check health fail for %s", role))
			return err
		}

		if !resp.IsHealthy {
4375
			log.Warn("check health fail")
4376 4377 4378 4379 4380 4381 4382 4383 4384 4385 4386 4387 4388 4389 4390 4391 4392 4393 4394 4395 4396 4397 4398
			errReasons = append(errReasons, resp.Reasons...)
		}
		return nil
	}

	group.Go(func() error {
		resp, err := node.rootCoord.CheckHealth(ctx, request)
		return fn("rootcoord", resp, err)
	})

	group.Go(func() error {
		resp, err := node.queryCoord.CheckHealth(ctx, request)
		return fn("querycoord", resp, err)
	})

	group.Go(func() error {
		resp, err := node.dataCoord.CheckHealth(ctx, request)
		return fn("datacoord", resp, err)
	})

	err := group.Wait()
	if err != nil || len(errReasons) != 0 {
		return &milvuspb.CheckHealthResponse{
4399 4400 4401
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_Success,
			},
4402 4403 4404 4405 4406
			IsHealthy: false,
			Reasons:   errReasons,
		}, nil
	}

4407
	states, reasons := node.multiRateLimiter.GetQuotaStates()
4408 4409 4410 4411 4412
	return &milvuspb.CheckHealthResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
			Reason:    "",
		},
4413 4414 4415
		QuotaStates: states,
		Reasons:     reasons,
		IsHealthy:   true,
4416
	}, nil
4417
}
W
wei liu 已提交
4418

J
jaime 已提交
4419
func (node *Proxy) RenameCollection(ctx context.Context, req *milvuspb.RenameCollectionRequest) (*commonpb.Status, error) {
4420
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-RenameCollection")
J
jaime 已提交
4421 4422 4423 4424 4425 4426 4427 4428 4429 4430 4431 4432 4433 4434 4435 4436 4437 4438 4439
	defer sp.End()

	log := log.Ctx(ctx).With(
		zap.String("role", typeutil.ProxyRole),
		zap.String("oldName", req.GetOldName()),
		zap.String("newName", req.GetNewName()))

	log.Info("received rename collection request")
	var err error

	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

	if err := validateCollectionName(req.GetNewName()); err != nil {
		log.Warn("validate new collection name fail", zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalCollectionName,
			Reason:    err.Error(),
J
jaime 已提交
4440
		}, nil
J
jaime 已提交
4441 4442 4443 4444 4445 4446 4447 4448 4449 4450 4451 4452 4453 4454 4455 4456 4457 4458 4459
	}

	req.Base = commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_RenameCollection),
		commonpbutil.WithMsgID(0),
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
	)
	resp, err := node.rootCoord.RenameCollection(ctx, req)
	if err != nil {
		log.Warn("failed to rename collection", zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, err
	}

	return resp, nil
}

W
wei liu 已提交
4460
func (node *Proxy) CreateResourceGroup(ctx context.Context, request *milvuspb.CreateResourceGroupRequest) (*commonpb.Status, error) {
W
wei liu 已提交
4461 4462 4463 4464
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

W
wei liu 已提交
4465 4466 4467 4468 4469 4470 4471 4472
	method := "CreateResourceGroup"
	if err := ValidateResourceGroupName(request.GetResourceGroup()); err != nil {
		log.Warn("CreateResourceGroup failed",
			zap.Error(err),
		)
		return getErrResponse(err, method), nil
	}

W
wei liu 已提交
4473 4474 4475 4476 4477 4478 4479 4480 4481 4482 4483 4484 4485 4486 4487 4488 4489 4490 4491 4492 4493
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CreateResourceGroup")
	defer sp.End()
	tr := timerecord.NewTimeRecorder(method)
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
	t := &CreateResourceGroupTask{
		ctx:                        ctx,
		Condition:                  NewTaskCondition(ctx),
		CreateResourceGroupRequest: request,
		queryCoord:                 node.queryCoord,
	}

	log := log.Ctx(ctx).With(
		zap.String("role", typeutil.ProxyRole),
	)

	log.Debug("CreateResourceGroup received")

	if err := node.sched.ddQueue.Enqueue(t); err != nil {
		log.Warn("CreateResourceGroup failed to enqueue",
			zap.Error(err))
W
wei liu 已提交
4494
		return getErrResponse(err, method), nil
W
wei liu 已提交
4495 4496 4497 4498 4499 4500 4501 4502 4503 4504 4505
	}

	log.Debug("CreateResourceGroup enqueued",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	if err := t.WaitToFinish(); err != nil {
		log.Warn("CreateResourceGroup failed to WaitToFinish",
			zap.Error(err),
			zap.Uint64("BeginTS", t.BeginTs()),
			zap.Uint64("EndTS", t.EndTs()))
W
wei liu 已提交
4506
		return getErrResponse(err, method), nil
W
wei liu 已提交
4507 4508 4509 4510 4511 4512 4513 4514 4515 4516
	}

	log.Debug("CreateResourceGroup done",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
	return t.result, nil
W
wei liu 已提交
4517 4518
}

W
wei liu 已提交
4519 4520 4521 4522
func getErrResponse(err error, method string) *commonpb.Status {
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()

	return &commonpb.Status{
W
wei liu 已提交
4523
		ErrorCode: commonpb.ErrorCode_IllegalArgument,
W
wei liu 已提交
4524 4525 4526 4527
		Reason:    err.Error(),
	}
}

W
wei liu 已提交
4528
func (node *Proxy) DropResourceGroup(ctx context.Context, request *milvuspb.DropResourceGroupRequest) (*commonpb.Status, error) {
W
wei liu 已提交
4529 4530 4531 4532
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

W
wei liu 已提交
4533
	method := "DropResourceGroup"
W
wei liu 已提交
4534 4535 4536 4537 4538 4539 4540 4541 4542 4543 4544 4545 4546 4547 4548 4549 4550 4551 4552 4553 4554 4555
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DropResourceGroup")
	defer sp.End()
	tr := timerecord.NewTimeRecorder(method)
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
	t := &DropResourceGroupTask{
		ctx:                      ctx,
		Condition:                NewTaskCondition(ctx),
		DropResourceGroupRequest: request,
		queryCoord:               node.queryCoord,
	}

	log := log.Ctx(ctx).With(
		zap.String("role", typeutil.ProxyRole),
	)

	log.Debug("DropResourceGroup received")

	if err := node.sched.ddQueue.Enqueue(t); err != nil {
		log.Warn("DropResourceGroup failed to enqueue",
			zap.Error(err))

W
wei liu 已提交
4556
		return getErrResponse(err, method), nil
W
wei liu 已提交
4557 4558 4559 4560 4561 4562 4563 4564 4565 4566 4567
	}

	log.Debug("DropResourceGroup enqueued",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	if err := t.WaitToFinish(); err != nil {
		log.Warn("DropResourceGroup failed to WaitToFinish",
			zap.Error(err),
			zap.Uint64("BeginTS", t.BeginTs()),
			zap.Uint64("EndTS", t.EndTs()))
W
wei liu 已提交
4568
		return getErrResponse(err, method), nil
W
wei liu 已提交
4569 4570 4571 4572 4573 4574 4575 4576 4577 4578
	}

	log.Debug("DropResourceGroup done",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
	return t.result, nil
W
wei liu 已提交
4579 4580 4581
}

func (node *Proxy) TransferNode(ctx context.Context, request *milvuspb.TransferNodeRequest) (*commonpb.Status, error) {
W
wei liu 已提交
4582 4583 4584 4585
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

W
wei liu 已提交
4586 4587 4588 4589 4590 4591 4592 4593 4594 4595 4596 4597 4598 4599 4600
	method := "TransferNode"
	if err := ValidateResourceGroupName(request.GetSourceResourceGroup()); err != nil {
		log.Warn("TransferNode failed",
			zap.Error(err),
		)
		return getErrResponse(err, method), nil
	}

	if err := ValidateResourceGroupName(request.GetTargetResourceGroup()); err != nil {
		log.Warn("TransferNode failed",
			zap.Error(err),
		)
		return getErrResponse(err, method), nil
	}

W
wei liu 已提交
4601 4602 4603 4604 4605 4606 4607 4608 4609 4610 4611 4612 4613 4614 4615 4616 4617 4618 4619 4620 4621 4622
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-TransferNode")
	defer sp.End()
	tr := timerecord.NewTimeRecorder(method)
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
	t := &TransferNodeTask{
		ctx:                 ctx,
		Condition:           NewTaskCondition(ctx),
		TransferNodeRequest: request,
		queryCoord:          node.queryCoord,
	}

	log := log.Ctx(ctx).With(
		zap.String("role", typeutil.ProxyRole),
	)

	log.Debug("TransferNode received")

	if err := node.sched.ddQueue.Enqueue(t); err != nil {
		log.Warn("TransferNode failed to enqueue",
			zap.Error(err))

W
wei liu 已提交
4623
		return getErrResponse(err, method), nil
W
wei liu 已提交
4624 4625 4626 4627 4628 4629 4630 4631 4632 4633 4634
	}

	log.Debug("TransferNode enqueued",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	if err := t.WaitToFinish(); err != nil {
		log.Warn("TransferNode failed to WaitToFinish",
			zap.Error(err),
			zap.Uint64("BeginTS", t.BeginTs()),
			zap.Uint64("EndTS", t.EndTs()))
W
wei liu 已提交
4635
		return getErrResponse(err, method), nil
W
wei liu 已提交
4636 4637 4638 4639 4640 4641 4642 4643 4644 4645
	}

	log.Debug("TransferNode done",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
	return t.result, nil
W
wei liu 已提交
4646 4647 4648
}

func (node *Proxy) TransferReplica(ctx context.Context, request *milvuspb.TransferReplicaRequest) (*commonpb.Status, error) {
W
wei liu 已提交
4649 4650 4651
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
W
wei liu 已提交
4652

W
wei liu 已提交
4653 4654 4655 4656 4657 4658 4659 4660 4661 4662 4663 4664 4665 4666 4667
	method := "TransferReplica"
	if err := ValidateResourceGroupName(request.GetSourceResourceGroup()); err != nil {
		log.Warn("TransferReplica failed",
			zap.Error(err),
		)
		return getErrResponse(err, method), nil
	}

	if err := ValidateResourceGroupName(request.GetTargetResourceGroup()); err != nil {
		log.Warn("TransferReplica failed",
			zap.Error(err),
		)
		return getErrResponse(err, method), nil
	}

W
wei liu 已提交
4668 4669 4670 4671 4672 4673 4674 4675 4676 4677 4678 4679 4680 4681 4682 4683 4684 4685 4686 4687 4688 4689
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-TransferReplica")
	defer sp.End()
	tr := timerecord.NewTimeRecorder(method)
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
	t := &TransferReplicaTask{
		ctx:                    ctx,
		Condition:              NewTaskCondition(ctx),
		TransferReplicaRequest: request,
		queryCoord:             node.queryCoord,
	}

	log := log.Ctx(ctx).With(
		zap.String("role", typeutil.ProxyRole),
	)

	log.Debug("TransferReplica received")

	if err := node.sched.ddQueue.Enqueue(t); err != nil {
		log.Warn("TransferReplica failed to enqueue",
			zap.Error(err))

W
wei liu 已提交
4690
		return getErrResponse(err, method), nil
W
wei liu 已提交
4691 4692 4693 4694 4695 4696 4697 4698 4699 4700 4701
	}

	log.Debug("TransferReplica enqueued",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	if err := t.WaitToFinish(); err != nil {
		log.Warn("TransferReplica failed to WaitToFinish",
			zap.Error(err),
			zap.Uint64("BeginTS", t.BeginTs()),
			zap.Uint64("EndTS", t.EndTs()))
W
wei liu 已提交
4702
		return getErrResponse(err, method), nil
W
wei liu 已提交
4703 4704 4705 4706 4707 4708 4709 4710 4711 4712
	}

	log.Debug("TransferReplica done",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
	return t.result, nil
W
wei liu 已提交
4713 4714
}

W
wei liu 已提交
4715 4716 4717 4718 4719 4720 4721 4722 4723 4724 4725 4726 4727 4728 4729 4730 4731 4732 4733 4734 4735 4736 4737 4738 4739 4740 4741 4742 4743 4744 4745 4746 4747 4748 4749 4750 4751 4752 4753 4754 4755 4756 4757 4758 4759 4760 4761 4762 4763 4764 4765 4766 4767 4768 4769 4770 4771 4772 4773 4774 4775 4776 4777 4778 4779 4780 4781
func (node *Proxy) ListResourceGroups(ctx context.Context, request *milvuspb.ListResourceGroupsRequest) (*milvuspb.ListResourceGroupsResponse, error) {
	if !node.checkHealthy() {
		return &milvuspb.ListResourceGroupsResponse{
			Status: unhealthyStatus(),
		}, nil
	}

	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ListResourceGroups")
	defer sp.End()
	method := "ListResourceGroups"
	tr := timerecord.NewTimeRecorder(method)
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
	t := &ListResourceGroupsTask{
		ctx:                       ctx,
		Condition:                 NewTaskCondition(ctx),
		ListResourceGroupsRequest: request,
		queryCoord:                node.queryCoord,
	}

	log := log.Ctx(ctx).With(
		zap.String("role", typeutil.ProxyRole),
	)

	log.Debug("ListResourceGroups received")

	if err := node.sched.ddQueue.Enqueue(t); err != nil {
		log.Warn("ListResourceGroups failed to enqueue",
			zap.Error(err))

		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
			metrics.AbandonLabel).Inc()
		return &milvuspb.ListResourceGroupsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

	log.Debug("ListResourceGroups enqueued",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	if err := t.WaitToFinish(); err != nil {
		log.Warn("ListResourceGroups failed to WaitToFinish",
			zap.Error(err),
			zap.Uint64("BeginTS", t.BeginTs()),
			zap.Uint64("EndTS", t.EndTs()))
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
			metrics.FailLabel).Inc()
		return &milvuspb.ListResourceGroupsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

	log.Debug("ListResourceGroups done",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
	return t.result, nil
W
wei liu 已提交
4782 4783 4784
}

func (node *Proxy) DescribeResourceGroup(ctx context.Context, request *milvuspb.DescribeResourceGroupRequest) (*milvuspb.DescribeResourceGroupResponse, error) {
W
wei liu 已提交
4785 4786 4787 4788 4789 4790
	if !node.checkHealthy() {
		return &milvuspb.DescribeResourceGroupResponse{
			Status: unhealthyStatus(),
		}, nil
	}

W
wei liu 已提交
4791 4792 4793 4794 4795 4796 4797 4798 4799 4800 4801 4802
	method := "DescribeResourceGroup"
	GetErrResponse := func(err error) *milvuspb.DescribeResourceGroupResponse {
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()

		return &milvuspb.DescribeResourceGroupResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}
	}

W
wei liu 已提交
4803 4804 4805 4806 4807 4808 4809 4810 4811 4812 4813 4814 4815 4816 4817 4818 4819 4820 4821 4822 4823 4824
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DescribeResourceGroup")
	defer sp.End()
	tr := timerecord.NewTimeRecorder(method)
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
	t := &DescribeResourceGroupTask{
		ctx:                          ctx,
		Condition:                    NewTaskCondition(ctx),
		DescribeResourceGroupRequest: request,
		queryCoord:                   node.queryCoord,
	}

	log := log.Ctx(ctx).With(
		zap.String("role", typeutil.ProxyRole),
	)

	log.Debug("DescribeResourceGroup received")

	if err := node.sched.ddQueue.Enqueue(t); err != nil {
		log.Warn("DescribeResourceGroup failed to enqueue",
			zap.Error(err))

W
wei liu 已提交
4825
		return GetErrResponse(err), nil
W
wei liu 已提交
4826 4827 4828 4829 4830 4831 4832 4833 4834 4835 4836
	}

	log.Debug("DescribeResourceGroup enqueued",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	if err := t.WaitToFinish(); err != nil {
		log.Warn("DescribeResourceGroup failed to WaitToFinish",
			zap.Error(err),
			zap.Uint64("BeginTS", t.BeginTs()),
			zap.Uint64("EndTS", t.EndTs()))
W
wei liu 已提交
4837
		return GetErrResponse(err), nil
W
wei liu 已提交
4838 4839 4840 4841 4842 4843 4844 4845 4846 4847
	}

	log.Debug("DescribeResourceGroup done",
		zap.Uint64("BeginTS", t.BeginTs()),
		zap.Uint64("EndTS", t.EndTs()))

	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
	return t.result, nil
W
wei liu 已提交
4848
}