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

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

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

	"golang.org/x/sync/errgroup"
27

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

53 54
const moduleName = "Proxy"

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

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

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

101
// InvalidateCollectionMetaCache invalidate the meta cache of specific collection.
C
Cai Yudong 已提交
102
func (node *Proxy) InvalidateCollectionMetaCache(ctx context.Context, request *proxypb.InvalidateCollMetaCacheRequest) (*commonpb.Status, error) {
103 104 105
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
106
	ctx = logutil.WithModule(ctx, moduleName)
107 108 109
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-InvalidateCollectionMetaCache")
	defer sp.Finish()
	log := log.Ctx(ctx).With(
110
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
111
		zap.String("db", request.DbName),
112 113
		zap.String("collectionName", request.CollectionName),
		zap.Int64("collectionID", request.CollectionID))
D
dragondriver 已提交
114

115 116
	log.Info("received request to invalidate collection meta cache")

117
	collectionName := request.CollectionName
118
	collectionID := request.CollectionID
X
Xiaofan 已提交
119 120

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

140
	return &commonpb.Status{
141
		ErrorCode: commonpb.ErrorCode_Success,
142 143
		Reason:    "",
	}, nil
144 145
}

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

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreateCollection")
	defer sp.Finish()
155 156 157
	method := "CreateCollection"
	tr := timerecord.NewTimeRecorder(method)

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

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

167 168 169
	// avoid data race
	lenOfSchema := len(request.Schema)

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

178 179
	log.Debug(rpcReceived(method))

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

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

192 193
	log.Debug(
		rpcEnqueued(method),
194 195
		zap.Uint64("BeginTs", cct.BeginTs()),
		zap.Uint64("EndTs", cct.EndTs()),
196
		zap.Uint64("timestamp", request.Base.Timestamp))
197

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

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

212 213
	log.Debug(
		rpcDone(method),
214
		zap.Uint64("BeginTs", cct.BeginTs()),
215
		zap.Uint64("EndTs", cct.EndTs()))
216

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

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

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DropCollection")
	defer sp.Finish()
230 231
	method := "DropCollection"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
232
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
233

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

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

248 249
	log.Debug("DropCollection received")

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

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

261 262
	log.Debug("DropCollection enqueued",
		zap.Uint64("BeginTs", dct.BeginTs()),
263
		zap.Uint64("EndTs", dct.EndTs()))
264 265 266

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

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

278 279
	log.Debug("DropCollection done",
		zap.Uint64("BeginTs", dct.BeginTs()),
280
		zap.Uint64("EndTs", dct.EndTs()))
281

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

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

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-HasCollection")
	defer sp.Finish()
297 298
	method := "HasCollection"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
299
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
300
		metrics.TotalLabel).Inc()
301

302
	log := log.Ctx(ctx).With(
303
		zap.String("role", typeutil.ProxyRole),
304 305 306
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

307 308
	log.Debug("HasCollection received")

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

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

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

330 331
	log.Debug("HasCollection enqueued",
		zap.Uint64("BeginTS", hct.BeginTs()),
332
		zap.Uint64("EndTS", hct.EndTs()))
333 334 335

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

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

350 351
	log.Debug("HasCollection done",
		zap.Uint64("BeginTS", hct.BeginTs()),
352
		zap.Uint64("EndTS", hct.EndTs()))
353

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

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

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-LoadCollection")
	defer sp.Finish()
368 369
	method := "LoadCollection"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
370
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
371
		metrics.TotalLabel).Inc()
372
	lct := &loadCollectionTask{
S
sunby 已提交
373
		ctx:                   ctx,
374 375
		Condition:             NewTaskCondition(ctx),
		LoadCollectionRequest: request,
376
		queryCoord:            node.queryCoord,
C
cai.zhang 已提交
377
		indexCoord:            node.indexCoord,
378 379
	}

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

385 386
	log.Debug("LoadCollection received")

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

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

399 400
	log.Debug("LoadCollection enqueued",
		zap.Uint64("BeginTS", lct.BeginTs()),
401
		zap.Uint64("EndTS", lct.EndTs()))
402 403 404

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

416 417
	log.Debug("LoadCollection done",
		zap.Uint64("BeginTS", lct.BeginTs()),
418
		zap.Uint64("EndTS", lct.EndTs()))
419

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

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

432
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-ReleaseCollection")
433
	defer sp.Finish()
434 435
	method := "ReleaseCollection"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
436
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
437
		metrics.TotalLabel).Inc()
438
	rct := &releaseCollectionTask{
S
sunby 已提交
439
		ctx:                      ctx,
440 441
		Condition:                NewTaskCondition(ctx),
		ReleaseCollectionRequest: request,
442
		queryCoord:               node.queryCoord,
443
		chMgr:                    node.chMgr,
444 445
	}

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

451 452
	log.Debug(rpcReceived(method))

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

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

466 467
	log.Debug(
		rpcEnqueued(method),
468
		zap.Uint64("BeginTS", rct.BeginTs()),
469
		zap.Uint64("EndTS", rct.EndTs()))
470 471

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

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

486 487
	log.Debug(
		rpcDone(method),
488
		zap.Uint64("BeginTS", rct.BeginTs()),
489
		zap.Uint64("EndTS", rct.EndTs()))
490

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

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

505
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DescribeCollection")
506
	defer sp.Finish()
507 508
	method := "DescribeCollection"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
509
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
510
		metrics.TotalLabel).Inc()
511

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

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

524 525
	log.Debug("DescribeCollection received")

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

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

540 541
	log.Debug("DescribeCollection enqueued",
		zap.Uint64("BeginTS", dct.BeginTs()),
542
		zap.Uint64("EndTS", dct.EndTs()))
543 544 545

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

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

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

561 562
	log.Debug("DescribeCollection done",
		zap.Uint64("BeginTS", dct.BeginTs()),
563
		zap.Uint64("EndTS", dct.EndTs()))
564

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

571 572 573 574 575 576 577 578 579
// 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
	}

580
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetStatistics")
581 582 583
	defer sp.Finish()
	method := "GetStatistics"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
584
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
585
		metrics.TotalLabel).Inc()
586 587 588 589 590 591 592 593 594 595
	g := &getStatisticsTask{
		request:   request,
		Condition: NewTaskCondition(ctx),
		ctx:       ctx,
		tr:        tr,
		dc:        node.dataCoord,
		qc:        node.queryCoord,
		shardMgr:  node.shardMgr,
	}

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

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

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

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

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

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

684 685
	log.Debug(rpcReceived(method))

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

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

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

702 703
	log.Debug(
		rpcEnqueued(method),
704
		zap.Uint64("BeginTS", g.BeginTs()),
705
		zap.Uint64("EndTS", g.EndTs()))
706 707

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

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

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

725 726
	log.Debug(
		rpcDone(method),
727
		zap.Uint64("BeginTS", g.BeginTs()),
728
		zap.Uint64("EndTS", g.EndTs()))
729

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

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

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

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

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

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

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

781
	log.Debug("ShowCollections enqueued",
782
		zap.Any("CollectionNames", request.CollectionNames))
D
dragondriver 已提交
783

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

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

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

800
	log.Debug("ShowCollections Done",
801 802
		zap.Int("len(CollectionNames)", len(request.CollectionNames)),
		zap.Int("num_collections", len(sct.result.CollectionNames)))
803

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

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

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-AlterCollection")
	defer sp.Finish()
	method := "AlterCollection"
	tr := timerecord.NewTimeRecorder(method)

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

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

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

833 834 835
	log.Debug(
		rpcReceived(method))

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

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

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

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

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

E
Enwei Jiao 已提交
873 874
	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 已提交
875 876 877
	return act.result, nil
}

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

884
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreatePartition")
885
	defer sp.Finish()
886 887
	method := "CreatePartition"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
888
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
889

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

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

904 905
	log.Debug(rpcReceived("CreatePartition"))

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

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

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

919 920 921
	log.Debug(
		rpcEnqueued("CreatePartition"),
		zap.Uint64("BeginTS", cpt.BeginTs()),
922
		zap.Uint64("EndTS", cpt.EndTs()))
923 924 925 926

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

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

933
		return &commonpb.Status{
934
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
935 936 937
			Reason:    err.Error(),
		}, nil
	}
938 939 940 941

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

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

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

955
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DropPartition")
956
	defer sp.Finish()
957 958
	method := "DropPartition"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
959
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
960

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

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

976 977
	log.Debug(rpcReceived(method))

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

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

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

991 992 993
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", dpt.BeginTs()),
994
		zap.Uint64("EndTS", dpt.EndTs()))
995 996 997 998

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

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

1005
		return &commonpb.Status{
1006
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1007 1008 1009
			Reason:    err.Error(),
		}, nil
	}
1010 1011 1012 1013

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

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

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

D
dragondriver 已提交
1029
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-HasPartition")
D
dragondriver 已提交
1030
	defer sp.Finish()
1031 1032 1033
	method := "HasPartition"
	tr := timerecord.NewTimeRecorder(method)
	//TODO: use collectionID instead of collectionName
E
Enwei Jiao 已提交
1034
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1035
		metrics.TotalLabel).Inc()
D
dragondriver 已提交
1036

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

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

1051 1052
	log.Debug(rpcReceived(method))

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

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

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

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

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

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

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

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

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

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

D
dragondriver 已提交
1111
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-LoadPartitions")
1112
	defer sp.Finish()
1113 1114
	method := "LoadPartitions"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1115
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1116
		metrics.TotalLabel).Inc()
1117
	lpt := &loadPartitionsTask{
G
godchen 已提交
1118 1119 1120
		ctx:                   ctx,
		Condition:             NewTaskCondition(ctx),
		LoadPartitionsRequest: request,
1121
		queryCoord:            node.queryCoord,
C
cai.zhang 已提交
1122
		indexCoord:            node.indexCoord,
1123 1124
	}

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

1131 1132
	log.Debug(rpcReceived(method))

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

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

1141
		return &commonpb.Status{
1142
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1143 1144 1145 1146
			Reason:    err.Error(),
		}, nil
	}

1147 1148 1149
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", lpt.BeginTs()),
1150
		zap.Uint64("EndTS", lpt.EndTs()))
1151 1152 1153 1154

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

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

1162
		return &commonpb.Status{
1163
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1164 1165 1166 1167
			Reason:    err.Error(),
		}, nil
	}

1168 1169 1170
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", lpt.BeginTs()),
1171
		zap.Uint64("EndTS", lpt.EndTs()))
1172

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

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

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-ReleasePartitions")
	defer sp.Finish()

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

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

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

1206 1207
	log.Debug(rpcReceived(method))

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

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

1216
		return &commonpb.Status{
1217
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1218 1219 1220 1221
			Reason:    err.Error(),
		}, nil
	}

1222 1223 1224
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", rpt.BeginTs()),
1225
		zap.Uint64("EndTS", rpt.EndTs()))
1226 1227 1228 1229

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

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

1237
		return &commonpb.Status{
1238
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1239 1240 1241 1242
			Reason:    err.Error(),
		}, nil
	}

1243 1244 1245
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", rpt.BeginTs()),
1246
		zap.Uint64("EndTS", rpt.EndTs()))
1247

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

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

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetPartitionStatistics")
	defer sp.Finish()
1264 1265
	method := "GetPartitionStatistics"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1266
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1267
		metrics.TotalLabel).Inc()
1268

1269
	g := &getPartitionStatisticsTask{
1270 1271 1272
		ctx:                           ctx,
		Condition:                     NewTaskCondition(ctx),
		GetPartitionStatisticsRequest: request,
1273
		dataCoord:                     node.dataCoord,
1274 1275
	}

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

1282 1283
	log.Debug(rpcReceived(method))

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

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

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

1300 1301 1302
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTS", g.BeginTs()),
1303
		zap.Uint64("EndTS", g.EndTs()))
1304 1305 1306 1307

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

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

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

1323 1324 1325
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTS", g.BeginTs()),
1326
		zap.Uint64("EndTS", g.EndTs()))
1327

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

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

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-ShowPartitions")
	defer sp.Finish()

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

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

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

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

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

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

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

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

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

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

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

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

S
SimFG 已提交
1426 1427 1428 1429 1430 1431
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)
1432
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetLoadingProgress")
S
SimFG 已提交
1433
	defer sp.Finish()
E
Enwei Jiao 已提交
1434
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
1435 1436 1437
	log := log.Ctx(ctx)

	log.Debug(
S
SimFG 已提交
1438 1439 1440 1441
		rpcReceived(method),
		zap.Any("request", request))

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

	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
	}

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

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

1506
func (node *Proxy) GetLoadState(ctx context.Context, request *milvuspb.GetLoadStateRequest) (*milvuspb.GetLoadStateResponse, error) {
S
SimFG 已提交
1507 1508 1509 1510 1511 1512 1513 1514 1515 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 1541 1542 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 1584 1585 1586 1587 1588 1589 1590 1591 1592 1593 1594 1595
	if !node.checkHealthy() {
		return &milvuspb.GetLoadStateResponse{Status: unhealthyStatus()}, nil
	}
	method := "GetLoadState"
	tr := timerecord.NewTimeRecorder(method)
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetLoadState")
	defer sp.Finish()
	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
	}

	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 {
			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 {
			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
1596 1597
}

1598
// CreateIndex create index for collection.
C
Cai Yudong 已提交
1599
func (node *Proxy) CreateIndex(ctx context.Context, request *milvuspb.CreateIndexRequest) (*commonpb.Status, error) {
1600 1601 1602
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
1603

1604
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreateIndex")
D
dragondriver 已提交
1605 1606
	defer sp.Finish()

1607
	cit := &createIndexTask{
Z
zhenshan.cao 已提交
1608 1609 1610 1611 1612
		ctx:        ctx,
		Condition:  NewTaskCondition(ctx),
		req:        request,
		rootCoord:  node.rootCoord,
		indexCoord: node.indexCoord,
1613
		queryCoord: node.queryCoord,
1614 1615
	}

D
dragondriver 已提交
1616
	method := "CreateIndex"
1617
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1618
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1619
		metrics.TotalLabel).Inc()
1620 1621

	log := log.Ctx(ctx).With(
1622
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1623 1624 1625 1626
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.Any("extra_params", request.ExtraParams))
D
dragondriver 已提交
1627

1628 1629
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
1630 1631 1632
	if err := node.sched.ddQueue.Enqueue(cit); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1633
			zap.Error(err))
D
dragondriver 已提交
1634

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

1638
		return &commonpb.Status{
1639
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1640 1641 1642 1643
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
1644 1645 1646
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", cit.BeginTs()),
1647
		zap.Uint64("EndTs", cit.EndTs()))
D
dragondriver 已提交
1648 1649 1650 1651

	if err := cit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1652
			zap.Error(err),
D
dragondriver 已提交
1653
			zap.Uint64("BeginTs", cit.BeginTs()),
1654
			zap.Uint64("EndTs", cit.EndTs()))
D
dragondriver 已提交
1655

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

1659
		return &commonpb.Status{
1660
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1661 1662 1663 1664
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
1665 1666 1667
	log.Debug(
		rpcDone(method),
		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.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1672
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1673 1674 1675
	return cit.result, nil
}

1676
// DescribeIndex get the meta information of index, such as index state, index id and etc.
C
Cai Yudong 已提交
1677
func (node *Proxy) DescribeIndex(ctx context.Context, request *milvuspb.DescribeIndexRequest) (*milvuspb.DescribeIndexResponse, error) {
1678 1679 1680 1681 1682
	if !node.checkHealthy() {
		return &milvuspb.DescribeIndexResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1683 1684 1685 1686

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DescribeIndex")
	defer sp.Finish()

1687
	dit := &describeIndexTask{
S
sunby 已提交
1688
		ctx:                  ctx,
1689 1690
		Condition:            NewTaskCondition(ctx),
		DescribeIndexRequest: request,
1691
		indexCoord:           node.indexCoord,
1692 1693
	}

1694 1695
	method := "DescribeIndex"
	// avoid data race
1696
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1697
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1698
		metrics.TotalLabel).Inc()
1699 1700

	log := log.Ctx(ctx).With(
1701
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1702 1703 1704
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
1705 1706 1707
		zap.String("index name", request.IndexName))

	log.Debug(rpcReceived(method))
1708 1709 1710 1711

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

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

1717 1718
		return &milvuspb.DescribeIndexResponse{
			Status: &commonpb.Status{
1719
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1720 1721 1722 1723 1724
				Reason:    err.Error(),
			},
		}, nil
	}

1725 1726 1727
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", dit.BeginTs()),
1728
		zap.Uint64("EndTs", dit.EndTs()))
1729 1730 1731 1732

	if err := dit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1733
			zap.Error(err),
1734
			zap.Uint64("BeginTs", dit.BeginTs()),
1735
			zap.Uint64("EndTs", dit.EndTs()))
D
dragondriver 已提交
1736

Z
zhenshan.cao 已提交
1737 1738 1739 1740
		errCode := commonpb.ErrorCode_UnexpectedError
		if dit.result != nil {
			errCode = dit.result.Status.GetErrorCode()
		}
E
Enwei Jiao 已提交
1741
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1742
			metrics.FailLabel).Inc()
1743

1744 1745
		return &milvuspb.DescribeIndexResponse{
			Status: &commonpb.Status{
Z
zhenshan.cao 已提交
1746
				ErrorCode: errCode,
1747 1748 1749 1750 1751
				Reason:    err.Error(),
			},
		}, nil
	}

1752 1753 1754
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", dit.BeginTs()),
1755
		zap.Uint64("EndTs", dit.EndTs()))
1756

E
Enwei Jiao 已提交
1757
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1758
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1759
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1760 1761 1762
	return dit.result, nil
}

1763
// DropIndex drop the index of collection.
C
Cai Yudong 已提交
1764
func (node *Proxy) DropIndex(ctx context.Context, request *milvuspb.DropIndexRequest) (*commonpb.Status, error) {
1765 1766 1767
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
1768 1769 1770 1771

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DropIndex")
	defer sp.Finish()

1772
	dit := &dropIndexTask{
S
sunby 已提交
1773
		ctx:              ctx,
B
BossZou 已提交
1774 1775
		Condition:        NewTaskCondition(ctx),
		DropIndexRequest: request,
1776
		indexCoord:       node.indexCoord,
1777
		queryCoord:       node.queryCoord,
B
BossZou 已提交
1778
	}
G
godchen 已提交
1779

D
dragondriver 已提交
1780
	method := "DropIndex"
1781
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1782
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1783
		metrics.TotalLabel).Inc()
D
dragondriver 已提交
1784

1785
	log := log.Ctx(ctx).With(
1786
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1787 1788 1789 1790 1791
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))

1792 1793
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
1794 1795 1796
	if err := node.sched.ddQueue.Enqueue(dit); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1797
			zap.Error(err))
E
Enwei Jiao 已提交
1798
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1799
			metrics.AbandonLabel).Inc()
D
dragondriver 已提交
1800

B
BossZou 已提交
1801
		return &commonpb.Status{
1802
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
B
BossZou 已提交
1803 1804 1805
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
1806

D
dragondriver 已提交
1807 1808 1809
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", dit.BeginTs()),
1810
		zap.Uint64("EndTs", dit.EndTs()))
D
dragondriver 已提交
1811 1812 1813 1814

	if err := dit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1815
			zap.Error(err),
D
dragondriver 已提交
1816
			zap.Uint64("BeginTs", dit.BeginTs()),
1817
			zap.Uint64("EndTs", dit.EndTs()))
D
dragondriver 已提交
1818

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

B
BossZou 已提交
1822
		return &commonpb.Status{
1823
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
B
BossZou 已提交
1824 1825 1826
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
1827 1828 1829 1830

	log.Debug(
		rpcDone(method),
		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.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1835
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
B
BossZou 已提交
1836 1837 1838
	return dit.result, nil
}

1839 1840
// 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.
1841
// Deprecated: use DescribeIndex instead
C
Cai Yudong 已提交
1842
func (node *Proxy) GetIndexBuildProgress(ctx context.Context, request *milvuspb.GetIndexBuildProgressRequest) (*milvuspb.GetIndexBuildProgressResponse, error) {
1843 1844 1845 1846 1847
	if !node.checkHealthy() {
		return &milvuspb.GetIndexBuildProgressResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1848 1849 1850 1851

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetIndexBuildProgress")
	defer sp.Finish()

1852
	gibpt := &getIndexBuildProgressTask{
1853 1854 1855
		ctx:                          ctx,
		Condition:                    NewTaskCondition(ctx),
		GetIndexBuildProgressRequest: request,
1856 1857
		indexCoord:                   node.indexCoord,
		rootCoord:                    node.rootCoord,
1858
		dataCoord:                    node.dataCoord,
1859 1860
	}

1861
	method := "GetIndexBuildProgress"
1862
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1863
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1864
		metrics.TotalLabel).Inc()
1865 1866

	log := log.Ctx(ctx).With(
1867
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1868 1869 1870 1871
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
1872

1873 1874
	log.Debug(rpcReceived(method))

1875 1876 1877
	if err := node.sched.ddQueue.Enqueue(gibpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1878
			zap.Error(err))
E
Enwei Jiao 已提交
1879
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1880
			metrics.AbandonLabel).Inc()
1881

1882 1883 1884 1885 1886 1887 1888 1889
		return &milvuspb.GetIndexBuildProgressResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

1890 1891 1892
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", gibpt.BeginTs()),
1893
		zap.Uint64("EndTs", gibpt.EndTs()))
1894 1895 1896 1897

	if err := gibpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
1898
			zap.Error(err),
1899
			zap.Uint64("BeginTs", gibpt.BeginTs()),
1900
			zap.Uint64("EndTs", gibpt.EndTs()))
E
Enwei Jiao 已提交
1901
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1902
			metrics.FailLabel).Inc()
1903 1904 1905 1906 1907 1908 1909 1910

		return &milvuspb.GetIndexBuildProgressResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}
1911 1912 1913 1914

	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", gibpt.BeginTs()),
1915
		zap.Uint64("EndTs", gibpt.EndTs()))
1916

E
Enwei Jiao 已提交
1917
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1918
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
1919
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1920
	return gibpt.result, nil
1921 1922
}

1923
// GetIndexState get the build-state of index.
1924
// Deprecated: use DescribeIndex instead
C
Cai Yudong 已提交
1925
func (node *Proxy) GetIndexState(ctx context.Context, request *milvuspb.GetIndexStateRequest) (*milvuspb.GetIndexStateResponse, error) {
1926 1927 1928 1929 1930
	if !node.checkHealthy() {
		return &milvuspb.GetIndexStateResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1931 1932 1933 1934

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Insert")
	defer sp.Finish()

1935
	dipt := &getIndexStateTask{
G
godchen 已提交
1936 1937 1938
		ctx:                  ctx,
		Condition:            NewTaskCondition(ctx),
		GetIndexStateRequest: request,
1939 1940
		indexCoord:           node.indexCoord,
		rootCoord:            node.rootCoord,
1941 1942
	}

1943
	method := "GetIndexState"
1944
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
1945
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1946
		metrics.TotalLabel).Inc()
1947 1948

	log := log.Ctx(ctx).With(
1949
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1950 1951 1952 1953
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
1954

1955 1956
	log.Debug(rpcReceived(method))

1957 1958 1959
	if err := node.sched.ddQueue.Enqueue(dipt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
1960
			zap.Error(err))
1961

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

G
godchen 已提交
1965
		return &milvuspb.GetIndexStateResponse{
1966
			Status: &commonpb.Status{
1967
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1968 1969 1970 1971 1972
				Reason:    err.Error(),
			},
		}, nil
	}

1973 1974 1975
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", dipt.BeginTs()),
1976
		zap.Uint64("EndTs", dipt.EndTs()))
1977 1978 1979 1980

	if err := dipt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1981
			zap.Error(err),
1982
			zap.Uint64("BeginTs", dipt.BeginTs()),
1983
			zap.Uint64("EndTs", dipt.EndTs()))
E
Enwei Jiao 已提交
1984
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
1985
			metrics.FailLabel).Inc()
1986

G
godchen 已提交
1987
		return &milvuspb.GetIndexStateResponse{
1988
			Status: &commonpb.Status{
1989
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1990 1991 1992 1993 1994
				Reason:    err.Error(),
			},
		}, nil
	}

1995 1996 1997
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", dipt.BeginTs()),
1998
		zap.Uint64("EndTs", dipt.EndTs()))
1999

E
Enwei Jiao 已提交
2000
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2001
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
2002
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
2003 2004 2005
	return dipt.result, nil
}

2006
// Insert insert records into collection.
C
Cai Yudong 已提交
2007
func (node *Proxy) Insert(ctx context.Context, request *milvuspb.InsertRequest) (*milvuspb.MutationResult, error) {
X
Xiangyu Wang 已提交
2008 2009
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Insert")
	defer sp.Finish()
2010 2011 2012
	log := log.Ctx(ctx)
	log.Debug("Start processing insert request in Proxy")
	defer log.Debug("Finish processing insert request in Proxy")
X
Xiangyu Wang 已提交
2013

2014 2015 2016 2017 2018
	if !node.checkHealthy() {
		return &milvuspb.MutationResult{
			Status: unhealthyStatus(),
		}, nil
	}
2019 2020
	method := "Insert"
	tr := timerecord.NewTimeRecorder(method)
2021
	receiveSize := proto.Size(request)
2022
	rateCol.Add(internalpb.RateType_DMLInsert.String(), float64(receiveSize))
E
Enwei Jiao 已提交
2023
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.InsertLabel).Add(float64(receiveSize))
D
dragondriver 已提交
2024

E
Enwei Jiao 已提交
2025
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
2026
	it := &insertTask{
2027 2028
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),
X
xige-16 已提交
2029
		// req:       request,
2030
		insertMsg: &msgstream.InsertMsg{
2031 2032 2033
			BaseMsg: msgstream.BaseMsg{
				HashValues: request.HashKeys,
			},
G
godchen 已提交
2034
			InsertRequest: internalpb.InsertRequest{
2035 2036 2037
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Insert),
					commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
2038
					commonpbutil.WithSourceID(paramtable.GetNodeID()),
2039
				),
2040 2041
				CollectionName: request.CollectionName,
				PartitionName:  request.PartitionName,
X
xige-16 已提交
2042 2043 2044
				FieldsData:     request.FieldsData,
				NumRows:        uint64(request.NumRows),
				Version:        internalpb.InsertDataVersion_ColumnBased,
2045
				// RowData: transfer column based request to this
2046 2047
			},
		},
2048
		idAllocator:   node.rowIDAllocator,
2049 2050 2051
		segIDAssigner: node.segAssigner,
		chMgr:         node.chMgr,
		chTicker:      node.chTicker,
2052
	}
2053

2054 2055
	if len(it.insertMsg.PartitionName) <= 0 {
		it.insertMsg.PartitionName = Params.CommonCfg.DefaultPartitionName.GetValue()
2056 2057
	}

X
Xiangyu Wang 已提交
2058
	constructFailedResponse := func(err error) *milvuspb.MutationResult {
X
xige-16 已提交
2059
		numRows := request.NumRows
2060 2061 2062 2063
		errIndex := make([]uint32, numRows)
		for i := uint32(0); i < numRows; i++ {
			errIndex[i] = i
		}
2064

X
Xiangyu Wang 已提交
2065 2066 2067 2068 2069 2070 2071
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
			ErrIndex: errIndex,
		}
2072 2073
	}

X
Xiangyu Wang 已提交
2074
	log.Debug("Enqueue insert request in Proxy",
2075
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2076 2077 2078 2079 2080
		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)),
2081
		zap.Uint32("NumRows", request.NumRows))
D
dragondriver 已提交
2082

X
Xiangyu Wang 已提交
2083
	if err := node.sched.dmQueue.Enqueue(it); err != nil {
J
Jiquan Long 已提交
2084
		log.Warn("Failed to enqueue insert task: " + err.Error())
E
Enwei Jiao 已提交
2085
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2086
			metrics.AbandonLabel).Inc()
X
Xiangyu Wang 已提交
2087
		return constructFailedResponse(err), nil
2088
	}
D
dragondriver 已提交
2089

X
Xiangyu Wang 已提交
2090
	log.Debug("Detail of insert request in Proxy",
2091
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2092 2093 2094 2095 2096
		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),
2097
		zap.Uint32("NumRows", request.NumRows))
X
Xiangyu Wang 已提交
2098 2099

	if err := it.WaitToFinish(); err != nil {
2100
		log.Warn("Failed to execute insert task in task scheduler: " + err.Error())
E
Enwei Jiao 已提交
2101
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2102
			metrics.FailLabel).Inc()
X
Xiangyu Wang 已提交
2103 2104 2105 2106 2107
		return constructFailedResponse(err), nil
	}

	if it.result.Status.ErrorCode != commonpb.ErrorCode_Success {
		setErrorIndex := func() {
X
xige-16 已提交
2108
			numRows := request.NumRows
X
Xiangyu Wang 已提交
2109 2110 2111 2112 2113 2114 2115 2116 2117 2118 2119
			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 已提交
2120
	it.result.InsertCnt = int64(request.NumRows)
D
dragondriver 已提交
2121

E
Enwei Jiao 已提交
2122
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2123
		metrics.SuccessLabel).Inc()
2124
	successCnt := it.result.InsertCnt - int64(len(it.result.ErrIndex))
E
Enwei Jiao 已提交
2125 2126 2127
	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()))
2128 2129 2130
	return it.result, nil
}

2131
// Delete delete records from collection, then these records cannot be searched.
G
groot 已提交
2132
func (node *Proxy) Delete(ctx context.Context, request *milvuspb.DeleteRequest) (*milvuspb.MutationResult, error) {
2133 2134
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Delete")
	defer sp.Finish()
2135 2136 2137
	log := log.Ctx(ctx)
	log.Debug("Start processing delete request in Proxy")
	defer log.Debug("Finish processing delete request in Proxy")
2138

2139
	receiveSize := proto.Size(request)
2140
	rateCol.Add(internalpb.RateType_DMLDelete.String(), float64(receiveSize))
E
Enwei Jiao 已提交
2141
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.DeleteLabel).Add(float64(receiveSize))
2142

G
groot 已提交
2143 2144 2145 2146 2147 2148
	if !node.checkHealthy() {
		return &milvuspb.MutationResult{
			Status: unhealthyStatus(),
		}, nil
	}

2149 2150 2151
	method := "Delete"
	tr := timerecord.NewTimeRecorder(method)

E
Enwei Jiao 已提交
2152
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2153
		metrics.TotalLabel).Inc()
2154
	dt := &deleteTask{
X
xige-16 已提交
2155 2156 2157
		ctx:        ctx,
		Condition:  NewTaskCondition(ctx),
		deleteExpr: request.Expr,
2158
		deleteMsg: &BaseDeleteTask{
G
godchen 已提交
2159 2160 2161
			BaseMsg: msgstream.BaseMsg{
				HashValues: request.HashKeys,
			},
G
godchen 已提交
2162
			DeleteRequest: internalpb.DeleteRequest{
2163 2164 2165 2166
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Delete),
					commonpbutil.WithMsgID(0),
				),
X
xige-16 已提交
2167
				DbName:         request.DbName,
G
godchen 已提交
2168 2169 2170
				CollectionName: request.CollectionName,
				PartitionName:  request.PartitionName,
				// RowData: transfer column based request to this
C
Cai Yudong 已提交
2171 2172 2173 2174
			},
		},
		chMgr:    node.chMgr,
		chTicker: node.chTicker,
G
groot 已提交
2175 2176
	}

2177
	log.Debug("Enqueue delete request in Proxy",
2178
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2179 2180 2181 2182
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
		zap.String("expr", request.Expr))
2183 2184 2185

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

G
groot 已提交
2190 2191 2192 2193 2194 2195 2196 2197
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

2198
	log.Debug("Detail of delete request in Proxy",
2199
		zap.String("role", typeutil.ProxyRole),
2200
		zap.Uint64("timestamp", dt.deleteMsg.Base.Timestamp),
G
groot 已提交
2201 2202 2203
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
2204
		zap.String("expr", request.Expr))
G
groot 已提交
2205

2206
	if err := dt.WaitToFinish(); err != nil {
2207
		log.Error("Failed to execute delete task in task scheduler: " + err.Error())
E
Enwei Jiao 已提交
2208
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2209
			metrics.FailLabel).Inc()
G
groot 已提交
2210 2211 2212 2213 2214 2215 2216 2217
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

E
Enwei Jiao 已提交
2218
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2219
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
2220 2221
	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 已提交
2222 2223 2224
	return dt.result, nil
}

2225
// Search search the most similar records of requests.
C
Cai Yudong 已提交
2226
func (node *Proxy) Search(ctx context.Context, request *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
2227
	receiveSize := proto.Size(request)
E
Enwei Jiao 已提交
2228
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.SearchLabel).Add(float64(receiveSize))
2229 2230 2231

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

2232 2233 2234 2235 2236
	if !node.checkHealthy() {
		return &milvuspb.SearchResults{
			Status: unhealthyStatus(),
		}, nil
	}
2237 2238
	method := "Search"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2239
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2240
		metrics.TotalLabel).Inc()
D
dragondriver 已提交
2241

C
cai.zhang 已提交
2242 2243
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Search")
	defer sp.Finish()
D
dragondriver 已提交
2244

2245
	qt := &searchTask{
S
sunby 已提交
2246
		ctx:       ctx,
2247
		Condition: NewTaskCondition(ctx),
G
godchen 已提交
2248
		SearchRequest: &internalpb.SearchRequest{
2249 2250
			Base: commonpbutil.NewMsgBase(
				commonpbutil.WithMsgType(commonpb.MsgType_Search),
E
Enwei Jiao 已提交
2251
				commonpbutil.WithSourceID(paramtable.GetNodeID()),
2252
			),
E
Enwei Jiao 已提交
2253
			ReqID: paramtable.GetNodeID(),
2254
		},
2255 2256 2257 2258
		request:  request,
		qc:       node.queryCoord,
		tr:       timerecord.NewTimeRecorder("search"),
		shardMgr: node.shardMgr,
2259 2260
	}

2261 2262 2263
	travelTs := request.TravelTimestamp
	guaranteeTs := request.GuaranteeTimestamp

2264
	log := log.Ctx(ctx).With(
2265
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2266 2267 2268 2269 2270
		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)),
2271 2272 2273 2274
		zap.Any("OutputFields", request.OutputFields),
		zap.Any("search_params", request.SearchParams),
		zap.Uint64("travel_timestamp", travelTs),
		zap.Uint64("guarantee_timestamp", guaranteeTs))
D
dragondriver 已提交
2275

2276 2277 2278
	log.Debug(
		rpcReceived(method))

2279
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
2280
		log.Warn(
2281
			rpcFailedToEnqueue(method),
2282
			zap.Error(err))
D
dragondriver 已提交
2283

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

2287 2288
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
2289
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2290 2291 2292 2293
				Reason:    err.Error(),
			},
		}, nil
	}
Z
Zach 已提交
2294
	tr.CtxRecord(ctx, "search request enqueue")
2295

2296
	log.Debug(
2297
		rpcEnqueued(method),
2298
		zap.Uint64("timestamp", qt.Base.Timestamp))
D
dragondriver 已提交
2299

2300
	if err := qt.WaitToFinish(); err != nil {
2301
		log.Warn(
2302
			rpcFailedToWaitToFinish(method),
2303
			zap.Error(err))
2304

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

2308 2309
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
2310
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2311 2312 2313 2314 2315
				Reason:    err.Error(),
			},
		}, nil
	}

Z
Zach 已提交
2316
	span := tr.CtxRecord(ctx, "wait search result")
E
Enwei Jiao 已提交
2317
	metrics.ProxyWaitForSearchResultLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2318
		metrics.SearchLabel).Observe(float64(span.Milliseconds()))
2319
	tr.CtxRecord(ctx, "wait search result")
2320
	log.Debug(rpcDone(method))
D
dragondriver 已提交
2321

E
Enwei Jiao 已提交
2322
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2323
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
2324
	metrics.ProxySearchVectors.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Add(float64(qt.result.GetResults().GetNumQueries()))
C
cai.zhang 已提交
2325
	searchDur := tr.ElapseSpan().Milliseconds()
E
Enwei Jiao 已提交
2326
	metrics.ProxySQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2327
		metrics.SearchLabel).Observe(float64(searchDur))
E
Enwei Jiao 已提交
2328
	metrics.ProxyCollectionSQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2329
		metrics.SearchLabel, request.CollectionName).Observe(float64(searchDur))
2330 2331
	if qt.result != nil {
		sentSize := proto.Size(qt.result)
E
Enwei Jiao 已提交
2332
		metrics.ProxyReadReqSendBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Add(float64(sentSize))
2333
		rateCol.Add(metricsinfo.ReadResultThroughput, float64(sentSize))
2334
	}
2335 2336 2337
	return qt.result, nil
}

2338
// Flush notify data nodes to persist the data of collection.
2339 2340 2341 2342 2343 2344 2345
func (node *Proxy) Flush(ctx context.Context, request *milvuspb.FlushRequest) (*milvuspb.FlushResponse, error) {
	resp := &milvuspb.FlushResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    "",
		},
	}
2346
	if !node.checkHealthy() {
2347 2348
		resp.Status.Reason = "proxy is not healthy"
		return resp, nil
2349
	}
D
dragondriver 已提交
2350 2351 2352 2353

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Flush")
	defer sp.Finish()

2354
	ft := &flushTask{
T
ThreadDao 已提交
2355 2356 2357
		ctx:          ctx,
		Condition:    NewTaskCondition(ctx),
		FlushRequest: request,
2358
		dataCoord:    node.dataCoord,
2359 2360
	}

D
dragondriver 已提交
2361
	method := "Flush"
2362
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2363
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2364

2365
	log := log.Ctx(ctx).With(
2366
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2367 2368
		zap.String("db", request.DbName),
		zap.Any("collections", request.CollectionNames))
D
dragondriver 已提交
2369

2370 2371
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2372 2373 2374
	if err := node.sched.ddQueue.Enqueue(ft); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2375
			zap.Error(err))
D
dragondriver 已提交
2376

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

2379 2380
		resp.Status.Reason = err.Error()
		return resp, nil
2381 2382
	}

D
dragondriver 已提交
2383 2384 2385
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", ft.BeginTs()),
2386
		zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2387 2388 2389 2390

	if err := ft.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
2391
			zap.Error(err),
D
dragondriver 已提交
2392
			zap.Uint64("BeginTs", ft.BeginTs()),
2393
			zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2394

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

D
dragondriver 已提交
2397
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
2398 2399
		resp.Status.Reason = err.Error()
		return resp, nil
2400 2401
	}

D
dragondriver 已提交
2402 2403 2404
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", ft.BeginTs()),
2405
		zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2406

E
Enwei Jiao 已提交
2407 2408
	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()))
2409
	return ft.result, nil
2410 2411
}

2412
// Query get the records by primary keys.
C
Cai Yudong 已提交
2413
func (node *Proxy) Query(ctx context.Context, request *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) {
2414
	receiveSize := proto.Size(request)
E
Enwei Jiao 已提交
2415
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.QueryLabel).Add(float64(receiveSize))
2416 2417 2418

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

2419 2420 2421 2422 2423
	if !node.checkHealthy() {
		return &milvuspb.QueryResults{
			Status: unhealthyStatus(),
		}, nil
	}
2424

D
dragondriver 已提交
2425 2426
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Query")
	defer sp.Finish()
2427
	tr := timerecord.NewTimeRecorder("Query")
D
dragondriver 已提交
2428

2429
	qt := &queryTask{
2430 2431 2432
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),
		RetrieveRequest: &internalpb.RetrieveRequest{
2433 2434
			Base: commonpbutil.NewMsgBase(
				commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
E
Enwei Jiao 已提交
2435
				commonpbutil.WithSourceID(paramtable.GetNodeID()),
2436
			),
E
Enwei Jiao 已提交
2437
			ReqID: paramtable.GetNodeID(),
2438
		},
2439 2440
		request:          request,
		qc:               node.queryCoord,
2441
		queryShardPolicy: mergeRoundRobinPolicy,
2442
		shardMgr:         node.shardMgr,
2443 2444
	}

D
dragondriver 已提交
2445 2446
	method := "Query"

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

2450
	log := log.Ctx(ctx).With(
2451
		zap.String("role", typeutil.ProxyRole),
2452 2453
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
2454 2455 2456 2457
		zap.Strings("partitions", request.PartitionNames))

	log.Debug(
		rpcReceived(method),
2458 2459 2460 2461
		zap.String("expr", request.Expr),
		zap.Strings("OutputFields", request.OutputFields),
		zap.Uint64("travel_timestamp", request.TravelTimestamp),
		zap.Uint64("guarantee_timestamp", request.GuaranteeTimestamp))
G
godchen 已提交
2462

D
dragondriver 已提交
2463
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
2464
		log.Warn(
D
dragondriver 已提交
2465
			rpcFailedToEnqueue(method),
2466
			zap.Error(err))
D
dragondriver 已提交
2467

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

2471 2472 2473 2474 2475 2476
		return &milvuspb.QueryResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
2477
	}
Z
Zach 已提交
2478
	tr.CtxRecord(ctx, "query request enqueue")
2479

2480
	log.Debug(rpcEnqueued(method))
D
dragondriver 已提交
2481 2482

	if err := qt.WaitToFinish(); err != nil {
2483
		log.Warn(
D
dragondriver 已提交
2484
			rpcFailedToWaitToFinish(method),
2485
			zap.Error(err))
2486

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

2490 2491 2492 2493 2494 2495 2496
		return &milvuspb.QueryResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}
Z
Zach 已提交
2497
	span := tr.CtxRecord(ctx, "wait query result")
E
Enwei Jiao 已提交
2498
	metrics.ProxyWaitForSearchResultLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2499
		metrics.QueryLabel).Observe(float64(span.Milliseconds()))
2500

2501
	log.Debug(rpcDone(method))
D
dragondriver 已提交
2502

E
Enwei Jiao 已提交
2503
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2504 2505
		metrics.SuccessLabel).Inc()

E
Enwei Jiao 已提交
2506
	metrics.ProxySQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2507
		metrics.QueryLabel).Observe(float64(tr.ElapseSpan().Milliseconds()))
E
Enwei Jiao 已提交
2508
	metrics.ProxyCollectionSQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2509
		metrics.QueryLabel, request.CollectionName).Observe(float64(tr.ElapseSpan().Milliseconds()))
2510 2511

	ret := &milvuspb.QueryResults{
2512 2513
		Status:     qt.result.Status,
		FieldsData: qt.result.FieldsData,
2514 2515
	}
	sentSize := proto.Size(qt.result)
2516
	rateCol.Add(metricsinfo.ReadResultThroughput, float64(sentSize))
E
Enwei Jiao 已提交
2517
	metrics.ProxyReadReqSendBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Add(float64(sentSize))
2518
	return ret, nil
2519
}
2520

2521
// CreateAlias create alias for collection, then you can search the collection with alias.
Y
Yusup 已提交
2522 2523 2524 2525
func (node *Proxy) CreateAlias(ctx context.Context, request *milvuspb.CreateAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2526 2527 2528 2529

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreateAlias")
	defer sp.Finish()

Y
Yusup 已提交
2530 2531 2532 2533 2534 2535 2536
	cat := &CreateAliasTask{
		ctx:                ctx,
		Condition:          NewTaskCondition(ctx),
		CreateAliasRequest: request,
		rootCoord:          node.rootCoord,
	}

D
dragondriver 已提交
2537
	method := "CreateAlias"
2538
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2539
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2540

2541
	log := log.Ctx(ctx).With(
D
dragondriver 已提交
2542 2543 2544 2545 2546
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

2547 2548
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2549 2550 2551
	if err := node.sched.ddQueue.Enqueue(cat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2552
			zap.Error(err))
D
dragondriver 已提交
2553

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

Y
Yusup 已提交
2556 2557 2558 2559 2560 2561
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2562 2563 2564
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", cat.BeginTs()),
2565
		zap.Uint64("EndTs", cat.EndTs()))
D
dragondriver 已提交
2566 2567 2568 2569

	if err := cat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2570
			zap.Error(err),
D
dragondriver 已提交
2571
			zap.Uint64("BeginTs", cat.BeginTs()),
2572
			zap.Uint64("EndTs", cat.EndTs()))
E
Enwei Jiao 已提交
2573
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
Y
Yusup 已提交
2574 2575 2576 2577 2578 2579 2580

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

D
dragondriver 已提交
2581 2582 2583
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", cat.BeginTs()),
2584
		zap.Uint64("EndTs", cat.EndTs()))
D
dragondriver 已提交
2585

E
Enwei Jiao 已提交
2586 2587
	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 已提交
2588 2589 2590
	return cat.result, nil
}

2591
// DropAlias alter the alias of collection.
Y
Yusup 已提交
2592 2593 2594 2595
func (node *Proxy) DropAlias(ctx context.Context, request *milvuspb.DropAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2596 2597 2598 2599

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DropAlias")
	defer sp.Finish()

Y
Yusup 已提交
2600 2601 2602 2603 2604 2605 2606
	dat := &DropAliasTask{
		ctx:              ctx,
		Condition:        NewTaskCondition(ctx),
		DropAliasRequest: request,
		rootCoord:        node.rootCoord,
	}

D
dragondriver 已提交
2607
	method := "DropAlias"
2608
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2609
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2610

2611
	log := log.Ctx(ctx).With(
D
dragondriver 已提交
2612 2613 2614 2615
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias))

2616 2617
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2618 2619 2620
	if err := node.sched.ddQueue.Enqueue(dat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2621
			zap.Error(err))
E
Enwei Jiao 已提交
2622
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
D
dragondriver 已提交
2623

Y
Yusup 已提交
2624 2625 2626 2627 2628 2629
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2630 2631 2632
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", dat.BeginTs()),
2633
		zap.Uint64("EndTs", dat.EndTs()))
D
dragondriver 已提交
2634 2635 2636 2637

	if err := dat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2638
			zap.Error(err),
D
dragondriver 已提交
2639
			zap.Uint64("BeginTs", dat.BeginTs()),
2640
			zap.Uint64("EndTs", dat.EndTs()))
Y
Yusup 已提交
2641

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

Y
Yusup 已提交
2644 2645 2646 2647 2648 2649
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2650 2651 2652
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", dat.BeginTs()),
2653
		zap.Uint64("EndTs", dat.EndTs()))
D
dragondriver 已提交
2654

E
Enwei Jiao 已提交
2655 2656
	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 已提交
2657 2658 2659
	return dat.result, nil
}

2660
// AlterAlias alter alias of collection.
Y
Yusup 已提交
2661 2662 2663 2664
func (node *Proxy) AlterAlias(ctx context.Context, request *milvuspb.AlterAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2665 2666 2667 2668

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-AlterAlias")
	defer sp.Finish()

Y
Yusup 已提交
2669 2670 2671 2672 2673 2674 2675
	aat := &AlterAliasTask{
		ctx:               ctx,
		Condition:         NewTaskCondition(ctx),
		AlterAliasRequest: request,
		rootCoord:         node.rootCoord,
	}

D
dragondriver 已提交
2676
	method := "AlterAlias"
2677
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2678
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2679

2680
	log := log.Ctx(ctx).With(
D
dragondriver 已提交
2681 2682 2683 2684 2685
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

2686 2687
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2688 2689 2690
	if err := node.sched.ddQueue.Enqueue(aat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2691
			zap.Error(err))
E
Enwei Jiao 已提交
2692
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
D
dragondriver 已提交
2693

Y
Yusup 已提交
2694 2695 2696 2697 2698 2699
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2700 2701 2702
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", aat.BeginTs()),
2703
		zap.Uint64("EndTs", aat.EndTs()))
D
dragondriver 已提交
2704 2705 2706 2707

	if err := aat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2708
			zap.Error(err),
D
dragondriver 已提交
2709
			zap.Uint64("BeginTs", aat.BeginTs()),
2710
			zap.Uint64("EndTs", aat.EndTs()))
Y
Yusup 已提交
2711

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

Y
Yusup 已提交
2714 2715 2716 2717 2718 2719
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2720 2721 2722
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", aat.BeginTs()),
2723
		zap.Uint64("EndTs", aat.EndTs()))
D
dragondriver 已提交
2724

E
Enwei Jiao 已提交
2725 2726
	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 已提交
2727 2728 2729
	return aat.result, nil
}

2730
// CalcDistance calculates the distances between vectors.
2731
func (node *Proxy) CalcDistance(ctx context.Context, request *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) {
2732 2733 2734 2735 2736
	if !node.checkHealthy() {
		return &milvuspb.CalcDistanceResults{
			Status: unhealthyStatus(),
		}, nil
	}
2737

2738 2739 2740 2741
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CalcDistance")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

2742 2743
	query := func(ids *milvuspb.VectorIDs) (*milvuspb.QueryResults, error) {
		outputFields := []string{ids.FieldName}
2744

2745 2746 2747 2748 2749
		queryRequest := &milvuspb.QueryRequest{
			DbName:         "",
			CollectionName: ids.CollectionName,
			PartitionNames: ids.PartitionNames,
			OutputFields:   outputFields,
2750 2751
		}

2752
		qt := &queryTask{
2753 2754 2755
			ctx:       ctx,
			Condition: NewTaskCondition(ctx),
			RetrieveRequest: &internalpb.RetrieveRequest{
2756 2757
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
E
Enwei Jiao 已提交
2758
					commonpbutil.WithSourceID(paramtable.GetNodeID()),
2759
				),
E
Enwei Jiao 已提交
2760
				ReqID: paramtable.GetNodeID(),
2761
			},
2762 2763 2764 2765
			request: queryRequest,
			qc:      node.queryCoord,
			ids:     ids.IdArray,

2766
			queryShardPolicy: mergeRoundRobinPolicy,
2767
			shardMgr:         node.shardMgr,
2768 2769
		}

2770
		log := log.Ctx(ctx).With(
G
groot 已提交
2771 2772
			zap.String("collection", queryRequest.CollectionName),
			zap.Any("partitions", queryRequest.PartitionNames),
2773
			zap.Any("OutputFields", queryRequest.OutputFields))
G
groot 已提交
2774

2775
		err := node.sched.dqQueue.Enqueue(qt)
2776
		if err != nil {
2777 2778
			log.Error("CalcDistance queryTask failed to enqueue",
				zap.Error(err))
2779

2780 2781 2782 2783 2784
			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2785
			}, err
2786
		}
2787

2788
		log.Debug("CalcDistance queryTask enqueued")
2789 2790 2791

		err = qt.WaitToFinish()
		if err != nil {
2792 2793
			log.Error("CalcDistance queryTask failed to WaitToFinish",
				zap.Error(err))
2794 2795 2796 2797 2798 2799

			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2800
			}, err
2801
		}
2802

2803
		log.Debug("CalcDistance queryTask Done")
2804 2805

		return &milvuspb.QueryResults{
2806 2807
			Status:     qt.result.Status,
			FieldsData: qt.result.FieldsData,
2808 2809 2810
		}, nil
	}

G
groot 已提交
2811 2812 2813 2814
	// calcDistanceTask is not a standard task, no need to enqueue
	task := &calcDistanceTask{
		traceID:   traceID,
		queryFunc: query,
2815 2816
	}

G
groot 已提交
2817
	return task.Execute(ctx, request)
2818 2819
}

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

2825
// GetPersistentSegmentInfo get the information of sealed segment.
C
Cai Yudong 已提交
2826
func (node *Proxy) GetPersistentSegmentInfo(ctx context.Context, req *milvuspb.GetPersistentSegmentInfoRequest) (*milvuspb.GetPersistentSegmentInfoResponse, error) {
2827 2828 2829 2830 2831
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetPersistentSegmentInfo")
	defer sp.Finish()

	log := log.Ctx(ctx)

D
dragondriver 已提交
2832
	log.Debug("GetPersistentSegmentInfo",
2833
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2834 2835 2836
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
2837
	resp := &milvuspb.GetPersistentSegmentInfoResponse{
X
XuanYang-cn 已提交
2838
		Status: &commonpb.Status{
2839
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
X
XuanYang-cn 已提交
2840 2841
		},
	}
2842 2843 2844 2845
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
2846 2847
	method := "GetPersistentSegmentInfo"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2848
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2849
		metrics.TotalLabel).Inc()
2850 2851 2852

	// list segments
	collectionID, err := globalMetaCache.GetCollectionID(ctx, req.GetCollectionName())
X
XuanYang-cn 已提交
2853
	if err != nil {
E
Enwei Jiao 已提交
2854
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
2855 2856 2857 2858 2859 2860 2861 2862 2863 2864 2865
		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 已提交
2866
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
2867
		resp.Status.Reason = fmt.Errorf("getSegmentsOfCollection, err:%w", err).Error()
X
XuanYang-cn 已提交
2868 2869
		return resp, nil
	}
2870 2871

	// get Segment info
2872
	infoResp, err := node.dataCoord.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{
2873 2874 2875
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_SegmentInfo),
			commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
2876
			commonpbutil.WithSourceID(paramtable.GetNodeID()),
2877
		),
2878
		SegmentIDs: getSegmentsByStatesResponse.Segments,
X
XuanYang-cn 已提交
2879 2880
	})
	if err != nil {
E
Enwei Jiao 已提交
2881
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2882
			metrics.FailLabel).Inc()
2883 2884
		log.Warn("GetPersistentSegmentInfo fail",
			zap.Error(err))
2885
		resp.Status.Reason = fmt.Errorf("dataCoord:GetSegmentInfo, err:%w", err).Error()
X
XuanYang-cn 已提交
2886 2887
		return resp, nil
	}
2888 2889 2890
	log.Debug("GetPersistentSegmentInfo",
		zap.Int("len(infos)", len(infoResp.Infos)),
		zap.Any("status", infoResp.Status))
2891
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
E
Enwei Jiao 已提交
2892
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2893
			metrics.FailLabel).Inc()
X
XuanYang-cn 已提交
2894 2895 2896 2897 2898 2899
		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 已提交
2900
			SegmentID:    info.ID,
X
XuanYang-cn 已提交
2901 2902
			CollectionID: info.CollectionID,
			PartitionID:  info.PartitionID,
S
sunby 已提交
2903
			NumRows:      info.NumOfRows,
X
XuanYang-cn 已提交
2904 2905 2906
			State:        info.State,
		}
	}
E
Enwei Jiao 已提交
2907
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2908
		metrics.SuccessLabel).Inc()
E
Enwei Jiao 已提交
2909
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
2910
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
X
XuanYang-cn 已提交
2911 2912 2913 2914
	resp.Infos = persistentInfos
	return resp, nil
}

J
jingkl 已提交
2915
// GetQuerySegmentInfo gets segment information from QueryCoord.
C
Cai Yudong 已提交
2916
func (node *Proxy) GetQuerySegmentInfo(ctx context.Context, req *milvuspb.GetQuerySegmentInfoRequest) (*milvuspb.GetQuerySegmentInfoResponse, error) {
2917 2918 2919 2920 2921
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetQuerySegmentInfo")
	defer sp.Finish()

	log := log.Ctx(ctx)

D
dragondriver 已提交
2922
	log.Debug("GetQuerySegmentInfo",
2923
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2924 2925 2926
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
2927
	resp := &milvuspb.GetQuerySegmentInfoResponse{
Z
zhenshan.cao 已提交
2928
		Status: &commonpb.Status{
2929
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
Z
zhenshan.cao 已提交
2930 2931
		},
	}
2932 2933 2934 2935
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
2936

2937 2938
	method := "GetQuerySegmentInfo"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2939
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2940 2941
		metrics.TotalLabel).Inc()

2942 2943
	collID, err := globalMetaCache.GetCollectionID(ctx, req.CollectionName)
	if err != nil {
E
Enwei Jiao 已提交
2944
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
2945 2946 2947
		resp.Status.Reason = err.Error()
		return resp, nil
	}
2948
	infoResp, err := node.queryCoord.GetSegmentInfo(ctx, &querypb.GetSegmentInfoRequest{
2949 2950 2951
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_SegmentInfo),
			commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
2952
			commonpbutil.WithSourceID(paramtable.GetNodeID()),
2953
		),
2954
		CollectionID: collID,
Z
zhenshan.cao 已提交
2955 2956
	})
	if err != nil {
E
Enwei Jiao 已提交
2957
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
2958 2959
		log.Error("Failed to get segment info from QueryCoord",
			zap.Error(err))
Z
zhenshan.cao 已提交
2960 2961 2962
		resp.Status.Reason = err.Error()
		return resp, nil
	}
2963 2964 2965
	log.Debug("GetQuerySegmentInfo",
		zap.Any("infos", infoResp.Infos),
		zap.Any("status", infoResp.Status))
2966
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
E
Enwei Jiao 已提交
2967
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
2968 2969
		log.Error("Failed to get segment info from QueryCoord",
			zap.String("errMsg", infoResp.Status.Reason))
Z
zhenshan.cao 已提交
2970 2971 2972 2973 2974 2975 2976 2977 2978 2979 2980 2981 2982
		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 已提交
2983
			State:        info.SegmentState,
2984
			NodeIds:      info.NodeIds,
Z
zhenshan.cao 已提交
2985 2986
		}
	}
2987

E
Enwei Jiao 已提交
2988 2989
	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()))
2990
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
Z
zhenshan.cao 已提交
2991 2992 2993 2994
	resp.Infos = queryInfos
	return resp, nil
}

J
jingkl 已提交
2995
// Dummy handles dummy request
C
Cai Yudong 已提交
2996
func (node *Proxy) Dummy(ctx context.Context, req *milvuspb.DummyRequest) (*milvuspb.DummyResponse, error) {
2997 2998 2999 3000 3001 3002
	failedResponse := &milvuspb.DummyResponse{
		Response: `{"status": "fail"}`,
	}

	// TODO(wxyu): change name RequestType to Request
	drt, err := parseDummyRequestType(req.RequestType)
3003 3004 3005 3006 3007 3008

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Dummy")
	defer sp.Finish()

	log := log.Ctx(ctx)

3009
	if err != nil {
3010 3011
		log.Warn("Failed to parse dummy request type",
			zap.Error(err))
3012 3013 3014
		return failedResponse, nil
	}

3015 3016
	if drt.RequestType == "query" {
		drr, err := parseDummyQueryRequest(req.RequestType)
3017
		if err != nil {
3018 3019
			log.Warn("Failed to parse dummy query request",
				zap.Error(err))
3020 3021 3022
			return failedResponse, nil
		}

3023
		request := &milvuspb.QueryRequest{
3024 3025 3026
			DbName:         drr.DbName,
			CollectionName: drr.CollectionName,
			PartitionNames: drr.PartitionNames,
3027
			OutputFields:   drr.OutputFields,
X
Xiangyu Wang 已提交
3028 3029
		}

3030
		_, err = node.Query(ctx, request)
3031
		if err != nil {
3032 3033
			log.Warn("Failed to execute dummy query",
				zap.Error(err))
3034 3035
			return failedResponse, err
		}
X
Xiangyu Wang 已提交
3036 3037 3038 3039 3040 3041

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

3042 3043
	log.Debug("cannot find specify dummy request type")
	return failedResponse, nil
X
Xiangyu Wang 已提交
3044 3045
}

J
jingkl 已提交
3046
// RegisterLink registers a link
C
Cai Yudong 已提交
3047
func (node *Proxy) RegisterLink(ctx context.Context, req *milvuspb.RegisterLinkRequest) (*milvuspb.RegisterLinkResponse, error) {
3048
	code := node.stateCode.Load().(commonpb.StateCode)
3049 3050 3051 3052 3053

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-RegisterLink")
	defer sp.Finish()

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

3057 3058
	log.Debug("RegisterLink")

3059
	if code != commonpb.StateCode_Healthy {
3060 3061 3062
		return &milvuspb.RegisterLinkResponse{
			Address: nil,
			Status: &commonpb.Status{
3063
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
C
Cai Yudong 已提交
3064
				Reason:    "proxy not healthy",
3065 3066 3067
			},
		}, nil
	}
E
Enwei Jiao 已提交
3068
	//metrics.ProxyLinkedSDKs.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Inc()
3069 3070 3071
	return &milvuspb.RegisterLinkResponse{
		Address: nil,
		Status: &commonpb.Status{
3072
			ErrorCode: commonpb.ErrorCode_Success,
3073
			Reason:    os.Getenv(metricsinfo.DeployModeEnvKey),
3074 3075 3076
		},
	}, nil
}
3077

3078
// GetMetrics gets the metrics of proxy
3079 3080
// 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) {
3081 3082 3083 3084 3085
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetMetrics")
	defer sp.Finish()

	log := log.Ctx(ctx)

3086
	log.Debug("Proxy.GetMetrics",
E
Enwei Jiao 已提交
3087
		zap.Int64("node_id", paramtable.GetNodeID()),
3088 3089 3090 3091
		zap.String("req", req.Request))

	if !node.checkHealthy() {
		log.Warn("Proxy.GetMetrics failed",
E
Enwei Jiao 已提交
3092
			zap.Int64("node_id", paramtable.GetNodeID()),
3093
			zap.String("req", req.Request),
E
Enwei Jiao 已提交
3094
			zap.Error(errProxyIsUnhealthy(paramtable.GetNodeID())))
3095 3096 3097 3098

		return &milvuspb.GetMetricsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
E
Enwei Jiao 已提交
3099
				Reason:    msgProxyIsUnhealthy(paramtable.GetNodeID()),
3100 3101 3102 3103 3104 3105 3106 3107
			},
			Response: "",
		}, nil
	}

	metricType, err := metricsinfo.ParseMetricType(req.Request)
	if err != nil {
		log.Warn("Proxy.GetMetrics failed to parse metric type",
E
Enwei Jiao 已提交
3108
			zap.Int64("node_id", paramtable.GetNodeID()),
3109 3110 3111 3112 3113 3114 3115 3116 3117 3118 3119 3120 3121 3122 3123
			zap.String("req", req.Request),
			zap.Error(err))

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

	log.Debug("Proxy.GetMetrics",
		zap.String("metric_type", metricType))

3124 3125 3126
	req.Base = commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_SystemInfo),
		commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
3127
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
3128
	)
3129
	if metricType == metricsinfo.SystemInfoMetrics {
3130 3131 3132 3133 3134 3135 3136
		ret, err := node.metricsCacheManager.GetSystemInfoMetrics()
		if err == nil && ret != nil {
			return ret, nil
		}
		log.Debug("failed to get system info metrics from cache, recompute instead",
			zap.Error(err))

3137
		metrics, err := getSystemInfoMetrics(ctx, req, node)
3138 3139

		log.Debug("Proxy.GetMetrics",
E
Enwei Jiao 已提交
3140
			zap.Int64("node_id", paramtable.GetNodeID()),
3141 3142 3143 3144 3145
			zap.String("req", req.Request),
			zap.String("metric_type", metricType),
			zap.Any("metrics", metrics), // TODO(dragondriver): necessary? may be very large
			zap.Error(err))

3146 3147
		node.metricsCacheManager.UpdateSystemInfoMetrics(metrics)

G
godchen 已提交
3148
		return metrics, nil
3149 3150
	}

J
Jiquan Long 已提交
3151
	log.Warn("Proxy.GetMetrics failed, request metric type is not implemented yet",
E
Enwei Jiao 已提交
3152
		zap.Int64("node_id", paramtable.GetNodeID()),
3153 3154 3155 3156 3157 3158 3159 3160 3161 3162 3163 3164
		zap.String("req", req.Request),
		zap.String("metric_type", metricType))

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

3165 3166 3167
// 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) {
3168 3169 3170 3171 3172 3173 3174
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetProxyMetrics")
	defer sp.Finish()

	log := log.Ctx(ctx).With(
		zap.Int64("node_id", paramtable.GetNodeID()),
		zap.String("req", req.Request))

3175 3176
	if !node.checkHealthy() {
		log.Warn("Proxy.GetProxyMetrics failed",
E
Enwei Jiao 已提交
3177
			zap.Error(errProxyIsUnhealthy(paramtable.GetNodeID())))
3178 3179 3180 3181

		return &milvuspb.GetMetricsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
E
Enwei Jiao 已提交
3182
				Reason:    msgProxyIsUnhealthy(paramtable.GetNodeID()),
3183 3184 3185 3186 3187 3188 3189 3190 3191 3192 3193 3194 3195 3196 3197 3198 3199
			},
		}, 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
	}

3200 3201 3202
	req.Base = commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_SystemInfo),
		commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
3203
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
3204
	)
3205 3206 3207 3208 3209 3210 3211 3212 3213 3214 3215 3216 3217 3218 3219 3220

	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",
3221
			zap.String("metric_type", metricType))
3222 3223 3224 3225

		return proxyMetrics, nil
	}

J
Jiquan Long 已提交
3226
	log.Warn("Proxy.GetProxyMetrics failed, request metric type is not implemented yet",
3227 3228 3229 3230 3231 3232 3233 3234 3235 3236
		zap.String("metric_type", metricType))

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

B
bigsheeper 已提交
3237 3238
// LoadBalance would do a load balancing operation between query nodes
func (node *Proxy) LoadBalance(ctx context.Context, req *milvuspb.LoadBalanceRequest) (*commonpb.Status, error) {
3239 3240 3241 3242 3243
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-LoadBalance")
	defer sp.Finish()

	log := log.Ctx(ctx)

B
bigsheeper 已提交
3244
	log.Debug("Proxy.LoadBalance",
E
Enwei Jiao 已提交
3245
		zap.Int64("proxy_id", paramtable.GetNodeID()),
B
bigsheeper 已提交
3246 3247 3248 3249 3250 3251 3252 3253 3254
		zap.Any("req", req))

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

	status := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
	}
3255 3256 3257

	collectionID, err := globalMetaCache.GetCollectionID(ctx, req.GetCollectionName())
	if err != nil {
J
Jiquan Long 已提交
3258
		log.Warn("failed to get collection id",
3259 3260
			zap.String("collection name", req.GetCollectionName()),
			zap.Error(err))
3261 3262 3263
		status.Reason = err.Error()
		return status, nil
	}
B
bigsheeper 已提交
3264
	infoResp, err := node.queryCoord.LoadBalance(ctx, &querypb.LoadBalanceRequest{
3265 3266 3267
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_LoadBalanceSegments),
			commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
3268
			commonpbutil.WithSourceID(paramtable.GetNodeID()),
3269
		),
B
bigsheeper 已提交
3270 3271
		SourceNodeIDs:    []int64{req.SrcNodeID},
		DstNodeIDs:       req.DstNodeIDs,
X
xige-16 已提交
3272
		BalanceReason:    querypb.TriggerCondition_GrpcRequest,
B
bigsheeper 已提交
3273
		SealedSegmentIDs: req.SealedSegmentIDs,
3274
		CollectionID:     collectionID,
B
bigsheeper 已提交
3275 3276
	})
	if err != nil {
J
Jiquan Long 已提交
3277
		log.Warn("Failed to LoadBalance from Query Coordinator",
3278 3279
			zap.Any("req", req),
			zap.Error(err))
B
bigsheeper 已提交
3280 3281 3282 3283
		status.Reason = err.Error()
		return status, nil
	}
	if infoResp.ErrorCode != commonpb.ErrorCode_Success {
J
Jiquan Long 已提交
3284
		log.Warn("Failed to LoadBalance from Query Coordinator",
3285
			zap.String("errMsg", infoResp.Reason))
B
bigsheeper 已提交
3286 3287 3288
		status.Reason = infoResp.Reason
		return status, nil
	}
3289 3290 3291
	log.Debug("LoadBalance Done",
		zap.Any("req", req),
		zap.Any("status", infoResp))
B
bigsheeper 已提交
3292 3293 3294 3295
	status.ErrorCode = commonpb.ErrorCode_Success
	return status, nil
}

3296 3297
// GetReplicas gets replica info
func (node *Proxy) GetReplicas(ctx context.Context, req *milvuspb.GetReplicasRequest) (*milvuspb.GetReplicasResponse, error) {
3298 3299 3300 3301 3302 3303 3304 3305
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetReplicas")
	defer sp.Finish()

	log := log.Ctx(ctx)

	log.Debug("received get replicas request",
		zap.Int64("collection", req.GetCollectionID()),
		zap.Bool("with shard nodes", req.GetWithShardNodes()))
3306 3307 3308 3309 3310 3311
	resp := &milvuspb.GetReplicasResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

S
smellthemoon 已提交
3312 3313
	req.Base = commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_GetReplicas),
E
Enwei Jiao 已提交
3314
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
S
smellthemoon 已提交
3315
	)
3316 3317 3318

	resp, err := node.queryCoord.GetReplicas(ctx, req)
	if err != nil {
3319 3320
		log.Error("Failed to get replicas from Query Coordinator",
			zap.Error(err))
3321 3322 3323 3324
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3325 3326 3327
	log.Debug("received get replicas response",
		zap.Any("resp", resp),
		zap.Error(err))
3328 3329 3330
	return resp, nil
}

3331
// GetCompactionState gets the compaction state of multiple segments
3332
func (node *Proxy) GetCompactionState(ctx context.Context, req *milvuspb.GetCompactionStateRequest) (*milvuspb.GetCompactionStateResponse, error) {
3333 3334 3335 3336 3337 3338 3339
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetCompactionState")
	defer sp.Finish()

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

	log.Debug("received GetCompactionState request")
3340 3341 3342 3343 3344 3345 3346
	resp := &milvuspb.GetCompactionStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionState(ctx, req)
3347 3348 3349
	log.Debug("received GetCompactionState response",
		zap.Any("resp", resp),
		zap.Error(err))
3350 3351 3352
	return resp, err
}

3353
// ManualCompaction invokes compaction on specified collection
3354
func (node *Proxy) ManualCompaction(ctx context.Context, req *milvuspb.ManualCompactionRequest) (*milvuspb.ManualCompactionResponse, error) {
3355 3356 3357 3358 3359 3360 3361
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-ManualCompaction")
	defer sp.Finish()

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

	log.Info("received ManualCompaction request")
3362 3363 3364 3365 3366 3367 3368
	resp := &milvuspb.ManualCompactionResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.ManualCompaction(ctx, req)
3369 3370 3371
	log.Info("received ManualCompaction response",
		zap.Any("resp", resp),
		zap.Error(err))
3372 3373 3374
	return resp, err
}

3375
// GetCompactionStateWithPlans returns the compactions states with the given plan ID
3376
func (node *Proxy) GetCompactionStateWithPlans(ctx context.Context, req *milvuspb.GetCompactionPlansRequest) (*milvuspb.GetCompactionPlansResponse, error) {
3377 3378 3379 3380 3381 3382 3383
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetCompactionStateWithPlans")
	defer sp.Finish()

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

	log.Debug("received GetCompactionStateWithPlans request")
3384 3385 3386 3387 3388 3389 3390
	resp := &milvuspb.GetCompactionPlansResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionStateWithPlans(ctx, req)
3391 3392 3393
	log.Debug("received GetCompactionStateWithPlans response",
		zap.Any("resp", resp),
		zap.Error(err))
3394 3395 3396
	return resp, err
}

B
Bingyi Sun 已提交
3397 3398
// GetFlushState gets the flush state of multiple segments
func (node *Proxy) GetFlushState(ctx context.Context, req *milvuspb.GetFlushStateRequest) (*milvuspb.GetFlushStateResponse, error) {
3399 3400 3401 3402 3403 3404 3405
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetFlushState")
	defer sp.Finish()

	log := log.Ctx(ctx)

	log.Debug("received get flush state request",
		zap.Any("request", req))
3406
	var err error
B
Bingyi Sun 已提交
3407 3408 3409
	resp := &milvuspb.GetFlushStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
J
Jiquan Long 已提交
3410
		log.Warn("unable to get flush state because of closed server")
B
Bingyi Sun 已提交
3411 3412 3413
		return resp, nil
	}

3414
	resp, err = node.dataCoord.GetFlushState(ctx, req)
X
Xiaofan 已提交
3415
	if err != nil {
3416 3417
		log.Warn("failed to get flush state response",
			zap.Error(err))
X
Xiaofan 已提交
3418 3419
		return nil, err
	}
3420 3421
	log.Debug("received get flush state response",
		zap.Any("response", resp))
B
Bingyi Sun 已提交
3422 3423 3424
	return resp, err
}

C
Cai Yudong 已提交
3425 3426
// checkHealthy checks proxy state is Healthy
func (node *Proxy) checkHealthy() bool {
3427 3428
	code := node.stateCode.Load().(commonpb.StateCode)
	return code == commonpb.StateCode_Healthy
3429 3430
}

3431 3432 3433
func (node *Proxy) checkHealthyAndReturnCode() (commonpb.StateCode, bool) {
	code := node.stateCode.Load().(commonpb.StateCode)
	return code, code == commonpb.StateCode_Healthy
3434 3435
}

3436
// unhealthyStatus returns the proxy not healthy status
3437 3438 3439
func unhealthyStatus() *commonpb.Status {
	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
C
Cai Yudong 已提交
3440
		Reason:    "proxy not healthy",
3441 3442
	}
}
G
groot 已提交
3443 3444 3445

// 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) {
3446 3447 3448 3449 3450
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Import")
	defer sp.Finish()

	log := log.Ctx(ctx)

3451 3452
	log.Info("received import request",
		zap.String("collection name", req.GetCollectionName()),
G
groot 已提交
3453 3454
		zap.String("partition name", req.GetPartitionName()),
		zap.Strings("files", req.GetFiles()))
3455 3456 3457 3458 3459 3460
	resp := &milvuspb.ImportResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
			Reason:    "",
		},
	}
G
groot 已提交
3461 3462 3463 3464
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3465

3466 3467
	err := importutil.ValidateOptions(req.GetOptions())
	if err != nil {
3468 3469
		log.Error("failed to execute import request",
			zap.Error(err))
3470 3471 3472 3473 3474
		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
	}

3475 3476
	method := "Import"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
3477
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3478 3479
		metrics.TotalLabel).Inc()

3480
	// Call rootCoord to finish import.
3481 3482
	respFromRC, err := node.rootCoord.Import(ctx, req)
	if err != nil {
E
Enwei Jiao 已提交
3483
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3484 3485
		log.Error("failed to execute bulk insert request",
			zap.Error(err))
3486 3487 3488 3489
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3490

E
Enwei Jiao 已提交
3491 3492
	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()))
3493
	return respFromRC, nil
G
groot 已提交
3494 3495
}

3496
// GetImportState checks import task state from RootCoord.
G
groot 已提交
3497
func (node *Proxy) GetImportState(ctx context.Context, req *milvuspb.GetImportStateRequest) (*milvuspb.GetImportStateResponse, error) {
3498 3499 3500 3501 3502 3503 3504
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetImportState")
	defer sp.Finish()

	log := log.Ctx(ctx)

	log.Debug("received get import state request",
		zap.Int64("taskID", req.GetTask()))
G
groot 已提交
3505 3506 3507 3508 3509
	resp := &milvuspb.GetImportStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3510 3511
	method := "GetImportState"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
3512
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3513
		metrics.TotalLabel).Inc()
G
groot 已提交
3514 3515

	resp, err := node.rootCoord.GetImportState(ctx, req)
3516
	if err != nil {
E
Enwei Jiao 已提交
3517
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3518 3519
		log.Error("failed to execute get import state",
			zap.Error(err))
3520 3521 3522 3523 3524
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
		return resp, nil
	}

3525 3526 3527
	log.Debug("successfully received get import state response",
		zap.Int64("taskID", req.GetTask()),
		zap.Any("resp", resp), zap.Error(err))
E
Enwei Jiao 已提交
3528 3529
	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()))
3530
	return resp, nil
G
groot 已提交
3531 3532 3533 3534
}

// ListImportTasks get id array of all import tasks from rootcoord
func (node *Proxy) ListImportTasks(ctx context.Context, req *milvuspb.ListImportTasksRequest) (*milvuspb.ListImportTasksResponse, error) {
3535 3536 3537 3538 3539
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-ListImportTasks")
	defer sp.Finish()

	log := log.Ctx(ctx)

J
Jiquan Long 已提交
3540
	log.Debug("received list import tasks request")
G
groot 已提交
3541 3542 3543 3544 3545
	resp := &milvuspb.ListImportTasksResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3546 3547
	method := "ListImportTasks"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
3548
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3549
		metrics.TotalLabel).Inc()
G
groot 已提交
3550
	resp, err := node.rootCoord.ListImportTasks(ctx, req)
3551
	if err != nil {
E
Enwei Jiao 已提交
3552
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3553 3554
		log.Error("failed to execute list import tasks",
			zap.Error(err))
3555 3556
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
X
XuanYang-cn 已提交
3557 3558 3559
		return resp, nil
	}

3560 3561 3562
	log.Debug("successfully received list import tasks response",
		zap.String("collection", req.CollectionName),
		zap.Any("tasks", resp.Tasks))
E
Enwei Jiao 已提交
3563 3564
	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 已提交
3565 3566 3567
	return resp, err
}

3568 3569 3570
// InvalidateCredentialCache invalidate the credential cache of specified username.
func (node *Proxy) InvalidateCredentialCache(ctx context.Context, request *proxypb.InvalidateCredCacheRequest) (*commonpb.Status, error) {
	ctx = logutil.WithModule(ctx, moduleName)
3571 3572 3573 3574 3575

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-InvalidateCredentialCache")
	defer sp.Finish()

	log := log.Ctx(ctx).With(
3576 3577
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))
3578 3579

	log.Debug("received request to invalidate credential cache")
3580
	if !node.checkHealthy() {
3581
		return unhealthyStatus(), nil
3582
	}
3583 3584 3585 3586 3587

	username := request.Username
	if globalMetaCache != nil {
		globalMetaCache.RemoveCredential(username) // no need to return error, though credential may be not cached
	}
3588
	log.Debug("complete to invalidate credential cache")
3589 3590 3591 3592 3593 3594 3595 3596 3597 3598

	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) {
	ctx = logutil.WithModule(ctx, moduleName)
3599 3600 3601 3602 3603

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-UpdateCredentialCache")
	defer sp.Finish()

	log := log.Ctx(ctx).With(
3604 3605
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))
3606 3607

	log.Debug("received request to update credential cache")
3608
	if !node.checkHealthy() {
3609
		return unhealthyStatus(), nil
3610
	}
3611 3612

	credInfo := &internalpb.CredentialInfo{
3613 3614
		Username:       request.Username,
		Sha256Password: request.Password,
3615 3616 3617 3618
	}
	if globalMetaCache != nil {
		globalMetaCache.UpdateCredential(credInfo) // no need to return error, though credential may be not cached
	}
3619
	log.Debug("complete to update credential cache")
3620 3621 3622 3623 3624 3625 3626 3627

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

func (node *Proxy) CreateCredential(ctx context.Context, req *milvuspb.CreateCredentialRequest) (*commonpb.Status, error) {
3628 3629 3630 3631 3632 3633 3634 3635
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreateCredential")
	defer sp.Finish()

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

	log.Debug("CreateCredential",
		zap.String("role", typeutil.ProxyRole))
3636
	if !node.checkHealthy() {
3637
		return unhealthyStatus(), nil
3638
	}
3639 3640 3641 3642 3643 3644 3645 3646 3647 3648
	// 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 {
3649 3650
		log.Error("decode password fail",
			zap.Error(err))
3651 3652 3653 3654 3655 3656
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_CreateCredentialFailure,
			Reason:    "decode password fail key:" + req.Username,
		}, nil
	}
	if err = ValidatePassword(rawPassword); err != nil {
3657 3658
		log.Error("illegal password",
			zap.Error(err))
3659 3660 3661 3662 3663 3664 3665
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
		}, nil
	}
	encryptedPassword, err := crypto.PasswordEncrypt(rawPassword)
	if err != nil {
3666 3667
		log.Error("encrypt password fail",
			zap.Error(err))
3668 3669 3670 3671 3672
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_CreateCredentialFailure,
			Reason:    "encrypt password fail key:" + req.Username,
		}, nil
	}
3673

3674 3675 3676
	credInfo := &internalpb.CredentialInfo{
		Username:          req.Username,
		EncryptedPassword: encryptedPassword,
3677
		Sha256Password:    crypto.SHA256(rawPassword, req.Username),
3678 3679 3680
	}
	result, err := node.rootCoord.CreateCredential(ctx, credInfo)
	if err != nil { // for error like conntext timeout etc.
3681 3682
		log.Error("create credential fail",
			zap.Error(err))
3683 3684 3685 3686 3687 3688 3689 3690
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}
	return result, err
}

C
codeman 已提交
3691
func (node *Proxy) UpdateCredential(ctx context.Context, req *milvuspb.UpdateCredentialRequest) (*commonpb.Status, error) {
3692 3693 3694 3695 3696 3697 3698 3699
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-UpdateCredential")
	defer sp.Finish()

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

	log.Debug("UpdateCredential",
		zap.String("role", typeutil.ProxyRole))
3700
	if !node.checkHealthy() {
3701
		return unhealthyStatus(), nil
3702
	}
C
codeman 已提交
3703 3704
	rawOldPassword, err := crypto.Base64Decode(req.OldPassword)
	if err != nil {
3705 3706
		log.Error("decode old password fail",
			zap.Error(err))
C
codeman 已提交
3707 3708 3709 3710 3711 3712
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "decode old password fail when updating:" + req.Username,
		}, nil
	}
	rawNewPassword, err := crypto.Base64Decode(req.NewPassword)
3713
	if err != nil {
3714 3715
		log.Error("decode password fail",
			zap.Error(err))
3716 3717 3718 3719 3720
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "decode password fail when updating:" + req.Username,
		}, nil
	}
C
codeman 已提交
3721 3722
	// valid new password
	if err = ValidatePassword(rawNewPassword); err != nil {
3723 3724
		log.Error("illegal password",
			zap.Error(err))
3725 3726 3727 3728 3729
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
		}, nil
	}
3730 3731

	if !passwordVerify(ctx, req.Username, rawOldPassword, globalMetaCache) {
C
codeman 已提交
3732 3733 3734 3735 3736 3737 3738
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "old password is not correct:" + req.Username,
		}, nil
	}
	// update meta data
	encryptedPassword, err := crypto.PasswordEncrypt(rawNewPassword)
3739
	if err != nil {
3740 3741
		log.Error("encrypt password fail",
			zap.Error(err))
3742 3743 3744 3745 3746
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "encrypt password fail when updating:" + req.Username,
		}, nil
	}
C
codeman 已提交
3747
	updateCredReq := &internalpb.CredentialInfo{
3748
		Username:          req.Username,
3749
		Sha256Password:    crypto.SHA256(rawNewPassword, req.Username),
3750 3751
		EncryptedPassword: encryptedPassword,
	}
C
codeman 已提交
3752
	result, err := node.rootCoord.UpdateCredential(ctx, updateCredReq)
3753
	if err != nil { // for error like conntext timeout etc.
3754 3755
		log.Error("update credential fail",
			zap.Error(err))
3756 3757 3758 3759 3760 3761 3762 3763 3764
		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) {
3765 3766 3767 3768 3769 3770 3771 3772
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DeleteCredential")
	defer sp.Finish()

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

	log.Debug("DeleteCredential",
		zap.String("role", typeutil.ProxyRole))
3773
	if !node.checkHealthy() {
3774
		return unhealthyStatus(), nil
3775 3776
	}

3777 3778 3779 3780 3781 3782
	if req.Username == util.UserRoot {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_DeleteCredentialFailure,
			Reason:    "user root cannot be deleted",
		}, nil
	}
3783 3784
	result, err := node.rootCoord.DeleteCredential(ctx, req)
	if err != nil { // for error like conntext timeout etc.
3785 3786
		log.Error("delete credential fail",
			zap.Error(err))
3787 3788 3789 3790 3791 3792 3793 3794 3795
		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) {
3796 3797 3798 3799 3800 3801 3802
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-ListCredUsers")
	defer sp.Finish()

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

	log.Debug("ListCredUsers")
3803
	if !node.checkHealthy() {
3804
		return &milvuspb.ListCredUsersResponse{Status: unhealthyStatus()}, nil
3805
	}
3806
	rootCoordReq := &milvuspb.ListCredUsersRequest{
3807 3808 3809
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_ListCredUsernames),
		),
3810 3811
	}
	resp, err := node.rootCoord.ListCredUsers(ctx, rootCoordReq)
3812 3813 3814 3815 3816 3817 3818 3819 3820 3821 3822 3823
	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,
		},
3824
		Usernames: resp.Usernames,
3825 3826
	}, nil
}
3827

3828
func (node *Proxy) CreateRole(ctx context.Context, req *milvuspb.CreateRoleRequest) (*commonpb.Status, error) {
3829 3830 3831 3832 3833 3834 3835
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreateRole")
	defer sp.Finish()

	log := log.Ctx(ctx)

	log.Debug("CreateRole",
		zap.Any("req", req))
3836
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
3837
		return errorutil.UnhealthyStatus(code), nil
3838 3839 3840 3841 3842 3843 3844 3845 3846 3847
	}

	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(),
3848
		}, nil
3849 3850 3851 3852
	}

	result, err := node.rootCoord.CreateRole(ctx, req)
	if err != nil {
3853 3854
		log.Error("fail to create role",
			zap.Error(err))
3855 3856 3857
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
3858
		}, nil
3859 3860
	}
	return result, nil
3861 3862
}

3863
func (node *Proxy) DropRole(ctx context.Context, req *milvuspb.DropRoleRequest) (*commonpb.Status, error) {
3864 3865 3866 3867 3868 3869 3870
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DropRole")
	defer sp.Finish()

	log := log.Ctx(ctx)

	log.Debug("DropRole",
		zap.Any("req", req))
3871
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
3872
		return errorutil.UnhealthyStatus(code), nil
3873 3874 3875 3876 3877
	}
	if err := ValidateRoleName(req.RoleName); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
3878
		}, nil
3879
	}
3880 3881 3882 3883 3884
	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,
3885
		}, nil
3886
	}
3887 3888
	result, err := node.rootCoord.DropRole(ctx, req)
	if err != nil {
3889 3890 3891
		log.Error("fail to drop role",
			zap.String("role_name", req.RoleName),
			zap.Error(err))
3892 3893 3894
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
3895
		}, nil
3896 3897
	}
	return result, nil
3898 3899
}

3900
func (node *Proxy) OperateUserRole(ctx context.Context, req *milvuspb.OperateUserRoleRequest) (*commonpb.Status, error) {
3901 3902 3903 3904 3905 3906 3907
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-OperateUserRole")
	defer sp.Finish()

	log := log.Ctx(ctx)

	log.Debug("OperateUserRole",
		zap.Any("req", req))
3908
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
3909
		return errorutil.UnhealthyStatus(code), nil
3910 3911 3912 3913 3914
	}
	if err := ValidateUsername(req.Username); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
3915
		}, nil
3916 3917 3918 3919 3920
	}
	if err := ValidateRoleName(req.RoleName); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
3921
		}, nil
3922 3923 3924 3925
	}

	result, err := node.rootCoord.OperateUserRole(ctx, req)
	if err != nil {
3926 3927
		logger.Error("fail to operate user role",
			zap.Error(err))
3928 3929 3930
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
3931
		}, nil
3932 3933
	}
	return result, nil
3934 3935
}

3936
func (node *Proxy) SelectRole(ctx context.Context, req *milvuspb.SelectRoleRequest) (*milvuspb.SelectRoleResponse, error) {
3937 3938 3939 3940 3941 3942
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-SelectRole")
	defer sp.Finish()

	log := log.Ctx(ctx)

	log.Debug("SelectRole", zap.Any("req", req))
3943
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
3944
		return &milvuspb.SelectRoleResponse{Status: errorutil.UnhealthyStatus(code)}, nil
3945 3946 3947 3948 3949 3950 3951 3952 3953
	}

	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(),
				},
3954
			}, nil
3955 3956 3957 3958 3959
		}
	}

	result, err := node.rootCoord.SelectRole(ctx, req)
	if err != nil {
3960 3961
		log.Error("fail to select role",
			zap.Error(err))
3962 3963 3964 3965 3966
		return &milvuspb.SelectRoleResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
3967
		}, nil
3968 3969
	}
	return result, nil
3970 3971
}

3972
func (node *Proxy) SelectUser(ctx context.Context, req *milvuspb.SelectUserRequest) (*milvuspb.SelectUserResponse, error) {
3973 3974 3975 3976 3977 3978 3979
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-SelectUser")
	defer sp.Finish()

	log := log.Ctx(ctx)

	log.Debug("SelectUser",
		zap.Any("req", req))
3980
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
3981
		return &milvuspb.SelectUserResponse{Status: errorutil.UnhealthyStatus(code)}, nil
3982 3983 3984 3985 3986 3987 3988 3989 3990
	}

	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(),
				},
3991
			}, nil
3992 3993 3994 3995 3996
		}
	}

	result, err := node.rootCoord.SelectUser(ctx, req)
	if err != nil {
3997 3998
		log.Error("fail to select user",
			zap.Error(err))
3999 4000 4001 4002 4003
		return &milvuspb.SelectUserResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
4004
		}, nil
4005 4006
	}
	return result, nil
4007 4008
}

4009 4010 4011 4012 4013 4014 4015 4016 4017 4018 4019 4020 4021 4022 4023 4024 4025 4026 4027 4028 4029 4030 4031 4032 4033 4034 4035 4036 4037 4038
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
4039 4040
}

4041
func (node *Proxy) OperatePrivilege(ctx context.Context, req *milvuspb.OperatePrivilegeRequest) (*commonpb.Status, error) {
4042 4043 4044 4045 4046 4047 4048
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-OperatePrivilege")
	defer sp.Finish()

	log := log.Ctx(ctx)

	log.Debug("OperatePrivilege",
		zap.Any("req", req))
4049
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4050
		return errorutil.UnhealthyStatus(code), nil
4051 4052 4053 4054 4055
	}
	if err := node.validPrivilegeParams(req); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4056
		}, nil
4057 4058 4059 4060 4061 4062
	}
	curUser, err := GetCurUserFromContext(ctx)
	if err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4063
		}, nil
4064 4065 4066 4067
	}
	req.Entity.Grantor.User = &milvuspb.UserEntity{Name: curUser}
	result, err := node.rootCoord.OperatePrivilege(ctx, req)
	if err != nil {
4068 4069
		log.Error("fail to operate privilege",
			zap.Error(err))
4070 4071 4072
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4073
		}, nil
4074 4075
	}
	return result, nil
4076 4077
}

4078 4079 4080 4081 4082 4083 4084 4085 4086 4087 4088 4089 4090 4091 4092 4093 4094 4095 4096 4097 4098 4099 4100 4101 4102 4103 4104
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) {
4105 4106 4107 4108 4109 4110 4111
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-SelectGrant")
	defer sp.Finish()

	log := log.Ctx(ctx)

	log.Debug("SelectGrant",
		zap.Any("req", req))
4112
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4113
		return &milvuspb.SelectGrantResponse{Status: errorutil.UnhealthyStatus(code)}, nil
4114 4115 4116 4117 4118 4119 4120 4121
	}

	if err := node.validGrantParams(req); err != nil {
		return &milvuspb.SelectGrantResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_IllegalArgument,
				Reason:    err.Error(),
			},
4122
		}, nil
4123 4124 4125 4126
	}

	result, err := node.rootCoord.SelectGrant(ctx, req)
	if err != nil {
4127 4128
		log.Error("fail to select grant",
			zap.Error(err))
4129 4130 4131 4132 4133
		return &milvuspb.SelectGrantResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
4134
		}, nil
4135 4136 4137 4138 4139
	}
	return result, nil
}

func (node *Proxy) RefreshPolicyInfoCache(ctx context.Context, req *proxypb.RefreshPolicyInfoCacheRequest) (*commonpb.Status, error) {
4140 4141 4142 4143 4144 4145 4146
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-RefreshPolicyInfoCache")
	defer sp.Finish()

	log := log.Ctx(ctx)

	log.Debug("RefreshPrivilegeInfoCache",
		zap.Any("req", req))
4147 4148 4149 4150 4151 4152 4153 4154 4155 4156
	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 {
4157 4158
			log.Error("fail to refresh policy info",
				zap.Error(err))
4159 4160 4161 4162 4163 4164
			return &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_RefreshPolicyInfoCacheFailure,
				Reason:    err.Error(),
			}, err
		}
	}
4165
	log.Debug("RefreshPrivilegeInfoCache success")
4166 4167 4168 4169

	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
	}, nil
4170
}
4171 4172 4173 4174 4175 4176 4177 4178 4179 4180 4181 4182 4183 4184 4185 4186 4187 4188 4189 4190

// 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
	}
	resp.ErrorCode = commonpb.ErrorCode_Success
	return resp, nil
}
4191 4192 4193 4194

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")
4195 4196 4197 4198
		return &milvuspb.CheckHealthResponse{
			Status:    unhealthyStatus(),
			IsHealthy: false,
			Reasons:   []string{reason}}, nil
4199 4200 4201 4202 4203 4204 4205 4206 4207 4208
	}

	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()

4209 4210 4211 4212 4213
		sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-RefreshPolicyInfoCache")
		defer sp.Finish()

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

4214
		if err != nil {
4215 4216
			log.Warn("check health fail",
				zap.Error(err))
4217 4218 4219 4220 4221
			errReasons = append(errReasons, fmt.Sprintf("check health fail for %s", role))
			return err
		}

		if !resp.IsHealthy {
4222
			log.Warn("check health fail")
4223 4224 4225 4226 4227 4228 4229 4230 4231 4232 4233 4234 4235 4236 4237 4238 4239 4240 4241 4242 4243 4244 4245 4246 4247 4248 4249 4250 4251 4252 4253 4254 4255
			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)
	})

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

	err := group.Wait()
	if err != nil || len(errReasons) != 0 {
		return &milvuspb.CheckHealthResponse{
			IsHealthy: false,
			Reasons:   errReasons,
		}, nil
	}

4256 4257 4258 4259 4260 4261 4262
	return &milvuspb.CheckHealthResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
			Reason:    "",
		},
		IsHealthy: true,
	}, nil
4263
}