impl.go 155.0 KB
Newer Older
1 2 3 4 5 6
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
7 8
// with the License. You may obtain a copy of the License at
//
9
//     http://www.apache.org/licenses/LICENSE-2.0
10
//
11 12 13 14 15
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
16

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	log := log.Ctx(ctx)

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

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

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

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

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

// 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 已提交
3588 3589
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Import")
	defer sp.End()
3590 3591 3592

	log := log.Ctx(ctx)

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

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

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

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

E
Enwei Jiao 已提交
3633 3634
	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()))
3635
	return respFromRC, nil
G
groot 已提交
3636 3637
}

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

	log := log.Ctx(ctx)

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

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

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

// 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 已提交
3677 3678
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-ListImportTasks")
	defer sp.End()
3679 3680 3681

	log := log.Ctx(ctx)

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

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

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

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

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

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

	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 已提交
3738 3739
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-UpdateCredentialCache")
	defer sp.End()
3740 3741

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	log := log.Ctx(ctx)

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

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

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

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

	log := log.Ctx(ctx)

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

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

	log := log.Ctx(ctx)

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

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

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

	log := log.Ctx(ctx)

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

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

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

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

	log := log.Ctx(ctx)

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

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

	result, err := node.rootCoord.SelectUser(ctx, req)
	if err != nil {
4135 4136
		log.Error("fail to select user",
			zap.Error(err))
4137 4138 4139 4140 4141
		return &milvuspb.SelectUserResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
4142
		}, nil
4143 4144
	}
	return result, nil
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 4176
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
4177 4178
}

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

	log := log.Ctx(ctx)

	log.Debug("OperatePrivilege",
		zap.Any("req", req))
4187
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4188
		return errorutil.UnhealthyStatus(code), nil
4189 4190 4191 4192 4193
	}
	if err := node.validPrivilegeParams(req); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4194
		}, nil
4195 4196 4197 4198 4199 4200
	}
	curUser, err := GetCurUserFromContext(ctx)
	if err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4201
		}, nil
4202 4203 4204 4205
	}
	req.Entity.Grantor.User = &milvuspb.UserEntity{Name: curUser}
	result, err := node.rootCoord.OperatePrivilege(ctx, req)
	if err != nil {
4206 4207
		log.Error("fail to operate privilege",
			zap.Error(err))
4208 4209 4210
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4211
		}, nil
4212 4213
	}
	return result, nil
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 4242
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 已提交
4243 4244
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-SelectGrant")
	defer sp.End()
4245 4246 4247 4248 4249

	log := log.Ctx(ctx)

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

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

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

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

	log := log.Ctx(ctx)

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

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

// 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
	}
4326
	node.multiRateLimiter.SetQuotaStates(request.GetStates(), request.GetCodes())
4327 4328 4329
	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() {
4330
			log.Warn("Proxy set quota states", zap.String("state", request.GetStates()[i].String()), zap.String("reason", request.GetCodes()[i].String()))
4331 4332
		}
	}
4333 4334 4335
	resp.ErrorCode = commonpb.ErrorCode_Success
	return resp, nil
}
4336 4337 4338 4339

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

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

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

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

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

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

J
jaime 已提交
4411 4412 4413 4414 4415 4416 4417 4418 4419 4420 4421 4422 4423 4424 4425 4426 4427 4428 4429 4430 4431
func (node *Proxy) RenameCollection(ctx context.Context, req *milvuspb.RenameCollectionRequest) (*commonpb.Status, error) {
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-RefreshPolicyInfoCache")
	defer sp.End()

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

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

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

	if err := validateCollectionName(req.GetNewName()); err != nil {
		log.Warn("validate new collection name fail", zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalCollectionName,
			Reason:    err.Error(),
J
jaime 已提交
4432
		}, nil
J
jaime 已提交
4433 4434 4435 4436 4437 4438 4439 4440 4441 4442 4443 4444 4445 4446 4447 4448 4449 4450 4451
	}

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

	return resp, nil
}

W
wei liu 已提交
4452
func (node *Proxy) CreateResourceGroup(ctx context.Context, request *milvuspb.CreateResourceGroupRequest) (*commonpb.Status, error) {
W
wei liu 已提交
4453 4454 4455 4456 4457 4458 4459 4460 4461 4462 4463 4464 4465 4466 4467 4468 4469 4470 4471 4472 4473 4474 4475 4476 4477 4478 4479 4480 4481 4482 4483 4484 4485 4486 4487 4488 4489 4490 4491 4492 4493 4494 4495 4496 4497 4498 4499 4500 4501 4502 4503 4504 4505 4506 4507 4508 4509 4510 4511 4512
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

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

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

	log.Debug("CreateResourceGroup received")

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

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

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

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

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

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

func (node *Proxy) DropResourceGroup(ctx context.Context, request *milvuspb.DropResourceGroupRequest) (*commonpb.Status, error) {
W
wei liu 已提交
4516 4517 4518 4519 4520 4521 4522 4523 4524 4525 4526 4527 4528 4529 4530 4531 4532 4533 4534 4535 4536 4537 4538 4539 4540 4541 4542 4543 4544 4545 4546 4547 4548 4549 4550 4551 4552 4553 4554 4555 4556 4557 4558 4559 4560 4561 4562 4563 4564 4565 4566 4567 4568 4569 4570 4571 4572 4573 4574 4575
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

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

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

	log.Debug("DropResourceGroup received")

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

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

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

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

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

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

func (node *Proxy) TransferNode(ctx context.Context, request *milvuspb.TransferNodeRequest) (*commonpb.Status, error) {
W
wei liu 已提交
4579 4580 4581 4582 4583 4584 4585 4586 4587 4588 4589 4590 4591 4592 4593 4594 4595 4596 4597 4598 4599 4600 4601 4602 4603 4604 4605 4606 4607 4608 4609 4610 4611 4612 4613 4614 4615 4616 4617 4618 4619 4620 4621 4622 4623 4624 4625 4626 4627 4628 4629 4630 4631 4632 4633 4634 4635 4636 4637 4638
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

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

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

	log.Debug("TransferNode received")

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

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

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

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

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

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

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

W
wei liu 已提交
4646 4647 4648 4649 4650 4651 4652 4653 4654 4655 4656 4657 4658 4659 4660 4661 4662 4663 4664 4665 4666 4667 4668 4669 4670 4671 4672 4673 4674 4675 4676 4677 4678 4679 4680 4681 4682 4683 4684 4685 4686 4687 4688 4689 4690 4691 4692 4693 4694 4695 4696 4697 4698 4699 4700 4701
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-TransferReplica")
	defer sp.End()
	method := "TransferReplica"
	tr := timerecord.NewTimeRecorder(method)
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
	t := &TransferReplicaTask{
		ctx:                    ctx,
		Condition:              NewTaskCondition(ctx),
		TransferReplicaRequest: request,
		queryCoord:             node.queryCoord,
	}

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

	log.Debug("TransferReplica received")

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

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

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

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

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

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

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

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

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

	log.Debug("ListResourceGroups received")

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

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

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

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

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

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

func (node *Proxy) DescribeResourceGroup(ctx context.Context, request *milvuspb.DescribeResourceGroupRequest) (*milvuspb.DescribeResourceGroupResponse, error) {
W
wei liu 已提交
4774 4775 4776 4777 4778 4779 4780 4781 4782 4783 4784 4785 4786 4787 4788 4789 4790 4791 4792 4793 4794 4795 4796 4797 4798 4799 4800 4801 4802 4803 4804 4805 4806 4807 4808 4809 4810 4811 4812 4813 4814 4815 4816 4817 4818 4819 4820 4821 4822 4823 4824 4825 4826 4827 4828 4829 4830 4831 4832 4833 4834 4835 4836 4837 4838 4839
	if !node.checkHealthy() {
		return &milvuspb.DescribeResourceGroupResponse{
			Status: unhealthyStatus(),
		}, nil
	}

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

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

	log.Debug("DescribeResourceGroup received")

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

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

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

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

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

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