impl.go 140.8 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
	"errors"
22
	"fmt"
C
cai.zhang 已提交
23
	"os"
24
	"strconv"
25 26
	"sync"

E
Enwei Jiao 已提交
27
	"go.opentelemetry.io/otel"
28
	"golang.org/x/sync/errgroup"
29

30
	"github.com/golang/protobuf/proto"
S
SimFG 已提交
31 32
	"github.com/milvus-io/milvus-proto/go-api/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/milvuspb"
S
smellthemoon 已提交
33
	"github.com/milvus-io/milvus-proto/go-api/schemapb"
34
	"github.com/milvus-io/milvus/internal/common"
X
Xiangyu Wang 已提交
35
	"github.com/milvus-io/milvus/internal/log"
36
	"github.com/milvus-io/milvus/internal/metrics"
J
jaime 已提交
37
	"github.com/milvus-io/milvus/internal/mq/msgstream"
X
Xiangyu Wang 已提交
38 39 40 41
	"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"
42
	"github.com/milvus-io/milvus/internal/util"
43
	"github.com/milvus-io/milvus/internal/util/commonpbutil"
44
	"github.com/milvus-io/milvus/internal/util/crypto"
45
	"github.com/milvus-io/milvus/internal/util/errorutil"
46
	"github.com/milvus-io/milvus/internal/util/importutil"
47 48
	"github.com/milvus-io/milvus/internal/util/logutil"
	"github.com/milvus-io/milvus/internal/util/metricsinfo"
E
Enwei Jiao 已提交
49
	"github.com/milvus-io/milvus/internal/util/paramtable"
50
	"github.com/milvus-io/milvus/internal/util/timerecord"
X
Xiangyu Wang 已提交
51
	"github.com/milvus-io/milvus/internal/util/typeutil"
52
	"go.uber.org/zap"
53 54
)

55 56
const moduleName = "Proxy"

57
// UpdateStateCode updates the state code of Proxy.
58
func (node *Proxy) UpdateStateCode(code commonpb.StateCode) {
59
	node.stateCode.Store(code)
Z
zhenshan.cao 已提交
60 61
}

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

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

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

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

118 119
	log.Info("received request to invalidate collection meta cache")

120
	collectionName := request.CollectionName
121
	collectionID := request.CollectionID
X
Xiaofan 已提交
122 123

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

143
	return &commonpb.Status{
144
		ErrorCode: commonpb.ErrorCode_Success,
145 146
		Reason:    "",
	}, nil
147 148
}

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

E
Enwei Jiao 已提交
156 157
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CreateCollection")
	defer sp.End()
158 159 160
	method := "CreateCollection"
	tr := timerecord.NewTimeRecorder(method)

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

163
	cct := &createCollectionTask{
S
sunby 已提交
164
		ctx:                     ctx,
165 166
		Condition:               NewTaskCondition(ctx),
		CreateCollectionRequest: request,
167
		rootCoord:               node.rootCoord,
168 169
	}

170 171 172
	// avoid data race
	lenOfSchema := len(request.Schema)

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

181 182
	log.Debug(rpcReceived(method))

183 184 185
	if err := node.sched.ddQueue.Enqueue(cct); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
186
			zap.Error(err))
187

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

195 196
	log.Debug(
		rpcEnqueued(method),
197 198
		zap.Uint64("BeginTs", cct.BeginTs()),
		zap.Uint64("EndTs", cct.EndTs()),
199
		zap.Uint64("timestamp", request.Base.Timestamp))
200

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

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

215 216
	log.Debug(
		rpcDone(method),
217
		zap.Uint64("BeginTs", cct.BeginTs()),
218
		zap.Uint64("EndTs", cct.EndTs()))
219

E
Enwei Jiao 已提交
220 221
	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()))
222 223 224
	return cct.result, nil
}

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

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

237
	dct := &dropCollectionTask{
S
sunby 已提交
238
		ctx:                   ctx,
239 240
		Condition:             NewTaskCondition(ctx),
		DropCollectionRequest: request,
241
		rootCoord:             node.rootCoord,
242
		chMgr:                 node.chMgr,
S
sunby 已提交
243
		chTicker:              node.chTicker,
244 245
	}

246
	log := log.Ctx(ctx).With(
247
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
248 249
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
250

251 252
	log.Debug("DropCollection received")

253 254
	if err := node.sched.ddQueue.Enqueue(dct); err != nil {
		log.Warn("DropCollection failed to enqueue",
255
			zap.Error(err))
256

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

264 265
	log.Debug("DropCollection enqueued",
		zap.Uint64("BeginTs", dct.BeginTs()),
266
		zap.Uint64("EndTs", dct.EndTs()))
267 268 269

	if err := dct.WaitToFinish(); err != nil {
		log.Warn("DropCollection failed to WaitToFinish",
D
dragondriver 已提交
270
			zap.Error(err),
271
			zap.Uint64("BeginTs", dct.BeginTs()),
272
			zap.Uint64("EndTs", dct.EndTs()))
D
dragondriver 已提交
273

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

281 282
	log.Debug("DropCollection done",
		zap.Uint64("BeginTs", dct.BeginTs()),
283
		zap.Uint64("EndTs", dct.EndTs()))
284

E
Enwei Jiao 已提交
285 286
	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()))
287 288 289
	return dct.result, nil
}

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

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

305
	log := log.Ctx(ctx).With(
306
		zap.String("role", typeutil.ProxyRole),
307 308 309
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

310 311
	log.Debug("HasCollection received")

312
	hct := &hasCollectionTask{
S
sunby 已提交
313
		ctx:                  ctx,
314 315
		Condition:            NewTaskCondition(ctx),
		HasCollectionRequest: request,
316
		rootCoord:            node.rootCoord,
317 318
	}

319 320
	if err := node.sched.ddQueue.Enqueue(hct); err != nil {
		log.Warn("HasCollection failed to enqueue",
321
			zap.Error(err))
322

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

333 334
	log.Debug("HasCollection enqueued",
		zap.Uint64("BeginTS", hct.BeginTs()),
335
		zap.Uint64("EndTS", hct.EndTs()))
336 337 338

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

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

353 354
	log.Debug("HasCollection done",
		zap.Uint64("BeginTS", hct.BeginTs()),
355
		zap.Uint64("EndTS", hct.EndTs()))
356

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

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

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

383
	log := log.Ctx(ctx).With(
384
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
385 386
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
387

388 389
	log.Debug("LoadCollection received")

390 391
	if err := node.sched.ddQueue.Enqueue(lct); err != nil {
		log.Warn("LoadCollection failed to enqueue",
392
			zap.Error(err))
393

E
Enwei Jiao 已提交
394
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
395
			metrics.AbandonLabel).Inc()
396
		return &commonpb.Status{
397
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
398 399 400
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
401

402 403
	log.Debug("LoadCollection enqueued",
		zap.Uint64("BeginTS", lct.BeginTs()),
404
		zap.Uint64("EndTS", lct.EndTs()))
405 406 407

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

419 420
	log.Debug("LoadCollection done",
		zap.Uint64("BeginTS", lct.BeginTs()),
421
		zap.Uint64("EndTS", lct.EndTs()))
422

E
Enwei Jiao 已提交
423
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
424
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
425
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
426
	return lct.result, nil
427 428
}

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

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

449
	log := log.Ctx(ctx).With(
450
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
451 452
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
453

454 455
	log.Debug(rpcReceived(method))

456
	if err := node.sched.ddQueue.Enqueue(rct); err != nil {
457 458
		log.Warn(
			rpcFailedToEnqueue(method),
459
			zap.Error(err))
460

E
Enwei Jiao 已提交
461
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
462
			metrics.AbandonLabel).Inc()
463
		return &commonpb.Status{
464
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
465 466 467 468
			Reason:    err.Error(),
		}, nil
	}

469 470
	log.Debug(
		rpcEnqueued(method),
471
		zap.Uint64("BeginTS", rct.BeginTs()),
472
		zap.Uint64("EndTS", rct.EndTs()))
473 474

	if err := rct.WaitToFinish(); err != nil {
475 476
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
477
			zap.Error(err),
478
			zap.Uint64("BeginTS", rct.BeginTs()),
479
			zap.Uint64("EndTS", rct.EndTs()))
D
dragondriver 已提交
480

E
Enwei Jiao 已提交
481
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
482
			metrics.FailLabel).Inc()
483
		return &commonpb.Status{
484
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
485 486 487 488
			Reason:    err.Error(),
		}, nil
	}

489 490
	log.Debug(
		rpcDone(method),
491
		zap.Uint64("BeginTS", rct.BeginTs()),
492
		zap.Uint64("EndTS", rct.EndTs()))
493

E
Enwei Jiao 已提交
494
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
495
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
496
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
497
	return rct.result, nil
498 499
}

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

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

515
	dct := &describeCollectionTask{
S
sunby 已提交
516
		ctx:                       ctx,
517 518
		Condition:                 NewTaskCondition(ctx),
		DescribeCollectionRequest: request,
519
		rootCoord:                 node.rootCoord,
520 521
	}

522
	log := log.Ctx(ctx).With(
523
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
524 525
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
526

527 528
	log.Debug("DescribeCollection received")

529 530
	if err := node.sched.ddQueue.Enqueue(dct); err != nil {
		log.Warn("DescribeCollection failed to enqueue",
531
			zap.Error(err))
532

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

543 544
	log.Debug("DescribeCollection enqueued",
		zap.Uint64("BeginTS", dct.BeginTs()),
545
		zap.Uint64("EndTS", dct.EndTs()))
546 547 548

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

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

556 557
		return &milvuspb.DescribeCollectionResponse{
			Status: &commonpb.Status{
558
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
559 560 561 562 563
				Reason:    err.Error(),
			},
		}, nil
	}

564 565
	log.Debug("DescribeCollection done",
		zap.Uint64("BeginTS", dct.BeginTs()),
566
		zap.Uint64("EndTS", dct.EndTs()))
567

E
Enwei Jiao 已提交
568
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
569
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
570
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
571 572 573
	return dct.result, nil
}

574 575 576 577 578 579 580 581 582
// 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 已提交
583 584
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetStatistics")
	defer sp.End()
585 586
	method := "GetStatistics"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
587
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
588
		metrics.TotalLabel).Inc()
589 590 591 592 593 594 595 596 597 598
	g := &getStatisticsTask{
		request:   request,
		Condition: NewTaskCondition(ctx),
		ctx:       ctx,
		tr:        tr,
		dc:        node.dataCoord,
		qc:        node.queryCoord,
		shardMgr:  node.shardMgr,
	}

599
	log := log.Ctx(ctx).With(
600 601
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
602 603 604 605
		zap.String("collection", request.CollectionName))

	log.Debug(
		rpcReceived(method),
606 607 608 609 610 611 612 613
		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 已提交
614
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638
			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 已提交
639
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
640 641 642 643 644 645 646 647 648 649 650 651 652
			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()),
653
		zap.Uint64("EndTS", g.EndTs()))
654

E
Enwei Jiao 已提交
655
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
656
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
657
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
658 659 660
	return g.result, nil
}

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

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

682
	log := log.Ctx(ctx).With(
683
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
684 685
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
686

687 688
	log.Debug(rpcReceived(method))

689
	if err := node.sched.ddQueue.Enqueue(g); err != nil {
690 691
		log.Warn(
			rpcFailedToEnqueue(method),
692
			zap.Error(err))
693

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

G
godchen 已提交
697
		return &milvuspb.GetCollectionStatisticsResponse{
698
			Status: &commonpb.Status{
699
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
700 701 702 703 704
				Reason:    err.Error(),
			},
		}, nil
	}

705 706
	log.Debug(
		rpcEnqueued(method),
707
		zap.Uint64("BeginTS", g.BeginTs()),
708
		zap.Uint64("EndTS", g.EndTs()))
709 710

	if err := g.WaitToFinish(); err != nil {
711 712
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
713
			zap.Error(err),
714
			zap.Uint64("BeginTS", g.BeginTs()),
715
			zap.Uint64("EndTS", g.EndTs()))
D
dragondriver 已提交
716

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

G
godchen 已提交
720
		return &milvuspb.GetCollectionStatisticsResponse{
721
			Status: &commonpb.Status{
722
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
723 724 725 726 727
				Reason:    err.Error(),
			},
		}, nil
	}

728 729
	log.Debug(
		rpcDone(method),
730
		zap.Uint64("BeginTS", g.BeginTs()),
731
		zap.Uint64("EndTS", g.EndTs()))
732

E
Enwei Jiao 已提交
733
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
734
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
735
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
736
	return g.result, nil
737 738
}

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

752
	sct := &showCollectionsTask{
G
godchen 已提交
753 754 755
		ctx:                    ctx,
		Condition:              NewTaskCondition(ctx),
		ShowCollectionsRequest: request,
756
		queryCoord:             node.queryCoord,
757
		rootCoord:              node.rootCoord,
758 759
	}

760
	log := log.Ctx(ctx).With(
761
		zap.String("role", typeutil.ProxyRole),
762 763
		zap.String("DbName", request.DbName),
		zap.Uint64("TimeStamp", request.TimeStamp),
764 765 766 767
		zap.String("ShowType", request.Type.String()))

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

769
	err := node.sched.ddQueue.Enqueue(sct)
770
	if err != nil {
771 772
		log.Warn("ShowCollections failed to enqueue",
			zap.Error(err),
773
			zap.Any("CollectionNames", request.CollectionNames))
774

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

784
	log.Debug("ShowCollections enqueued",
785
		zap.Any("CollectionNames", request.CollectionNames))
D
dragondriver 已提交
786

787 788
	err = sct.WaitToFinish()
	if err != nil {
789 790
		log.Warn("ShowCollections failed to WaitToFinish",
			zap.Error(err),
791
			zap.Any("CollectionNames", request.CollectionNames))
792

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

G
godchen 已提交
795
		return &milvuspb.ShowCollectionsResponse{
796
			Status: &commonpb.Status{
797
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
798 799 800 801 802
				Reason:    err.Error(),
			},
		}, nil
	}

803
	log.Debug("ShowCollections Done",
804 805
		zap.Int("len(CollectionNames)", len(request.CollectionNames)),
		zap.Int("num_collections", len(sct.result.CollectionNames)))
806

E
Enwei Jiao 已提交
807 808
	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()))
809 810 811
	return sct.result, nil
}

J
jaime 已提交
812 813 814 815 816
func (node *Proxy) AlterCollection(ctx context.Context, request *milvuspb.AlterCollectionRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

E
Enwei Jiao 已提交
817 818
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-AlterCollection")
	defer sp.End()
J
jaime 已提交
819 820 821
	method := "AlterCollection"
	tr := timerecord.NewTimeRecorder(method)

E
Enwei Jiao 已提交
822
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
J
jaime 已提交
823 824 825 826 827 828 829 830

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

831
	log := log.Ctx(ctx).With(
J
jaime 已提交
832 833 834 835
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

836 837 838
	log.Debug(
		rpcReceived(method))

J
jaime 已提交
839 840 841
	if err := node.sched.ddQueue.Enqueue(act); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
842
			zap.Error(err))
J
jaime 已提交
843

E
Enwei Jiao 已提交
844
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
J
jaime 已提交
845 846 847 848 849 850 851 852 853 854
		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()),
855
		zap.Uint64("timestamp", request.Base.Timestamp))
J
jaime 已提交
856 857 858 859 860 861

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

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

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

E
Enwei Jiao 已提交
876 877
	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 已提交
878 879 880
	return act.result, nil
}

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

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

893
	cpt := &createPartitionTask{
S
sunby 已提交
894
		ctx:                    ctx,
895 896
		Condition:              NewTaskCondition(ctx),
		CreatePartitionRequest: request,
897
		rootCoord:              node.rootCoord,
898 899 900
		result:                 nil,
	}

901
	log := log.Ctx(ctx).With(
902
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
903 904 905
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
906

907 908
	log.Debug(rpcReceived("CreatePartition"))

909 910 911
	if err := node.sched.ddQueue.Enqueue(cpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue("CreatePartition"),
912
			zap.Error(err))
913

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

916
		return &commonpb.Status{
917
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
918 919 920
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
921

922 923 924
	log.Debug(
		rpcEnqueued("CreatePartition"),
		zap.Uint64("BeginTS", cpt.BeginTs()),
925
		zap.Uint64("EndTS", cpt.EndTs()))
926 927 928 929

	if err := cpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish("CreatePartition"),
D
dragondriver 已提交
930
			zap.Error(err),
931
			zap.Uint64("BeginTS", cpt.BeginTs()),
932
			zap.Uint64("EndTS", cpt.EndTs()))
D
dragondriver 已提交
933

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

936
		return &commonpb.Status{
937
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
938 939 940
			Reason:    err.Error(),
		}, nil
	}
941 942 943 944

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

E
Enwei Jiao 已提交
947 948
	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()))
949 950 951
	return cpt.result, nil
}

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

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

964
	dpt := &dropPartitionTask{
S
sunby 已提交
965
		ctx:                  ctx,
966 967
		Condition:            NewTaskCondition(ctx),
		DropPartitionRequest: request,
968
		rootCoord:            node.rootCoord,
C
cai.zhang 已提交
969
		queryCoord:           node.queryCoord,
970 971 972
		result:               nil,
	}

973
	log := log.Ctx(ctx).With(
974
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
975 976 977
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
978

979 980
	log.Debug(rpcReceived(method))

981 982 983
	if err := node.sched.ddQueue.Enqueue(dpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
984
			zap.Error(err))
985

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

988
		return &commonpb.Status{
989
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
990 991 992
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
993

994 995 996
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", dpt.BeginTs()),
997
		zap.Uint64("EndTS", dpt.EndTs()))
998 999 1000 1001

	if err := dpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1002
			zap.Error(err),
1003
			zap.Uint64("BeginTS", dpt.BeginTs()),
1004
			zap.Uint64("EndTS", dpt.EndTs()))
D
dragondriver 已提交
1005

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

1008
		return &commonpb.Status{
1009
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1010 1011 1012
			Reason:    err.Error(),
		}, nil
	}
1013 1014 1015 1016

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

E
Enwei Jiao 已提交
1019 1020
	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()))
1021 1022 1023
	return dpt.result, nil
}

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

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

1040
	hpt := &hasPartitionTask{
S
sunby 已提交
1041
		ctx:                 ctx,
1042 1043
		Condition:           NewTaskCondition(ctx),
		HasPartitionRequest: request,
1044
		rootCoord:           node.rootCoord,
1045 1046 1047
		result:              nil,
	}

1048
	log := log.Ctx(ctx).With(
1049
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1050 1051 1052
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
D
dragondriver 已提交
1053

1054 1055
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
1056 1057 1058
	if err := node.sched.ddQueue.Enqueue(hpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1059
			zap.Error(err))
D
dragondriver 已提交
1060

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

1064 1065
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
1066
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1067 1068 1069 1070 1071
				Reason:    err.Error(),
			},
			Value: false,
		}, nil
	}
D
dragondriver 已提交
1072

D
dragondriver 已提交
1073 1074 1075
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", hpt.BeginTs()),
1076
		zap.Uint64("EndTS", hpt.EndTs()))
D
dragondriver 已提交
1077 1078 1079 1080

	if err := hpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1081
			zap.Error(err),
D
dragondriver 已提交
1082
			zap.Uint64("BeginTS", hpt.BeginTs()),
1083
			zap.Uint64("EndTS", hpt.EndTs()))
D
dragondriver 已提交
1084

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

1088 1089
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
1090
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1091 1092 1093 1094 1095
				Reason:    err.Error(),
			},
			Value: false,
		}, nil
	}
D
dragondriver 已提交
1096 1097 1098 1099

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

E
Enwei Jiao 已提交
1102
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1103
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1104
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1105 1106 1107
	return hpt.result, nil
}

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

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

1128
	log := log.Ctx(ctx).With(
1129
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1130 1131 1132
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
1133

1134 1135
	log.Debug(rpcReceived(method))

1136 1137 1138
	if err := node.sched.ddQueue.Enqueue(lpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1139
			zap.Error(err))
1140

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

1144
		return &commonpb.Status{
1145
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1146 1147 1148 1149
			Reason:    err.Error(),
		}, nil
	}

1150 1151 1152
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", lpt.BeginTs()),
1153
		zap.Uint64("EndTS", lpt.EndTs()))
1154 1155 1156 1157

	if err := lpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1158
			zap.Error(err),
1159
			zap.Uint64("BeginTS", lpt.BeginTs()),
1160
			zap.Uint64("EndTS", lpt.EndTs()))
D
dragondriver 已提交
1161

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

1165
		return &commonpb.Status{
1166
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1167 1168 1169 1170
			Reason:    err.Error(),
		}, nil
	}

1171 1172 1173
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", lpt.BeginTs()),
1174
		zap.Uint64("EndTS", lpt.EndTs()))
1175

E
Enwei Jiao 已提交
1176
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1177
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1178
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1179
	return lpt.result, nil
1180 1181
}

1182
// ReleasePartitions release specific partitions from query nodes.
C
Cai Yudong 已提交
1183
func (node *Proxy) ReleasePartitions(ctx context.Context, request *milvuspb.ReleasePartitionsRequest) (*commonpb.Status, error) {
1184 1185 1186
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
1187

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

1191
	rpt := &releasePartitionsTask{
G
godchen 已提交
1192 1193 1194
		ctx:                      ctx,
		Condition:                NewTaskCondition(ctx),
		ReleasePartitionsRequest: request,
1195
		queryCoord:               node.queryCoord,
1196 1197
	}

1198
	method := "ReleasePartitions"
1199
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1200
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1201
		metrics.TotalLabel).Inc()
1202 1203

	log := log.Ctx(ctx).With(
1204
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1205 1206 1207
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
1208

1209 1210
	log.Debug(rpcReceived(method))

1211 1212 1213
	if err := node.sched.ddQueue.Enqueue(rpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1214
			zap.Error(err))
1215

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

1219
		return &commonpb.Status{
1220
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1221 1222 1223 1224
			Reason:    err.Error(),
		}, nil
	}

1225 1226 1227
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", rpt.BeginTs()),
1228
		zap.Uint64("EndTS", rpt.EndTs()))
1229 1230 1231 1232

	if err := rpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1233
			zap.Error(err),
1234
			zap.Uint64("BeginTS", rpt.BeginTs()),
1235
			zap.Uint64("EndTS", rpt.EndTs()))
D
dragondriver 已提交
1236

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

1240
		return &commonpb.Status{
1241
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1242 1243 1244 1245
			Reason:    err.Error(),
		}, nil
	}

1246 1247 1248
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", rpt.BeginTs()),
1249
		zap.Uint64("EndTS", rpt.EndTs()))
1250

E
Enwei Jiao 已提交
1251
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1252
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1253
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1254
	return rpt.result, nil
1255 1256
}

1257
// GetPartitionStatistics get the statistics of partition, such as num_rows.
C
Cai Yudong 已提交
1258
func (node *Proxy) GetPartitionStatistics(ctx context.Context, request *milvuspb.GetPartitionStatisticsRequest) (*milvuspb.GetPartitionStatisticsResponse, error) {
1259 1260 1261 1262 1263
	if !node.checkHealthy() {
		return &milvuspb.GetPartitionStatisticsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1264

E
Enwei Jiao 已提交
1265 1266
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetPartitionStatistics")
	defer sp.End()
1267 1268
	method := "GetPartitionStatistics"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1269
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1270
		metrics.TotalLabel).Inc()
1271

1272
	g := &getPartitionStatisticsTask{
1273 1274 1275
		ctx:                           ctx,
		Condition:                     NewTaskCondition(ctx),
		GetPartitionStatisticsRequest: request,
1276
		dataCoord:                     node.dataCoord,
1277 1278
	}

1279
	log := log.Ctx(ctx).With(
1280
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1281 1282 1283
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
1284

1285 1286
	log.Debug(rpcReceived(method))

1287 1288 1289
	if err := node.sched.ddQueue.Enqueue(g); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1290
			zap.Error(err))
1291

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

1295 1296 1297 1298 1299 1300 1301 1302
		return &milvuspb.GetPartitionStatisticsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

1303 1304 1305
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", g.BeginTs()),
1306
		zap.Uint64("EndTS", g.EndTs()))
1307 1308 1309 1310

	if err := g.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
1311
			zap.Error(err),
1312
			zap.Uint64("BeginTS", g.BeginTs()),
1313
			zap.Uint64("EndTS", g.EndTs()))
1314

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

1318 1319 1320 1321 1322 1323 1324 1325
		return &milvuspb.GetPartitionStatisticsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

1326 1327 1328
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", g.BeginTs()),
1329
		zap.Uint64("EndTS", g.EndTs()))
1330

E
Enwei Jiao 已提交
1331
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1332
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1333
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1334
	return g.result, nil
1335 1336
}

1337
// ShowPartitions list all partitions in the specific collection.
C
Cai Yudong 已提交
1338
func (node *Proxy) ShowPartitions(ctx context.Context, request *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
1339 1340 1341 1342 1343
	if !node.checkHealthy() {
		return &milvuspb.ShowPartitionsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1344

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

1348
	spt := &showPartitionsTask{
G
godchen 已提交
1349 1350 1351
		ctx:                   ctx,
		Condition:             NewTaskCondition(ctx),
		ShowPartitionsRequest: request,
1352
		rootCoord:             node.rootCoord,
1353
		queryCoord:            node.queryCoord,
G
godchen 已提交
1354
		result:                nil,
1355 1356
	}

1357
	method := "ShowPartitions"
1358 1359
	tr := timerecord.NewTimeRecorder(method)
	//TODO: use collectionID instead of collectionName
E
Enwei Jiao 已提交
1360
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1361
		metrics.TotalLabel).Inc()
1362

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

1365 1366
	log.Debug(
		rpcReceived(method),
G
godchen 已提交
1367
		zap.Any("request", request))
1368 1369 1370 1371 1372 1373 1374

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

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

G
godchen 已提交
1378
		return &milvuspb.ShowPartitionsResponse{
1379
			Status: &commonpb.Status{
1380
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1381 1382 1383 1384 1385
				Reason:    err.Error(),
			},
		}, nil
	}

1386 1387 1388 1389
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", spt.BeginTs()),
		zap.Uint64("EndTS", spt.EndTs()),
1390 1391
		zap.String("db", spt.ShowPartitionsRequest.DbName),
		zap.String("collection", spt.ShowPartitionsRequest.CollectionName),
1392 1393 1394 1395 1396
		zap.Any("partitions", spt.ShowPartitionsRequest.PartitionNames))

	if err := spt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1397
			zap.Error(err),
1398 1399 1400 1401 1402
			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 已提交
1403

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

G
godchen 已提交
1407
		return &milvuspb.ShowPartitionsResponse{
1408
			Status: &commonpb.Status{
1409
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1410 1411 1412 1413
				Reason:    err.Error(),
			},
		}, nil
	}
1414 1415 1416 1417 1418 1419 1420 1421 1422

	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 已提交
1423
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1424
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1425
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1426 1427 1428
	return spt.result, nil
}

S
SimFG 已提交
1429 1430 1431 1432 1433 1434
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 已提交
1435 1436
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetLoadingProgress")
	defer sp.End()
E
Enwei Jiao 已提交
1437
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
1438 1439 1440
	log := log.Ctx(ctx)

	log.Debug(
S
SimFG 已提交
1441 1442 1443 1444
		rpcReceived(method),
		zap.Any("request", request))

	getErrResponse := func(err error) *milvuspb.GetLoadingProgressResponse {
J
Jiquan Long 已提交
1445
		log.Warn("fail to get loading progress",
1446
			zap.String("collection_name", request.CollectionName),
S
SimFG 已提交
1447 1448
			zap.Strings("partition_name", request.PartitionNames),
			zap.Error(err))
E
Enwei Jiao 已提交
1449
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
1450 1451 1452 1453 1454
		if errors.Is(err, ErrInsufficientMemory) {
			return &milvuspb.GetLoadingProgressResponse{
				Status: InSufficientMemoryStatus(request.GetCollectionName()),
			}
		}
S
SimFG 已提交
1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468
		return &milvuspb.GetLoadingProgressResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}
	}
	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 已提交
1469

1470 1471 1472
	msgBase := commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_SystemInfo),
		commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
1473
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
1474
	)
S
SimFG 已提交
1475 1476 1477 1478 1479 1480 1481 1482 1483 1484
	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 已提交
1485
		if progress, err = getCollectionProgress(ctx, node.queryCoord, request.GetBase(), collectionID); err != nil {
S
SimFG 已提交
1486 1487 1488
			return getErrResponse(err), nil
		}
	} else {
S
SimFG 已提交
1489 1490
		if progress, err = getPartitionProgress(ctx, node.queryCoord, request.GetBase(),
			request.GetPartitionNames(), request.GetCollectionName(), collectionID); err != nil {
S
SimFG 已提交
1491 1492 1493 1494
			return getErrResponse(err), nil
		}
	}

1495
	log.Debug(
S
SimFG 已提交
1496 1497
		rpcDone(method),
		zap.Any("request", request))
E
Enwei Jiao 已提交
1498 1499
	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 已提交
1500 1501 1502 1503 1504 1505 1506 1507
	return &milvuspb.GetLoadingProgressResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
		},
		Progress: progress,
	}, nil
}

1508
func (node *Proxy) GetLoadState(ctx context.Context, request *milvuspb.GetLoadStateRequest) (*milvuspb.GetLoadStateResponse, error) {
S
SimFG 已提交
1509 1510 1511 1512 1513
	if !node.checkHealthy() {
		return &milvuspb.GetLoadStateResponse{Status: unhealthyStatus()}, nil
	}
	method := "GetLoadState"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1514 1515
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetLoadState")
	defer sp.End()
S
SimFG 已提交
1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526 1527 1528 1529 1530 1531 1532 1533 1534 1535 1536 1537 1538 1539 1540
	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{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}
	}

	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,
			},
G
godchen 已提交
2042
			InsertRequest: internalpb.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 2052
				FieldsData:     request.FieldsData,
				NumRows:        uint64(request.NumRows),
				Version:        internalpb.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,
			},
G
godchen 已提交
2171
			DeleteRequest: internalpb.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 2182 2183
			},
		},
		chMgr:    node.chMgr,
		chTicker: node.chTicker,
G
groot 已提交
2184 2185
	}

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

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

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

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

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

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

E
Enwei Jiao 已提交
2230
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2231
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
2232 2233
	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 已提交
2234 2235 2236
	return dt.result, nil
}

S
smellthemoon 已提交
2237 2238
// Upsert upsert records into collection.
func (node *Proxy) Upsert(ctx context.Context, request *milvuspb.UpsertRequest) (*milvuspb.MutationResult, error) {
E
Enwei Jiao 已提交
2239 2240
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Upsert")
	defer sp.End()
S
smellthemoon 已提交
2241 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 2272 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 2342 2343 2344 2345 2346 2347 2348 2349 2350 2351 2352 2353 2354 2355 2356 2357 2358 2359 2360 2361 2362 2363 2364 2365 2366 2367 2368 2369 2370

	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(
				commonpbutil.WithMsgType(commonpb.MsgType(commonpb.MsgType_Upsert)),
				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()
		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
}

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

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

2378 2379 2380 2381 2382
	if !node.checkHealthy() {
		return &milvuspb.SearchResults{
			Status: unhealthyStatus(),
		}, nil
	}
2383 2384
	method := "Search"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2385
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2386
		metrics.TotalLabel).Inc()
D
dragondriver 已提交
2387

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

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

2407 2408 2409
	travelTs := request.TravelTimestamp
	guaranteeTs := request.GuaranteeTimestamp

2410
	log := log.Ctx(ctx).With(
2411
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2412 2413 2414 2415 2416
		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)),
2417 2418 2419 2420
		zap.Any("OutputFields", request.OutputFields),
		zap.Any("search_params", request.SearchParams),
		zap.Uint64("travel_timestamp", travelTs),
		zap.Uint64("guarantee_timestamp", guaranteeTs))
D
dragondriver 已提交
2421

2422 2423 2424
	log.Debug(
		rpcReceived(method))

2425
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
2426
		log.Warn(
2427
			rpcFailedToEnqueue(method),
2428
			zap.Error(err))
D
dragondriver 已提交
2429

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

2433 2434
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
2435
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2436 2437 2438 2439
				Reason:    err.Error(),
			},
		}, nil
	}
Z
Zach 已提交
2440
	tr.CtxRecord(ctx, "search request enqueue")
2441

2442
	log.Debug(
2443
		rpcEnqueued(method),
2444
		zap.Uint64("timestamp", qt.Base.Timestamp))
D
dragondriver 已提交
2445

2446
	if err := qt.WaitToFinish(); err != nil {
2447
		log.Warn(
2448
			rpcFailedToWaitToFinish(method),
2449
			zap.Error(err))
2450

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

2454 2455
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
2456
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2457 2458 2459 2460 2461
				Reason:    err.Error(),
			},
		}, nil
	}

Z
Zach 已提交
2462
	span := tr.CtxRecord(ctx, "wait search result")
E
Enwei Jiao 已提交
2463
	metrics.ProxyWaitForSearchResultLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2464
		metrics.SearchLabel).Observe(float64(span.Milliseconds()))
2465
	tr.CtxRecord(ctx, "wait search result")
2466
	log.Debug(rpcDone(method))
D
dragondriver 已提交
2467

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

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

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

2500
	ft := &flushTask{
T
ThreadDao 已提交
2501 2502 2503
		ctx:          ctx,
		Condition:    NewTaskCondition(ctx),
		FlushRequest: request,
2504
		dataCoord:    node.dataCoord,
2505 2506
	}

D
dragondriver 已提交
2507
	method := "Flush"
2508
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2509
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2510

2511
	log := log.Ctx(ctx).With(
2512
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2513 2514
		zap.String("db", request.DbName),
		zap.Any("collections", request.CollectionNames))
D
dragondriver 已提交
2515

2516 2517
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2518 2519 2520
	if err := node.sched.ddQueue.Enqueue(ft); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2521
			zap.Error(err))
D
dragondriver 已提交
2522

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

2525 2526
		resp.Status.Reason = err.Error()
		return resp, nil
2527 2528
	}

D
dragondriver 已提交
2529 2530 2531
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", ft.BeginTs()),
2532
		zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2533 2534 2535 2536

	if err := ft.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
2537
			zap.Error(err),
D
dragondriver 已提交
2538
			zap.Uint64("BeginTs", ft.BeginTs()),
2539
			zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2540

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

D
dragondriver 已提交
2543
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
2544 2545
		resp.Status.Reason = err.Error()
		return resp, nil
2546 2547
	}

D
dragondriver 已提交
2548 2549 2550
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", ft.BeginTs()),
2551
		zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2552

E
Enwei Jiao 已提交
2553 2554
	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()))
2555
	return ft.result, nil
2556 2557
}

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

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

2565 2566 2567 2568 2569
	if !node.checkHealthy() {
		return &milvuspb.QueryResults{
			Status: unhealthyStatus(),
		}, nil
	}
2570

E
Enwei Jiao 已提交
2571 2572
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Query")
	defer sp.End()
2573
	tr := timerecord.NewTimeRecorder("Query")
D
dragondriver 已提交
2574

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

D
dragondriver 已提交
2591 2592
	method := "Query"

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

2596
	log := log.Ctx(ctx).With(
2597
		zap.String("role", typeutil.ProxyRole),
2598 2599
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
2600 2601 2602 2603
		zap.Strings("partitions", request.PartitionNames))

	log.Debug(
		rpcReceived(method),
2604 2605 2606 2607
		zap.String("expr", request.Expr),
		zap.Strings("OutputFields", request.OutputFields),
		zap.Uint64("travel_timestamp", request.TravelTimestamp),
		zap.Uint64("guarantee_timestamp", request.GuaranteeTimestamp))
G
godchen 已提交
2608

D
dragondriver 已提交
2609
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
2610
		log.Warn(
D
dragondriver 已提交
2611
			rpcFailedToEnqueue(method),
2612
			zap.Error(err))
D
dragondriver 已提交
2613

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

2617 2618 2619 2620 2621 2622
		return &milvuspb.QueryResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
2623
	}
Z
Zach 已提交
2624
	tr.CtxRecord(ctx, "query request enqueue")
2625

2626
	log.Debug(rpcEnqueued(method))
D
dragondriver 已提交
2627 2628

	if err := qt.WaitToFinish(); err != nil {
2629
		log.Warn(
D
dragondriver 已提交
2630
			rpcFailedToWaitToFinish(method),
2631
			zap.Error(err))
2632

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

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

2647
	log.Debug(rpcDone(method))
D
dragondriver 已提交
2648

E
Enwei Jiao 已提交
2649
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2650 2651
		metrics.SuccessLabel).Inc()

E
Enwei Jiao 已提交
2652
	metrics.ProxySQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2653
		metrics.QueryLabel).Observe(float64(tr.ElapseSpan().Milliseconds()))
E
Enwei Jiao 已提交
2654
	metrics.ProxyCollectionSQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2655
		metrics.QueryLabel, request.CollectionName).Observe(float64(tr.ElapseSpan().Milliseconds()))
2656 2657

	ret := &milvuspb.QueryResults{
2658 2659
		Status:     qt.result.Status,
		FieldsData: qt.result.FieldsData,
2660 2661
	}
	sentSize := proto.Size(qt.result)
2662
	rateCol.Add(metricsinfo.ReadResultThroughput, float64(sentSize))
E
Enwei Jiao 已提交
2663
	metrics.ProxyReadReqSendBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Add(float64(sentSize))
2664
	return ret, nil
2665
}
2666

2667
// CreateAlias create alias for collection, then you can search the collection with alias.
Y
Yusup 已提交
2668 2669 2670 2671
func (node *Proxy) CreateAlias(ctx context.Context, request *milvuspb.CreateAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2672

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

Y
Yusup 已提交
2676 2677 2678 2679 2680 2681 2682
	cat := &CreateAliasTask{
		ctx:                ctx,
		Condition:          NewTaskCondition(ctx),
		CreateAliasRequest: request,
		rootCoord:          node.rootCoord,
	}

D
dragondriver 已提交
2683
	method := "CreateAlias"
2684
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2685
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2686

2687
	log := log.Ctx(ctx).With(
D
dragondriver 已提交
2688 2689 2690 2691 2692
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

2693 2694
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2695 2696 2697
	if err := node.sched.ddQueue.Enqueue(cat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2698
			zap.Error(err))
D
dragondriver 已提交
2699

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

Y
Yusup 已提交
2702 2703 2704 2705 2706 2707
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2708 2709 2710
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", cat.BeginTs()),
2711
		zap.Uint64("EndTs", cat.EndTs()))
D
dragondriver 已提交
2712 2713 2714 2715

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

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

D
dragondriver 已提交
2727 2728 2729
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", cat.BeginTs()),
2730
		zap.Uint64("EndTs", cat.EndTs()))
D
dragondriver 已提交
2731

E
Enwei Jiao 已提交
2732 2733
	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 已提交
2734 2735 2736
	return cat.result, nil
}

2737
// DropAlias alter the alias of collection.
Y
Yusup 已提交
2738 2739 2740 2741
func (node *Proxy) DropAlias(ctx context.Context, request *milvuspb.DropAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2742

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

Y
Yusup 已提交
2746 2747 2748 2749 2750 2751 2752
	dat := &DropAliasTask{
		ctx:              ctx,
		Condition:        NewTaskCondition(ctx),
		DropAliasRequest: request,
		rootCoord:        node.rootCoord,
	}

D
dragondriver 已提交
2753
	method := "DropAlias"
2754
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2755
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2756

2757
	log := log.Ctx(ctx).With(
D
dragondriver 已提交
2758 2759 2760 2761
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias))

2762 2763
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2764 2765 2766
	if err := node.sched.ddQueue.Enqueue(dat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2767
			zap.Error(err))
E
Enwei Jiao 已提交
2768
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
D
dragondriver 已提交
2769

Y
Yusup 已提交
2770 2771 2772 2773 2774 2775
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2776 2777 2778
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", dat.BeginTs()),
2779
		zap.Uint64("EndTs", dat.EndTs()))
D
dragondriver 已提交
2780 2781 2782 2783

	if err := dat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2784
			zap.Error(err),
D
dragondriver 已提交
2785
			zap.Uint64("BeginTs", dat.BeginTs()),
2786
			zap.Uint64("EndTs", dat.EndTs()))
Y
Yusup 已提交
2787

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

Y
Yusup 已提交
2790 2791 2792 2793 2794 2795
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2796 2797 2798
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", dat.BeginTs()),
2799
		zap.Uint64("EndTs", dat.EndTs()))
D
dragondriver 已提交
2800

E
Enwei Jiao 已提交
2801 2802
	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 已提交
2803 2804 2805
	return dat.result, nil
}

2806
// AlterAlias alter alias of collection.
Y
Yusup 已提交
2807 2808 2809 2810
func (node *Proxy) AlterAlias(ctx context.Context, request *milvuspb.AlterAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2811

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

Y
Yusup 已提交
2815 2816 2817 2818 2819 2820 2821
	aat := &AlterAliasTask{
		ctx:               ctx,
		Condition:         NewTaskCondition(ctx),
		AlterAliasRequest: request,
		rootCoord:         node.rootCoord,
	}

D
dragondriver 已提交
2822
	method := "AlterAlias"
2823
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2824
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2825

2826
	log := log.Ctx(ctx).With(
D
dragondriver 已提交
2827 2828 2829 2830 2831
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

2832 2833
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2834 2835 2836
	if err := node.sched.ddQueue.Enqueue(aat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2837
			zap.Error(err))
E
Enwei Jiao 已提交
2838
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
D
dragondriver 已提交
2839

Y
Yusup 已提交
2840 2841 2842 2843 2844 2845
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2846 2847 2848
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", aat.BeginTs()),
2849
		zap.Uint64("EndTs", aat.EndTs()))
D
dragondriver 已提交
2850 2851 2852 2853

	if err := aat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2854
			zap.Error(err),
D
dragondriver 已提交
2855
			zap.Uint64("BeginTs", aat.BeginTs()),
2856
			zap.Uint64("EndTs", aat.EndTs()))
Y
Yusup 已提交
2857

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

Y
Yusup 已提交
2860 2861 2862 2863 2864 2865
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2866 2867 2868
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", aat.BeginTs()),
2869
		zap.Uint64("EndTs", aat.EndTs()))
D
dragondriver 已提交
2870

E
Enwei Jiao 已提交
2871 2872
	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 已提交
2873 2874 2875
	return aat.result, nil
}

2876
// CalcDistance calculates the distances between vectors.
2877
func (node *Proxy) CalcDistance(ctx context.Context, request *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) {
2878 2879 2880 2881 2882
	if !node.checkHealthy() {
		return &milvuspb.CalcDistanceResults{
			Status: unhealthyStatus(),
		}, nil
	}
2883

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

2887 2888
	query := func(ids *milvuspb.VectorIDs) (*milvuspb.QueryResults, error) {
		outputFields := []string{ids.FieldName}
2889

2890 2891 2892 2893 2894
		queryRequest := &milvuspb.QueryRequest{
			DbName:         "",
			CollectionName: ids.CollectionName,
			PartitionNames: ids.PartitionNames,
			OutputFields:   outputFields,
2895 2896
		}

2897
		qt := &queryTask{
2898 2899 2900
			ctx:       ctx,
			Condition: NewTaskCondition(ctx),
			RetrieveRequest: &internalpb.RetrieveRequest{
2901 2902
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
E
Enwei Jiao 已提交
2903
					commonpbutil.WithSourceID(paramtable.GetNodeID()),
2904
				),
E
Enwei Jiao 已提交
2905
				ReqID: paramtable.GetNodeID(),
2906
			},
2907 2908 2909 2910
			request: queryRequest,
			qc:      node.queryCoord,
			ids:     ids.IdArray,

2911
			queryShardPolicy: mergeRoundRobinPolicy,
2912
			shardMgr:         node.shardMgr,
2913 2914
		}

2915
		log := log.Ctx(ctx).With(
G
groot 已提交
2916 2917
			zap.String("collection", queryRequest.CollectionName),
			zap.Any("partitions", queryRequest.PartitionNames),
2918
			zap.Any("OutputFields", queryRequest.OutputFields))
G
groot 已提交
2919

2920
		err := node.sched.dqQueue.Enqueue(qt)
2921
		if err != nil {
2922 2923
			log.Error("CalcDistance queryTask failed to enqueue",
				zap.Error(err))
2924

2925 2926 2927 2928 2929
			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2930
			}, err
2931
		}
2932

2933
		log.Debug("CalcDistance queryTask enqueued")
2934 2935 2936

		err = qt.WaitToFinish()
		if err != nil {
2937 2938
			log.Error("CalcDistance queryTask failed to WaitToFinish",
				zap.Error(err))
2939 2940 2941 2942 2943 2944

			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2945
			}, err
2946
		}
2947

2948
		log.Debug("CalcDistance queryTask Done")
2949 2950

		return &milvuspb.QueryResults{
2951 2952
			Status:     qt.result.Status,
			FieldsData: qt.result.FieldsData,
2953 2954 2955
		}, nil
	}

G
groot 已提交
2956 2957
	// calcDistanceTask is not a standard task, no need to enqueue
	task := &calcDistanceTask{
E
Enwei Jiao 已提交
2958
		traceID:   sp.SpanContext().TraceID().String(),
G
groot 已提交
2959
		queryFunc: query,
2960 2961
	}

G
groot 已提交
2962
	return task.Execute(ctx, request)
2963 2964
}

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

2970
// GetPersistentSegmentInfo get the information of sealed segment.
C
Cai Yudong 已提交
2971
func (node *Proxy) GetPersistentSegmentInfo(ctx context.Context, req *milvuspb.GetPersistentSegmentInfoRequest) (*milvuspb.GetPersistentSegmentInfoResponse, error) {
E
Enwei Jiao 已提交
2972 2973
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetPersistentSegmentInfo")
	defer sp.End()
2974 2975 2976

	log := log.Ctx(ctx)

D
dragondriver 已提交
2977
	log.Debug("GetPersistentSegmentInfo",
2978
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2979 2980 2981
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

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

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

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

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

	log := log.Ctx(ctx)

D
dragondriver 已提交
3067
	log.Debug("GetQuerySegmentInfo",
3068
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3069 3070 3071
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
3072
	resp := &milvuspb.GetQuerySegmentInfoResponse{
Z
zhenshan.cao 已提交
3073
		Status: &commonpb.Status{
3074
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
Z
zhenshan.cao 已提交
3075 3076
		},
	}
3077 3078 3079 3080
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3081

3082 3083
	method := "GetQuerySegmentInfo"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
3084
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3085 3086
		metrics.TotalLabel).Inc()

3087 3088
	collID, err := globalMetaCache.GetCollectionID(ctx, req.CollectionName)
	if err != nil {
E
Enwei Jiao 已提交
3089
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3090 3091 3092
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3093
	infoResp, err := node.queryCoord.GetSegmentInfo(ctx, &querypb.GetSegmentInfoRequest{
3094 3095 3096
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_SegmentInfo),
			commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
3097
			commonpbutil.WithSourceID(paramtable.GetNodeID()),
3098
		),
3099
		CollectionID: collID,
Z
zhenshan.cao 已提交
3100 3101
	})
	if err != nil {
E
Enwei Jiao 已提交
3102
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3103 3104
		log.Error("Failed to get segment info from QueryCoord",
			zap.Error(err))
Z
zhenshan.cao 已提交
3105 3106 3107
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3108 3109 3110
	log.Debug("GetQuerySegmentInfo",
		zap.Any("infos", infoResp.Infos),
		zap.Any("status", infoResp.Status))
3111
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
E
Enwei Jiao 已提交
3112
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3113 3114
		log.Error("Failed to get segment info from QueryCoord",
			zap.String("errMsg", infoResp.Status.Reason))
Z
zhenshan.cao 已提交
3115 3116 3117 3118 3119 3120 3121 3122 3123 3124 3125 3126 3127
		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 已提交
3128
			State:        info.SegmentState,
3129
			NodeIds:      info.NodeIds,
Z
zhenshan.cao 已提交
3130 3131
		}
	}
3132

E
Enwei Jiao 已提交
3133 3134
	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()))
3135
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
Z
zhenshan.cao 已提交
3136 3137 3138 3139
	resp.Infos = queryInfos
	return resp, nil
}

J
jingkl 已提交
3140
// Dummy handles dummy request
C
Cai Yudong 已提交
3141
func (node *Proxy) Dummy(ctx context.Context, req *milvuspb.DummyRequest) (*milvuspb.DummyResponse, error) {
3142 3143 3144 3145 3146 3147
	failedResponse := &milvuspb.DummyResponse{
		Response: `{"status": "fail"}`,
	}

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

E
Enwei Jiao 已提交
3149 3150
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Dummy")
	defer sp.End()
3151 3152 3153

	log := log.Ctx(ctx)

3154
	if err != nil {
3155 3156
		log.Warn("Failed to parse dummy request type",
			zap.Error(err))
3157 3158 3159
		return failedResponse, nil
	}

3160 3161
	if drt.RequestType == "query" {
		drr, err := parseDummyQueryRequest(req.RequestType)
3162
		if err != nil {
3163 3164
			log.Warn("Failed to parse dummy query request",
				zap.Error(err))
3165 3166 3167
			return failedResponse, nil
		}

3168
		request := &milvuspb.QueryRequest{
3169 3170 3171
			DbName:         drr.DbName,
			CollectionName: drr.CollectionName,
			PartitionNames: drr.PartitionNames,
3172
			OutputFields:   drr.OutputFields,
X
Xiangyu Wang 已提交
3173 3174
		}

3175
		_, err = node.Query(ctx, request)
3176
		if err != nil {
3177 3178
			log.Warn("Failed to execute dummy query",
				zap.Error(err))
3179 3180
			return failedResponse, err
		}
X
Xiangyu Wang 已提交
3181 3182 3183 3184 3185 3186

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

3187 3188
	log.Debug("cannot find specify dummy request type")
	return failedResponse, nil
X
Xiangyu Wang 已提交
3189 3190
}

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

E
Enwei Jiao 已提交
3195 3196
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-RegisterLink")
	defer sp.End()
3197 3198

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

3202 3203
	log.Debug("RegisterLink")

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

3223
// GetMetrics gets the metrics of proxy
3224 3225
// 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 已提交
3226 3227
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetMetrics")
	defer sp.End()
3228 3229 3230

	log := log.Ctx(ctx)

3231 3232
	log.RatedDebug(60, "Proxy.GetMetrics",
		zap.Int64("nodeID", paramtable.GetNodeID()),
3233 3234 3235 3236
		zap.String("req", req.Request))

	if !node.checkHealthy() {
		log.Warn("Proxy.GetMetrics failed",
3237
			zap.Int64("nodeID", paramtable.GetNodeID()),
3238
			zap.String("req", req.Request),
E
Enwei Jiao 已提交
3239
			zap.Error(errProxyIsUnhealthy(paramtable.GetNodeID())))
3240 3241 3242 3243

		return &milvuspb.GetMetricsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
E
Enwei Jiao 已提交
3244
				Reason:    msgProxyIsUnhealthy(paramtable.GetNodeID()),
3245 3246 3247 3248 3249 3250 3251 3252
			},
			Response: "",
		}, nil
	}

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

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

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

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

3284 3285
		node.metricsCacheManager.UpdateSystemInfoMetrics(metrics)

G
godchen 已提交
3286
		return metrics, nil
3287 3288
	}

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

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

3303 3304 3305
// 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 已提交
3306 3307
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetProxyMetrics")
	defer sp.End()
3308 3309

	log := log.Ctx(ctx).With(
3310
		zap.Int64("nodeID", paramtable.GetNodeID()),
3311 3312
		zap.String("req", req.Request))

3313 3314
	if !node.checkHealthy() {
		log.Warn("Proxy.GetProxyMetrics failed",
E
Enwei Jiao 已提交
3315
			zap.Error(errProxyIsUnhealthy(paramtable.GetNodeID())))
3316 3317 3318 3319

		return &milvuspb.GetMetricsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
E
Enwei Jiao 已提交
3320
				Reason:    msgProxyIsUnhealthy(paramtable.GetNodeID()),
3321 3322 3323 3324 3325 3326 3327 3328 3329 3330 3331 3332 3333 3334 3335 3336 3337
			},
		}, 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
	}

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

	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
		}

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

		return proxyMetrics, nil
	}

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

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

B
bigsheeper 已提交
3375 3376
// 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 已提交
3377 3378
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-LoadBalance")
	defer sp.End()
3379 3380 3381

	log := log.Ctx(ctx)

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

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

	status := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
	}
3393 3394 3395

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

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

	log := log.Ctx(ctx)

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

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

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

3469
// GetCompactionState gets the compaction state of multiple segments
3470
func (node *Proxy) GetCompactionState(ctx context.Context, req *milvuspb.GetCompactionStateRequest) (*milvuspb.GetCompactionStateResponse, error) {
E
Enwei Jiao 已提交
3471 3472
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetCompactionState")
	defer sp.End()
3473 3474 3475 3476 3477

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

	log.Debug("received GetCompactionState request")
3478 3479 3480 3481 3482 3483 3484
	resp := &milvuspb.GetCompactionStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionState(ctx, req)
3485 3486 3487
	log.Debug("received GetCompactionState response",
		zap.Any("resp", resp),
		zap.Error(err))
3488 3489 3490
	return resp, err
}

3491
// ManualCompaction invokes compaction on specified collection
3492
func (node *Proxy) ManualCompaction(ctx context.Context, req *milvuspb.ManualCompactionRequest) (*milvuspb.ManualCompactionResponse, error) {
E
Enwei Jiao 已提交
3493 3494
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ManualCompaction")
	defer sp.End()
3495 3496 3497 3498 3499

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

	log.Info("received ManualCompaction request")
3500 3501 3502 3503 3504 3505 3506
	resp := &milvuspb.ManualCompactionResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.ManualCompaction(ctx, req)
3507 3508 3509
	log.Info("received ManualCompaction response",
		zap.Any("resp", resp),
		zap.Error(err))
3510 3511 3512
	return resp, err
}

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

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

	log.Debug("received GetCompactionStateWithPlans request")
3522 3523 3524 3525 3526 3527 3528
	resp := &milvuspb.GetCompactionPlansResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionStateWithPlans(ctx, req)
3529 3530 3531
	log.Debug("received GetCompactionStateWithPlans response",
		zap.Any("resp", resp),
		zap.Error(err))
3532 3533 3534
	return resp, err
}

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

	log := log.Ctx(ctx)

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

3552
	resp, err = node.dataCoord.GetFlushState(ctx, req)
X
Xiaofan 已提交
3553
	if err != nil {
3554 3555
		log.Warn("failed to get flush state response",
			zap.Error(err))
X
Xiaofan 已提交
3556 3557
		return nil, err
	}
3558 3559
	log.Debug("received get flush state response",
		zap.Any("response", resp))
B
Bingyi Sun 已提交
3560 3561 3562
	return resp, err
}

C
Cai Yudong 已提交
3563 3564
// checkHealthy checks proxy state is Healthy
func (node *Proxy) checkHealthy() bool {
3565 3566
	code := node.stateCode.Load().(commonpb.StateCode)
	return code == commonpb.StateCode_Healthy
3567 3568
}

3569 3570 3571
func (node *Proxy) checkHealthyAndReturnCode() (commonpb.StateCode, bool) {
	code := node.stateCode.Load().(commonpb.StateCode)
	return code, code == commonpb.StateCode_Healthy
3572 3573
}

3574
// unhealthyStatus returns the proxy not healthy status
3575 3576 3577
func unhealthyStatus() *commonpb.Status {
	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
C
Cai Yudong 已提交
3578
		Reason:    "proxy not healthy",
3579 3580
	}
}
G
groot 已提交
3581 3582 3583

// 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 已提交
3584 3585
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Import")
	defer sp.End()
3586 3587 3588

	log := log.Ctx(ctx)

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

3604 3605
	err := importutil.ValidateOptions(req.GetOptions())
	if err != nil {
3606 3607
		log.Error("failed to execute import request",
			zap.Error(err))
3608 3609 3610 3611 3612
		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
	}

3613 3614
	method := "Import"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
3615
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3616 3617
		metrics.TotalLabel).Inc()

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

E
Enwei Jiao 已提交
3629 3630
	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()))
3631
	return respFromRC, nil
G
groot 已提交
3632 3633
}

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

	log := log.Ctx(ctx)

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

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

3663 3664 3665
	log.Debug("successfully received get import state response",
		zap.Int64("taskID", req.GetTask()),
		zap.Any("resp", resp), zap.Error(err))
E
Enwei Jiao 已提交
3666 3667
	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()))
3668
	return resp, nil
G
groot 已提交
3669 3670 3671 3672
}

// 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 已提交
3673 3674
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ListImportTasks")
	defer sp.End()
3675 3676 3677

	log := log.Ctx(ctx)

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

3698 3699 3700
	log.Debug("successfully received list import tasks response",
		zap.String("collection", req.CollectionName),
		zap.Any("tasks", resp.Tasks))
E
Enwei Jiao 已提交
3701 3702
	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 已提交
3703 3704 3705
	return resp, err
}

3706 3707
// InvalidateCredentialCache invalidate the credential cache of specified username.
func (node *Proxy) InvalidateCredentialCache(ctx context.Context, request *proxypb.InvalidateCredCacheRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
3708 3709
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-InvalidateCredentialCache")
	defer sp.End()
3710 3711

	log := log.Ctx(ctx).With(
3712 3713
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))
3714 3715

	log.Debug("received request to invalidate credential cache")
3716
	if !node.checkHealthy() {
3717
		return unhealthyStatus(), nil
3718
	}
3719 3720 3721 3722 3723

	username := request.Username
	if globalMetaCache != nil {
		globalMetaCache.RemoveCredential(username) // no need to return error, though credential may be not cached
	}
3724
	log.Debug("complete to invalidate credential cache")
3725 3726 3727 3728 3729 3730 3731 3732 3733

	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 已提交
3734 3735
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-UpdateCredentialCache")
	defer sp.End()
3736 3737

	log := log.Ctx(ctx).With(
3738 3739
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))
3740 3741

	log.Debug("received request to update credential cache")
3742
	if !node.checkHealthy() {
3743
		return unhealthyStatus(), nil
3744
	}
3745 3746

	credInfo := &internalpb.CredentialInfo{
3747 3748
		Username:       request.Username,
		Sha256Password: request.Password,
3749 3750 3751 3752
	}
	if globalMetaCache != nil {
		globalMetaCache.UpdateCredential(credInfo) // no need to return error, though credential may be not cached
	}
3753
	log.Debug("complete to update credential cache")
3754 3755 3756 3757 3758 3759 3760

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

E
Enwei Jiao 已提交
3761
//
3762
func (node *Proxy) CreateCredential(ctx context.Context, req *milvuspb.CreateCredentialRequest) (*commonpb.Status, error) {
E
Enwei Jiao 已提交
3763 3764
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CreateCredential")
	defer sp.End()
3765 3766 3767 3768 3769 3770

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

	log.Debug("CreateCredential",
		zap.String("role", typeutil.ProxyRole))
3771
	if !node.checkHealthy() {
3772
		return unhealthyStatus(), nil
3773
	}
3774 3775 3776 3777 3778 3779 3780 3781 3782 3783
	// 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 {
3784 3785
		log.Error("decode password fail",
			zap.Error(err))
3786 3787 3788 3789 3790 3791
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_CreateCredentialFailure,
			Reason:    "decode password fail key:" + req.Username,
		}, nil
	}
	if err = ValidatePassword(rawPassword); err != nil {
3792 3793
		log.Error("illegal password",
			zap.Error(err))
3794 3795 3796 3797 3798 3799 3800
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
		}, nil
	}
	encryptedPassword, err := crypto.PasswordEncrypt(rawPassword)
	if err != nil {
3801 3802
		log.Error("encrypt password fail",
			zap.Error(err))
3803 3804 3805 3806 3807
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_CreateCredentialFailure,
			Reason:    "encrypt password fail key:" + req.Username,
		}, nil
	}
3808

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

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

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

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

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

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

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

	log.Debug("DeleteCredential",
		zap.String("role", typeutil.ProxyRole))
3910
	if !node.checkHealthy() {
3911
		return unhealthyStatus(), nil
3912 3913
	}

3914 3915 3916 3917 3918 3919
	if req.Username == util.UserRoot {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_DeleteCredentialFailure,
			Reason:    "user root cannot be deleted",
		}, nil
	}
3920 3921
	result, err := node.rootCoord.DeleteCredential(ctx, req)
	if err != nil { // for error like conntext timeout etc.
3922 3923
		log.Error("delete credential fail",
			zap.Error(err))
3924 3925 3926 3927 3928 3929 3930 3931 3932
		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 已提交
3933 3934
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ListCredUsers")
	defer sp.End()
3935 3936 3937 3938 3939

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

	log.Debug("ListCredUsers")
3940
	if !node.checkHealthy() {
3941
		return &milvuspb.ListCredUsersResponse{Status: unhealthyStatus()}, nil
3942
	}
3943
	rootCoordReq := &milvuspb.ListCredUsersRequest{
3944 3945 3946
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_ListCredUsernames),
		),
3947 3948
	}
	resp, err := node.rootCoord.ListCredUsers(ctx, rootCoordReq)
3949 3950 3951 3952 3953 3954 3955 3956 3957 3958 3959 3960
	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,
		},
3961
		Usernames: resp.Usernames,
3962 3963
	}, nil
}
3964

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

	log := log.Ctx(ctx)

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

	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(),
3985
		}, nil
3986 3987 3988 3989
	}

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

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

	log := log.Ctx(ctx)

	log.Debug("DropRole",
		zap.Any("req", req))
4008
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4009
		return errorutil.UnhealthyStatus(code), nil
4010 4011 4012 4013 4014
	}
	if err := ValidateRoleName(req.RoleName); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4015
		}, nil
4016
	}
4017 4018 4019 4020 4021
	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,
4022
		}, nil
4023
	}
4024 4025
	result, err := node.rootCoord.DropRole(ctx, req)
	if err != nil {
4026 4027 4028
		log.Error("fail to drop role",
			zap.String("role_name", req.RoleName),
			zap.Error(err))
4029 4030 4031
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4032
		}, nil
4033 4034
	}
	return result, nil
4035 4036
}

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

	log := log.Ctx(ctx)

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

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

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

	log := log.Ctx(ctx)

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

	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(),
				},
4091
			}, nil
4092 4093 4094 4095 4096
		}
	}

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

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

	log := log.Ctx(ctx)

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

	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(),
				},
4128
			}, nil
4129 4130 4131 4132 4133
		}
	}

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

4146 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
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
4176 4177
}

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

	log := log.Ctx(ctx)

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

4215 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
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 已提交
4242 4243
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-SelectGrant")
	defer sp.End()
4244 4245 4246 4247 4248

	log := log.Ctx(ctx)

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

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

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

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

	log := log.Ctx(ctx)

	log.Debug("RefreshPrivilegeInfoCache",
		zap.Any("req", req))
4284 4285 4286 4287 4288 4289 4290 4291 4292 4293
	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 {
4294 4295
			log.Error("fail to refresh policy info",
				zap.Error(err))
4296 4297 4298 4299 4300 4301
			return &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_RefreshPolicyInfoCacheFailure,
				Reason:    err.Error(),
			}, err
		}
	}
4302
	log.Debug("RefreshPrivilegeInfoCache success")
4303 4304 4305 4306

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

// 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
	}
4325
	node.multiRateLimiter.SetQuotaStates(request.GetStates(), request.GetCodes())
4326 4327 4328
	log.Info("current rates in proxy", zap.Int64("proxyNodeID", paramtable.GetNodeID()), zap.Any("rates", request.GetRates()))
	if len(request.GetStates()) != 0 {
		for i := range request.GetStates() {
4329
			log.Warn("Proxy set quota states", zap.String("state", request.GetStates()[i].String()), zap.String("reason", request.GetCodes()[i].String()))
4330 4331
		}
	}
4332 4333 4334
	resp.ErrorCode = commonpb.ErrorCode_Success
	return resp, nil
}
4335 4336 4337 4338

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")
4339 4340 4341 4342
		return &milvuspb.CheckHealthResponse{
			Status:    unhealthyStatus(),
			IsHealthy: false,
			Reasons:   []string{reason}}, nil
4343 4344 4345 4346 4347 4348 4349 4350 4351 4352
	}

	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 已提交
4353 4354
		ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-RefreshPolicyInfoCache")
		defer sp.End()
4355 4356 4357

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

4358
		if err != nil {
4359 4360
			log.Warn("check health fail",
				zap.Error(err))
4361 4362 4363 4364 4365
			errReasons = append(errReasons, fmt.Sprintf("check health fail for %s", role))
			return err
		}

		if !resp.IsHealthy {
4366
			log.Warn("check health fail")
4367 4368 4369 4370 4371 4372 4373 4374 4375 4376 4377 4378 4379 4380 4381 4382 4383 4384 4385 4386 4387 4388 4389
			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{
4390 4391 4392
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_Success,
			},
4393 4394 4395 4396 4397
			IsHealthy: false,
			Reasons:   errReasons,
		}, nil
	}

4398
	states, reasons := node.multiRateLimiter.GetQuotaStates()
4399 4400 4401 4402 4403
	return &milvuspb.CheckHealthResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
			Reason:    "",
		},
4404 4405 4406
		QuotaStates: states,
		Reasons:     reasons,
		IsHealthy:   true,
4407
	}, nil
4408
}