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

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

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

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

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

54 55
const moduleName = "Proxy"

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

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

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

102
// InvalidateCollectionMetaCache invalidate the meta cache of specific collection.
C
Cai Yudong 已提交
103
func (node *Proxy) InvalidateCollectionMetaCache(ctx context.Context, request *proxypb.InvalidateCollMetaCacheRequest) (*commonpb.Status, error) {
104
	ctx = logutil.WithModule(ctx, moduleName)
105
	logutil.Logger(ctx).Info("received request to invalidate collection meta cache",
106
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
107
		zap.String("db", request.DbName),
108 109
		zap.String("collectionName", request.CollectionName),
		zap.Int64("collectionID", request.CollectionID))
D
dragondriver 已提交
110

111
	collectionName := request.CollectionName
112
	collectionID := request.CollectionID
N
neza2017 已提交
113
	if globalMetaCache != nil {
114 115 116 117 118 119
		if collectionName != "" {
			globalMetaCache.RemoveCollection(ctx, collectionName) // no need to return error, though collection may be not cached
		}
		if request.CollectionID != UniqueID(0) {
			globalMetaCache.RemoveCollectionsByID(ctx, collectionID)
		}
N
neza2017 已提交
120
	}
121 122
	if request.GetBase().GetMsgType() == commonpb.MsgType_DropCollection {
		// no need to handle error, since this Proxy may not create dml stream for the collection.
123 124 125
		node.chMgr.removeDMLStream(request.GetCollectionID())
		// clean up collection level metrics
		metrics.CleanupCollectionMetrics(Params.ProxyCfg.GetNodeID(), collectionName)
126
	}
127
	logutil.Logger(ctx).Info("complete to invalidate collection meta cache",
128
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
129
		zap.String("db", request.DbName),
130 131
		zap.String("collection", collectionName),
		zap.Int64("collectionID", collectionID))
D
dragondriver 已提交
132

133
	return &commonpb.Status{
134
		ErrorCode: commonpb.ErrorCode_Success,
135 136
		Reason:    "",
	}, nil
137 138
}

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

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreateCollection")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
149 150 151
	method := "CreateCollection"
	tr := timerecord.NewTimeRecorder(method)

152
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
153

154
	cct := &createCollectionTask{
S
sunby 已提交
155
		ctx:                     ctx,
156 157
		Condition:               NewTaskCondition(ctx),
		CreateCollectionRequest: request,
158
		rootCoord:               node.rootCoord,
159 160
	}

161 162 163
	// avoid data race
	lenOfSchema := len(request.Schema)

164 165
	log.Debug(
		rpcReceived(method),
166
		zap.String("traceID", traceID),
167
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
168 169
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
170
		zap.Int("len(schema)", lenOfSchema),
171 172
		zap.Int32("shards_num", request.ShardsNum),
		zap.String("consistency_level", request.ConsistencyLevel.String()))
173

174 175 176
	if err := node.sched.ddQueue.Enqueue(cct); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
177 178
			zap.Error(err),
			zap.String("traceID", traceID),
179
			zap.String("role", typeutil.ProxyRole),
180 181 182
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Int("len(schema)", lenOfSchema),
183 184
			zap.Int32("shards_num", request.ShardsNum),
			zap.String("consistency_level", request.ConsistencyLevel.String()))
185

186
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
187
		return &commonpb.Status{
188
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
189 190 191 192
			Reason:    err.Error(),
		}, nil
	}

193 194
	log.Debug(
		rpcEnqueued(method),
195
		zap.String("traceID", traceID),
196
		zap.String("role", typeutil.ProxyRole),
197 198 199
		zap.Int64("MsgID", cct.ID()),
		zap.Uint64("BeginTs", cct.BeginTs()),
		zap.Uint64("EndTs", cct.EndTs()),
D
dragondriver 已提交
200 201
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
202 203
		zap.String("collection", request.CollectionName),
		zap.Int("len(schema)", lenOfSchema),
204 205
		zap.Int32("shards_num", request.ShardsNum),
		zap.String("consistency_level", request.ConsistencyLevel.String()))
206

207 208 209
	if err := cct.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
210
			zap.Error(err),
211
			zap.String("traceID", traceID),
212
			zap.String("role", typeutil.ProxyRole),
213 214 215
			zap.Int64("MsgID", cct.ID()),
			zap.Uint64("BeginTs", cct.BeginTs()),
			zap.Uint64("EndTs", cct.EndTs()),
D
dragondriver 已提交
216 217
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
218
			zap.Int("len(schema)", lenOfSchema),
219 220
			zap.Int32("shards_num", request.ShardsNum),
			zap.String("consistency_level", request.ConsistencyLevel.String()))
D
dragondriver 已提交
221

222
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
223
		return &commonpb.Status{
224
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
225 226 227 228
			Reason:    err.Error(),
		}, nil
	}

229 230
	log.Debug(
		rpcDone(method),
231
		zap.String("traceID", traceID),
232
		zap.String("role", typeutil.ProxyRole),
233 234 235 236 237 238
		zap.Int64("MsgID", cct.ID()),
		zap.Uint64("BeginTs", cct.BeginTs()),
		zap.Uint64("EndTs", cct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Int("len(schema)", lenOfSchema),
239 240
		zap.Int32("shards_num", request.ShardsNum),
		zap.String("consistency_level", request.ConsistencyLevel.String()))
241

242 243
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
244 245 246
	return cct.result, nil
}

247
// DropCollection drop a collection.
C
Cai Yudong 已提交
248
func (node *Proxy) DropCollection(ctx context.Context, request *milvuspb.DropCollectionRequest) (*commonpb.Status, error) {
249 250 251
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
252 253 254 255

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DropCollection")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
256 257
	method := "DropCollection"
	tr := timerecord.NewTimeRecorder(method)
258
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
259

260
	dct := &dropCollectionTask{
S
sunby 已提交
261
		ctx:                   ctx,
262 263
		Condition:             NewTaskCondition(ctx),
		DropCollectionRequest: request,
264
		rootCoord:             node.rootCoord,
265
		chMgr:                 node.chMgr,
S
sunby 已提交
266
		chTicker:              node.chTicker,
267 268
	}

269 270
	log.Debug("DropCollection received",
		zap.String("traceID", traceID),
271
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
272 273
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
274 275 276 277 278

	if err := node.sched.ddQueue.Enqueue(dct); err != nil {
		log.Warn("DropCollection failed to enqueue",
			zap.Error(err),
			zap.String("traceID", traceID),
279
			zap.String("role", typeutil.ProxyRole),
280 281 282
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

283
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
284
		return &commonpb.Status{
285
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
286 287 288 289
			Reason:    err.Error(),
		}, nil
	}

290 291
	log.Debug("DropCollection enqueued",
		zap.String("traceID", traceID),
292
		zap.String("role", typeutil.ProxyRole),
293 294 295
		zap.Int64("MsgID", dct.ID()),
		zap.Uint64("BeginTs", dct.BeginTs()),
		zap.Uint64("EndTs", dct.EndTs()),
D
dragondriver 已提交
296 297
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
298 299 300

	if err := dct.WaitToFinish(); err != nil {
		log.Warn("DropCollection failed to WaitToFinish",
D
dragondriver 已提交
301
			zap.Error(err),
302
			zap.String("traceID", traceID),
303
			zap.String("role", typeutil.ProxyRole),
304 305 306
			zap.Int64("MsgID", dct.ID()),
			zap.Uint64("BeginTs", dct.BeginTs()),
			zap.Uint64("EndTs", dct.EndTs()),
D
dragondriver 已提交
307 308 309
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

310
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
311
		return &commonpb.Status{
312
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
313 314 315 316
			Reason:    err.Error(),
		}, nil
	}

317 318
	log.Debug("DropCollection done",
		zap.String("traceID", traceID),
319
		zap.String("role", typeutil.ProxyRole),
320 321 322 323 324 325
		zap.Int64("MsgID", dct.ID()),
		zap.Uint64("BeginTs", dct.BeginTs()),
		zap.Uint64("EndTs", dct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

326 327
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
328 329 330
	return dct.result, nil
}

331
// HasCollection check if the specific collection exists in Milvus.
C
Cai Yudong 已提交
332
func (node *Proxy) HasCollection(ctx context.Context, request *milvuspb.HasCollectionRequest) (*milvuspb.BoolResponse, error) {
333 334 335 336 337
	if !node.checkHealthy() {
		return &milvuspb.BoolResponse{
			Status: unhealthyStatus(),
		}, nil
	}
338 339 340 341

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-HasCollection")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
342 343
	method := "HasCollection"
	tr := timerecord.NewTimeRecorder(method)
344
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
345
		metrics.TotalLabel).Inc()
346 347 348

	log.Debug("HasCollection received",
		zap.String("traceID", traceID),
349
		zap.String("role", typeutil.ProxyRole),
350 351 352
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

353
	hct := &hasCollectionTask{
S
sunby 已提交
354
		ctx:                  ctx,
355 356
		Condition:            NewTaskCondition(ctx),
		HasCollectionRequest: request,
357
		rootCoord:            node.rootCoord,
358 359
	}

360 361 362 363
	if err := node.sched.ddQueue.Enqueue(hct); err != nil {
		log.Warn("HasCollection failed to enqueue",
			zap.Error(err),
			zap.String("traceID", traceID),
364
			zap.String("role", typeutil.ProxyRole),
365 366 367
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

368
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
369
			metrics.AbandonLabel).Inc()
370 371
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
372
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
373 374 375 376 377
				Reason:    err.Error(),
			},
		}, nil
	}

378 379
	log.Debug("HasCollection enqueued",
		zap.String("traceID", traceID),
380
		zap.String("role", typeutil.ProxyRole),
381 382 383
		zap.Int64("MsgID", hct.ID()),
		zap.Uint64("BeginTS", hct.BeginTs()),
		zap.Uint64("EndTS", hct.EndTs()),
D
dragondriver 已提交
384 385
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
386 387 388

	if err := hct.WaitToFinish(); err != nil {
		log.Warn("HasCollection failed to WaitToFinish",
D
dragondriver 已提交
389
			zap.Error(err),
390
			zap.String("traceID", traceID),
391
			zap.String("role", typeutil.ProxyRole),
392 393 394
			zap.Int64("MsgID", hct.ID()),
			zap.Uint64("BeginTS", hct.BeginTs()),
			zap.Uint64("EndTS", hct.EndTs()),
D
dragondriver 已提交
395 396 397
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

398
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
399
			metrics.FailLabel).Inc()
400 401
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
402
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
403 404 405 406 407
				Reason:    err.Error(),
			},
		}, nil
	}

408 409
	log.Debug("HasCollection done",
		zap.String("traceID", traceID),
410
		zap.String("role", typeutil.ProxyRole),
411 412 413 414 415 416
		zap.Int64("MsgID", hct.ID()),
		zap.Uint64("BeginTS", hct.BeginTs()),
		zap.Uint64("EndTS", hct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

417
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
418
		metrics.SuccessLabel).Inc()
419
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
420 421 422
	return hct.result, nil
}

423
// LoadCollection load a collection into query nodes.
C
Cai Yudong 已提交
424
func (node *Proxy) LoadCollection(ctx context.Context, request *milvuspb.LoadCollectionRequest) (*commonpb.Status, error) {
425 426 427
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
428 429 430 431

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-LoadCollection")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
432 433
	method := "LoadCollection"
	tr := timerecord.NewTimeRecorder(method)
434 435
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
436
	lct := &loadCollectionTask{
S
sunby 已提交
437
		ctx:                   ctx,
438 439
		Condition:             NewTaskCondition(ctx),
		LoadCollectionRequest: request,
440
		queryCoord:            node.queryCoord,
C
cai.zhang 已提交
441
		indexCoord:            node.indexCoord,
442 443
	}

444 445
	log.Debug("LoadCollection received",
		zap.String("traceID", traceID),
446
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
447 448
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
449 450 451 452 453

	if err := node.sched.ddQueue.Enqueue(lct); err != nil {
		log.Warn("LoadCollection failed to enqueue",
			zap.Error(err),
			zap.String("traceID", traceID),
454
			zap.String("role", typeutil.ProxyRole),
455 456 457
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

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

466 467
	log.Debug("LoadCollection enqueued",
		zap.String("traceID", traceID),
468
		zap.String("role", typeutil.ProxyRole),
469 470 471
		zap.Int64("MsgID", lct.ID()),
		zap.Uint64("BeginTS", lct.BeginTs()),
		zap.Uint64("EndTS", lct.EndTs()),
D
dragondriver 已提交
472 473
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
474 475 476

	if err := lct.WaitToFinish(); err != nil {
		log.Warn("LoadCollection failed to WaitToFinish",
D
dragondriver 已提交
477
			zap.Error(err),
478
			zap.String("traceID", traceID),
479
			zap.String("role", typeutil.ProxyRole),
480 481 482
			zap.Int64("MsgID", lct.ID()),
			zap.Uint64("BeginTS", lct.BeginTs()),
			zap.Uint64("EndTS", lct.EndTs()),
D
dragondriver 已提交
483 484
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))
485
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
486
			metrics.FailLabel).Inc()
487
		return &commonpb.Status{
488
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
489 490 491 492
			Reason:    err.Error(),
		}, nil
	}

493 494
	log.Debug("LoadCollection done",
		zap.String("traceID", traceID),
495
		zap.String("role", typeutil.ProxyRole),
496 497 498 499 500 501
		zap.Int64("MsgID", lct.ID()),
		zap.Uint64("BeginTS", lct.BeginTs()),
		zap.Uint64("EndTS", lct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

502
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
503
		metrics.SuccessLabel).Inc()
504
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
505
	return lct.result, nil
506 507
}

508
// ReleaseCollection remove the loaded collection from query nodes.
C
Cai Yudong 已提交
509
func (node *Proxy) ReleaseCollection(ctx context.Context, request *milvuspb.ReleaseCollectionRequest) (*commonpb.Status, error) {
510 511 512
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
513

514
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-ReleaseCollection")
515 516
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
517 518
	method := "ReleaseCollection"
	tr := timerecord.NewTimeRecorder(method)
519 520
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
521
	rct := &releaseCollectionTask{
S
sunby 已提交
522
		ctx:                      ctx,
523 524
		Condition:                NewTaskCondition(ctx),
		ReleaseCollectionRequest: request,
525
		queryCoord:               node.queryCoord,
526
		chMgr:                    node.chMgr,
527 528
	}

529 530
	log.Debug(
		rpcReceived(method),
531
		zap.String("traceID", traceID),
532
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
533 534
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
535 536

	if err := node.sched.ddQueue.Enqueue(rct); err != nil {
537 538
		log.Warn(
			rpcFailedToEnqueue(method),
539 540
			zap.Error(err),
			zap.String("traceID", traceID),
541
			zap.String("role", typeutil.ProxyRole),
542 543 544
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

545
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
546
			metrics.AbandonLabel).Inc()
547
		return &commonpb.Status{
548
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
549 550 551 552
			Reason:    err.Error(),
		}, nil
	}

553 554
	log.Debug(
		rpcEnqueued(method),
555
		zap.String("traceID", traceID),
556
		zap.String("role", typeutil.ProxyRole),
557 558 559
		zap.Int64("MsgID", rct.ID()),
		zap.Uint64("BeginTS", rct.BeginTs()),
		zap.Uint64("EndTS", rct.EndTs()),
D
dragondriver 已提交
560 561
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
562 563

	if err := rct.WaitToFinish(); err != nil {
564 565
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
566
			zap.Error(err),
567
			zap.String("traceID", traceID),
568
			zap.String("role", typeutil.ProxyRole),
569 570 571
			zap.Int64("MsgID", rct.ID()),
			zap.Uint64("BeginTS", rct.BeginTs()),
			zap.Uint64("EndTS", rct.EndTs()),
D
dragondriver 已提交
572 573 574
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

575
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
576
			metrics.FailLabel).Inc()
577
		return &commonpb.Status{
578
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
579 580 581 582
			Reason:    err.Error(),
		}, nil
	}

583 584
	log.Debug(
		rpcDone(method),
585
		zap.String("traceID", traceID),
586
		zap.String("role", typeutil.ProxyRole),
587 588 589 590 591 592
		zap.Int64("MsgID", rct.ID()),
		zap.Uint64("BeginTS", rct.BeginTs()),
		zap.Uint64("EndTS", rct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

593
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
594
		metrics.SuccessLabel).Inc()
595
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
596
	return rct.result, nil
597 598
}

599
// DescribeCollection get the meta information of specific collection, such as schema, created timestamp and etc.
C
Cai Yudong 已提交
600
func (node *Proxy) DescribeCollection(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
601 602 603 604 605
	if !node.checkHealthy() {
		return &milvuspb.DescribeCollectionResponse{
			Status: unhealthyStatus(),
		}, nil
	}
606

607
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DescribeCollection")
608 609
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
610 611
	method := "DescribeCollection"
	tr := timerecord.NewTimeRecorder(method)
612 613
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
614

615
	dct := &describeCollectionTask{
S
sunby 已提交
616
		ctx:                       ctx,
617 618
		Condition:                 NewTaskCondition(ctx),
		DescribeCollectionRequest: request,
619
		rootCoord:                 node.rootCoord,
620 621
	}

622 623
	log.Debug("DescribeCollection received",
		zap.String("traceID", traceID),
624
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
625 626
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
627 628 629 630 631

	if err := node.sched.ddQueue.Enqueue(dct); err != nil {
		log.Warn("DescribeCollection failed to enqueue",
			zap.Error(err),
			zap.String("traceID", traceID),
632
			zap.String("role", typeutil.ProxyRole),
633 634 635
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

636
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
637
			metrics.AbandonLabel).Inc()
638 639
		return &milvuspb.DescribeCollectionResponse{
			Status: &commonpb.Status{
640
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
641 642 643 644 645
				Reason:    err.Error(),
			},
		}, nil
	}

646 647
	log.Debug("DescribeCollection enqueued",
		zap.String("traceID", traceID),
648
		zap.String("role", typeutil.ProxyRole),
649 650 651
		zap.Int64("MsgID", dct.ID()),
		zap.Uint64("BeginTS", dct.BeginTs()),
		zap.Uint64("EndTS", dct.EndTs()),
D
dragondriver 已提交
652 653
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
654 655 656

	if err := dct.WaitToFinish(); err != nil {
		log.Warn("DescribeCollection failed to WaitToFinish",
D
dragondriver 已提交
657
			zap.Error(err),
658
			zap.String("traceID", traceID),
659
			zap.String("role", typeutil.ProxyRole),
660 661 662
			zap.Int64("MsgID", dct.ID()),
			zap.Uint64("BeginTS", dct.BeginTs()),
			zap.Uint64("EndTS", dct.EndTs()),
D
dragondriver 已提交
663 664 665
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

666
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
667
			metrics.FailLabel).Inc()
668

669 670
		return &milvuspb.DescribeCollectionResponse{
			Status: &commonpb.Status{
671
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
672 673 674 675 676
				Reason:    err.Error(),
			},
		}, nil
	}

677 678
	log.Debug("DescribeCollection done",
		zap.String("traceID", traceID),
679
		zap.String("role", typeutil.ProxyRole),
680 681 682 683 684 685
		zap.Int64("MsgID", dct.ID()),
		zap.Uint64("BeginTS", dct.BeginTs()),
		zap.Uint64("EndTS", dct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

686
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
687
		metrics.SuccessLabel).Inc()
688
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
689 690 691
	return dct.result, nil
}

692 693 694 695 696 697 698 699 700
// 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
	}

701
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetStatistics")
702 703 704 705
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
	method := "GetStatistics"
	tr := timerecord.NewTimeRecorder(method)
706 707
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735
	g := &getStatisticsTask{
		request:   request,
		Condition: NewTaskCondition(ctx),
		ctx:       ctx,
		tr:        tr,
		dc:        node.dataCoord,
		qc:        node.queryCoord,
		shardMgr:  node.shardMgr,
	}

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Strings("partitions", request.PartitionNames))

	if err := node.sched.ddQueue.Enqueue(g); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
			zap.String("role", typeutil.ProxyRole),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Strings("partitions", request.PartitionNames))

736
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770
			metrics.AbandonLabel).Inc()

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

	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.Int64("msgID", g.ID()),
		zap.Uint64("BeginTS", g.BeginTs()),
		zap.Uint64("EndTS", g.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Strings("partitions", request.PartitionNames))

	if err := g.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
			zap.Error(err),
			zap.String("traceID", traceID),
			zap.String("role", typeutil.ProxyRole),
			zap.Int64("MsgID", g.ID()),
			zap.Uint64("BeginTS", g.BeginTs()),
			zap.Uint64("EndTS", g.EndTs()),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Strings("partitions", request.PartitionNames))

771
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791
			metrics.FailLabel).Inc()

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

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.Int64("msgID", g.ID()),
		zap.Uint64("BeginTS", g.BeginTs()),
		zap.Uint64("EndTS", g.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

792
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
793
		metrics.SuccessLabel).Inc()
794
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
795 796 797
	return g.result, nil
}

798
// GetCollectionStatistics get the collection statistics, such as `num_rows`.
C
Cai Yudong 已提交
799
func (node *Proxy) GetCollectionStatistics(ctx context.Context, request *milvuspb.GetCollectionStatisticsRequest) (*milvuspb.GetCollectionStatisticsResponse, error) {
800 801 802 803 804
	if !node.checkHealthy() {
		return &milvuspb.GetCollectionStatisticsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
805 806 807 808

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetCollectionStatistics")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
809 810
	method := "GetCollectionStatistics"
	tr := timerecord.NewTimeRecorder(method)
811 812
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
813
	g := &getCollectionStatisticsTask{
G
godchen 已提交
814 815 816
		ctx:                            ctx,
		Condition:                      NewTaskCondition(ctx),
		GetCollectionStatisticsRequest: request,
817
		dataCoord:                      node.dataCoord,
818 819
	}

820 821
	log.Debug(
		rpcReceived(method),
822
		zap.String("traceID", traceID),
823
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
824 825
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
826 827

	if err := node.sched.ddQueue.Enqueue(g); err != nil {
828 829
		log.Warn(
			rpcFailedToEnqueue(method),
830 831
			zap.Error(err),
			zap.String("traceID", traceID),
832
			zap.String("role", typeutil.ProxyRole),
833 834 835
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

836
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
837
			metrics.AbandonLabel).Inc()
838

G
godchen 已提交
839
		return &milvuspb.GetCollectionStatisticsResponse{
840
			Status: &commonpb.Status{
841
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
842 843 844 845 846
				Reason:    err.Error(),
			},
		}, nil
	}

847 848
	log.Debug(
		rpcEnqueued(method),
849
		zap.String("traceID", traceID),
850
		zap.String("role", typeutil.ProxyRole),
851
		zap.Int64("msgID", g.ID()),
852 853
		zap.Uint64("BeginTS", g.BeginTs()),
		zap.Uint64("EndTS", g.EndTs()),
D
dragondriver 已提交
854 855
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
856 857

	if err := g.WaitToFinish(); err != nil {
858 859
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
860
			zap.Error(err),
861
			zap.String("traceID", traceID),
862
			zap.String("role", typeutil.ProxyRole),
863 864 865
			zap.Int64("MsgID", g.ID()),
			zap.Uint64("BeginTS", g.BeginTs()),
			zap.Uint64("EndTS", g.EndTs()),
D
dragondriver 已提交
866 867 868
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

869
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
870
			metrics.FailLabel).Inc()
871

G
godchen 已提交
872
		return &milvuspb.GetCollectionStatisticsResponse{
873
			Status: &commonpb.Status{
874
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
875 876 877 878 879
				Reason:    err.Error(),
			},
		}, nil
	}

880 881
	log.Debug(
		rpcDone(method),
882
		zap.String("traceID", traceID),
883
		zap.String("role", typeutil.ProxyRole),
884
		zap.Int64("msgID", g.ID()),
885 886 887 888 889
		zap.Uint64("BeginTS", g.BeginTs()),
		zap.Uint64("EndTS", g.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

890
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
891
		metrics.SuccessLabel).Inc()
892
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
893
	return g.result, nil
894 895
}

896
// ShowCollections list all collections in Milvus.
C
Cai Yudong 已提交
897
func (node *Proxy) ShowCollections(ctx context.Context, request *milvuspb.ShowCollectionsRequest) (*milvuspb.ShowCollectionsResponse, error) {
898 899 900 901 902
	if !node.checkHealthy() {
		return &milvuspb.ShowCollectionsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
903 904
	method := "ShowCollections"
	tr := timerecord.NewTimeRecorder(method)
905
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
906

907
	sct := &showCollectionsTask{
G
godchen 已提交
908 909 910
		ctx:                    ctx,
		Condition:              NewTaskCondition(ctx),
		ShowCollectionsRequest: request,
911
		queryCoord:             node.queryCoord,
912
		rootCoord:              node.rootCoord,
913 914
	}

915
	log.Debug("ShowCollections received",
916
		zap.String("role", typeutil.ProxyRole),
917 918 919 920 921 922
		zap.String("DbName", request.DbName),
		zap.Uint64("TimeStamp", request.TimeStamp),
		zap.String("ShowType", request.Type.String()),
		zap.Any("CollectionNames", request.CollectionNames),
	)

923
	err := node.sched.ddQueue.Enqueue(sct)
924
	if err != nil {
925 926
		log.Warn("ShowCollections failed to enqueue",
			zap.Error(err),
927
			zap.String("role", typeutil.ProxyRole),
928 929 930 931 932 933
			zap.String("DbName", request.DbName),
			zap.Uint64("TimeStamp", request.TimeStamp),
			zap.String("ShowType", request.Type.String()),
			zap.Any("CollectionNames", request.CollectionNames),
		)

934
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
G
godchen 已提交
935
		return &milvuspb.ShowCollectionsResponse{
936
			Status: &commonpb.Status{
937
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
938 939 940 941 942
				Reason:    err.Error(),
			},
		}, nil
	}

943
	log.Debug("ShowCollections enqueued",
944
		zap.String("role", typeutil.ProxyRole),
945
		zap.Int64("MsgID", sct.ID()),
946
		zap.String("DbName", sct.ShowCollectionsRequest.DbName),
947
		zap.Uint64("TimeStamp", request.TimeStamp),
948 949 950
		zap.String("ShowType", sct.ShowCollectionsRequest.Type.String()),
		zap.Any("CollectionNames", sct.ShowCollectionsRequest.CollectionNames),
	)
D
dragondriver 已提交
951

952 953
	err = sct.WaitToFinish()
	if err != nil {
954 955
		log.Warn("ShowCollections failed to WaitToFinish",
			zap.Error(err),
956
			zap.String("role", typeutil.ProxyRole),
957 958 959 960 961 962 963
			zap.Int64("MsgID", sct.ID()),
			zap.String("DbName", request.DbName),
			zap.Uint64("TimeStamp", request.TimeStamp),
			zap.String("ShowType", request.Type.String()),
			zap.Any("CollectionNames", request.CollectionNames),
		)

964
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
965

G
godchen 已提交
966
		return &milvuspb.ShowCollectionsResponse{
967
			Status: &commonpb.Status{
968
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
969 970 971 972 973
				Reason:    err.Error(),
			},
		}, nil
	}

974
	log.Debug("ShowCollections Done",
975
		zap.String("role", typeutil.ProxyRole),
976 977 978 979
		zap.Int64("MsgID", sct.ID()),
		zap.String("DbName", request.DbName),
		zap.Uint64("TimeStamp", request.TimeStamp),
		zap.String("ShowType", request.Type.String()),
980 981
		zap.Int("len(CollectionNames)", len(request.CollectionNames)),
		zap.Int("num_collections", len(sct.result.CollectionNames)))
982

983 984
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
985 986 987
	return sct.result, nil
}

J
jaime 已提交
988 989 990 991 992 993 994 995 996 997 998
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()
	traceID, _, _ := trace.InfoFromSpan(sp)
	method := "AlterCollection"
	tr := timerecord.NewTimeRecorder(method)

999
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
J
jaime 已提交
1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023

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

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

	if err := node.sched.ddQueue.Enqueue(act); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
			zap.String("role", typeutil.ProxyRole),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

1024
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
J
jaime 已提交
1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.Int64("MsgID", act.ID()),
		zap.Uint64("BeginTs", act.BeginTs()),
		zap.Uint64("EndTs", act.EndTs()),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

	if err := act.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
			zap.Error(err),
			zap.String("traceID", traceID),
			zap.String("role", typeutil.ProxyRole),
			zap.Int64("MsgID", act.ID()),
			zap.Uint64("BeginTs", act.BeginTs()),
			zap.Uint64("EndTs", act.EndTs()),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

1054
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
J
jaime 已提交
1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.Int64("MsgID", act.ID()),
		zap.Uint64("BeginTs", act.BeginTs()),
		zap.Uint64("EndTs", act.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

1071 1072
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
J
jaime 已提交
1073 1074 1075
	return act.result, nil
}

1076
// CreatePartition create a partition in specific collection.
C
Cai Yudong 已提交
1077
func (node *Proxy) CreatePartition(ctx context.Context, request *milvuspb.CreatePartitionRequest) (*commonpb.Status, error) {
1078 1079 1080
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
1081

1082
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreatePartition")
1083 1084
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
1085 1086
	method := "CreatePartition"
	tr := timerecord.NewTimeRecorder(method)
1087
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
1088

1089
	cpt := &createPartitionTask{
S
sunby 已提交
1090
		ctx:                    ctx,
1091 1092
		Condition:              NewTaskCondition(ctx),
		CreatePartitionRequest: request,
1093
		rootCoord:              node.rootCoord,
1094 1095 1096
		result:                 nil,
	}

1097 1098 1099
	log.Debug(
		rpcReceived("CreatePartition"),
		zap.String("traceID", traceID),
1100
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1101 1102 1103
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
1104 1105 1106 1107 1108 1109

	if err := node.sched.ddQueue.Enqueue(cpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue("CreatePartition"),
			zap.Error(err),
			zap.String("traceID", traceID),
1110
			zap.String("role", typeutil.ProxyRole),
1111 1112 1113 1114
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

1115
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
1116

1117
		return &commonpb.Status{
1118
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1119 1120 1121
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
1122

1123 1124 1125
	log.Debug(
		rpcEnqueued("CreatePartition"),
		zap.String("traceID", traceID),
1126
		zap.String("role", typeutil.ProxyRole),
1127 1128 1129
		zap.Int64("MsgID", cpt.ID()),
		zap.Uint64("BeginTS", cpt.BeginTs()),
		zap.Uint64("EndTS", cpt.EndTs()),
D
dragondriver 已提交
1130 1131 1132
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
1133 1134 1135 1136

	if err := cpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish("CreatePartition"),
D
dragondriver 已提交
1137
			zap.Error(err),
1138
			zap.String("traceID", traceID),
1139
			zap.String("role", typeutil.ProxyRole),
1140 1141 1142
			zap.Int64("MsgID", cpt.ID()),
			zap.Uint64("BeginTS", cpt.BeginTs()),
			zap.Uint64("EndTS", cpt.EndTs()),
D
dragondriver 已提交
1143 1144 1145 1146
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

1147
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
1148

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

	log.Debug(
		rpcDone("CreatePartition"),
		zap.String("traceID", traceID),
1158
		zap.String("role", typeutil.ProxyRole),
1159 1160 1161 1162 1163 1164 1165
		zap.Int64("MsgID", cpt.ID()),
		zap.Uint64("BeginTS", cpt.BeginTs()),
		zap.Uint64("EndTS", cpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))

1166 1167
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1168 1169 1170
	return cpt.result, nil
}

1171
// DropPartition drop a partition in specific collection.
C
Cai Yudong 已提交
1172
func (node *Proxy) DropPartition(ctx context.Context, request *milvuspb.DropPartitionRequest) (*commonpb.Status, error) {
1173 1174 1175
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
1176

1177
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DropPartition")
1178 1179
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
1180 1181
	method := "DropPartition"
	tr := timerecord.NewTimeRecorder(method)
1182
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
1183

1184
	dpt := &dropPartitionTask{
S
sunby 已提交
1185
		ctx:                  ctx,
1186 1187
		Condition:            NewTaskCondition(ctx),
		DropPartitionRequest: request,
1188
		rootCoord:            node.rootCoord,
1189 1190 1191
		result:               nil,
	}

1192 1193 1194
	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1195
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1196 1197 1198
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
1199 1200 1201 1202 1203 1204

	if err := node.sched.ddQueue.Enqueue(dpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1205
			zap.String("role", typeutil.ProxyRole),
1206 1207 1208 1209
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

1210
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
1211

1212
		return &commonpb.Status{
1213
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1214 1215 1216
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
1217

1218 1219 1220
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1221
		zap.String("role", typeutil.ProxyRole),
1222 1223 1224
		zap.Int64("MsgID", dpt.ID()),
		zap.Uint64("BeginTS", dpt.BeginTs()),
		zap.Uint64("EndTS", dpt.EndTs()),
D
dragondriver 已提交
1225 1226 1227
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
1228 1229 1230 1231

	if err := dpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1232
			zap.Error(err),
1233
			zap.String("traceID", traceID),
1234
			zap.String("role", typeutil.ProxyRole),
1235 1236 1237
			zap.Int64("MsgID", dpt.ID()),
			zap.Uint64("BeginTS", dpt.BeginTs()),
			zap.Uint64("EndTS", dpt.EndTs()),
D
dragondriver 已提交
1238 1239 1240 1241
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

1242
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
1243

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

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1253
		zap.String("role", typeutil.ProxyRole),
1254 1255 1256 1257 1258 1259 1260
		zap.Int64("MsgID", dpt.ID()),
		zap.Uint64("BeginTS", dpt.BeginTs()),
		zap.Uint64("EndTS", dpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))

1261 1262
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1263 1264 1265
	return dpt.result, nil
}

1266
// HasPartition check if partition exist.
C
Cai Yudong 已提交
1267
func (node *Proxy) HasPartition(ctx context.Context, request *milvuspb.HasPartitionRequest) (*milvuspb.BoolResponse, error) {
1268 1269 1270 1271 1272
	if !node.checkHealthy() {
		return &milvuspb.BoolResponse{
			Status: unhealthyStatus(),
		}, nil
	}
D
dragondriver 已提交
1273

D
dragondriver 已提交
1274
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-HasPartition")
D
dragondriver 已提交
1275 1276
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
1277 1278 1279
	method := "HasPartition"
	tr := timerecord.NewTimeRecorder(method)
	//TODO: use collectionID instead of collectionName
1280
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1281
		metrics.TotalLabel).Inc()
D
dragondriver 已提交
1282

1283
	hpt := &hasPartitionTask{
S
sunby 已提交
1284
		ctx:                 ctx,
1285 1286
		Condition:           NewTaskCondition(ctx),
		HasPartitionRequest: request,
1287
		rootCoord:           node.rootCoord,
1288 1289 1290
		result:              nil,
	}

D
dragondriver 已提交
1291 1292 1293
	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1294
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1295 1296 1297
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
D
dragondriver 已提交
1298 1299 1300 1301 1302 1303

	if err := node.sched.ddQueue.Enqueue(hpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1304
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1305 1306 1307 1308
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

1309
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1310
			metrics.AbandonLabel).Inc()
1311

1312 1313
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
1314
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1315 1316 1317 1318 1319
				Reason:    err.Error(),
			},
			Value: false,
		}, nil
	}
D
dragondriver 已提交
1320

D
dragondriver 已提交
1321 1322 1323
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1324
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1325 1326 1327
		zap.Int64("MsgID", hpt.ID()),
		zap.Uint64("BeginTS", hpt.BeginTs()),
		zap.Uint64("EndTS", hpt.EndTs()),
D
dragondriver 已提交
1328 1329 1330
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
D
dragondriver 已提交
1331 1332 1333 1334

	if err := hpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1335
			zap.Error(err),
D
dragondriver 已提交
1336
			zap.String("traceID", traceID),
1337
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1338 1339 1340
			zap.Int64("MsgID", hpt.ID()),
			zap.Uint64("BeginTS", hpt.BeginTs()),
			zap.Uint64("EndTS", hpt.EndTs()),
D
dragondriver 已提交
1341 1342 1343 1344
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

1345
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1346
			metrics.FailLabel).Inc()
1347

1348 1349
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
1350
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1351 1352 1353 1354 1355
				Reason:    err.Error(),
			},
			Value: false,
		}, nil
	}
D
dragondriver 已提交
1356 1357 1358 1359

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1360
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1361 1362 1363 1364 1365 1366 1367
		zap.Int64("MsgID", hpt.ID()),
		zap.Uint64("BeginTS", hpt.BeginTs()),
		zap.Uint64("EndTS", hpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))

1368
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1369
		metrics.SuccessLabel).Inc()
1370
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1371 1372 1373
	return hpt.result, nil
}

1374
// LoadPartitions load specific partitions into query nodes.
C
Cai Yudong 已提交
1375
func (node *Proxy) LoadPartitions(ctx context.Context, request *milvuspb.LoadPartitionsRequest) (*commonpb.Status, error) {
1376 1377 1378
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
1379

D
dragondriver 已提交
1380
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-LoadPartitions")
1381 1382
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
1383 1384
	method := "LoadPartitions"
	tr := timerecord.NewTimeRecorder(method)
1385 1386
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
1387
	lpt := &loadPartitionsTask{
G
godchen 已提交
1388 1389 1390
		ctx:                   ctx,
		Condition:             NewTaskCondition(ctx),
		LoadPartitionsRequest: request,
1391
		queryCoord:            node.queryCoord,
C
cai.zhang 已提交
1392
		indexCoord:            node.indexCoord,
1393 1394
	}

1395 1396 1397
	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1398
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1399 1400 1401
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
1402 1403 1404 1405 1406 1407

	if err := node.sched.ddQueue.Enqueue(lpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1408
			zap.String("role", typeutil.ProxyRole),
1409 1410 1411 1412
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))

1413
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1414
			metrics.AbandonLabel).Inc()
1415

1416
		return &commonpb.Status{
1417
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1418 1419 1420 1421
			Reason:    err.Error(),
		}, nil
	}

1422 1423 1424
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1425
		zap.String("role", typeutil.ProxyRole),
1426 1427 1428
		zap.Int64("MsgID", lpt.ID()),
		zap.Uint64("BeginTS", lpt.BeginTs()),
		zap.Uint64("EndTS", lpt.EndTs()),
D
dragondriver 已提交
1429 1430 1431
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
1432 1433 1434 1435

	if err := lpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1436
			zap.Error(err),
1437
			zap.String("traceID", traceID),
1438
			zap.String("role", typeutil.ProxyRole),
1439 1440 1441
			zap.Int64("MsgID", lpt.ID()),
			zap.Uint64("BeginTS", lpt.BeginTs()),
			zap.Uint64("EndTS", lpt.EndTs()),
D
dragondriver 已提交
1442 1443 1444 1445
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))

1446
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1447
			metrics.FailLabel).Inc()
1448

1449
		return &commonpb.Status{
1450
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1451 1452 1453 1454
			Reason:    err.Error(),
		}, nil
	}

1455 1456 1457
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1458
		zap.String("role", typeutil.ProxyRole),
1459 1460 1461 1462 1463 1464 1465
		zap.Int64("MsgID", lpt.ID()),
		zap.Uint64("BeginTS", lpt.BeginTs()),
		zap.Uint64("EndTS", lpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))

1466
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1467
		metrics.SuccessLabel).Inc()
1468
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1469
	return lpt.result, nil
1470 1471
}

1472
// ReleasePartitions release specific partitions from query nodes.
C
Cai Yudong 已提交
1473
func (node *Proxy) ReleasePartitions(ctx context.Context, request *milvuspb.ReleasePartitionsRequest) (*commonpb.Status, error) {
1474 1475 1476
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
1477 1478 1479 1480 1481

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

1482
	rpt := &releasePartitionsTask{
G
godchen 已提交
1483 1484 1485
		ctx:                      ctx,
		Condition:                NewTaskCondition(ctx),
		ReleasePartitionsRequest: request,
1486
		queryCoord:               node.queryCoord,
1487 1488
	}

1489
	method := "ReleasePartitions"
1490
	tr := timerecord.NewTimeRecorder(method)
1491 1492
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
1493 1494 1495
	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1496
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1497 1498 1499
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
1500 1501 1502 1503 1504 1505

	if err := node.sched.ddQueue.Enqueue(rpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1506
			zap.String("role", typeutil.ProxyRole),
1507 1508 1509 1510
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))

1511
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1512
			metrics.AbandonLabel).Inc()
1513

1514
		return &commonpb.Status{
1515
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1516 1517 1518 1519
			Reason:    err.Error(),
		}, nil
	}

1520 1521 1522
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1523
		zap.String("role", typeutil.ProxyRole),
1524 1525 1526
		zap.Int64("msgID", rpt.Base.MsgID),
		zap.Uint64("BeginTS", rpt.BeginTs()),
		zap.Uint64("EndTS", rpt.EndTs()),
D
dragondriver 已提交
1527 1528 1529
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
1530 1531 1532 1533

	if err := rpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1534
			zap.Error(err),
1535
			zap.String("traceID", traceID),
1536
			zap.String("role", typeutil.ProxyRole),
1537 1538 1539
			zap.Int64("msgID", rpt.Base.MsgID),
			zap.Uint64("BeginTS", rpt.BeginTs()),
			zap.Uint64("EndTS", rpt.EndTs()),
D
dragondriver 已提交
1540 1541 1542 1543
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))

1544
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1545
			metrics.FailLabel).Inc()
1546

1547
		return &commonpb.Status{
1548
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1549 1550 1551 1552
			Reason:    err.Error(),
		}, nil
	}

1553 1554 1555
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1556
		zap.String("role", typeutil.ProxyRole),
1557 1558 1559 1560 1561 1562 1563
		zap.Int64("msgID", rpt.Base.MsgID),
		zap.Uint64("BeginTS", rpt.BeginTs()),
		zap.Uint64("EndTS", rpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))

1564
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1565
		metrics.SuccessLabel).Inc()
1566
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1567
	return rpt.result, nil
1568 1569
}

1570
// GetPartitionStatistics get the statistics of partition, such as num_rows.
C
Cai Yudong 已提交
1571
func (node *Proxy) GetPartitionStatistics(ctx context.Context, request *milvuspb.GetPartitionStatisticsRequest) (*milvuspb.GetPartitionStatisticsResponse, error) {
1572 1573 1574 1575 1576
	if !node.checkHealthy() {
		return &milvuspb.GetPartitionStatisticsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1577 1578 1579 1580

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetPartitionStatistics")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
1581 1582
	method := "GetPartitionStatistics"
	tr := timerecord.NewTimeRecorder(method)
1583 1584
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
1585

1586
	g := &getPartitionStatisticsTask{
1587 1588 1589
		ctx:                           ctx,
		Condition:                     NewTaskCondition(ctx),
		GetPartitionStatisticsRequest: request,
1590
		dataCoord:                     node.dataCoord,
1591 1592
	}

1593 1594 1595
	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1596
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1597 1598 1599
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
1600 1601 1602 1603 1604 1605

	if err := node.sched.ddQueue.Enqueue(g); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1606
			zap.String("role", typeutil.ProxyRole),
1607 1608 1609 1610
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

1611
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1612
			metrics.AbandonLabel).Inc()
1613

1614 1615 1616 1617 1618 1619 1620 1621
		return &milvuspb.GetPartitionStatisticsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

1622 1623 1624
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1625
		zap.String("role", typeutil.ProxyRole),
1626 1627 1628
		zap.Int64("msgID", g.ID()),
		zap.Uint64("BeginTS", g.BeginTs()),
		zap.Uint64("EndTS", g.EndTs()),
1629 1630 1631
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
1632 1633 1634 1635

	if err := g.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
1636
			zap.Error(err),
1637
			zap.String("traceID", traceID),
1638
			zap.String("role", typeutil.ProxyRole),
1639 1640 1641
			zap.Int64("msgID", g.ID()),
			zap.Uint64("BeginTS", g.BeginTs()),
			zap.Uint64("EndTS", g.EndTs()),
1642 1643 1644 1645
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

1646
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1647
			metrics.FailLabel).Inc()
1648

1649 1650 1651 1652 1653 1654 1655 1656
		return &milvuspb.GetPartitionStatisticsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

1657 1658 1659
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1660
		zap.String("role", typeutil.ProxyRole),
1661 1662 1663 1664 1665 1666 1667
		zap.Int64("msgID", g.ID()),
		zap.Uint64("BeginTS", g.BeginTs()),
		zap.Uint64("EndTS", g.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))

1668
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1669
		metrics.SuccessLabel).Inc()
1670
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1671
	return g.result, nil
1672 1673
}

1674
// ShowPartitions list all partitions in the specific collection.
C
Cai Yudong 已提交
1675
func (node *Proxy) ShowPartitions(ctx context.Context, request *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
1676 1677 1678 1679 1680
	if !node.checkHealthy() {
		return &milvuspb.ShowPartitionsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1681 1682 1683 1684 1685

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

1686
	spt := &showPartitionsTask{
G
godchen 已提交
1687 1688 1689
		ctx:                   ctx,
		Condition:             NewTaskCondition(ctx),
		ShowPartitionsRequest: request,
1690
		rootCoord:             node.rootCoord,
1691
		queryCoord:            node.queryCoord,
G
godchen 已提交
1692
		result:                nil,
1693 1694
	}

1695
	method := "ShowPartitions"
1696 1697
	tr := timerecord.NewTimeRecorder(method)
	//TODO: use collectionID instead of collectionName
1698
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1699
		metrics.TotalLabel).Inc()
1700 1701 1702 1703

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1704
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1705
		zap.Any("request", request))
1706 1707 1708 1709 1710 1711

	if err := node.sched.ddQueue.Enqueue(spt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1712
			zap.String("role", typeutil.ProxyRole),
1713 1714
			zap.Any("request", request))

1715
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1716
			metrics.AbandonLabel).Inc()
1717

G
godchen 已提交
1718
		return &milvuspb.ShowPartitionsResponse{
1719
			Status: &commonpb.Status{
1720
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1721 1722 1723 1724 1725
				Reason:    err.Error(),
			},
		}, nil
	}

1726 1727 1728
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1729
		zap.String("role", typeutil.ProxyRole),
1730 1731 1732
		zap.Int64("msgID", spt.ID()),
		zap.Uint64("BeginTS", spt.BeginTs()),
		zap.Uint64("EndTS", spt.EndTs()),
1733 1734
		zap.String("db", spt.ShowPartitionsRequest.DbName),
		zap.String("collection", spt.ShowPartitionsRequest.CollectionName),
1735 1736 1737 1738 1739
		zap.Any("partitions", spt.ShowPartitionsRequest.PartitionNames))

	if err := spt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1740
			zap.Error(err),
1741
			zap.String("traceID", traceID),
1742
			zap.String("role", typeutil.ProxyRole),
1743 1744 1745 1746 1747 1748
			zap.Int64("msgID", spt.ID()),
			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 已提交
1749

1750
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1751
			metrics.FailLabel).Inc()
1752

G
godchen 已提交
1753
		return &milvuspb.ShowPartitionsResponse{
1754
			Status: &commonpb.Status{
1755
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1756 1757 1758 1759
				Reason:    err.Error(),
			},
		}, nil
	}
1760 1761 1762 1763

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1764
		zap.String("role", typeutil.ProxyRole),
1765 1766 1767 1768 1769 1770 1771
		zap.Int64("msgID", spt.ID()),
		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))

1772
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1773
		metrics.SuccessLabel).Inc()
1774
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
1775 1776 1777
	return spt.result, nil
}

S
SimFG 已提交
1778 1779
func (node *Proxy) getCollectionProgress(ctx context.Context, request *milvuspb.GetLoadingProgressRequest, collectionID int64) (int64, error) {
	resp, err := node.queryCoord.ShowCollections(ctx, &querypb.ShowCollectionsRequest{
1780 1781 1782 1783
		Base: commonpbutil.NewMsgBaseCopy(
			request.Base,
			commonpbutil.WithMsgType(commonpb.MsgType_DescribeCollection),
		),
S
SimFG 已提交
1784 1785 1786 1787 1788
		CollectionIDs: []int64{collectionID},
	})
	if err != nil {
		return 0, err
	}
X
xige-16 已提交
1789 1790 1791 1792 1793

	if resp.Status.ErrorCode != commonpb.ErrorCode_Success {
		return 0, errors.New(resp.Status.Reason)
	}

S
SimFG 已提交
1794 1795 1796 1797 1798 1799 1800 1801 1802 1803 1804 1805 1806 1807 1808 1809 1810 1811
	if len(resp.InMemoryPercentages) == 0 {
		return 0, errors.New("fail to show collections from the querycoord, no data")
	}
	return resp.InMemoryPercentages[0], nil
}

func (node *Proxy) getPartitionProgress(ctx context.Context, request *milvuspb.GetLoadingProgressRequest, collectionID int64) (int64, error) {
	IDs2Names := make(map[int64]string)
	partitionIDs := make([]int64, 0)
	for _, partitionName := range request.PartitionNames {
		partitionID, err := globalMetaCache.GetPartitionID(ctx, request.CollectionName, partitionName)
		if err != nil {
			return 0, err
		}
		IDs2Names[partitionID] = partitionName
		partitionIDs = append(partitionIDs, partitionID)
	}
	resp, err := node.queryCoord.ShowPartitions(ctx, &querypb.ShowPartitionsRequest{
1812 1813 1814 1815
		Base: commonpbutil.NewMsgBaseCopy(
			request.Base,
			commonpbutil.WithMsgType(commonpb.MsgType_ShowPartitions),
		),
S
SimFG 已提交
1816 1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829 1830 1831 1832 1833 1834 1835 1836 1837 1838
		CollectionID: collectionID,
		PartitionIDs: partitionIDs,
	})
	if err != nil {
		return 0, err
	}
	if len(resp.InMemoryPercentages) != len(partitionIDs) {
		return 0, errors.New("fail to show partitions from the querycoord, invalid data num")
	}
	var progress int64
	for _, p := range resp.InMemoryPercentages {
		progress += p
	}
	progress /= int64(len(partitionIDs))
	return progress, nil
}

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)
1839
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetLoadingProgress")
S
SimFG 已提交
1840 1841
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
1842
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
S
SimFG 已提交
1843 1844 1845 1846 1847 1848 1849 1850
	logger.Info(
		rpcReceived(method),
		zap.String("traceID", traceID),
		zap.Any("request", request))

	getErrResponse := func(err error) *milvuspb.GetLoadingProgressResponse {
		logger.Error("fail to get loading progress", zap.String("collection_name", request.CollectionName),
			zap.Strings("partition_name", request.PartitionNames), zap.Error(err))
1851
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
S
SimFG 已提交
1852 1853 1854 1855 1856 1857 1858 1859 1860 1861 1862 1863 1864 1865
		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
	}
1866 1867 1868 1869 1870 1871
	msgBase := commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_SystemInfo),
		commonpbutil.WithMsgID(0),
		commonpbutil.WithTimeStamp(0),
		commonpbutil.WithSourceID(Params.ProxyCfg.GetNodeID()),
	)
S
SimFG 已提交
1872 1873 1874 1875 1876 1877 1878 1879 1880 1881 1882 1883 1884 1885 1886 1887 1888 1889 1890 1891 1892 1893 1894
	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 = node.getCollectionProgress(ctx, request, collectionID); err != nil {
			return getErrResponse(err), nil
		}
	} else {
		if progress, err = node.getPartitionProgress(ctx, request, collectionID); err != nil {
			return getErrResponse(err), nil
		}
	}

	logger.Info(
		rpcDone(method),
		zap.String("traceID", traceID),
		zap.Any("request", request))
1895 1896
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
S
SimFG 已提交
1897 1898 1899 1900 1901 1902 1903 1904
	return &milvuspb.GetLoadingProgressResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
		},
		Progress: progress,
	}, nil
}

1905
// CreateIndex create index for collection.
C
Cai Yudong 已提交
1906
func (node *Proxy) CreateIndex(ctx context.Context, request *milvuspb.CreateIndexRequest) (*commonpb.Status, error) {
1907 1908 1909
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
1910

1911
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreateIndex")
D
dragondriver 已提交
1912 1913 1914
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

1915
	cit := &createIndexTask{
Z
zhenshan.cao 已提交
1916 1917 1918 1919 1920
		ctx:        ctx,
		Condition:  NewTaskCondition(ctx),
		req:        request,
		rootCoord:  node.rootCoord,
		indexCoord: node.indexCoord,
1921 1922
	}

D
dragondriver 已提交
1923
	method := "CreateIndex"
1924
	tr := timerecord.NewTimeRecorder(method)
1925 1926
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
D
dragondriver 已提交
1927 1928 1929
	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1930
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1931 1932 1933 1934
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.Any("extra_params", request.ExtraParams))
D
dragondriver 已提交
1935 1936 1937 1938 1939 1940

	if err := node.sched.ddQueue.Enqueue(cit); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1941
			zap.String("role", typeutil.ProxyRole),
Z
zhenshan.cao 已提交
1942 1943 1944 1945
			zap.String("db", request.GetDbName()),
			zap.String("collection", request.GetCollectionName()),
			zap.String("field", request.GetFieldName()),
			zap.Any("extra_params", request.GetExtraParams()))
D
dragondriver 已提交
1946

1947
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1948
			metrics.AbandonLabel).Inc()
1949

1950
		return &commonpb.Status{
1951
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1952 1953 1954 1955
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
1956 1957 1958
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1959
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1960 1961 1962
		zap.Int64("MsgID", cit.ID()),
		zap.Uint64("BeginTs", cit.BeginTs()),
		zap.Uint64("EndTs", cit.EndTs()),
D
dragondriver 已提交
1963 1964 1965 1966
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.Any("extra_params", request.ExtraParams))
D
dragondriver 已提交
1967 1968 1969 1970

	if err := cit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1971
			zap.Error(err),
D
dragondriver 已提交
1972
			zap.String("traceID", traceID),
1973
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1974 1975 1976
			zap.Int64("MsgID", cit.ID()),
			zap.Uint64("BeginTs", cit.BeginTs()),
			zap.Uint64("EndTs", cit.EndTs()),
D
dragondriver 已提交
1977 1978 1979 1980 1981
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.Any("extra_params", request.ExtraParams))

1982
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
1983
			metrics.FailLabel).Inc()
1984

1985
		return &commonpb.Status{
1986
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1987 1988 1989 1990
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
1991 1992 1993
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1994
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1995 1996 1997 1998 1999 2000 2001 2002
		zap.Int64("MsgID", cit.ID()),
		zap.Uint64("BeginTs", cit.BeginTs()),
		zap.Uint64("EndTs", cit.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.Any("extra_params", request.ExtraParams))

2003
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2004
		metrics.SuccessLabel).Inc()
2005
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
2006 2007 2008
	return cit.result, nil
}

2009
// DescribeIndex get the meta information of index, such as index state, index id and etc.
C
Cai Yudong 已提交
2010
func (node *Proxy) DescribeIndex(ctx context.Context, request *milvuspb.DescribeIndexRequest) (*milvuspb.DescribeIndexResponse, error) {
2011 2012 2013 2014 2015
	if !node.checkHealthy() {
		return &milvuspb.DescribeIndexResponse{
			Status: unhealthyStatus(),
		}, nil
	}
2016 2017 2018 2019 2020

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

2021
	dit := &describeIndexTask{
S
sunby 已提交
2022
		ctx:                  ctx,
2023 2024
		Condition:            NewTaskCondition(ctx),
		DescribeIndexRequest: request,
2025
		indexCoord:           node.indexCoord,
2026 2027
	}

2028 2029 2030
	method := "DescribeIndex"
	// avoid data race
	indexName := request.IndexName
2031
	tr := timerecord.NewTimeRecorder(method)
2032 2033
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
2034 2035 2036
	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
2037
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2038 2039 2040
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
2041 2042 2043 2044 2045 2046 2047
		zap.String("index name", indexName))

	if err := node.sched.ddQueue.Enqueue(dit); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
2048
			zap.String("role", typeutil.ProxyRole),
2049 2050 2051 2052 2053
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", indexName))

2054
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2055
			metrics.AbandonLabel).Inc()
2056

2057 2058
		return &milvuspb.DescribeIndexResponse{
			Status: &commonpb.Status{
2059
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2060 2061 2062 2063 2064
				Reason:    err.Error(),
			},
		}, nil
	}

2065 2066 2067
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2068
		zap.String("role", typeutil.ProxyRole),
2069 2070 2071
		zap.Int64("MsgID", dit.ID()),
		zap.Uint64("BeginTs", dit.BeginTs()),
		zap.Uint64("EndTs", dit.EndTs()),
D
dragondriver 已提交
2072 2073 2074
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
2075 2076 2077 2078 2079
		zap.String("index name", indexName))

	if err := dit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
2080
			zap.Error(err),
2081
			zap.String("traceID", traceID),
2082
			zap.String("role", typeutil.ProxyRole),
2083 2084 2085
			zap.Int64("MsgID", dit.ID()),
			zap.Uint64("BeginTs", dit.BeginTs()),
			zap.Uint64("EndTs", dit.EndTs()),
D
dragondriver 已提交
2086 2087 2088
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
2089
			zap.String("index name", indexName))
D
dragondriver 已提交
2090

Z
zhenshan.cao 已提交
2091 2092 2093 2094
		errCode := commonpb.ErrorCode_UnexpectedError
		if dit.result != nil {
			errCode = dit.result.Status.GetErrorCode()
		}
2095
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2096
			metrics.FailLabel).Inc()
2097

2098 2099
		return &milvuspb.DescribeIndexResponse{
			Status: &commonpb.Status{
Z
zhenshan.cao 已提交
2100
				ErrorCode: errCode,
2101 2102 2103 2104 2105
				Reason:    err.Error(),
			},
		}, nil
	}

2106 2107 2108
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
2109
		zap.String("role", typeutil.ProxyRole),
2110 2111 2112 2113 2114 2115 2116 2117
		zap.Int64("MsgID", dit.ID()),
		zap.Uint64("BeginTs", dit.BeginTs()),
		zap.Uint64("EndTs", dit.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", indexName))

2118
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2119
		metrics.SuccessLabel).Inc()
2120
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
2121 2122 2123
	return dit.result, nil
}

2124
// DropIndex drop the index of collection.
C
Cai Yudong 已提交
2125
func (node *Proxy) DropIndex(ctx context.Context, request *milvuspb.DropIndexRequest) (*commonpb.Status, error) {
2126 2127 2128
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2129 2130 2131 2132 2133

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

2134
	dit := &dropIndexTask{
S
sunby 已提交
2135
		ctx:              ctx,
B
BossZou 已提交
2136 2137
		Condition:        NewTaskCondition(ctx),
		DropIndexRequest: request,
2138
		indexCoord:       node.indexCoord,
2139
		queryCoord:       node.queryCoord,
B
BossZou 已提交
2140
	}
G
godchen 已提交
2141

D
dragondriver 已提交
2142
	method := "DropIndex"
2143
	tr := timerecord.NewTimeRecorder(method)
2144 2145
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
D
dragondriver 已提交
2146 2147 2148 2149

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
2150
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2151 2152 2153 2154 2155
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))

D
dragondriver 已提交
2156 2157 2158 2159 2160
	if err := node.sched.ddQueue.Enqueue(dit); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
2161
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2162 2163 2164 2165
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))
2166
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2167
			metrics.AbandonLabel).Inc()
D
dragondriver 已提交
2168

B
BossZou 已提交
2169
		return &commonpb.Status{
2170
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
B
BossZou 已提交
2171 2172 2173
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
2174

D
dragondriver 已提交
2175 2176 2177
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2178
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2179 2180 2181
		zap.Int64("MsgID", dit.ID()),
		zap.Uint64("BeginTs", dit.BeginTs()),
		zap.Uint64("EndTs", dit.EndTs()),
D
dragondriver 已提交
2182 2183 2184 2185
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
D
dragondriver 已提交
2186 2187 2188 2189

	if err := dit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
2190
			zap.Error(err),
D
dragondriver 已提交
2191
			zap.String("traceID", traceID),
2192
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2193 2194 2195
			zap.Int64("MsgID", dit.ID()),
			zap.Uint64("BeginTs", dit.BeginTs()),
			zap.Uint64("EndTs", dit.EndTs()),
D
dragondriver 已提交
2196 2197 2198 2199 2200
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))

2201
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2202
			metrics.FailLabel).Inc()
2203

B
BossZou 已提交
2204
		return &commonpb.Status{
2205
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
B
BossZou 已提交
2206 2207 2208
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
2209 2210 2211 2212

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
2213
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2214 2215 2216 2217 2218 2219 2220 2221
		zap.Int64("MsgID", dit.ID()),
		zap.Uint64("BeginTs", dit.BeginTs()),
		zap.Uint64("EndTs", dit.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))

2222
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2223
		metrics.SuccessLabel).Inc()
2224
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
B
BossZou 已提交
2225 2226 2227
	return dit.result, nil
}

2228 2229
// 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.
2230
// Deprecated: use DescribeIndex instead
C
Cai Yudong 已提交
2231
func (node *Proxy) GetIndexBuildProgress(ctx context.Context, request *milvuspb.GetIndexBuildProgressRequest) (*milvuspb.GetIndexBuildProgressResponse, error) {
2232 2233 2234 2235 2236
	if !node.checkHealthy() {
		return &milvuspb.GetIndexBuildProgressResponse{
			Status: unhealthyStatus(),
		}, nil
	}
2237 2238 2239 2240 2241

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

2242
	gibpt := &getIndexBuildProgressTask{
2243 2244 2245
		ctx:                          ctx,
		Condition:                    NewTaskCondition(ctx),
		GetIndexBuildProgressRequest: request,
2246 2247
		indexCoord:                   node.indexCoord,
		rootCoord:                    node.rootCoord,
2248
		dataCoord:                    node.dataCoord,
2249 2250
	}

2251
	method := "GetIndexBuildProgress"
2252
	tr := timerecord.NewTimeRecorder(method)
2253 2254
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
2255 2256 2257
	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
2258
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2259 2260 2261 2262
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
2263 2264 2265 2266 2267 2268

	if err := node.sched.ddQueue.Enqueue(gibpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
2269
			zap.String("role", typeutil.ProxyRole),
2270 2271 2272 2273
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))
2274
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2275
			metrics.AbandonLabel).Inc()
2276

2277 2278 2279 2280 2281 2282 2283 2284
		return &milvuspb.GetIndexBuildProgressResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

2285 2286 2287
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2288
		zap.String("role", typeutil.ProxyRole),
2289 2290 2291
		zap.Int64("MsgID", gibpt.ID()),
		zap.Uint64("BeginTs", gibpt.BeginTs()),
		zap.Uint64("EndTs", gibpt.EndTs()),
2292 2293 2294 2295
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
2296 2297 2298 2299

	if err := gibpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
2300
			zap.Error(err),
2301
			zap.String("traceID", traceID),
2302
			zap.String("role", typeutil.ProxyRole),
2303 2304 2305
			zap.Int64("MsgID", gibpt.ID()),
			zap.Uint64("BeginTs", gibpt.BeginTs()),
			zap.Uint64("EndTs", gibpt.EndTs()),
2306 2307 2308 2309
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))
2310
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2311
			metrics.FailLabel).Inc()
2312 2313 2314 2315 2316 2317 2318 2319

		return &milvuspb.GetIndexBuildProgressResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}
2320 2321 2322 2323

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
2324
		zap.String("role", typeutil.ProxyRole),
2325 2326 2327 2328 2329 2330 2331 2332
		zap.Int64("MsgID", gibpt.ID()),
		zap.Uint64("BeginTs", gibpt.BeginTs()),
		zap.Uint64("EndTs", gibpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName),
		zap.Any("result", gibpt.result))
2333

2334
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2335
		metrics.SuccessLabel).Inc()
2336
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
2337
	return gibpt.result, nil
2338 2339
}

2340
// GetIndexState get the build-state of index.
2341
// Deprecated: use DescribeIndex instead
C
Cai Yudong 已提交
2342
func (node *Proxy) GetIndexState(ctx context.Context, request *milvuspb.GetIndexStateRequest) (*milvuspb.GetIndexStateResponse, error) {
2343 2344 2345 2346 2347
	if !node.checkHealthy() {
		return &milvuspb.GetIndexStateResponse{
			Status: unhealthyStatus(),
		}, nil
	}
2348 2349 2350 2351 2352

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

2353
	dipt := &getIndexStateTask{
G
godchen 已提交
2354 2355 2356
		ctx:                  ctx,
		Condition:            NewTaskCondition(ctx),
		GetIndexStateRequest: request,
2357 2358
		indexCoord:           node.indexCoord,
		rootCoord:            node.rootCoord,
2359 2360
	}

2361
	method := "GetIndexState"
2362
	tr := timerecord.NewTimeRecorder(method)
2363 2364
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
2365 2366 2367
	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
2368
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2369 2370 2371 2372
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
2373 2374 2375 2376 2377 2378

	if err := node.sched.ddQueue.Enqueue(dipt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
2379
			zap.String("role", typeutil.ProxyRole),
2380 2381 2382 2383 2384
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))

2385
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2386
			metrics.AbandonLabel).Inc()
2387

G
godchen 已提交
2388
		return &milvuspb.GetIndexStateResponse{
2389
			Status: &commonpb.Status{
2390
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2391 2392 2393 2394 2395
				Reason:    err.Error(),
			},
		}, nil
	}

2396 2397 2398
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2399
		zap.String("role", typeutil.ProxyRole),
2400 2401 2402
		zap.Int64("MsgID", dipt.ID()),
		zap.Uint64("BeginTs", dipt.BeginTs()),
		zap.Uint64("EndTs", dipt.EndTs()),
D
dragondriver 已提交
2403 2404 2405 2406
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
2407 2408 2409 2410

	if err := dipt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
2411
			zap.Error(err),
2412
			zap.String("traceID", traceID),
2413
			zap.String("role", typeutil.ProxyRole),
2414 2415 2416
			zap.Int64("MsgID", dipt.ID()),
			zap.Uint64("BeginTs", dipt.BeginTs()),
			zap.Uint64("EndTs", dipt.EndTs()),
D
dragondriver 已提交
2417 2418 2419 2420
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))
2421
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2422
			metrics.FailLabel).Inc()
2423

G
godchen 已提交
2424
		return &milvuspb.GetIndexStateResponse{
2425
			Status: &commonpb.Status{
2426
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2427 2428 2429 2430 2431
				Reason:    err.Error(),
			},
		}, nil
	}

2432 2433 2434
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
2435
		zap.String("role", typeutil.ProxyRole),
2436 2437 2438 2439 2440 2441 2442 2443
		zap.Int64("MsgID", dipt.ID()),
		zap.Uint64("BeginTs", dipt.BeginTs()),
		zap.Uint64("EndTs", dipt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))

2444
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2445
		metrics.SuccessLabel).Inc()
2446
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
2447 2448 2449
	return dipt.result, nil
}

2450
// Insert insert records into collection.
C
Cai Yudong 已提交
2451
func (node *Proxy) Insert(ctx context.Context, request *milvuspb.InsertRequest) (*milvuspb.MutationResult, error) {
X
Xiangyu Wang 已提交
2452 2453 2454 2455 2456 2457
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Insert")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
	log.Info("Start processing insert request in Proxy", zap.String("traceID", traceID))
	defer log.Info("Finish processing insert request in Proxy", zap.String("traceID", traceID))

2458 2459 2460 2461 2462
	if !node.checkHealthy() {
		return &milvuspb.MutationResult{
			Status: unhealthyStatus(),
		}, nil
	}
2463 2464
	method := "Insert"
	tr := timerecord.NewTimeRecorder(method)
2465
	receiveSize := proto.Size(request)
2466 2467
	rateCol.Add(internalpb.RateType_DMLInsert.String(), float64(receiveSize))
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), metrics.InsertLabel).Add(float64(receiveSize))
D
dragondriver 已提交
2468

2469
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
2470
	it := &insertTask{
2471 2472
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),
X
xige-16 已提交
2473
		// req:       request,
2474 2475 2476 2477
		BaseInsertTask: BaseInsertTask{
			BaseMsg: msgstream.BaseMsg{
				HashValues: request.HashKeys,
			},
G
godchen 已提交
2478
			InsertRequest: internalpb.InsertRequest{
2479 2480 2481 2482 2483
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Insert),
					commonpbutil.WithMsgID(0),
					commonpbutil.WithSourceID(Params.ProxyCfg.GetNodeID()),
				),
2484 2485
				CollectionName: request.CollectionName,
				PartitionName:  request.PartitionName,
X
xige-16 已提交
2486 2487 2488
				FieldsData:     request.FieldsData,
				NumRows:        uint64(request.NumRows),
				Version:        internalpb.InsertDataVersion_ColumnBased,
2489
				// RowData: transfer column based request to this
2490 2491
			},
		},
2492
		idAllocator:   node.rowIDAllocator,
2493 2494 2495
		segIDAssigner: node.segAssigner,
		chMgr:         node.chMgr,
		chTicker:      node.chTicker,
2496
	}
2497 2498

	if len(it.PartitionName) <= 0 {
2499
		it.PartitionName = Params.CommonCfg.DefaultPartitionName
2500 2501
	}

X
Xiangyu Wang 已提交
2502
	constructFailedResponse := func(err error) *milvuspb.MutationResult {
X
xige-16 已提交
2503
		numRows := request.NumRows
2504 2505 2506 2507
		errIndex := make([]uint32, numRows)
		for i := uint32(0); i < numRows; i++ {
			errIndex[i] = i
		}
2508

X
Xiangyu Wang 已提交
2509 2510 2511 2512 2513 2514 2515
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
			ErrIndex: errIndex,
		}
2516 2517
	}

X
Xiangyu Wang 已提交
2518
	log.Debug("Enqueue insert request in Proxy",
2519
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2520 2521 2522 2523 2524
		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)),
2525 2526
		zap.Uint32("NumRows", request.NumRows),
		zap.String("traceID", traceID))
D
dragondriver 已提交
2527

X
Xiangyu Wang 已提交
2528 2529
	if err := node.sched.dmQueue.Enqueue(it); err != nil {
		log.Debug("Failed to enqueue insert task: " + err.Error())
2530
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2531
			metrics.AbandonLabel).Inc()
X
Xiangyu Wang 已提交
2532
		return constructFailedResponse(err), nil
2533
	}
D
dragondriver 已提交
2534

X
Xiangyu Wang 已提交
2535
	log.Debug("Detail of insert request in Proxy",
2536
		zap.String("role", typeutil.ProxyRole),
X
Xiangyu Wang 已提交
2537
		zap.Int64("msgID", it.Base.MsgID),
D
dragondriver 已提交
2538 2539 2540 2541 2542
		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),
X
Xiangyu Wang 已提交
2543 2544 2545 2546 2547
		zap.Uint32("NumRows", request.NumRows),
		zap.String("traceID", traceID))

	if err := it.WaitToFinish(); err != nil {
		log.Debug("Failed to execute insert task in task scheduler: "+err.Error(), zap.String("traceID", traceID))
2548
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2549
			metrics.FailLabel).Inc()
X
Xiangyu Wang 已提交
2550 2551 2552 2553 2554
		return constructFailedResponse(err), nil
	}

	if it.result.Status.ErrorCode != commonpb.ErrorCode_Success {
		setErrorIndex := func() {
X
xige-16 已提交
2555
			numRows := request.NumRows
X
Xiangyu Wang 已提交
2556 2557 2558 2559 2560 2561 2562 2563 2564 2565 2566
			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 已提交
2567
	it.result.InsertCnt = int64(request.NumRows)
D
dragondriver 已提交
2568

2569
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2570
		metrics.SuccessLabel).Inc()
2571 2572
	successCnt := it.result.InsertCnt - int64(len(it.result.ErrIndex))
	metrics.ProxyInsertVectors.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10)).Add(float64(successCnt))
2573
	metrics.ProxyMutationLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), metrics.InsertLabel).Observe(float64(tr.ElapseSpan().Milliseconds()))
2574
	metrics.ProxyCollectionMutationLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), metrics.InsertLabel, request.CollectionName).Observe(float64(tr.ElapseSpan().Milliseconds()))
2575 2576 2577
	return it.result, nil
}

2578
// Delete delete records from collection, then these records cannot be searched.
G
groot 已提交
2579
func (node *Proxy) Delete(ctx context.Context, request *milvuspb.DeleteRequest) (*milvuspb.MutationResult, error) {
2580 2581 2582
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Delete")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
2583 2584
	log.Info("Start processing delete request in Proxy", zap.String("traceID", traceID))
	defer log.Info("Finish processing delete request in Proxy", zap.String("traceID", traceID))
2585

2586
	receiveSize := proto.Size(request)
2587 2588
	rateCol.Add(internalpb.RateType_DMLDelete.String(), float64(receiveSize))
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), metrics.DeleteLabel).Add(float64(receiveSize))
2589

G
groot 已提交
2590 2591 2592 2593 2594 2595
	if !node.checkHealthy() {
		return &milvuspb.MutationResult{
			Status: unhealthyStatus(),
		}, nil
	}

2596 2597 2598
	method := "Delete"
	tr := timerecord.NewTimeRecorder(method)

2599
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2600
		metrics.TotalLabel).Inc()
2601
	dt := &deleteTask{
X
xige-16 已提交
2602 2603 2604
		ctx:        ctx,
		Condition:  NewTaskCondition(ctx),
		deleteExpr: request.Expr,
G
godchen 已提交
2605
		BaseDeleteTask: BaseDeleteTask{
G
godchen 已提交
2606 2607 2608
			BaseMsg: msgstream.BaseMsg{
				HashValues: request.HashKeys,
			},
G
godchen 已提交
2609
			DeleteRequest: internalpb.DeleteRequest{
2610 2611 2612 2613
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Delete),
					commonpbutil.WithMsgID(0),
				),
X
xige-16 已提交
2614
				DbName:         request.DbName,
G
godchen 已提交
2615 2616 2617
				CollectionName: request.CollectionName,
				PartitionName:  request.PartitionName,
				// RowData: transfer column based request to this
C
Cai Yudong 已提交
2618 2619 2620 2621
			},
		},
		chMgr:    node.chMgr,
		chTicker: node.chTicker,
G
groot 已提交
2622 2623
	}

2624
	log.Debug("Enqueue delete request in Proxy",
2625
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2626 2627 2628 2629
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
		zap.String("expr", request.Expr))
2630 2631 2632 2633

	// MsgID will be set by Enqueue()
	if err := node.sched.dmQueue.Enqueue(dt); err != nil {
		log.Error("Failed to enqueue delete task: "+err.Error(), zap.String("traceID", traceID))
2634 2635
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
			metrics.AbandonLabel).Inc()
2636

G
groot 已提交
2637 2638 2639 2640 2641 2642 2643 2644
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

2645
	log.Debug("Detail of delete request in Proxy",
2646
		zap.String("role", typeutil.ProxyRole),
G
groot 已提交
2647 2648 2649 2650 2651
		zap.Int64("msgID", dt.Base.MsgID),
		zap.Uint64("timestamp", dt.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
2652 2653
		zap.String("expr", request.Expr),
		zap.String("traceID", traceID))
G
groot 已提交
2654

2655 2656
	if err := dt.WaitToFinish(); err != nil {
		log.Error("Failed to execute delete task in task scheduler: "+err.Error(), zap.String("traceID", traceID))
2657
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2658
			metrics.FailLabel).Inc()
G
groot 已提交
2659 2660 2661 2662 2663 2664 2665 2666
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

2667
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2668
		metrics.SuccessLabel).Inc()
2669
	metrics.ProxyMutationLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), metrics.DeleteLabel).Observe(float64(tr.ElapseSpan().Milliseconds()))
2670
	metrics.ProxyCollectionMutationLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), metrics.DeleteLabel, request.CollectionName).Observe(float64(tr.ElapseSpan().Milliseconds()))
G
groot 已提交
2671 2672 2673
	return dt.result, nil
}

2674
// Search search the most similar records of requests.
C
Cai Yudong 已提交
2675
func (node *Proxy) Search(ctx context.Context, request *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
2676 2677 2678 2679 2680
	receiveSize := proto.Size(request)
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), metrics.SearchLabel).Add(float64(receiveSize))

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

2681 2682 2683 2684 2685
	if !node.checkHealthy() {
		return &milvuspb.SearchResults{
			Status: unhealthyStatus(),
		}, nil
	}
2686 2687
	method := "Search"
	tr := timerecord.NewTimeRecorder(method)
2688
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2689
		metrics.TotalLabel).Inc()
D
dragondriver 已提交
2690

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

2694
	qt := &searchTask{
S
sunby 已提交
2695
		ctx:       ctx,
2696
		Condition: NewTaskCondition(ctx),
G
godchen 已提交
2697
		SearchRequest: &internalpb.SearchRequest{
2698 2699 2700 2701
			Base: commonpbutil.NewMsgBase(
				commonpbutil.WithMsgType(commonpb.MsgType_Search),
				commonpbutil.WithSourceID(Params.ProxyCfg.GetNodeID()),
			),
2702
			ReqID: Params.ProxyCfg.GetNodeID(),
2703
		},
2704 2705 2706 2707
		request:  request,
		qc:       node.queryCoord,
		tr:       timerecord.NewTimeRecorder("search"),
		shardMgr: node.shardMgr,
2708 2709
	}

2710 2711 2712
	travelTs := request.TravelTimestamp
	guaranteeTs := request.GuaranteeTimestamp

Z
Zach 已提交
2713
	log.Ctx(ctx).Info(
2714
		rpcReceived(method),
2715
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2716 2717 2718 2719 2720
		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)),
2721 2722 2723 2724
		zap.Any("OutputFields", request.OutputFields),
		zap.Any("search_params", request.SearchParams),
		zap.Uint64("travel_timestamp", travelTs),
		zap.Uint64("guarantee_timestamp", guaranteeTs))
D
dragondriver 已提交
2725

2726
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
Z
Zach 已提交
2727
		log.Ctx(ctx).Warn(
2728
			rpcFailedToEnqueue(method),
D
dragondriver 已提交
2729
			zap.Error(err),
2730
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2731 2732 2733 2734 2735 2736
			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)),
			zap.Any("OutputFields", request.OutputFields),
2737 2738 2739
			zap.Any("search_params", request.SearchParams),
			zap.Uint64("travel_timestamp", travelTs),
			zap.Uint64("guarantee_timestamp", guaranteeTs))
D
dragondriver 已提交
2740

2741
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2742
			metrics.AbandonLabel).Inc()
2743

2744 2745
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
2746
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2747 2748 2749 2750
				Reason:    err.Error(),
			},
		}, nil
	}
Z
Zach 已提交
2751
	tr.CtxRecord(ctx, "search request enqueue")
2752

Z
Zach 已提交
2753
	log.Ctx(ctx).Debug(
2754
		rpcEnqueued(method),
2755
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2756
		zap.Int64("msgID", qt.ID()),
D
dragondriver 已提交
2757 2758 2759 2760 2761
		zap.Uint64("timestamp", qt.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames),
		zap.Any("dsl", request.Dsl),
2762
		zap.Any("len(PlaceholderGroup)", len(request.PlaceholderGroup)),
2763 2764 2765 2766
		zap.Any("OutputFields", request.OutputFields),
		zap.Any("search_params", request.SearchParams),
		zap.Uint64("travel_timestamp", travelTs),
		zap.Uint64("guarantee_timestamp", guaranteeTs))
D
dragondriver 已提交
2767

2768
	if err := qt.WaitToFinish(); err != nil {
Z
Zach 已提交
2769
		log.Ctx(ctx).Warn(
2770
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
2771
			zap.Error(err),
2772
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2773
			zap.Int64("msgID", qt.ID()),
D
dragondriver 已提交
2774 2775 2776 2777
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames),
			zap.Any("dsl", request.Dsl),
2778
			zap.Any("len(PlaceholderGroup)", len(request.PlaceholderGroup)),
2779 2780 2781 2782
			zap.Any("OutputFields", request.OutputFields),
			zap.Any("search_params", request.SearchParams),
			zap.Uint64("travel_timestamp", travelTs),
			zap.Uint64("guarantee_timestamp", guaranteeTs))
2783

2784
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2785
			metrics.FailLabel).Inc()
2786

2787 2788
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
2789
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2790 2791 2792 2793 2794
				Reason:    err.Error(),
			},
		}, nil
	}

Z
Zach 已提交
2795
	span := tr.CtxRecord(ctx, "wait search result")
2796 2797
	metrics.ProxyWaitForSearchResultLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10),
		metrics.SearchLabel).Observe(float64(span.Milliseconds()))
2798
	tr.CtxRecord(ctx, "wait search result")
Z
Zach 已提交
2799
	log.Ctx(ctx).Debug(
2800
		rpcDone(method),
2801
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2802 2803 2804 2805 2806 2807
		zap.Int64("msgID", qt.ID()),
		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)),
2808 2809 2810 2811
		zap.Any("OutputFields", request.OutputFields),
		zap.Any("search_params", request.SearchParams),
		zap.Uint64("travel_timestamp", travelTs),
		zap.Uint64("guarantee_timestamp", guaranteeTs))
D
dragondriver 已提交
2812

2813
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2814 2815
		metrics.SuccessLabel).Inc()
	metrics.ProxySearchVectors.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10)).Add(float64(qt.result.GetResults().GetNumQueries()))
C
cai.zhang 已提交
2816
	searchDur := tr.ElapseSpan().Milliseconds()
2817
	metrics.ProxySQLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10),
2818
		metrics.SearchLabel).Observe(float64(searchDur))
2819 2820
	metrics.ProxyCollectionSQLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10),
		metrics.SearchLabel, request.CollectionName).Observe(float64(searchDur))
2821 2822 2823
	if qt.result != nil {
		sentSize := proto.Size(qt.result)
		metrics.ProxyReadReqSendBytes.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10)).Add(float64(sentSize))
2824
		rateCol.Add(metricsinfo.ReadResultThroughput, float64(sentSize))
2825
	}
2826 2827 2828
	return qt.result, nil
}

2829
// Flush notify data nodes to persist the data of collection.
2830 2831 2832 2833 2834 2835 2836
func (node *Proxy) Flush(ctx context.Context, request *milvuspb.FlushRequest) (*milvuspb.FlushResponse, error) {
	resp := &milvuspb.FlushResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    "",
		},
	}
2837
	if !node.checkHealthy() {
2838 2839
		resp.Status.Reason = "proxy is not healthy"
		return resp, nil
2840
	}
D
dragondriver 已提交
2841 2842 2843 2844 2845

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

2846
	ft := &flushTask{
T
ThreadDao 已提交
2847 2848 2849
		ctx:          ctx,
		Condition:    NewTaskCondition(ctx),
		FlushRequest: request,
2850
		dataCoord:    node.dataCoord,
2851 2852
	}

D
dragondriver 已提交
2853
	method := "Flush"
2854
	tr := timerecord.NewTimeRecorder(method)
2855
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
2856 2857 2858 2859

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
2860
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2861 2862
		zap.String("db", request.DbName),
		zap.Any("collections", request.CollectionNames))
D
dragondriver 已提交
2863 2864 2865 2866 2867 2868

	if err := node.sched.ddQueue.Enqueue(ft); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
2869
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2870 2871 2872
			zap.String("db", request.DbName),
			zap.Any("collections", request.CollectionNames))

2873
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
2874

2875 2876
		resp.Status.Reason = err.Error()
		return resp, nil
2877 2878
	}

D
dragondriver 已提交
2879 2880 2881
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2882
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2883 2884 2885
		zap.Int64("MsgID", ft.ID()),
		zap.Uint64("BeginTs", ft.BeginTs()),
		zap.Uint64("EndTs", ft.EndTs()),
D
dragondriver 已提交
2886 2887
		zap.String("db", request.DbName),
		zap.Any("collections", request.CollectionNames))
D
dragondriver 已提交
2888 2889 2890 2891

	if err := ft.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
2892
			zap.Error(err),
D
dragondriver 已提交
2893
			zap.String("traceID", traceID),
2894
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2895 2896 2897
			zap.Int64("MsgID", ft.ID()),
			zap.Uint64("BeginTs", ft.BeginTs()),
			zap.Uint64("EndTs", ft.EndTs()),
D
dragondriver 已提交
2898 2899 2900
			zap.String("db", request.DbName),
			zap.Any("collections", request.CollectionNames))

2901
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
2902

D
dragondriver 已提交
2903
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
2904 2905
		resp.Status.Reason = err.Error()
		return resp, nil
2906 2907
	}

D
dragondriver 已提交
2908 2909 2910
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
2911
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2912 2913 2914 2915 2916 2917
		zap.Int64("MsgID", ft.ID()),
		zap.Uint64("BeginTs", ft.BeginTs()),
		zap.Uint64("EndTs", ft.EndTs()),
		zap.String("db", request.DbName),
		zap.Any("collections", request.CollectionNames))

2918 2919
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
2920
	return ft.result, nil
2921 2922
}

2923
// Query get the records by primary keys.
C
Cai Yudong 已提交
2924
func (node *Proxy) Query(ctx context.Context, request *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) {
2925 2926 2927 2928 2929
	receiveSize := proto.Size(request)
	metrics.ProxyReceiveBytes.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), metrics.QueryLabel).Add(float64(receiveSize))

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

2930 2931 2932 2933 2934
	if !node.checkHealthy() {
		return &milvuspb.QueryResults{
			Status: unhealthyStatus(),
		}, nil
	}
2935

D
dragondriver 已提交
2936 2937
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Query")
	defer sp.Finish()
2938
	tr := timerecord.NewTimeRecorder("Query")
D
dragondriver 已提交
2939

2940
	qt := &queryTask{
2941 2942 2943
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),
		RetrieveRequest: &internalpb.RetrieveRequest{
2944 2945 2946 2947
			Base: commonpbutil.NewMsgBase(
				commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
				commonpbutil.WithSourceID(Params.ProxyCfg.GetNodeID()),
			),
2948
			ReqID: Params.ProxyCfg.GetNodeID(),
2949
		},
2950 2951
		request:          request,
		qc:               node.queryCoord,
2952
		queryShardPolicy: mergeRoundRobinPolicy,
2953
		shardMgr:         node.shardMgr,
2954 2955
	}

D
dragondriver 已提交
2956 2957
	method := "Query"

2958
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
2959 2960
		metrics.TotalLabel).Inc()

Z
Zach 已提交
2961
	log.Ctx(ctx).Info(
D
dragondriver 已提交
2962
		rpcReceived(method),
2963
		zap.String("role", typeutil.ProxyRole),
2964 2965
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
2966 2967 2968 2969 2970
		zap.Strings("partitions", request.PartitionNames),
		zap.String("expr", request.Expr),
		zap.Strings("OutputFields", request.OutputFields),
		zap.Uint64("travel_timestamp", request.TravelTimestamp),
		zap.Uint64("guarantee_timestamp", request.GuaranteeTimestamp))
G
godchen 已提交
2971

D
dragondriver 已提交
2972
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
Z
Zach 已提交
2973
		log.Ctx(ctx).Warn(
D
dragondriver 已提交
2974 2975 2976
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("role", typeutil.ProxyRole),
2977 2978 2979
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))
D
dragondriver 已提交
2980

2981 2982
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
			metrics.AbandonLabel).Inc()
2983

2984 2985 2986 2987 2988 2989
		return &milvuspb.QueryResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
2990
	}
Z
Zach 已提交
2991
	tr.CtxRecord(ctx, "query request enqueue")
2992

Z
Zach 已提交
2993
	log.Ctx(ctx).Debug(
D
dragondriver 已提交
2994
		rpcEnqueued(method),
2995
		zap.String("role", typeutil.ProxyRole),
2996
		zap.Int64("msgID", qt.ID()),
2997 2998
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
2999
		zap.Strings("partitions", request.PartitionNames))
D
dragondriver 已提交
3000 3001

	if err := qt.WaitToFinish(); err != nil {
Z
Zach 已提交
3002
		log.Ctx(ctx).Warn(
D
dragondriver 已提交
3003 3004
			rpcFailedToWaitToFinish(method),
			zap.Error(err),
3005
			zap.String("role", typeutil.ProxyRole),
3006
			zap.Int64("msgID", qt.ID()),
3007 3008 3009
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))
3010

3011
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
3012
			metrics.FailLabel).Inc()
3013

3014 3015 3016 3017 3018 3019 3020
		return &milvuspb.QueryResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}
Z
Zach 已提交
3021
	span := tr.CtxRecord(ctx, "wait query result")
3022 3023
	metrics.ProxyWaitForSearchResultLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10),
		metrics.QueryLabel).Observe(float64(span.Milliseconds()))
3024

Z
Zach 已提交
3025
	log.Ctx(ctx).Debug(
D
dragondriver 已提交
3026 3027
		rpcDone(method),
		zap.String("role", typeutil.ProxyRole),
3028
		zap.Int64("msgID", qt.ID()),
3029 3030 3031
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
D
dragondriver 已提交
3032

3033
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
3034 3035
		metrics.SuccessLabel).Inc()

3036
	metrics.ProxySQLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10),
3037
		metrics.QueryLabel).Observe(float64(tr.ElapseSpan().Milliseconds()))
3038 3039
	metrics.ProxyCollectionSQLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10),
		metrics.QueryLabel, request.CollectionName).Observe(float64(tr.ElapseSpan().Milliseconds()))
3040 3041

	ret := &milvuspb.QueryResults{
3042 3043
		Status:     qt.result.Status,
		FieldsData: qt.result.FieldsData,
3044 3045
	}
	sentSize := proto.Size(qt.result)
3046
	rateCol.Add(metricsinfo.ReadResultThroughput, float64(sentSize))
3047 3048
	metrics.ProxyReadReqSendBytes.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10)).Add(float64(sentSize))
	return ret, nil
3049
}
3050

3051
// CreateAlias create alias for collection, then you can search the collection with alias.
Y
Yusup 已提交
3052 3053 3054 3055
func (node *Proxy) CreateAlias(ctx context.Context, request *milvuspb.CreateAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
3056 3057 3058 3059 3060

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

Y
Yusup 已提交
3061 3062 3063 3064 3065 3066 3067
	cat := &CreateAliasTask{
		ctx:                ctx,
		Condition:          NewTaskCondition(ctx),
		CreateAliasRequest: request,
		rootCoord:          node.rootCoord,
	}

D
dragondriver 已提交
3068
	method := "CreateAlias"
3069
	tr := timerecord.NewTimeRecorder(method)
3070
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
3071 3072 3073 3074 3075 3076 3077 3078 3079 3080 3081 3082 3083 3084 3085 3086 3087 3088 3089

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

	if err := node.sched.ddQueue.Enqueue(cat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
			zap.String("role", typeutil.ProxyRole),
			zap.String("db", request.DbName),
			zap.String("alias", request.Alias),
			zap.String("collection", request.CollectionName))

3090
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
3091

Y
Yusup 已提交
3092 3093 3094 3095 3096 3097
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
3098 3099 3100
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
3101
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3102 3103 3104 3105
		zap.Int64("MsgID", cat.ID()),
		zap.Uint64("BeginTs", cat.BeginTs()),
		zap.Uint64("EndTs", cat.EndTs()),
		zap.String("db", request.DbName),
Y
Yusup 已提交
3106 3107
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))
D
dragondriver 已提交
3108 3109 3110 3111

	if err := cat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
3112
			zap.Error(err),
D
dragondriver 已提交
3113
			zap.String("traceID", traceID),
3114
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3115 3116 3117 3118
			zap.Int64("MsgID", cat.ID()),
			zap.Uint64("BeginTs", cat.BeginTs()),
			zap.Uint64("EndTs", cat.EndTs()),
			zap.String("db", request.DbName),
Y
Yusup 已提交
3119 3120
			zap.String("alias", request.Alias),
			zap.String("collection", request.CollectionName))
3121
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
Y
Yusup 已提交
3122 3123 3124 3125 3126 3127 3128

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

D
dragondriver 已提交
3129 3130 3131 3132 3133 3134 3135 3136 3137 3138 3139
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.Int64("MsgID", cat.ID()),
		zap.Uint64("BeginTs", cat.BeginTs()),
		zap.Uint64("EndTs", cat.EndTs()),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

3140 3141
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
Y
Yusup 已提交
3142 3143 3144
	return cat.result, nil
}

3145
// DropAlias alter the alias of collection.
Y
Yusup 已提交
3146 3147 3148 3149
func (node *Proxy) DropAlias(ctx context.Context, request *milvuspb.DropAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
3150 3151 3152 3153 3154

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

Y
Yusup 已提交
3155 3156 3157 3158 3159 3160 3161
	dat := &DropAliasTask{
		ctx:              ctx,
		Condition:        NewTaskCondition(ctx),
		DropAliasRequest: request,
		rootCoord:        node.rootCoord,
	}

D
dragondriver 已提交
3162
	method := "DropAlias"
3163
	tr := timerecord.NewTimeRecorder(method)
3164
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
3165 3166 3167 3168 3169 3170 3171 3172 3173 3174 3175 3176 3177 3178 3179 3180

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias))

	if err := node.sched.ddQueue.Enqueue(dat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
			zap.String("role", typeutil.ProxyRole),
			zap.String("db", request.DbName),
			zap.String("alias", request.Alias))
3181
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
D
dragondriver 已提交
3182

Y
Yusup 已提交
3183 3184 3185 3186 3187 3188
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
3189 3190 3191
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
3192
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3193 3194 3195 3196
		zap.Int64("MsgID", dat.ID()),
		zap.Uint64("BeginTs", dat.BeginTs()),
		zap.Uint64("EndTs", dat.EndTs()),
		zap.String("db", request.DbName),
Y
Yusup 已提交
3197
		zap.String("alias", request.Alias))
D
dragondriver 已提交
3198 3199 3200 3201

	if err := dat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
3202
			zap.Error(err),
D
dragondriver 已提交
3203
			zap.String("traceID", traceID),
3204
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3205 3206 3207 3208
			zap.Int64("MsgID", dat.ID()),
			zap.Uint64("BeginTs", dat.BeginTs()),
			zap.Uint64("EndTs", dat.EndTs()),
			zap.String("db", request.DbName),
Y
Yusup 已提交
3209 3210
			zap.String("alias", request.Alias))

3211
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3212

Y
Yusup 已提交
3213 3214 3215 3216 3217 3218
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
3219 3220 3221 3222 3223 3224 3225 3226 3227 3228
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.Int64("MsgID", dat.ID()),
		zap.Uint64("BeginTs", dat.BeginTs()),
		zap.Uint64("EndTs", dat.EndTs()),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias))

3229 3230
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
Y
Yusup 已提交
3231 3232 3233
	return dat.result, nil
}

3234
// AlterAlias alter alias of collection.
Y
Yusup 已提交
3235 3236 3237 3238
func (node *Proxy) AlterAlias(ctx context.Context, request *milvuspb.AlterAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
3239 3240 3241 3242 3243

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

Y
Yusup 已提交
3244 3245 3246 3247 3248 3249 3250
	aat := &AlterAliasTask{
		ctx:               ctx,
		Condition:         NewTaskCondition(ctx),
		AlterAliasRequest: request,
		rootCoord:         node.rootCoord,
	}

D
dragondriver 已提交
3251
	method := "AlterAlias"
3252
	tr := timerecord.NewTimeRecorder(method)
3253
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.TotalLabel).Inc()
D
dragondriver 已提交
3254 3255 3256 3257 3258 3259 3260 3261 3262 3263 3264 3265 3266 3267 3268 3269 3270 3271

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

	if err := node.sched.ddQueue.Enqueue(aat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
			zap.String("role", typeutil.ProxyRole),
			zap.String("db", request.DbName),
			zap.String("alias", request.Alias),
			zap.String("collection", request.CollectionName))
3272
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.AbandonLabel).Inc()
D
dragondriver 已提交
3273

Y
Yusup 已提交
3274 3275 3276 3277 3278 3279
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
3280 3281 3282
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
3283
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3284 3285 3286 3287
		zap.Int64("MsgID", aat.ID()),
		zap.Uint64("BeginTs", aat.BeginTs()),
		zap.Uint64("EndTs", aat.EndTs()),
		zap.String("db", request.DbName),
Y
Yusup 已提交
3288 3289
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))
D
dragondriver 已提交
3290 3291 3292 3293

	if err := aat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
3294
			zap.Error(err),
D
dragondriver 已提交
3295
			zap.String("traceID", traceID),
3296
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3297 3298 3299 3300
			zap.Int64("MsgID", aat.ID()),
			zap.Uint64("BeginTs", aat.BeginTs()),
			zap.Uint64("EndTs", aat.EndTs()),
			zap.String("db", request.DbName),
Y
Yusup 已提交
3301 3302 3303
			zap.String("alias", request.Alias),
			zap.String("collection", request.CollectionName))

3304
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3305

Y
Yusup 已提交
3306 3307 3308 3309 3310 3311
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
3312 3313 3314 3315 3316 3317 3318 3319 3320 3321 3322
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.Int64("MsgID", aat.ID()),
		zap.Uint64("BeginTs", aat.BeginTs()),
		zap.Uint64("EndTs", aat.EndTs()),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

3323 3324
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
Y
Yusup 已提交
3325 3326 3327
	return aat.result, nil
}

3328
// CalcDistance calculates the distances between vectors.
3329
func (node *Proxy) CalcDistance(ctx context.Context, request *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) {
3330 3331 3332 3333 3334
	if !node.checkHealthy() {
		return &milvuspb.CalcDistanceResults{
			Status: unhealthyStatus(),
		}, nil
	}
3335

3336 3337 3338 3339
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CalcDistance")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

3340 3341
	query := func(ids *milvuspb.VectorIDs) (*milvuspb.QueryResults, error) {
		outputFields := []string{ids.FieldName}
3342

3343 3344 3345 3346 3347
		queryRequest := &milvuspb.QueryRequest{
			DbName:         "",
			CollectionName: ids.CollectionName,
			PartitionNames: ids.PartitionNames,
			OutputFields:   outputFields,
3348 3349
		}

3350
		qt := &queryTask{
3351 3352 3353
			ctx:       ctx,
			Condition: NewTaskCondition(ctx),
			RetrieveRequest: &internalpb.RetrieveRequest{
3354 3355 3356 3357
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
					commonpbutil.WithSourceID(Params.ProxyCfg.GetNodeID()),
				),
3358
				ReqID: Params.ProxyCfg.GetNodeID(),
3359
			},
3360 3361 3362 3363
			request: queryRequest,
			qc:      node.queryCoord,
			ids:     ids.IdArray,

3364
			queryShardPolicy: mergeRoundRobinPolicy,
3365
			shardMgr:         node.shardMgr,
3366 3367
		}

G
groot 已提交
3368 3369 3370 3371 3372 3373
		items := []zapcore.Field{
			zap.String("collection", queryRequest.CollectionName),
			zap.Any("partitions", queryRequest.PartitionNames),
			zap.Any("OutputFields", queryRequest.OutputFields),
		}

3374
		err := node.sched.dqQueue.Enqueue(qt)
3375
		if err != nil {
G
groot 已提交
3376
			log.Error("CalcDistance queryTask failed to enqueue", append(items, zap.Error(err))...)
3377

3378 3379 3380 3381 3382
			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
3383
			}, err
3384
		}
3385

G
groot 已提交
3386
		log.Debug("CalcDistance queryTask enqueued", items...)
3387 3388 3389

		err = qt.WaitToFinish()
		if err != nil {
G
groot 已提交
3390
			log.Error("CalcDistance queryTask failed to WaitToFinish", append(items, zap.Error(err))...)
3391 3392 3393 3394 3395 3396

			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
3397
			}, err
3398
		}
3399

G
groot 已提交
3400
		log.Debug("CalcDistance queryTask Done", items...)
3401 3402

		return &milvuspb.QueryResults{
3403 3404
			Status:     qt.result.Status,
			FieldsData: qt.result.FieldsData,
3405 3406 3407
		}, nil
	}

G
groot 已提交
3408 3409 3410 3411
	// calcDistanceTask is not a standard task, no need to enqueue
	task := &calcDistanceTask{
		traceID:   traceID,
		queryFunc: query,
3412 3413
	}

G
groot 已提交
3414
	return task.Execute(ctx, request)
3415 3416
}

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

3422
// GetPersistentSegmentInfo get the information of sealed segment.
C
Cai Yudong 已提交
3423
func (node *Proxy) GetPersistentSegmentInfo(ctx context.Context, req *milvuspb.GetPersistentSegmentInfoRequest) (*milvuspb.GetPersistentSegmentInfoResponse, error) {
D
dragondriver 已提交
3424
	log.Debug("GetPersistentSegmentInfo",
3425
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3426 3427 3428
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
3429
	resp := &milvuspb.GetPersistentSegmentInfoResponse{
X
XuanYang-cn 已提交
3430
		Status: &commonpb.Status{
3431
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
X
XuanYang-cn 已提交
3432 3433
		},
	}
3434 3435 3436 3437
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3438 3439
	method := "GetPersistentSegmentInfo"
	tr := timerecord.NewTimeRecorder(method)
3440
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
3441
		metrics.TotalLabel).Inc()
3442 3443 3444

	// list segments
	collectionID, err := globalMetaCache.GetCollectionID(ctx, req.GetCollectionName())
X
XuanYang-cn 已提交
3445
	if err != nil {
3446 3447 3448 3449 3450 3451 3452 3453 3454 3455 3456 3457 3458
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
		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 {
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3459
		resp.Status.Reason = fmt.Errorf("getSegmentsOfCollection, err:%w", err).Error()
X
XuanYang-cn 已提交
3460 3461
		return resp, nil
	}
3462 3463

	// get Segment info
3464
	infoResp, err := node.dataCoord.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{
3465 3466 3467 3468 3469 3470
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_SegmentInfo),
			commonpbutil.WithMsgID(0),
			commonpbutil.WithTimeStamp(0),
			commonpbutil.WithSourceID(Params.ProxyCfg.GetNodeID()),
		),
3471
		SegmentIDs: getSegmentsByStatesResponse.Segments,
X
XuanYang-cn 已提交
3472 3473
	})
	if err != nil {
3474 3475 3476
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
			metrics.FailLabel).Inc()
		log.Warn("GetPersistentSegmentInfo fail", zap.Error(err))
3477
		resp.Status.Reason = fmt.Errorf("dataCoord:GetSegmentInfo, err:%w", err).Error()
X
XuanYang-cn 已提交
3478 3479
		return resp, nil
	}
3480
	log.Debug("GetPersistentSegmentInfo ", zap.Int("len(infos)", len(infoResp.Infos)), zap.Any("status", infoResp.Status))
3481
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
3482 3483
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
			metrics.FailLabel).Inc()
X
XuanYang-cn 已提交
3484 3485 3486 3487 3488 3489
		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 已提交
3490
			SegmentID:    info.ID,
X
XuanYang-cn 已提交
3491 3492
			CollectionID: info.CollectionID,
			PartitionID:  info.PartitionID,
S
sunby 已提交
3493
			NumRows:      info.NumOfRows,
X
XuanYang-cn 已提交
3494 3495 3496
			State:        info.State,
		}
	}
3497
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
3498
		metrics.SuccessLabel).Inc()
3499
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
3500
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
X
XuanYang-cn 已提交
3501 3502 3503 3504
	resp.Infos = persistentInfos
	return resp, nil
}

J
jingkl 已提交
3505
// GetQuerySegmentInfo gets segment information from QueryCoord.
C
Cai Yudong 已提交
3506
func (node *Proxy) GetQuerySegmentInfo(ctx context.Context, req *milvuspb.GetQuerySegmentInfoRequest) (*milvuspb.GetQuerySegmentInfoResponse, error) {
D
dragondriver 已提交
3507
	log.Debug("GetQuerySegmentInfo",
3508
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3509 3510 3511
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
3512
	resp := &milvuspb.GetQuerySegmentInfoResponse{
Z
zhenshan.cao 已提交
3513
		Status: &commonpb.Status{
3514
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
Z
zhenshan.cao 已提交
3515 3516
		},
	}
3517 3518 3519 3520
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3521

3522 3523 3524 3525 3526
	method := "GetQuerySegmentInfo"
	tr := timerecord.NewTimeRecorder(method)
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()

3527 3528
	collID, err := globalMetaCache.GetCollectionID(ctx, req.CollectionName)
	if err != nil {
3529
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3530 3531 3532
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3533
	infoResp, err := node.queryCoord.GetSegmentInfo(ctx, &querypb.GetSegmentInfoRequest{
3534 3535 3536 3537 3538 3539
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_SegmentInfo),
			commonpbutil.WithMsgID(0),
			commonpbutil.WithTimeStamp(0),
			commonpbutil.WithSourceID(Params.ProxyCfg.GetNodeID()),
		),
3540
		CollectionID: collID,
Z
zhenshan.cao 已提交
3541 3542
	})
	if err != nil {
3543 3544
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
		log.Error("Failed to get segment info from QueryCoord", zap.Error(err))
Z
zhenshan.cao 已提交
3545 3546 3547
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3548
	log.Debug("GetQuerySegmentInfo ", zap.Any("infos", infoResp.Infos), zap.Any("status", infoResp.Status))
3549
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
3550
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
3551
		log.Error("Failed to get segment info from QueryCoord", zap.String("errMsg", infoResp.Status.Reason))
Z
zhenshan.cao 已提交
3552 3553 3554 3555 3556 3557 3558 3559 3560 3561 3562 3563 3564
		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 已提交
3565
			State:        info.SegmentState,
3566
			NodeIds:      info.NodeIds,
Z
zhenshan.cao 已提交
3567 3568
		}
	}
3569 3570 3571

	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
3572
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
Z
zhenshan.cao 已提交
3573 3574 3575 3576
	resp.Infos = queryInfos
	return resp, nil
}

J
jingkl 已提交
3577
// Dummy handles dummy request
C
Cai Yudong 已提交
3578
func (node *Proxy) Dummy(ctx context.Context, req *milvuspb.DummyRequest) (*milvuspb.DummyResponse, error) {
3579 3580 3581 3582 3583 3584 3585 3586 3587 3588 3589
	failedResponse := &milvuspb.DummyResponse{
		Response: `{"status": "fail"}`,
	}

	// TODO(wxyu): change name RequestType to Request
	drt, err := parseDummyRequestType(req.RequestType)
	if err != nil {
		log.Debug("Failed to parse dummy request type")
		return failedResponse, nil
	}

3590 3591
	if drt.RequestType == "query" {
		drr, err := parseDummyQueryRequest(req.RequestType)
3592
		if err != nil {
3593
			log.Debug("Failed to parse dummy query request")
3594 3595 3596
			return failedResponse, nil
		}

3597
		request := &milvuspb.QueryRequest{
3598 3599 3600
			DbName:         drr.DbName,
			CollectionName: drr.CollectionName,
			PartitionNames: drr.PartitionNames,
3601
			OutputFields:   drr.OutputFields,
X
Xiangyu Wang 已提交
3602 3603
		}

3604
		_, err = node.Query(ctx, request)
3605
		if err != nil {
3606
			log.Debug("Failed to execute dummy query")
3607 3608
			return failedResponse, err
		}
X
Xiangyu Wang 已提交
3609 3610 3611 3612 3613 3614

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

3615 3616
	log.Debug("cannot find specify dummy request type")
	return failedResponse, nil
X
Xiangyu Wang 已提交
3617 3618
}

J
jingkl 已提交
3619
// RegisterLink registers a link
C
Cai Yudong 已提交
3620
func (node *Proxy) RegisterLink(ctx context.Context, req *milvuspb.RegisterLinkRequest) (*milvuspb.RegisterLinkResponse, error) {
3621
	code := node.stateCode.Load().(commonpb.StateCode)
D
dragondriver 已提交
3622
	log.Debug("RegisterLink",
3623
		zap.String("role", typeutil.ProxyRole),
C
Cai Yudong 已提交
3624
		zap.Any("state code of proxy", code))
D
dragondriver 已提交
3625

3626
	if code != commonpb.StateCode_Healthy {
3627 3628 3629
		return &milvuspb.RegisterLinkResponse{
			Address: nil,
			Status: &commonpb.Status{
3630
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
C
Cai Yudong 已提交
3631
				Reason:    "proxy not healthy",
3632 3633 3634
			},
		}, nil
	}
X
Xiaofan 已提交
3635
	//metrics.ProxyLinkedSDKs.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10)).Inc()
3636 3637 3638
	return &milvuspb.RegisterLinkResponse{
		Address: nil,
		Status: &commonpb.Status{
3639
			ErrorCode: commonpb.ErrorCode_Success,
3640
			Reason:    os.Getenv(metricsinfo.DeployModeEnvKey),
3641 3642 3643
		},
	}, nil
}
3644

3645
// GetMetrics gets the metrics of proxy
3646 3647 3648
// 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) {
	log.Debug("Proxy.GetMetrics",
X
Xiaofan 已提交
3649
		zap.Int64("node_id", Params.ProxyCfg.GetNodeID()),
3650 3651 3652 3653
		zap.String("req", req.Request))

	if !node.checkHealthy() {
		log.Warn("Proxy.GetMetrics failed",
X
Xiaofan 已提交
3654
			zap.Int64("node_id", Params.ProxyCfg.GetNodeID()),
3655
			zap.String("req", req.Request),
X
Xiaofan 已提交
3656
			zap.Error(errProxyIsUnhealthy(Params.ProxyCfg.GetNodeID())))
3657 3658 3659 3660

		return &milvuspb.GetMetricsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
X
Xiaofan 已提交
3661
				Reason:    msgProxyIsUnhealthy(Params.ProxyCfg.GetNodeID()),
3662 3663 3664 3665 3666 3667 3668 3669
			},
			Response: "",
		}, nil
	}

	metricType, err := metricsinfo.ParseMetricType(req.Request)
	if err != nil {
		log.Warn("Proxy.GetMetrics failed to parse metric type",
X
Xiaofan 已提交
3670
			zap.Int64("node_id", Params.ProxyCfg.GetNodeID()),
3671 3672 3673 3674 3675 3676 3677 3678 3679 3680 3681 3682 3683 3684 3685
			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))

3686 3687 3688 3689 3690 3691
	req.Base = commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_SystemInfo),
		commonpbutil.WithMsgID(0),
		commonpbutil.WithTimeStamp(0),
		commonpbutil.WithSourceID(Params.ProxyCfg.GetNodeID()),
	)
3692
	if metricType == metricsinfo.SystemInfoMetrics {
3693 3694 3695 3696 3697 3698 3699
		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))

3700
		metrics, err := getSystemInfoMetrics(ctx, req, node)
3701 3702

		log.Debug("Proxy.GetMetrics",
X
Xiaofan 已提交
3703
			zap.Int64("node_id", Params.ProxyCfg.GetNodeID()),
3704 3705 3706 3707 3708
			zap.String("req", req.Request),
			zap.String("metric_type", metricType),
			zap.Any("metrics", metrics), // TODO(dragondriver): necessary? may be very large
			zap.Error(err))

3709 3710
		node.metricsCacheManager.UpdateSystemInfoMetrics(metrics)

G
godchen 已提交
3711
		return metrics, nil
3712 3713 3714
	}

	log.Debug("Proxy.GetMetrics failed, request metric type is not implemented yet",
X
Xiaofan 已提交
3715
		zap.Int64("node_id", Params.ProxyCfg.GetNodeID()),
3716 3717 3718 3719 3720 3721 3722 3723 3724 3725 3726 3727
		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
}

3728 3729 3730 3731 3732 3733 3734 3735 3736 3737 3738 3739 3740 3741 3742 3743 3744 3745 3746 3747 3748 3749 3750 3751 3752 3753 3754 3755 3756 3757 3758 3759 3760 3761 3762 3763 3764 3765 3766
// 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) {
	log.Debug("Proxy.GetProxyMetrics",
		zap.Int64("node_id", Params.ProxyCfg.GetNodeID()),
		zap.String("req", req.Request))

	if !node.checkHealthy() {
		log.Warn("Proxy.GetProxyMetrics failed",
			zap.Int64("node_id", Params.ProxyCfg.GetNodeID()),
			zap.String("req", req.Request),
			zap.Error(errProxyIsUnhealthy(Params.ProxyCfg.GetNodeID())))

		return &milvuspb.GetMetricsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    msgProxyIsUnhealthy(Params.ProxyCfg.GetNodeID()),
			},
		}, nil
	}

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

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

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

3767 3768 3769 3770 3771 3772
	req.Base = commonpbutil.NewMsgBase(
		commonpbutil.WithMsgType(commonpb.MsgType_SystemInfo),
		commonpbutil.WithMsgID(0),
		commonpbutil.WithTimeStamp(0),
		commonpbutil.WithSourceID(Params.ProxyCfg.GetNodeID()),
	)
3773 3774 3775 3776 3777 3778 3779 3780 3781 3782 3783 3784 3785 3786 3787 3788 3789 3790 3791 3792 3793 3794 3795 3796 3797 3798 3799 3800 3801 3802 3803 3804 3805 3806 3807 3808 3809 3810 3811

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

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

		log.Debug("Proxy.GetProxyMetrics",
			zap.Int64("node_id", Params.ProxyCfg.GetNodeID()),
			zap.String("req", req.Request),
			zap.String("metric_type", metricType),
			zap.Error(err))

		return proxyMetrics, nil
	}

	log.Debug("Proxy.GetProxyMetrics failed, request metric type is not implemented yet",
		zap.Int64("node_id", Params.ProxyCfg.GetNodeID()),
		zap.String("req", req.Request),
		zap.String("metric_type", metricType))

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

B
bigsheeper 已提交
3812 3813 3814
// LoadBalance would do a load balancing operation between query nodes
func (node *Proxy) LoadBalance(ctx context.Context, req *milvuspb.LoadBalanceRequest) (*commonpb.Status, error) {
	log.Debug("Proxy.LoadBalance",
X
Xiaofan 已提交
3815
		zap.Int64("proxy_id", Params.ProxyCfg.GetNodeID()),
B
bigsheeper 已提交
3816 3817 3818 3819 3820 3821 3822 3823 3824
		zap.Any("req", req))

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

	status := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
	}
3825 3826 3827 3828 3829 3830 3831

	collectionID, err := globalMetaCache.GetCollectionID(ctx, req.GetCollectionName())
	if err != nil {
		log.Error("failed to get collection id", zap.String("collection name", req.GetCollectionName()), zap.Error(err))
		status.Reason = err.Error()
		return status, nil
	}
B
bigsheeper 已提交
3832
	infoResp, err := node.queryCoord.LoadBalance(ctx, &querypb.LoadBalanceRequest{
3833 3834 3835 3836 3837 3838
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_LoadBalanceSegments),
			commonpbutil.WithMsgID(0),
			commonpbutil.WithTimeStamp(0),
			commonpbutil.WithSourceID(Params.ProxyCfg.GetNodeID()),
		),
B
bigsheeper 已提交
3839 3840
		SourceNodeIDs:    []int64{req.SrcNodeID},
		DstNodeIDs:       req.DstNodeIDs,
X
xige-16 已提交
3841
		BalanceReason:    querypb.TriggerCondition_GrpcRequest,
B
bigsheeper 已提交
3842
		SealedSegmentIDs: req.SealedSegmentIDs,
3843
		CollectionID:     collectionID,
B
bigsheeper 已提交
3844 3845 3846 3847 3848 3849 3850 3851 3852 3853 3854 3855 3856 3857 3858 3859 3860
	})
	if err != nil {
		log.Error("Failed to LoadBalance from Query Coordinator",
			zap.Any("req", req), zap.Error(err))
		status.Reason = err.Error()
		return status, nil
	}
	if infoResp.ErrorCode != commonpb.ErrorCode_Success {
		log.Error("Failed to LoadBalance from Query Coordinator", zap.String("errMsg", infoResp.Reason))
		status.Reason = infoResp.Reason
		return status, nil
	}
	log.Debug("LoadBalance Done", zap.Any("req", req), zap.Any("status", infoResp))
	status.ErrorCode = commonpb.ErrorCode_Success
	return status, nil
}

3861 3862 3863 3864 3865 3866 3867 3868 3869 3870 3871 3872 3873 3874 3875 3876 3877 3878 3879 3880 3881 3882 3883 3884 3885
// GetReplicas gets replica info
func (node *Proxy) GetReplicas(ctx context.Context, req *milvuspb.GetReplicasRequest) (*milvuspb.GetReplicasResponse, error) {
	log.Info("received get replicas request")
	resp := &milvuspb.GetReplicasResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	req.Base = &commonpb.MsgBase{
		MsgType:  commonpb.MsgType_GetReplicas,
		SourceID: Params.ProxyCfg.GetNodeID(),
	}

	resp, err := node.queryCoord.GetReplicas(ctx, req)
	if err != nil {
		log.Error("Failed to get replicas from Query Coordinator", zap.Error(err))
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
		return resp, nil
	}
	log.Info("received get replicas response", zap.Any("resp", resp), zap.Error(err))
	return resp, nil
}

3886
// GetCompactionState gets the compaction state of multiple segments
3887 3888 3889 3890 3891 3892 3893 3894 3895 3896 3897 3898 3899
func (node *Proxy) GetCompactionState(ctx context.Context, req *milvuspb.GetCompactionStateRequest) (*milvuspb.GetCompactionStateResponse, error) {
	log.Info("received GetCompactionState request", zap.Int64("compactionID", req.GetCompactionID()))
	resp := &milvuspb.GetCompactionStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionState(ctx, req)
	log.Info("received GetCompactionState response", zap.Int64("compactionID", req.GetCompactionID()), zap.Any("resp", resp), zap.Error(err))
	return resp, err
}

3900
// ManualCompaction invokes compaction on specified collection
3901 3902 3903 3904 3905 3906 3907 3908 3909 3910 3911 3912 3913
func (node *Proxy) ManualCompaction(ctx context.Context, req *milvuspb.ManualCompactionRequest) (*milvuspb.ManualCompactionResponse, error) {
	log.Info("received ManualCompaction request", zap.Int64("collectionID", req.GetCollectionID()))
	resp := &milvuspb.ManualCompactionResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.ManualCompaction(ctx, req)
	log.Info("received ManualCompaction response", zap.Int64("collectionID", req.GetCollectionID()), zap.Any("resp", resp), zap.Error(err))
	return resp, err
}

3914
// GetCompactionStateWithPlans returns the compactions states with the given plan ID
3915 3916 3917 3918 3919 3920 3921 3922 3923 3924 3925 3926 3927
func (node *Proxy) GetCompactionStateWithPlans(ctx context.Context, req *milvuspb.GetCompactionPlansRequest) (*milvuspb.GetCompactionPlansResponse, error) {
	log.Info("received GetCompactionStateWithPlans request", zap.Int64("compactionID", req.GetCompactionID()))
	resp := &milvuspb.GetCompactionPlansResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionStateWithPlans(ctx, req)
	log.Info("received GetCompactionStateWithPlans response", zap.Int64("compactionID", req.GetCompactionID()), zap.Any("resp", resp), zap.Error(err))
	return resp, err
}

B
Bingyi Sun 已提交
3928 3929 3930
// GetFlushState gets the flush state of multiple segments
func (node *Proxy) GetFlushState(ctx context.Context, req *milvuspb.GetFlushStateRequest) (*milvuspb.GetFlushStateResponse, error) {
	log.Info("received get flush state request", zap.Any("request", req))
3931
	var err error
B
Bingyi Sun 已提交
3932 3933 3934 3935 3936 3937 3938
	resp := &milvuspb.GetFlushStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		log.Info("unable to get flush state because of closed server")
		return resp, nil
	}

3939
	resp, err = node.dataCoord.GetFlushState(ctx, req)
X
Xiaofan 已提交
3940 3941 3942 3943
	if err != nil {
		log.Info("failed to get flush state response", zap.Error(err))
		return nil, err
	}
B
Bingyi Sun 已提交
3944 3945 3946 3947
	log.Info("received get flush state response", zap.Any("response", resp))
	return resp, err
}

C
Cai Yudong 已提交
3948 3949
// checkHealthy checks proxy state is Healthy
func (node *Proxy) checkHealthy() bool {
3950 3951
	code := node.stateCode.Load().(commonpb.StateCode)
	return code == commonpb.StateCode_Healthy
3952 3953
}

3954 3955 3956
func (node *Proxy) checkHealthyAndReturnCode() (commonpb.StateCode, bool) {
	code := node.stateCode.Load().(commonpb.StateCode)
	return code, code == commonpb.StateCode_Healthy
3957 3958
}

3959
// unhealthyStatus returns the proxy not healthy status
3960 3961 3962
func unhealthyStatus() *commonpb.Status {
	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
C
Cai Yudong 已提交
3963
		Reason:    "proxy not healthy",
3964 3965
	}
}
G
groot 已提交
3966 3967 3968

// 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) {
3969 3970 3971
	log.Info("received import request",
		zap.String("collection name", req.GetCollectionName()),
		zap.Bool("row-based", req.GetRowBased()))
3972 3973 3974 3975 3976 3977
	resp := &milvuspb.ImportResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
			Reason:    "",
		},
	}
G
groot 已提交
3978 3979 3980 3981
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
3982

3983 3984 3985 3986 3987 3988 3989 3990
	err := importutil.ValidateOptions(req.GetOptions())
	if err != nil {
		log.Error("failed to execute import request", zap.Error(err))
		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
	}

3991 3992 3993 3994 3995
	method := "Import"
	tr := timerecord.NewTimeRecorder(method)
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()

3996
	// Call rootCoord to finish import.
3997 3998
	respFromRC, err := node.rootCoord.Import(ctx, req)
	if err != nil {
3999
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
4000 4001 4002 4003 4004
		log.Error("failed to execute bulk load request", zap.Error(err))
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
		return resp, nil
	}
4005 4006 4007

	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
4008
	return respFromRC, nil
G
groot 已提交
4009 4010
}

4011
// GetImportState checks import task state from RootCoord.
G
groot 已提交
4012 4013 4014 4015 4016 4017 4018
func (node *Proxy) GetImportState(ctx context.Context, req *milvuspb.GetImportStateRequest) (*milvuspb.GetImportStateResponse, error) {
	log.Info("received get import state request", zap.Int64("taskID", req.GetTask()))
	resp := &milvuspb.GetImportStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
4019 4020 4021 4022
	method := "GetImportState"
	tr := timerecord.NewTimeRecorder(method)
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
G
groot 已提交
4023 4024

	resp, err := node.rootCoord.GetImportState(ctx, req)
4025 4026 4027 4028 4029 4030 4031 4032 4033 4034 4035 4036
	if err != nil {
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
		log.Error("failed to execute get import state", zap.Error(err))
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
		return resp, nil
	}

	log.Info("successfully received get import state response", zap.Int64("taskID", req.GetTask()), zap.Any("resp", resp), zap.Error(err))
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
	return resp, nil
G
groot 已提交
4037 4038 4039 4040 4041 4042 4043 4044 4045 4046
}

// ListImportTasks get id array of all import tasks from rootcoord
func (node *Proxy) ListImportTasks(ctx context.Context, req *milvuspb.ListImportTasksRequest) (*milvuspb.ListImportTasksResponse, error) {
	log.Info("received list import tasks request")
	resp := &milvuspb.ListImportTasksResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
4047 4048 4049 4050
	method := "ListImportTasks"
	tr := timerecord.NewTimeRecorder(method)
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method,
		metrics.TotalLabel).Inc()
G
groot 已提交
4051
	resp, err := node.rootCoord.ListImportTasks(ctx, req)
4052 4053 4054 4055 4056
	if err != nil {
		metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.FailLabel).Inc()
		log.Error("failed to execute list import tasks", zap.Error(err))
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
		resp.Status.Reason = err.Error()
X
XuanYang-cn 已提交
4057 4058 4059
		return resp, nil
	}

4060 4061 4062
	log.Info("successfully received list import tasks response", zap.String("collection", req.CollectionName), zap.Any("tasks", resp.Tasks))
	metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method, metrics.SuccessLabel).Inc()
	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(Params.ProxyCfg.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
X
XuanYang-cn 已提交
4063 4064 4065
	return resp, err
}

4066 4067 4068 4069 4070 4071
// 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)
	logutil.Logger(ctx).Debug("received request to invalidate credential cache",
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))
4072
	if !node.checkHealthy() {
4073
		return unhealthyStatus(), nil
4074
	}
4075 4076 4077 4078 4079 4080 4081 4082 4083 4084 4085 4086 4087 4088 4089 4090 4091 4092 4093 4094 4095

	username := request.Username
	if globalMetaCache != nil {
		globalMetaCache.RemoveCredential(username) // no need to return error, though credential may be not cached
	}
	logutil.Logger(ctx).Debug("complete to invalidate credential cache",
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))

	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)
	logutil.Logger(ctx).Debug("received request to update credential cache",
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))
4096
	if !node.checkHealthy() {
4097
		return unhealthyStatus(), nil
4098
	}
4099 4100

	credInfo := &internalpb.CredentialInfo{
4101 4102
		Username:       request.Username,
		Sha256Password: request.Password,
4103 4104 4105 4106 4107 4108 4109 4110 4111 4112 4113 4114 4115 4116 4117
	}
	if globalMetaCache != nil {
		globalMetaCache.UpdateCredential(credInfo) // no need to return error, though credential may be not cached
	}
	logutil.Logger(ctx).Debug("complete to update credential cache",
		zap.String("role", typeutil.ProxyRole),
		zap.String("username", request.Username))

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

func (node *Proxy) CreateCredential(ctx context.Context, req *milvuspb.CreateCredentialRequest) (*commonpb.Status, error) {
4118 4119
	log.Debug("CreateCredential", zap.String("role", typeutil.ProxyRole), zap.String("username", req.Username))
	if !node.checkHealthy() {
4120
		return unhealthyStatus(), nil
4121
	}
4122 4123 4124 4125 4126 4127 4128 4129 4130 4131 4132 4133 4134 4135 4136 4137 4138 4139 4140 4141 4142 4143 4144 4145 4146 4147 4148 4149 4150 4151 4152
	// 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 {
		log.Error("decode password fail", zap.String("username", req.Username), zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_CreateCredentialFailure,
			Reason:    "decode password fail key:" + req.Username,
		}, nil
	}
	if err = ValidatePassword(rawPassword); err != nil {
		log.Error("illegal password", zap.String("username", req.Username), zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
		}, nil
	}
	encryptedPassword, err := crypto.PasswordEncrypt(rawPassword)
	if err != nil {
		log.Error("encrypt password fail", zap.String("username", req.Username), zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_CreateCredentialFailure,
			Reason:    "encrypt password fail key:" + req.Username,
		}, nil
	}
4153

4154 4155 4156
	credInfo := &internalpb.CredentialInfo{
		Username:          req.Username,
		EncryptedPassword: encryptedPassword,
4157
		Sha256Password:    crypto.SHA256(rawPassword, req.Username),
4158 4159 4160 4161 4162 4163 4164 4165 4166 4167 4168 4169
	}
	result, err := node.rootCoord.CreateCredential(ctx, credInfo)
	if err != nil { // for error like conntext timeout etc.
		log.Error("create credential fail", zap.String("username", req.Username), zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}
	return result, err
}

C
codeman 已提交
4170
func (node *Proxy) UpdateCredential(ctx context.Context, req *milvuspb.UpdateCredentialRequest) (*commonpb.Status, error) {
4171 4172
	log.Debug("UpdateCredential", zap.String("role", typeutil.ProxyRole), zap.String("username", req.Username))
	if !node.checkHealthy() {
4173
		return unhealthyStatus(), nil
4174
	}
C
codeman 已提交
4175 4176 4177 4178 4179 4180 4181 4182 4183
	rawOldPassword, err := crypto.Base64Decode(req.OldPassword)
	if err != nil {
		log.Error("decode old password fail", zap.String("username", req.Username), zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "decode old password fail when updating:" + req.Username,
		}, nil
	}
	rawNewPassword, err := crypto.Base64Decode(req.NewPassword)
4184 4185 4186 4187 4188 4189 4190
	if err != nil {
		log.Error("decode password fail", zap.String("username", req.Username), zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "decode password fail when updating:" + req.Username,
		}, nil
	}
C
codeman 已提交
4191 4192
	// valid new password
	if err = ValidatePassword(rawNewPassword); err != nil {
4193 4194 4195 4196 4197 4198
		log.Error("illegal password", zap.String("username", req.Username), zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
		}, nil
	}
4199 4200

	if !passwordVerify(ctx, req.Username, rawOldPassword, globalMetaCache) {
C
codeman 已提交
4201 4202 4203 4204 4205 4206 4207
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "old password is not correct:" + req.Username,
		}, nil
	}
	// update meta data
	encryptedPassword, err := crypto.PasswordEncrypt(rawNewPassword)
4208 4209 4210 4211 4212 4213 4214
	if err != nil {
		log.Error("encrypt password fail", zap.String("username", req.Username), zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UpdateCredentialFailure,
			Reason:    "encrypt password fail when updating:" + req.Username,
		}, nil
	}
C
codeman 已提交
4215
	updateCredReq := &internalpb.CredentialInfo{
4216
		Username:          req.Username,
4217
		Sha256Password:    crypto.SHA256(rawNewPassword, req.Username),
4218 4219
		EncryptedPassword: encryptedPassword,
	}
C
codeman 已提交
4220
	result, err := node.rootCoord.UpdateCredential(ctx, updateCredReq)
4221 4222 4223 4224 4225 4226 4227 4228 4229 4230 4231
	if err != nil { // for error like conntext timeout etc.
		log.Error("update credential fail", zap.String("username", req.Username), zap.Error(err))
		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) {
4232 4233
	log.Debug("DeleteCredential", zap.String("role", typeutil.ProxyRole), zap.String("username", req.Username))
	if !node.checkHealthy() {
4234
		return unhealthyStatus(), nil
4235 4236
	}

4237 4238 4239 4240 4241 4242
	if req.Username == util.UserRoot {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_DeleteCredentialFailure,
			Reason:    "user root cannot be deleted",
		}, nil
	}
4243 4244 4245 4246 4247 4248 4249 4250 4251 4252 4253 4254
	result, err := node.rootCoord.DeleteCredential(ctx, req)
	if err != nil { // for error like conntext timeout etc.
		log.Error("delete credential fail", zap.String("username", req.Username), zap.Error(err))
		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) {
4255 4256
	log.Debug("ListCredUsers", zap.String("role", typeutil.ProxyRole))
	if !node.checkHealthy() {
4257
		return &milvuspb.ListCredUsersResponse{Status: unhealthyStatus()}, nil
4258
	}
4259
	rootCoordReq := &milvuspb.ListCredUsersRequest{
4260 4261 4262
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_ListCredUsernames),
		),
4263 4264
	}
	resp, err := node.rootCoord.ListCredUsers(ctx, rootCoordReq)
4265 4266 4267 4268 4269 4270 4271 4272 4273 4274 4275 4276
	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,
		},
4277
		Usernames: resp.Usernames,
4278 4279
	}, nil
}
4280

4281 4282 4283
func (node *Proxy) CreateRole(ctx context.Context, req *milvuspb.CreateRoleRequest) (*commonpb.Status, error) {
	logger.Debug("CreateRole", zap.Any("req", req))
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4284
		return errorutil.UnhealthyStatus(code), nil
4285 4286 4287 4288 4289 4290 4291 4292 4293 4294
	}

	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(),
4295
		}, nil
4296 4297 4298 4299 4300 4301 4302 4303
	}

	result, err := node.rootCoord.CreateRole(ctx, req)
	if err != nil {
		logger.Error("fail to create role", zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4304
		}, nil
4305 4306
	}
	return result, nil
4307 4308
}

4309 4310 4311
func (node *Proxy) DropRole(ctx context.Context, req *milvuspb.DropRoleRequest) (*commonpb.Status, error) {
	logger.Debug("DropRole", zap.Any("req", req))
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4312
		return errorutil.UnhealthyStatus(code), nil
4313 4314 4315 4316 4317
	}
	if err := ValidateRoleName(req.RoleName); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4318
		}, nil
4319
	}
4320 4321 4322 4323 4324
	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,
4325
		}, nil
4326
	}
4327 4328 4329 4330 4331 4332
	result, err := node.rootCoord.DropRole(ctx, req)
	if err != nil {
		logger.Error("fail to drop role", zap.String("role_name", req.RoleName), zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4333
		}, nil
4334 4335
	}
	return result, nil
4336 4337
}

4338 4339 4340
func (node *Proxy) OperateUserRole(ctx context.Context, req *milvuspb.OperateUserRoleRequest) (*commonpb.Status, error) {
	logger.Debug("OperateUserRole", zap.Any("req", req))
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4341
		return errorutil.UnhealthyStatus(code), nil
4342 4343 4344 4345 4346
	}
	if err := ValidateUsername(req.Username); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4347
		}, nil
4348 4349 4350 4351 4352
	}
	if err := ValidateRoleName(req.RoleName); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4353
		}, nil
4354 4355 4356 4357 4358 4359 4360 4361
	}

	result, err := node.rootCoord.OperateUserRole(ctx, req)
	if err != nil {
		logger.Error("fail to operate user role", zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4362
		}, nil
4363 4364
	}
	return result, nil
4365 4366
}

4367 4368 4369
func (node *Proxy) SelectRole(ctx context.Context, req *milvuspb.SelectRoleRequest) (*milvuspb.SelectRoleResponse, error) {
	logger.Debug("SelectRole", zap.Any("req", req))
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4370
		return &milvuspb.SelectRoleResponse{Status: errorutil.UnhealthyStatus(code)}, nil
4371 4372 4373 4374 4375 4376 4377 4378 4379
	}

	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(),
				},
4380
			}, nil
4381 4382 4383 4384 4385 4386 4387 4388 4389 4390 4391
		}
	}

	result, err := node.rootCoord.SelectRole(ctx, req)
	if err != nil {
		logger.Error("fail to select role", zap.Error(err))
		return &milvuspb.SelectRoleResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
4392
		}, nil
4393 4394
	}
	return result, nil
4395 4396
}

4397 4398 4399
func (node *Proxy) SelectUser(ctx context.Context, req *milvuspb.SelectUserRequest) (*milvuspb.SelectUserResponse, error) {
	logger.Debug("SelectUser", zap.Any("req", req))
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4400
		return &milvuspb.SelectUserResponse{Status: errorutil.UnhealthyStatus(code)}, nil
4401 4402 4403 4404 4405 4406 4407 4408 4409
	}

	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(),
				},
4410
			}, nil
4411 4412 4413 4414 4415 4416 4417 4418 4419 4420 4421
		}
	}

	result, err := node.rootCoord.SelectUser(ctx, req)
	if err != nil {
		logger.Error("fail to select user", zap.Error(err))
		return &milvuspb.SelectUserResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
4422
		}, nil
4423 4424
	}
	return result, nil
4425 4426
}

4427 4428 4429 4430 4431 4432 4433 4434 4435 4436 4437 4438 4439 4440 4441 4442 4443 4444 4445 4446 4447 4448 4449 4450 4451 4452 4453 4454 4455 4456
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
4457 4458
}

4459 4460 4461
func (node *Proxy) OperatePrivilege(ctx context.Context, req *milvuspb.OperatePrivilegeRequest) (*commonpb.Status, error) {
	logger.Debug("OperatePrivilege", zap.Any("req", req))
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4462
		return errorutil.UnhealthyStatus(code), nil
4463 4464 4465 4466 4467
	}
	if err := node.validPrivilegeParams(req); err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_IllegalArgument,
			Reason:    err.Error(),
4468
		}, nil
4469 4470 4471 4472 4473 4474
	}
	curUser, err := GetCurUserFromContext(ctx)
	if err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4475
		}, nil
4476 4477 4478 4479 4480 4481 4482 4483
	}
	req.Entity.Grantor.User = &milvuspb.UserEntity{Name: curUser}
	result, err := node.rootCoord.OperatePrivilege(ctx, req)
	if err != nil {
		logger.Error("fail to operate privilege", zap.Error(err))
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
4484
		}, nil
4485 4486
	}
	return result, nil
4487 4488
}

4489 4490 4491 4492 4493 4494 4495 4496 4497 4498 4499 4500 4501 4502 4503 4504 4505 4506 4507 4508 4509 4510 4511 4512 4513 4514 4515 4516 4517
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) {
	logger.Debug("SelectGrant", zap.Any("req", req))
	if code, ok := node.checkHealthyAndReturnCode(); !ok {
4518
		return &milvuspb.SelectGrantResponse{Status: errorutil.UnhealthyStatus(code)}, nil
4519 4520 4521 4522 4523 4524 4525 4526
	}

	if err := node.validGrantParams(req); err != nil {
		return &milvuspb.SelectGrantResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_IllegalArgument,
				Reason:    err.Error(),
			},
4527
		}, nil
4528 4529 4530 4531 4532 4533 4534 4535 4536 4537
	}

	result, err := node.rootCoord.SelectGrant(ctx, req)
	if err != nil {
		logger.Error("fail to select grant", zap.Error(err))
		return &milvuspb.SelectGrantResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
4538
		}, nil
4539 4540 4541 4542 4543 4544 4545 4546 4547 4548 4549 4550 4551 4552 4553 4554 4555 4556 4557 4558 4559 4560 4561 4562 4563 4564 4565 4566
	}
	return result, nil
}

func (node *Proxy) RefreshPolicyInfoCache(ctx context.Context, req *proxypb.RefreshPolicyInfoCacheRequest) (*commonpb.Status, error) {
	logger.Debug("RefreshPrivilegeInfoCache", zap.Any("req", req))
	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 {
			log.Error("fail to refresh policy info", zap.Error(err))
			return &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_RefreshPolicyInfoCacheFailure,
				Reason:    err.Error(),
			}, err
		}
	}
	logger.Debug("RefreshPrivilegeInfoCache success")

	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
	}, nil
4567
}
4568 4569 4570 4571 4572 4573 4574 4575 4576 4577 4578 4579 4580 4581 4582 4583 4584 4585 4586 4587 4588

// SetRates limits the rates of requests.
func (node *Proxy) SetRates(ctx context.Context, request *proxypb.SetRatesRequest) (*commonpb.Status, error) {
	log.Debug("SetRates", zap.String("role", typeutil.ProxyRole), zap.Any("rates", request.GetRates()))
	resp := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
	}
	if !node.checkHealthy() {
		resp = unhealthyStatus()
		return resp, nil
	}

	err := node.multiRateLimiter.globalRateLimiter.setRates(request.GetRates())
	// TODO: set multiple rate limiter rates
	if err != nil {
		resp.Reason = err.Error()
		return resp, nil
	}
	resp.ErrorCode = commonpb.ErrorCode_Success
	return resp, nil
}
4589 4590 4591 4592

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")
4593 4594 4595 4596
		return &milvuspb.CheckHealthResponse{
			Status:    unhealthyStatus(),
			IsHealthy: false,
			Reasons:   []string{reason}}, nil
4597 4598 4599 4600 4601 4602 4603 4604 4605 4606 4607
	}

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

		if err != nil {
4608
			log.Warn("check health fail", zap.String("role", role), zap.Error(err))
4609 4610 4611 4612 4613
			errReasons = append(errReasons, fmt.Sprintf("check health fail for %s", role))
			return err
		}

		if !resp.IsHealthy {
4614
			log.Warn("check health fail", zap.String("role", role))
4615 4616 4617 4618 4619 4620 4621 4622 4623 4624 4625 4626 4627 4628 4629 4630 4631 4632 4633 4634 4635 4636 4637 4638 4639 4640 4641 4642 4643 4644 4645 4646 4647
			errReasons = append(errReasons, resp.Reasons...)
		}
		return nil
	}

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

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

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

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

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

4648 4649 4650 4651 4652 4653 4654
	return &milvuspb.CheckHealthResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
			Reason:    "",
		},
		IsHealthy: true,
	}, nil
4655
}