impl.go 137.3 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 2132 2133 2134
func (node *Proxy) Upsert(ctx context.Context, request *milvuspb.UpsertRequest) (*milvuspb.MutationResult, error) {
	panic("TODO: not implement")
}

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

2143
	receiveSize := proto.Size(request)
2144
	rateCol.Add(internalpb.RateType_DMLDelete.String(), float64(receiveSize))
E
Enwei Jiao 已提交
2145
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.DeleteLabel).Add(float64(receiveSize))
2146

G
groot 已提交
2147 2148 2149 2150 2151 2152
	if !node.checkHealthy() {
		return &milvuspb.MutationResult{
			Status: unhealthyStatus(),
		}, nil
	}

2153 2154 2155
	method := "Delete"
	tr := timerecord.NewTimeRecorder(method)

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

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

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

G
groot 已提交
2194 2195 2196 2197 2198 2199 2200 2201
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

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

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

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

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

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

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

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

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

2265 2266 2267
	travelTs := request.TravelTimestamp
	guaranteeTs := request.GuaranteeTimestamp

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

2280 2281 2282
	log.Debug(
		rpcReceived(method))

2283
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
2284
		log.Warn(
2285
			rpcFailedToEnqueue(method),
2286
			zap.Error(err))
D
dragondriver 已提交
2287

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

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

2300
	log.Debug(
2301
		rpcEnqueued(method),
2302
		zap.Uint64("timestamp", qt.Base.Timestamp))
D
dragondriver 已提交
2303

2304
	if err := qt.WaitToFinish(); err != nil {
2305
		log.Warn(
2306
			rpcFailedToWaitToFinish(method),
2307
			zap.Error(err))
2308

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

2312 2313
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
2314
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2315 2316 2317 2318 2319
				Reason:    err.Error(),
			},
		}, nil
	}

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

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

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

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

2358
	ft := &flushTask{
T
ThreadDao 已提交
2359 2360 2361
		ctx:          ctx,
		Condition:    NewTaskCondition(ctx),
		FlushRequest: request,
2362
		dataCoord:    node.dataCoord,
2363 2364
	}

D
dragondriver 已提交
2365
	method := "Flush"
2366
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2367
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2368

2369
	log := log.Ctx(ctx).With(
2370
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2371 2372
		zap.String("db", request.DbName),
		zap.Any("collections", request.CollectionNames))
D
dragondriver 已提交
2373

2374 2375
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2376 2377 2378
	if err := node.sched.ddQueue.Enqueue(ft); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2379
			zap.Error(err))
D
dragondriver 已提交
2380

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

2383 2384
		resp.Status.Reason = err.Error()
		return resp, nil
2385 2386
	}

D
dragondriver 已提交
2387 2388 2389
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", ft.BeginTs()),
2390
		zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2391 2392 2393 2394

	if err := ft.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
2395
			zap.Error(err),
D
dragondriver 已提交
2396
			zap.Uint64("BeginTs", ft.BeginTs()),
2397
			zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2398

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

D
dragondriver 已提交
2401
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
2402 2403
		resp.Status.Reason = err.Error()
		return resp, nil
2404 2405
	}

D
dragondriver 已提交
2406 2407 2408
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", ft.BeginTs()),
2409
		zap.Uint64("EndTs", ft.EndTs()))
D
dragondriver 已提交
2410

E
Enwei Jiao 已提交
2411 2412
	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()))
2413
	return ft.result, nil
2414 2415
}

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

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

2423 2424 2425 2426 2427
	if !node.checkHealthy() {
		return &milvuspb.QueryResults{
			Status: unhealthyStatus(),
		}, nil
	}
2428

D
dragondriver 已提交
2429 2430
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Query")
	defer sp.Finish()
2431
	tr := timerecord.NewTimeRecorder("Query")
D
dragondriver 已提交
2432

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

D
dragondriver 已提交
2449 2450
	method := "Query"

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

2454
	log := log.Ctx(ctx).With(
2455
		zap.String("role", typeutil.ProxyRole),
2456 2457
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
2458 2459 2460 2461
		zap.Strings("partitions", request.PartitionNames))

	log.Debug(
		rpcReceived(method),
2462 2463 2464 2465
		zap.String("expr", request.Expr),
		zap.Strings("OutputFields", request.OutputFields),
		zap.Uint64("travel_timestamp", request.TravelTimestamp),
		zap.Uint64("guarantee_timestamp", request.GuaranteeTimestamp))
G
godchen 已提交
2466

D
dragondriver 已提交
2467
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
2468
		log.Warn(
D
dragondriver 已提交
2469
			rpcFailedToEnqueue(method),
2470
			zap.Error(err))
D
dragondriver 已提交
2471

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

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

2484
	log.Debug(rpcEnqueued(method))
D
dragondriver 已提交
2485 2486

	if err := qt.WaitToFinish(); err != nil {
2487
		log.Warn(
D
dragondriver 已提交
2488
			rpcFailedToWaitToFinish(method),
2489
			zap.Error(err))
2490

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

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

2505
	log.Debug(rpcDone(method))
D
dragondriver 已提交
2506

E
Enwei Jiao 已提交
2507
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2508 2509
		metrics.SuccessLabel).Inc()

E
Enwei Jiao 已提交
2510
	metrics.ProxySQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2511
		metrics.QueryLabel).Observe(float64(tr.ElapseSpan().Milliseconds()))
E
Enwei Jiao 已提交
2512
	metrics.ProxyCollectionSQLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10),
2513
		metrics.QueryLabel, request.CollectionName).Observe(float64(tr.ElapseSpan().Milliseconds()))
2514 2515

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

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

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

Y
Yusup 已提交
2534 2535 2536 2537 2538 2539 2540
	cat := &CreateAliasTask{
		ctx:                ctx,
		Condition:          NewTaskCondition(ctx),
		CreateAliasRequest: request,
		rootCoord:          node.rootCoord,
	}

D
dragondriver 已提交
2541
	method := "CreateAlias"
2542
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2543
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2544

2545
	log := log.Ctx(ctx).With(
D
dragondriver 已提交
2546 2547 2548 2549 2550
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

2551 2552
	log.Debug(rpcReceived(method))

D
dragondriver 已提交
2553 2554 2555
	if err := node.sched.ddQueue.Enqueue(cat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
2556
			zap.Error(err))
D
dragondriver 已提交
2557

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

Y
Yusup 已提交
2560 2561 2562 2563 2564 2565
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2566 2567 2568
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", cat.BeginTs()),
2569
		zap.Uint64("EndTs", cat.EndTs()))
D
dragondriver 已提交
2570 2571 2572 2573

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

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

D
dragondriver 已提交
2585 2586 2587
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", cat.BeginTs()),
2588
		zap.Uint64("EndTs", cat.EndTs()))
D
dragondriver 已提交
2589

E
Enwei Jiao 已提交
2590 2591
	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 已提交
2592 2593 2594
	return cat.result, nil
}

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

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

Y
Yusup 已提交
2604 2605 2606 2607 2608 2609 2610
	dat := &DropAliasTask{
		ctx:              ctx,
		Condition:        NewTaskCondition(ctx),
		DropAliasRequest: request,
		rootCoord:        node.rootCoord,
	}

D
dragondriver 已提交
2611
	method := "DropAlias"
2612
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2613
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2614

2615
	log := log.Ctx(ctx).With(
D
dragondriver 已提交
2616 2617 2618 2619
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias))

2620 2621
	log.Debug(rpcReceived(method))

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

Y
Yusup 已提交
2628 2629 2630 2631 2632 2633
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2634 2635 2636
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", dat.BeginTs()),
2637
		zap.Uint64("EndTs", dat.EndTs()))
D
dragondriver 已提交
2638 2639 2640 2641

	if err := dat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2642
			zap.Error(err),
D
dragondriver 已提交
2643
			zap.Uint64("BeginTs", dat.BeginTs()),
2644
			zap.Uint64("EndTs", dat.EndTs()))
Y
Yusup 已提交
2645

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

Y
Yusup 已提交
2648 2649 2650 2651 2652 2653
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2654 2655 2656
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", dat.BeginTs()),
2657
		zap.Uint64("EndTs", dat.EndTs()))
D
dragondriver 已提交
2658

E
Enwei Jiao 已提交
2659 2660
	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 已提交
2661 2662 2663
	return dat.result, nil
}

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

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

Y
Yusup 已提交
2673 2674 2675 2676 2677 2678 2679
	aat := &AlterAliasTask{
		ctx:               ctx,
		Condition:         NewTaskCondition(ctx),
		AlterAliasRequest: request,
		rootCoord:         node.rootCoord,
	}

D
dragondriver 已提交
2680
	method := "AlterAlias"
2681
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2682
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2683

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

2690 2691
	log.Debug(rpcReceived(method))

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

Y
Yusup 已提交
2698 2699 2700 2701 2702 2703
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2704 2705 2706
	log.Debug(
		rpcEnqueued(method),
		zap.Uint64("BeginTs", aat.BeginTs()),
2707
		zap.Uint64("EndTs", aat.EndTs()))
D
dragondriver 已提交
2708 2709 2710 2711

	if err := aat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2712
			zap.Error(err),
D
dragondriver 已提交
2713
			zap.Uint64("BeginTs", aat.BeginTs()),
2714
			zap.Uint64("EndTs", aat.EndTs()))
Y
Yusup 已提交
2715

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

Y
Yusup 已提交
2718 2719 2720 2721 2722 2723
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2724 2725 2726
	log.Debug(
		rpcDone(method),
		zap.Uint64("BeginTs", aat.BeginTs()),
2727
		zap.Uint64("EndTs", aat.EndTs()))
D
dragondriver 已提交
2728

E
Enwei Jiao 已提交
2729 2730
	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 已提交
2731 2732 2733
	return aat.result, nil
}

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

2742 2743 2744 2745
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CalcDistance")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

2746 2747
	query := func(ids *milvuspb.VectorIDs) (*milvuspb.QueryResults, error) {
		outputFields := []string{ids.FieldName}
2748

2749 2750 2751 2752 2753
		queryRequest := &milvuspb.QueryRequest{
			DbName:         "",
			CollectionName: ids.CollectionName,
			PartitionNames: ids.PartitionNames,
			OutputFields:   outputFields,
2754 2755
		}

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

2770
			queryShardPolicy: mergeRoundRobinPolicy,
2771
			shardMgr:         node.shardMgr,
2772 2773
		}

2774
		log := log.Ctx(ctx).With(
G
groot 已提交
2775 2776
			zap.String("collection", queryRequest.CollectionName),
			zap.Any("partitions", queryRequest.PartitionNames),
2777
			zap.Any("OutputFields", queryRequest.OutputFields))
G
groot 已提交
2778

2779
		err := node.sched.dqQueue.Enqueue(qt)
2780
		if err != nil {
2781 2782
			log.Error("CalcDistance queryTask failed to enqueue",
				zap.Error(err))
2783

2784 2785 2786 2787 2788
			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2789
			}, err
2790
		}
2791

2792
		log.Debug("CalcDistance queryTask enqueued")
2793 2794 2795

		err = qt.WaitToFinish()
		if err != nil {
2796 2797
			log.Error("CalcDistance queryTask failed to WaitToFinish",
				zap.Error(err))
2798 2799 2800 2801 2802 2803

			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2804
			}, err
2805
		}
2806

2807
		log.Debug("CalcDistance queryTask Done")
2808 2809

		return &milvuspb.QueryResults{
2810 2811
			Status:     qt.result.Status,
			FieldsData: qt.result.FieldsData,
2812 2813 2814
		}, nil
	}

G
groot 已提交
2815 2816 2817 2818
	// calcDistanceTask is not a standard task, no need to enqueue
	task := &calcDistanceTask{
		traceID:   traceID,
		queryFunc: query,
2819 2820
	}

G
groot 已提交
2821
	return task.Execute(ctx, request)
2822 2823
}

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

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

	log := log.Ctx(ctx)

D
dragondriver 已提交
2836
	log.Debug("GetPersistentSegmentInfo",
2837
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2838 2839 2840
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

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

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

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

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

	log := log.Ctx(ctx)

D
dragondriver 已提交
2926
	log.Debug("GetQuerySegmentInfo",
2927
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2928 2929 2930
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

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

2941 2942
	method := "GetQuerySegmentInfo"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
2943
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
2944 2945
		metrics.TotalLabel).Inc()

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

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

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

	// TODO(wxyu): change name RequestType to Request
	drt, err := parseDummyRequestType(req.RequestType)
3007 3008 3009 3010 3011 3012

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

	log := log.Ctx(ctx)

3013
	if err != nil {
3014 3015
		log.Warn("Failed to parse dummy request type",
			zap.Error(err))
3016 3017 3018
		return failedResponse, nil
	}

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

3027
		request := &milvuspb.QueryRequest{
3028 3029 3030
			DbName:         drr.DbName,
			CollectionName: drr.CollectionName,
			PartitionNames: drr.PartitionNames,
3031
			OutputFields:   drr.OutputFields,
X
Xiangyu Wang 已提交
3032 3033
		}

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

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

3046 3047
	log.Debug("cannot find specify dummy request type")
	return failedResponse, nil
X
Xiangyu Wang 已提交
3048 3049
}

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

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

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

3061 3062
	log.Debug("RegisterLink")

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

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

	log := log.Ctx(ctx)

3090
	log.Debug("Proxy.GetMetrics",
E
Enwei Jiao 已提交
3091
		zap.Int64("node_id", paramtable.GetNodeID()),
3092 3093 3094 3095
		zap.String("req", req.Request))

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

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

	metricType, err := metricsinfo.ParseMetricType(req.Request)
	if err != nil {
		log.Warn("Proxy.GetMetrics failed to parse metric type",
E
Enwei Jiao 已提交
3112
			zap.Int64("node_id", paramtable.GetNodeID()),
3113 3114 3115 3116 3117 3118 3119 3120 3121 3122 3123 3124 3125 3126 3127
			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))

3128 3129 3130
	req.Base = commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_SystemInfo),
		commonpbutil.WithMsgID(0),
E
Enwei Jiao 已提交
3131
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
3132
	)
3133
	if metricType == metricsinfo.SystemInfoMetrics {
3134 3135 3136 3137 3138 3139 3140
		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))

3141
		metrics, err := getSystemInfoMetrics(ctx, req, node)
3142 3143

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

3150 3151
		node.metricsCacheManager.UpdateSystemInfoMetrics(metrics)

G
godchen 已提交
3152
		return metrics, nil
3153 3154
	}

J
Jiquan Long 已提交
3155
	log.Warn("Proxy.GetMetrics failed, request metric type is not implemented yet",
E
Enwei Jiao 已提交
3156
		zap.Int64("node_id", paramtable.GetNodeID()),
3157 3158 3159 3160 3161 3162 3163 3164 3165 3166 3167 3168
		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
}

3169 3170 3171
// 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) {
3172 3173 3174 3175 3176 3177 3178
	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))

3179 3180
	if !node.checkHealthy() {
		log.Warn("Proxy.GetProxyMetrics failed",
E
Enwei Jiao 已提交
3181
			zap.Error(errProxyIsUnhealthy(paramtable.GetNodeID())))
3182 3183 3184 3185

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

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

	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",
3225
			zap.String("metric_type", metricType))
3226 3227 3228 3229

		return proxyMetrics, nil
	}

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

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

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

	log := log.Ctx(ctx)

B
bigsheeper 已提交
3248
	log.Debug("Proxy.LoadBalance",
E
Enwei Jiao 已提交
3249
		zap.Int64("proxy_id", paramtable.GetNodeID()),
B
bigsheeper 已提交
3250 3251 3252 3253 3254 3255 3256 3257 3258
		zap.Any("req", req))

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

	status := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
	}
3259 3260 3261

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

3300 3301
// GetReplicas gets replica info
func (node *Proxy) GetReplicas(ctx context.Context, req *milvuspb.GetReplicasRequest) (*milvuspb.GetReplicasResponse, error) {
3302 3303 3304 3305 3306 3307 3308 3309
	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()))
3310 3311 3312 3313 3314 3315
	resp := &milvuspb.GetReplicasResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

S
smellthemoon 已提交
3316 3317
	req.Base = commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_GetReplicas),
E
Enwei Jiao 已提交
3318
		commonpbutil.WithSourceID(paramtable.GetNodeID()),
S
smellthemoon 已提交
3319
	)
3320 3321 3322

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

3335
// GetCompactionState gets the compaction state of multiple segments
3336
func (node *Proxy) GetCompactionState(ctx context.Context, req *milvuspb.GetCompactionStateRequest) (*milvuspb.GetCompactionStateResponse, error) {
3337 3338 3339 3340 3341 3342 3343
	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")
3344 3345 3346 3347 3348 3349 3350
	resp := &milvuspb.GetCompactionStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionState(ctx, req)
3351 3352 3353
	log.Debug("received GetCompactionState response",
		zap.Any("resp", resp),
		zap.Error(err))
3354 3355 3356
	return resp, err
}

3357
// ManualCompaction invokes compaction on specified collection
3358
func (node *Proxy) ManualCompaction(ctx context.Context, req *milvuspb.ManualCompactionRequest) (*milvuspb.ManualCompactionResponse, error) {
3359 3360 3361 3362 3363 3364 3365
	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")
3366 3367 3368 3369 3370 3371 3372
	resp := &milvuspb.ManualCompactionResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.ManualCompaction(ctx, req)
3373 3374 3375
	log.Info("received ManualCompaction response",
		zap.Any("resp", resp),
		zap.Error(err))
3376 3377 3378
	return resp, err
}

3379
// GetCompactionStateWithPlans returns the compactions states with the given plan ID
3380
func (node *Proxy) GetCompactionStateWithPlans(ctx context.Context, req *milvuspb.GetCompactionPlansRequest) (*milvuspb.GetCompactionPlansResponse, error) {
3381 3382 3383 3384 3385 3386 3387
	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")
3388 3389 3390 3391 3392 3393 3394
	resp := &milvuspb.GetCompactionPlansResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionStateWithPlans(ctx, req)
3395 3396 3397
	log.Debug("received GetCompactionStateWithPlans response",
		zap.Any("resp", resp),
		zap.Error(err))
3398 3399 3400
	return resp, err
}

B
Bingyi Sun 已提交
3401 3402
// GetFlushState gets the flush state of multiple segments
func (node *Proxy) GetFlushState(ctx context.Context, req *milvuspb.GetFlushStateRequest) (*milvuspb.GetFlushStateResponse, error) {
3403 3404 3405 3406 3407 3408 3409
	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))
3410
	var err error
B
Bingyi Sun 已提交
3411 3412 3413
	resp := &milvuspb.GetFlushStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
J
Jiquan Long 已提交
3414
		log.Warn("unable to get flush state because of closed server")
B
Bingyi Sun 已提交
3415 3416 3417
		return resp, nil
	}

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

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

3435 3436 3437
func (node *Proxy) checkHealthyAndReturnCode() (commonpb.StateCode, bool) {
	code := node.stateCode.Load().(commonpb.StateCode)
	return code, code == commonpb.StateCode_Healthy
3438 3439
}

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

// 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) {
3450 3451 3452 3453 3454
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Import")
	defer sp.Finish()

	log := log.Ctx(ctx)

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

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

3479 3480
	method := "Import"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
3481
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3482 3483
		metrics.TotalLabel).Inc()

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

E
Enwei Jiao 已提交
3495 3496
	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()))
3497
	return respFromRC, nil
G
groot 已提交
3498 3499
}

3500
// GetImportState checks import task state from RootCoord.
G
groot 已提交
3501
func (node *Proxy) GetImportState(ctx context.Context, req *milvuspb.GetImportStateRequest) (*milvuspb.GetImportStateResponse, error) {
3502 3503 3504 3505 3506 3507 3508
	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 已提交
3509 3510 3511 3512 3513
	resp := &milvuspb.GetImportStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3514 3515
	method := "GetImportState"
	tr := timerecord.NewTimeRecorder(method)
E
Enwei Jiao 已提交
3516
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method,
3517
		metrics.TotalLabel).Inc()
G
groot 已提交
3518 3519

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

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

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

	log := log.Ctx(ctx)

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

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

3572 3573 3574
// 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)
3575 3576 3577 3578 3579

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

	log := log.Ctx(ctx).With(
3580 3581
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))
3582 3583

	log.Debug("received request to invalidate credential cache")
3584
	if !node.checkHealthy() {
3585
		return unhealthyStatus(), nil
3586
	}
3587 3588 3589 3590 3591

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

	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)
3603 3604 3605 3606 3607

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

	log := log.Ctx(ctx).With(
3608 3609
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))
3610 3611

	log.Debug("received request to update credential cache")
3612
	if !node.checkHealthy() {
3613
		return unhealthyStatus(), nil
3614
	}
3615 3616

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

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

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

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

C
codeman 已提交
3695
func (node *Proxy) UpdateCredential(ctx context.Context, req *milvuspb.UpdateCredentialRequest) (*commonpb.Status, error) {
3696 3697 3698 3699 3700 3701 3702 3703
	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))
3704
	if !node.checkHealthy() {
3705
		return unhealthyStatus(), nil
3706
	}
C
codeman 已提交
3707 3708
	rawOldPassword, err := crypto.Base64Decode(req.OldPassword)
	if err != nil {
3709 3710
		log.Error("decode old password fail",
			zap.Error(err))
C
codeman 已提交
3711 3712 3713 3714 3715 3716
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "decode old password fail when updating:" + req.Username,
		}, nil
	}
	rawNewPassword, err := crypto.Base64Decode(req.NewPassword)
3717
	if err != nil {
3718 3719
		log.Error("decode password fail",
			zap.Error(err))
3720 3721 3722 3723 3724
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "decode password fail when updating:" + req.Username,
		}, nil
	}
C
codeman 已提交
3725 3726
	// valid new password
	if err = ValidatePassword(rawNewPassword); err != nil {
3727 3728
		log.Error("illegal password",
			zap.Error(err))
3729 3730 3731 3732 3733
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
		}, nil
	}
3734 3735

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

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

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

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

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

	log := log.Ctx(ctx)

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

	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(),
3852
		}, nil
3853 3854 3855 3856
	}

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

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

	log := log.Ctx(ctx)

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

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

	log := log.Ctx(ctx)

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

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

3940
func (node *Proxy) SelectRole(ctx context.Context, req *milvuspb.SelectRoleRequest) (*milvuspb.SelectRoleResponse, error) {
3941 3942 3943 3944 3945 3946
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-SelectRole")
	defer sp.Finish()

	log := log.Ctx(ctx)

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

	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(),
				},
3958
			}, nil
3959 3960 3961 3962 3963
		}
	}

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

3976
func (node *Proxy) SelectUser(ctx context.Context, req *milvuspb.SelectUserRequest) (*milvuspb.SelectUserResponse, error) {
3977 3978 3979 3980 3981 3982 3983
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-SelectUser")
	defer sp.Finish()

	log := log.Ctx(ctx)

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

	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(),
				},
3995
			}, nil
3996 3997 3998 3999 4000
		}
	}

	result, err := node.rootCoord.SelectUser(ctx, req)
	if err != nil {
4001 4002
		log.Error("fail to select user",
			zap.Error(err))
4003 4004 4005 4006 4007
		return &milvuspb.SelectUserResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
4008
		}, nil
4009 4010
	}
	return result, nil
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 4039 4040 4041 4042
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
4043 4044
}

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

	log := log.Ctx(ctx)

	log.Debug("OperatePrivilege",
		zap.Any("req", req))
4053
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4054
		return errorutil.UnhealthyStatus(code), nil
4055 4056 4057 4058 4059
	}
	if err := node.validPrivilegeParams(req); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4060
		}, nil
4061 4062 4063 4064 4065 4066
	}
	curUser, err := GetCurUserFromContext(ctx)
	if err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4067
		}, nil
4068 4069 4070 4071
	}
	req.Entity.Grantor.User = &milvuspb.UserEntity{Name: curUser}
	result, err := node.rootCoord.OperatePrivilege(ctx, req)
	if err != nil {
4072 4073
		log.Error("fail to operate privilege",
			zap.Error(err))
4074 4075 4076
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4077
		}, nil
4078 4079
	}
	return result, nil
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 4105 4106 4107 4108
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) {
4109 4110 4111 4112 4113 4114 4115
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-SelectGrant")
	defer sp.Finish()

	log := log.Ctx(ctx)

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

	if err := node.validGrantParams(req); err != nil {
		return &milvuspb.SelectGrantResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_IllegalArgument,
				Reason:    err.Error(),
			},
4126
		}, nil
4127 4128 4129 4130
	}

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

func (node *Proxy) RefreshPolicyInfoCache(ctx context.Context, req *proxypb.RefreshPolicyInfoCacheRequest) (*commonpb.Status, error) {
4144 4145 4146 4147 4148 4149 4150
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-RefreshPolicyInfoCache")
	defer sp.Finish()

	log := log.Ctx(ctx)

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

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

// 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
	}
4192 4193 4194 4195 4196 4197 4198
	node.multiRateLimiter.SetQuotaStates(request.GetStates(), request.GetStateReasons())
	log.Info("current rates in proxy", zap.Int64("proxyNodeID", paramtable.GetNodeID()), zap.Any("rates", request.GetRates()))
	if len(request.GetStates()) != 0 {
		for i := range request.GetStates() {
			log.Warn("Proxy set quota states", zap.String("state", request.GetStates()[i].String()), zap.String("reason", request.GetStateReasons()[i]))
		}
	}
4199 4200 4201
	resp.ErrorCode = commonpb.ErrorCode_Success
	return resp, nil
}
4202 4203 4204 4205

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")
4206 4207 4208 4209
		return &milvuspb.CheckHealthResponse{
			Status:    unhealthyStatus(),
			IsHealthy: false,
			Reasons:   []string{reason}}, nil
4210 4211 4212 4213 4214 4215 4216 4217 4218 4219
	}

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

4220 4221 4222 4223 4224
		sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-RefreshPolicyInfoCache")
		defer sp.Finish()

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

4225
		if err != nil {
4226 4227
			log.Warn("check health fail",
				zap.Error(err))
4228 4229 4230 4231 4232
			errReasons = append(errReasons, fmt.Sprintf("check health fail for %s", role))
			return err
		}

		if !resp.IsHealthy {
4233
			log.Warn("check health fail")
4234 4235 4236 4237 4238 4239 4240 4241 4242 4243 4244 4245 4246 4247 4248 4249 4250 4251 4252 4253 4254 4255 4256 4257 4258 4259 4260 4261
			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{
4262 4263 4264
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_Success,
			},
4265 4266 4267 4268 4269
			IsHealthy: false,
			Reasons:   errReasons,
		}, nil
	}

4270
	states, reasons := node.multiRateLimiter.GetQuotaStates()
4271 4272 4273 4274 4275
	return &milvuspb.CheckHealthResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
			Reason:    "",
		},
4276 4277 4278
		QuotaStates: states,
		Reasons:     reasons,
		IsHealthy:   true,
4279
	}, nil
4280
}