impl.go 37.7 KB
Newer Older
1 2 3 4
package proxynode

import (
	"context"
5
	"errors"
6
	"fmt"
C
cai.zhang 已提交
7
	"os"
8 9
	"strconv"

10
	"go.uber.org/zap"
S
sunby 已提交
11

12
	"github.com/zilliztech/milvus-distributed/internal/log"
13 14
	"github.com/zilliztech/milvus-distributed/internal/msgstream"
	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
G
godchen 已提交
15
	"github.com/zilliztech/milvus-distributed/internal/proto/datapb"
G
godchen 已提交
16
	"github.com/zilliztech/milvus-distributed/internal/proto/internalpb"
17 18
	"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
	"github.com/zilliztech/milvus-distributed/internal/proto/proxypb"
G
godchen 已提交
19
	"github.com/zilliztech/milvus-distributed/internal/proto/querypb"
G
godchen 已提交
20
	"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
21 22
)

G
godchen 已提交
23
func (node *ProxyNode) UpdateStateCode(code internalpb.StateCode) {
24
	node.stateCode.Store(code)
Z
zhenshan.cao 已提交
25 26
}

G
godchen 已提交
27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
func (node *ProxyNode) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
	stats := &internalpb.ComponentStates{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
		},
	}
	code, ok := node.stateCode.Load().(internalpb.StateCode)
	if !ok {
		errMsg := "unexpected error in type assertion"
		stats.Status = &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    errMsg,
		}
		return stats, errors.New(errMsg)
	}
	info := &internalpb.ComponentInfo{
		NodeID:    Params.ProxyID,
		Role:      typeutil.ProxyNodeRole,
		StateCode: code,
	}
	stats.State = info
	return stats, nil
}

func (node *ProxyNode) GetStatisticsChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
	return &milvuspb.StringResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
			Reason:    "",
		},
		Value: "",
	}, nil
}

61
func (node *ProxyNode) InvalidateCollectionMetaCache(ctx context.Context, request *proxypb.InvalidateCollMetaCacheRequest) (*commonpb.Status, error) {
D
dragondriver 已提交
62 63 64 65 66
	log.Debug("InvalidateCollectionMetaCache",
		zap.String("role", Params.RoleName),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

67
	collectionName := request.CollectionName
G
godchen 已提交
68
	globalMetaCache.RemoveCollection(ctx, collectionName) // no need to return error, though collection may be not cached
D
dragondriver 已提交
69 70 71 72 73 74

	log.Debug("InvalidateCollectionMetaCache Done",
		zap.String("role", Params.RoleName),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

75
	return &commonpb.Status{
76
		ErrorCode: commonpb.ErrorCode_Success,
77 78
		Reason:    "",
	}, nil
79 80
}

81
func (node *ProxyNode) CreateCollection(ctx context.Context, request *milvuspb.CreateCollectionRequest) (*commonpb.Status, error) {
82
	cct := &CreateCollectionTask{
S
sunby 已提交
83
		ctx:                     ctx,
84 85
		Condition:               NewTaskCondition(ctx),
		CreateCollectionRequest: request,
T
ThreadDao 已提交
86 87
		masterService:           node.masterService,
		dataServiceClient:       node.dataService,
88 89
	}

90
	err := node.sched.DdQueue.Enqueue(cct)
91 92
	if err != nil {
		return &commonpb.Status{
93
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
94 95 96 97
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115
	log.Debug("CreateCollection",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("schema", request.Schema))
	defer func() {
		log.Debug("CreateCollection Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("schema", request.Schema))
	}()

116 117 118
	err = cct.WaitToFinish()
	if err != nil {
		return &commonpb.Status{
119
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
120 121 122 123 124 125 126
			Reason:    err.Error(),
		}, nil
	}

	return cct.result, nil
}

127
func (node *ProxyNode) DropCollection(ctx context.Context, request *milvuspb.DropCollectionRequest) (*commonpb.Status, error) {
128
	dct := &DropCollectionTask{
S
sunby 已提交
129
		ctx:                   ctx,
130 131
		Condition:             NewTaskCondition(ctx),
		DropCollectionRequest: request,
T
ThreadDao 已提交
132
		masterService:         node.masterService,
133 134
	}

135
	err := node.sched.DdQueue.Enqueue(dct)
136 137
	if err != nil {
		return &commonpb.Status{
138
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
139 140 141 142
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158
	log.Debug("DropCollection",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
	defer func() {
		log.Debug("DropCollection Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))
	}()

159 160 161
	err = dct.WaitToFinish()
	if err != nil {
		return &commonpb.Status{
162
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
163 164 165 166 167 168 169
			Reason:    err.Error(),
		}, nil
	}

	return dct.result, nil
}

170
func (node *ProxyNode) HasCollection(ctx context.Context, request *milvuspb.HasCollectionRequest) (*milvuspb.BoolResponse, error) {
171
	hct := &HasCollectionTask{
S
sunby 已提交
172
		ctx:                  ctx,
173 174
		Condition:            NewTaskCondition(ctx),
		HasCollectionRequest: request,
T
ThreadDao 已提交
175
		masterService:        node.masterService,
176 177
	}

178
	err := node.sched.DdQueue.Enqueue(hct)
179 180 181
	if err != nil {
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
182
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
183 184 185 186 187
				Reason:    err.Error(),
			},
		}, nil
	}

D
dragondriver 已提交
188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203
	log.Debug("HasCollection",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
	defer func() {
		log.Debug("HasCollection Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))
	}()

204 205 206 207
	err = hct.WaitToFinish()
	if err != nil {
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
208
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
209 210 211 212 213 214 215 216
				Reason:    err.Error(),
			},
		}, nil
	}

	return hct.result, nil
}

217
func (node *ProxyNode) LoadCollection(ctx context.Context, request *milvuspb.LoadCollectionRequest) (*commonpb.Status, error) {
218 219

	lct := &LoadCollectionTask{
S
sunby 已提交
220
		ctx:                   ctx,
221 222
		Condition:             NewTaskCondition(ctx),
		LoadCollectionRequest: request,
T
ThreadDao 已提交
223
		queryService:          node.queryService,
224 225 226 227 228
	}

	err := node.sched.DdQueue.Enqueue(lct)
	if err != nil {
		return &commonpb.Status{
229
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
230 231 232
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249

	log.Debug("LoadCollection",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
	defer func() {
		log.Debug("LoadCollection Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))
	}()

250 251 252
	err = lct.WaitToFinish()
	if err != nil {
		return &commonpb.Status{
253
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
254 255 256 257 258
			Reason:    err.Error(),
		}, nil
	}

	return lct.result, nil
259 260
}

261
func (node *ProxyNode) ReleaseCollection(ctx context.Context, request *milvuspb.ReleaseCollectionRequest) (*commonpb.Status, error) {
262
	rct := &ReleaseCollectionTask{
S
sunby 已提交
263
		ctx:                      ctx,
264 265
		Condition:                NewTaskCondition(ctx),
		ReleaseCollectionRequest: request,
T
ThreadDao 已提交
266
		queryService:             node.queryService,
267 268 269 270 271
	}

	err := node.sched.DdQueue.Enqueue(rct)
	if err != nil {
		return &commonpb.Status{
272
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
273 274 275 276
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292
	log.Debug("ReleaseCollection",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
	defer func() {
		log.Debug("ReleaseCollection Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))
	}()

293 294 295
	err = rct.WaitToFinish()
	if err != nil {
		return &commonpb.Status{
296
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
297 298 299 300 301
			Reason:    err.Error(),
		}, nil
	}

	return rct.result, nil
302 303
}

304
func (node *ProxyNode) DescribeCollection(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
305
	dct := &DescribeCollectionTask{
S
sunby 已提交
306
		ctx:                       ctx,
307 308
		Condition:                 NewTaskCondition(ctx),
		DescribeCollectionRequest: request,
T
ThreadDao 已提交
309
		masterService:             node.masterService,
310 311
	}

312
	err := node.sched.DdQueue.Enqueue(dct)
313 314 315
	if err != nil {
		return &milvuspb.DescribeCollectionResponse{
			Status: &commonpb.Status{
316
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
317 318 319 320 321
				Reason:    err.Error(),
			},
		}, nil
	}

D
dragondriver 已提交
322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337
	log.Debug("DescribeCollection",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
	defer func() {
		log.Debug("DescribeCollection Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))
	}()

338 339 340 341
	err = dct.WaitToFinish()
	if err != nil {
		return &milvuspb.DescribeCollectionResponse{
			Status: &commonpb.Status{
342
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
343 344 345 346 347 348 349 350
				Reason:    err.Error(),
			},
		}, nil
	}

	return dct.result, nil
}

G
godchen 已提交
351
func (node *ProxyNode) GetCollectionStatistics(ctx context.Context, request *milvuspb.GetCollectionStatisticsRequest) (*milvuspb.GetCollectionStatisticsResponse, error) {
352
	g := &GetCollectionsStatisticsTask{
G
godchen 已提交
353 354 355 356
		ctx:                            ctx,
		Condition:                      NewTaskCondition(ctx),
		GetCollectionStatisticsRequest: request,
		dataService:                    node.dataService,
357 358
	}

359
	err := node.sched.DdQueue.Enqueue(g)
360
	if err != nil {
G
godchen 已提交
361
		return &milvuspb.GetCollectionStatisticsResponse{
362
			Status: &commonpb.Status{
363
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
364 365 366 367 368
				Reason:    err.Error(),
			},
		}, nil
	}

D
dragondriver 已提交
369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384
	log.Debug("GetCollectionStatistics",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
	defer func() {
		log.Debug("GetCollectionStatistics Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))
	}()

385 386
	err = g.WaitToFinish()
	if err != nil {
G
godchen 已提交
387
		return &milvuspb.GetCollectionStatisticsResponse{
388
			Status: &commonpb.Status{
389
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
390 391 392 393 394 395
				Reason:    err.Error(),
			},
		}, nil
	}

	return g.result, nil
396 397
}

G
godchen 已提交
398
func (node *ProxyNode) ShowCollections(ctx context.Context, request *milvuspb.ShowCollectionsRequest) (*milvuspb.ShowCollectionsResponse, error) {
399
	sct := &ShowCollectionsTask{
G
godchen 已提交
400 401 402 403
		ctx:                    ctx,
		Condition:              NewTaskCondition(ctx),
		ShowCollectionsRequest: request,
		masterService:          node.masterService,
404 405
	}

406
	err := node.sched.DdQueue.Enqueue(sct)
407
	if err != nil {
G
godchen 已提交
408
		return &milvuspb.ShowCollectionsResponse{
409
			Status: &commonpb.Status{
410
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
411 412 413 414 415
				Reason:    err.Error(),
			},
		}, nil
	}

D
dragondriver 已提交
416 417 418 419 420 421 422 423 424 425 426 427 428 429
	log.Debug("ShowCollections",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName))
	defer func() {
		log.Debug("ShowCollections Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName))
	}()

430 431
	err = sct.WaitToFinish()
	if err != nil {
G
godchen 已提交
432
		return &milvuspb.ShowCollectionsResponse{
433
			Status: &commonpb.Status{
434
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
435 436 437 438 439 440 441 442
				Reason:    err.Error(),
			},
		}, nil
	}

	return sct.result, nil
}

443
func (node *ProxyNode) CreatePartition(ctx context.Context, request *milvuspb.CreatePartitionRequest) (*commonpb.Status, error) {
444
	cpt := &CreatePartitionTask{
S
sunby 已提交
445
		ctx:                    ctx,
446 447
		Condition:              NewTaskCondition(ctx),
		CreatePartitionRequest: request,
T
ThreadDao 已提交
448
		masterService:          node.masterService,
449 450 451
		result:                 nil,
	}

452
	err := node.sched.DdQueue.Enqueue(cpt)
453 454
	if err != nil {
		return &commonpb.Status{
455
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
456 457 458
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477

	log.Debug("CreatePartition",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
	defer func() {
		log.Debug("CreatePartition Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))
	}()

478 479 480
	err = cpt.WaitToFinish()
	if err != nil {
		return &commonpb.Status{
481
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
482 483 484 485 486 487
			Reason:    err.Error(),
		}, nil
	}
	return cpt.result, nil
}

488
func (node *ProxyNode) DropPartition(ctx context.Context, request *milvuspb.DropPartitionRequest) (*commonpb.Status, error) {
489
	dpt := &DropPartitionTask{
S
sunby 已提交
490
		ctx:                  ctx,
491 492
		Condition:            NewTaskCondition(ctx),
		DropPartitionRequest: request,
T
ThreadDao 已提交
493
		masterService:        node.masterService,
494 495 496
		result:               nil,
	}

497
	err := node.sched.DdQueue.Enqueue(dpt)
498 499 500

	if err != nil {
		return &commonpb.Status{
501
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
502 503 504
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523

	log.Debug("DropPartition",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
	defer func() {
		log.Debug("DropPartition Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))
	}()

524 525 526
	err = dpt.WaitToFinish()
	if err != nil {
		return &commonpb.Status{
527
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
528 529 530 531 532 533
			Reason:    err.Error(),
		}, nil
	}
	return dpt.result, nil
}

534
func (node *ProxyNode) HasPartition(ctx context.Context, request *milvuspb.HasPartitionRequest) (*milvuspb.BoolResponse, error) {
535
	hpt := &HasPartitionTask{
S
sunby 已提交
536
		ctx:                 ctx,
537 538
		Condition:           NewTaskCondition(ctx),
		HasPartitionRequest: request,
T
ThreadDao 已提交
539
		masterService:       node.masterService,
540 541 542
		result:              nil,
	}

543
	err := node.sched.DdQueue.Enqueue(hpt)
544 545 546 547

	if err != nil {
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
548
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
549 550 551 552 553
				Reason:    err.Error(),
			},
			Value: false,
		}, nil
	}
D
dragondriver 已提交
554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572

	log.Debug("HasPartition",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
	defer func() {
		log.Debug("HasPartition Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))
	}()

573 574 575 576
	err = hpt.WaitToFinish()
	if err != nil {
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
577
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
578 579 580 581 582 583 584 585
				Reason:    err.Error(),
			},
			Value: false,
		}, nil
	}
	return hpt.result, nil
}

G
godchen 已提交
586
func (node *ProxyNode) LoadPartitions(ctx context.Context, request *milvuspb.LoadPartitionsRequest) (*commonpb.Status, error) {
587
	lpt := &LoadPartitionTask{
G
godchen 已提交
588 589 590 591
		ctx:                   ctx,
		Condition:             NewTaskCondition(ctx),
		LoadPartitionsRequest: request,
		queryService:          node.queryService,
592 593 594 595 596
	}

	err := node.sched.DdQueue.Enqueue(lpt)
	if err != nil {
		return &commonpb.Status{
597
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
598 599 600 601
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619
	log.Debug("LoadPartitions",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
	defer func() {
		log.Debug("LoadPartitions Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))
	}()

620 621 622
	err = lpt.WaitToFinish()
	if err != nil {
		return &commonpb.Status{
623
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
624 625 626 627 628
			Reason:    err.Error(),
		}, nil
	}

	return lpt.result, nil
629 630
}

G
godchen 已提交
631
func (node *ProxyNode) ReleasePartitions(ctx context.Context, request *milvuspb.ReleasePartitionsRequest) (*commonpb.Status, error) {
632
	rpt := &ReleasePartitionTask{
G
godchen 已提交
633 634 635 636
		ctx:                      ctx,
		Condition:                NewTaskCondition(ctx),
		ReleasePartitionsRequest: request,
		queryService:             node.queryService,
637 638 639 640 641
	}

	err := node.sched.DdQueue.Enqueue(rpt)
	if err != nil {
		return &commonpb.Status{
642
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
643 644 645 646
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664
	log.Debug("ReleasePartitions",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
	defer func() {
		log.Debug("ReleasePartitions Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))
	}()

665 666 667
	err = rpt.WaitToFinish()
	if err != nil {
		return &commonpb.Status{
668
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
669 670 671 672 673
			Reason:    err.Error(),
		}, nil
	}

	return rpt.result, nil
674 675
}

G
godchen 已提交
676
func (node *ProxyNode) GetPartitionStatistics(ctx context.Context, request *milvuspb.GetPartitionStatisticsRequest) (*milvuspb.GetPartitionStatisticsResponse, error) {
677 678 679
	panic("implement me")
}

G
godchen 已提交
680
func (node *ProxyNode) ShowPartitions(ctx context.Context, request *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
681
	spt := &ShowPartitionsTask{
G
godchen 已提交
682 683 684 685 686
		ctx:                   ctx,
		Condition:             NewTaskCondition(ctx),
		ShowPartitionsRequest: request,
		masterService:         node.masterService,
		result:                nil,
687 688
	}

689
	err := node.sched.DdQueue.Enqueue(spt)
690 691

	if err != nil {
G
godchen 已提交
692
		return &milvuspb.ShowPartitionsResponse{
693
			Status: &commonpb.Status{
694
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
695 696 697 698 699
				Reason:    err.Error(),
			},
		}, nil
	}

D
dragondriver 已提交
700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715
	log.Debug("ShowPartitions",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
	defer func() {
		log.Debug("ShowPartitions Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))
	}()

716 717
	err = spt.WaitToFinish()
	if err != nil {
G
godchen 已提交
718
		return &milvuspb.ShowPartitionsResponse{
719
			Status: &commonpb.Status{
720
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
721 722 723 724 725 726 727
				Reason:    err.Error(),
			},
		}, nil
	}
	return spt.result, nil
}

728
func (node *ProxyNode) CreateIndex(ctx context.Context, request *milvuspb.CreateIndexRequest) (*commonpb.Status, error) {
729
	cit := &CreateIndexTask{
S
sunby 已提交
730
		ctx:                ctx,
731 732
		Condition:          NewTaskCondition(ctx),
		CreateIndexRequest: request,
T
ThreadDao 已提交
733
		masterService:      node.masterService,
734 735
	}

736
	err := node.sched.DdQueue.Enqueue(cit)
737 738
	if err != nil {
		return &commonpb.Status{
739
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
740 741 742 743
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763
	log.Debug("CreateIndex",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.Any("extra_params", request.ExtraParams))
	defer func() {
		log.Debug("CreateIndex Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.Any("extra_params", request.ExtraParams))
	}()

764 765 766
	err = cit.WaitToFinish()
	if err != nil {
		return &commonpb.Status{
767
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
768 769 770 771 772 773 774
			Reason:    err.Error(),
		}, nil
	}

	return cit.result, nil
}

775
func (node *ProxyNode) DescribeIndex(ctx context.Context, request *milvuspb.DescribeIndexRequest) (*milvuspb.DescribeIndexResponse, error) {
776
	dit := &DescribeIndexTask{
S
sunby 已提交
777
		ctx:                  ctx,
778 779
		Condition:            NewTaskCondition(ctx),
		DescribeIndexRequest: request,
T
ThreadDao 已提交
780
		masterService:        node.masterService,
781 782
	}

783
	err := node.sched.DdQueue.Enqueue(dit)
784 785 786
	if err != nil {
		return &milvuspb.DescribeIndexResponse{
			Status: &commonpb.Status{
787
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
788 789 790 791 792
				Reason:    err.Error(),
			},
		}, nil
	}

D
dragondriver 已提交
793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812
	log.Debug("DescribeIndex",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
	defer func() {
		log.Debug("DescribeIndex Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))
	}()

813 814
	err = dit.WaitToFinish()
	if err != nil {
Z
zhenshan.cao 已提交
815 816 817 818
		errCode := commonpb.ErrorCode_UnexpectedError
		if dit.result != nil {
			errCode = dit.result.Status.GetErrorCode()
		}
819 820
		return &milvuspb.DescribeIndexResponse{
			Status: &commonpb.Status{
Z
zhenshan.cao 已提交
821
				ErrorCode: errCode,
822 823 824 825 826 827 828 829
				Reason:    err.Error(),
			},
		}, nil
	}

	return dit.result, nil
}

830
func (node *ProxyNode) DropIndex(ctx context.Context, request *milvuspb.DropIndexRequest) (*commonpb.Status, error) {
B
BossZou 已提交
831
	dit := &DropIndexTask{
S
sunby 已提交
832
		ctx:              ctx,
B
BossZou 已提交
833 834
		Condition:        NewTaskCondition(ctx),
		DropIndexRequest: request,
T
ThreadDao 已提交
835
		masterService:    node.masterService,
B
BossZou 已提交
836 837 838 839
	}
	err := node.sched.DdQueue.Enqueue(dit)
	if err != nil {
		return &commonpb.Status{
840
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
B
BossZou 已提交
841 842 843
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864

	log.Debug("DropIndex",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
	defer func() {
		log.Debug("DropIndex Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))
	}()

B
BossZou 已提交
865 866 867
	err = dit.WaitToFinish()
	if err != nil {
		return &commonpb.Status{
868
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
B
BossZou 已提交
869 870 871 872 873 874
			Reason:    err.Error(),
		}, nil
	}
	return dit.result, nil
}

G
godchen 已提交
875
func (node *ProxyNode) GetIndexState(ctx context.Context, request *milvuspb.GetIndexStateRequest) (*milvuspb.GetIndexStateResponse, error) {
876
	dipt := &GetIndexStateTask{
G
godchen 已提交
877 878 879 880 881
		ctx:                  ctx,
		Condition:            NewTaskCondition(ctx),
		GetIndexStateRequest: request,
		indexService:         node.indexService,
		masterService:        node.masterService,
882 883
	}

884
	err := node.sched.DdQueue.Enqueue(dipt)
885
	if err != nil {
G
godchen 已提交
886
		return &milvuspb.GetIndexStateResponse{
887
			Status: &commonpb.Status{
888
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
889 890 891 892 893
				Reason:    err.Error(),
			},
		}, nil
	}

D
dragondriver 已提交
894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913
	log.Debug("GetIndexState",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
	defer func() {
		log.Debug("GetIndexState Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))
	}()

914 915
	err = dipt.WaitToFinish()
	if err != nil {
G
godchen 已提交
916
		return &milvuspb.GetIndexStateResponse{
917
			Status: &commonpb.Status{
918
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
919 920 921 922 923 924 925 926
				Reason:    err.Error(),
			},
		}, nil
	}

	return dipt.result, nil
}

927
func (node *ProxyNode) Insert(ctx context.Context, request *milvuspb.InsertRequest) (*milvuspb.InsertResponse, error) {
928
	it := &InsertTask{
T
ThreadDao 已提交
929 930 931
		ctx:         ctx,
		Condition:   NewTaskCondition(ctx),
		dataService: node.dataService,
932 933 934 935
		BaseInsertTask: BaseInsertTask{
			BaseMsg: msgstream.BaseMsg{
				HashValues: request.HashKeys,
			},
G
godchen 已提交
936
			InsertRequest: internalpb.InsertRequest{
937
				Base: &commonpb.MsgBase{
938
					MsgType: commonpb.MsgType_Insert,
939 940 941 942 943 944 945
					MsgID:   0,
				},
				CollectionName: request.CollectionName,
				PartitionName:  request.PartitionName,
				RowData:        request.RowData,
			},
		},
946
		rowIDAllocator: node.idAllocator,
947 948
	}
	if len(it.PartitionName) <= 0 {
949
		it.PartitionName = Params.DefaultPartitionName
950 951
	}

952
	err := node.sched.DmQueue.Enqueue(it)
953 954 955 956

	if err != nil {
		return &milvuspb.InsertResponse{
			Status: &commonpb.Status{
957
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
958 959 960 961 962
				Reason:    err.Error(),
			},
		}, nil
	}

D
dragondriver 已提交
963 964 965 966 967 968 969
	log.Debug("Insert",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", it.BaseInsertTask.InsertRequest.Base.MsgID),
		zap.Uint64("timestamp", it.BaseInsertTask.InsertRequest.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
970 971
		zap.Any("len(RowData)", len(it.RowData)),
		zap.Any("len(RowIDs)", len(it.RowIDs)))
D
dragondriver 已提交
972 973 974 975 976 977 978 979 980
	defer func() {
		log.Debug("Insert Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", it.BaseInsertTask.InsertRequest.Base.MsgID),
			zap.Uint64("timestamp", it.BaseInsertTask.InsertRequest.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName),
981 982
			zap.Any("len(RowData)", len(it.RowData)),
			zap.Any("len(RowIDs)", len(it.RowIDs)))
D
dragondriver 已提交
983 984
	}()

985 986 987 988
	err = it.WaitToFinish()
	if err != nil {
		return &milvuspb.InsertResponse{
			Status: &commonpb.Status{
989
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
990 991 992 993 994 995 996 997
				Reason:    err.Error(),
			},
		}, nil
	}

	return it.result, nil
}

998
func (node *ProxyNode) Search(ctx context.Context, request *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
999
	qt := &SearchTask{
S
sunby 已提交
1000
		ctx:       ctx,
1001
		Condition: NewTaskCondition(ctx),
G
godchen 已提交
1002
		SearchRequest: &internalpb.SearchRequest{
1003
			Base: &commonpb.MsgBase{
1004
				MsgType:  commonpb.MsgType_Search,
1005
				SourceID: Params.ProxyID,
1006
			},
1007
			ResultChannelID: strconv.FormatInt(Params.ProxyID, 10),
1008 1009
		},
		queryMsgStream: node.queryMsgStream,
G
godchen 已提交
1010
		resultBuf:      make(chan []*internalpb.SearchResults),
1011 1012 1013
		query:          request,
	}

1014
	err := node.sched.DqQueue.Enqueue(qt)
1015 1016 1017
	if err != nil {
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
1018
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1019 1020 1021 1022 1023
				Reason:    err.Error(),
			},
		}, nil
	}

D
dragondriver 已提交
1024 1025 1026 1027 1028 1029 1030 1031
	log.Debug("Search",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", qt.Base.MsgID),
		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),
1032
		zap.Any("len(PlaceholderGroup)", len(request.PlaceholderGroup)))
D
dragondriver 已提交
1033 1034 1035 1036 1037 1038 1039 1040 1041 1042
	defer func() {
		log.Debug("Search Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", qt.Base.MsgID),
			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),
1043
			zap.Any("len(PlaceholderGroup)", len(request.PlaceholderGroup)))
D
dragondriver 已提交
1044 1045
	}()

1046 1047 1048 1049
	err = qt.WaitToFinish()
	if err != nil {
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
1050
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1051 1052 1053 1054 1055 1056 1057 1058
				Reason:    err.Error(),
			},
		}, nil
	}

	return qt.result, nil
}

1059
func (node *ProxyNode) Flush(ctx context.Context, request *milvuspb.FlushRequest) (*commonpb.Status, error) {
1060
	ft := &FlushTask{
T
ThreadDao 已提交
1061 1062 1063 1064
		ctx:          ctx,
		Condition:    NewTaskCondition(ctx),
		FlushRequest: request,
		dataService:  node.dataService,
1065 1066 1067 1068 1069
	}

	err := node.sched.DdQueue.Enqueue(ft)
	if err != nil {
		return &commonpb.Status{
1070
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1071 1072 1073 1074
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090
	log.Debug("Flush",
		zap.String("role", Params.RoleName),
		zap.Int64("msgID", request.Base.MsgID),
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.Any("collections", request.CollectionNames))
	defer func() {
		log.Debug("Flush Done",
			zap.Error(err),
			zap.String("role", Params.RoleName),
			zap.Int64("msgID", request.Base.MsgID),
			zap.Uint64("timestamp", request.Base.Timestamp),
			zap.String("db", request.DbName),
			zap.Any("collections", request.CollectionNames))
	}()

1091 1092 1093
	err = ft.WaitToFinish()
	if err != nil {
		return &commonpb.Status{
1094
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1095 1096 1097 1098 1099
			Reason:    err.Error(),
		}, nil
	}

	return ft.result, nil
1100 1101
}

G
godchen 已提交
1102
func (node *ProxyNode) GetDdChannel(ctx context.Context, request *internalpb.GetDdChannelRequest) (*milvuspb.StringResponse, error) {
1103 1104
	panic("implement me")
}
X
XuanYang-cn 已提交
1105

G
godchen 已提交
1106
func (node *ProxyNode) GetPersistentSegmentInfo(ctx context.Context, req *milvuspb.GetPersistentSegmentInfoRequest) (*milvuspb.GetPersistentSegmentInfoResponse, error) {
D
dragondriver 已提交
1107 1108 1109 1110 1111
	log.Debug("GetPersistentSegmentInfo",
		zap.String("role", Params.RoleName),
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
1112
	resp := &milvuspb.GetPersistentSegmentInfoResponse{
X
XuanYang-cn 已提交
1113
		Status: &commonpb.Status{
1114
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
X
XuanYang-cn 已提交
1115 1116
		},
	}
G
godchen 已提交
1117
	segments, err := node.getSegmentsOfCollection(ctx, req.DbName, req.CollectionName)
X
XuanYang-cn 已提交
1118
	if err != nil {
1119
		resp.Status.Reason = fmt.Errorf("getSegmentsOfCollection, err:%w", err).Error()
X
XuanYang-cn 已提交
1120 1121
		return resp, nil
	}
G
godchen 已提交
1122
	infoResp, err := node.dataService.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{
X
XuanYang-cn 已提交
1123
		Base: &commonpb.MsgBase{
1124
			MsgType:   commonpb.MsgType_SegmentInfo,
X
XuanYang-cn 已提交
1125 1126 1127 1128 1129 1130
			MsgID:     0,
			Timestamp: 0,
			SourceID:  Params.ProxyID,
		},
		SegmentIDs: segments,
	})
1131
	log.Debug("GetPersistentSegmentInfo ", zap.Any("infos", infoResp.Infos), zap.Any("status", infoResp.Status))
X
XuanYang-cn 已提交
1132
	if err != nil {
1133
		resp.Status.Reason = fmt.Errorf("dataService:GetSegmentInfo, err:%w", err).Error()
X
XuanYang-cn 已提交
1134 1135
		return resp, nil
	}
1136
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
X
XuanYang-cn 已提交
1137 1138 1139 1140 1141 1142
		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 已提交
1143
			SegmentID:    info.ID,
X
XuanYang-cn 已提交
1144 1145 1146 1147 1148 1149 1150 1151 1152 1153
			CollectionID: info.CollectionID,
			PartitionID:  info.PartitionID,
			OpenTime:     info.OpenTime,
			SealedTime:   info.SealedTime,
			FlushedTime:  info.FlushedTime,
			NumRows:      info.NumRows,
			MemSize:      info.MemSize,
			State:        info.State,
		}
	}
1154
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
X
XuanYang-cn 已提交
1155 1156 1157 1158
	resp.Infos = persistentInfos
	return resp, nil
}

G
godchen 已提交
1159
func (node *ProxyNode) GetQuerySegmentInfo(ctx context.Context, req *milvuspb.GetQuerySegmentInfoRequest) (*milvuspb.GetQuerySegmentInfoResponse, error) {
D
dragondriver 已提交
1160 1161 1162 1163 1164
	log.Debug("GetQuerySegmentInfo",
		zap.String("role", Params.RoleName),
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
1165
	resp := &milvuspb.GetQuerySegmentInfoResponse{
Z
zhenshan.cao 已提交
1166
		Status: &commonpb.Status{
1167
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
Z
zhenshan.cao 已提交
1168 1169
		},
	}
G
godchen 已提交
1170
	segments, err := node.getSegmentsOfCollection(ctx, req.DbName, req.CollectionName)
Z
zhenshan.cao 已提交
1171 1172 1173 1174
	if err != nil {
		resp.Status.Reason = err.Error()
		return resp, nil
	}
G
godchen 已提交
1175
	infoResp, err := node.queryService.GetSegmentInfo(ctx, &querypb.GetSegmentInfoRequest{
Z
zhenshan.cao 已提交
1176
		Base: &commonpb.MsgBase{
1177
			MsgType:   commonpb.MsgType_SegmentInfo,
Z
zhenshan.cao 已提交
1178 1179 1180 1181 1182 1183
			MsgID:     0,
			Timestamp: 0,
			SourceID:  Params.ProxyID,
		},
		SegmentIDs: segments,
	})
1184
	log.Debug("GetQuerySegmentInfo ", zap.Any("infos", infoResp.Infos), zap.Any("status", infoResp.Status))
Z
zhenshan.cao 已提交
1185 1186 1187 1188
	if err != nil {
		resp.Status.Reason = err.Error()
		return resp, nil
	}
1189
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
Z
zhenshan.cao 已提交
1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204
		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,
		}
	}
1205
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
Z
zhenshan.cao 已提交
1206 1207 1208 1209
	resp.Infos = queryInfos
	return resp, nil
}

1210
func (node *ProxyNode) getSegmentsOfCollection(ctx context.Context, dbName string, collectionName string) ([]UniqueID, error) {
T
ThreadDao 已提交
1211
	describeCollectionResponse, err := node.masterService.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{
X
XuanYang-cn 已提交
1212
		Base: &commonpb.MsgBase{
1213
			MsgType:   commonpb.MsgType_DescribeCollection,
X
XuanYang-cn 已提交
1214 1215 1216 1217 1218 1219 1220 1221 1222 1223
			MsgID:     0,
			Timestamp: 0,
			SourceID:  Params.ProxyID,
		},
		DbName:         dbName,
		CollectionName: collectionName,
	})
	if err != nil {
		return nil, err
	}
1224
	if describeCollectionResponse.Status.ErrorCode != commonpb.ErrorCode_Success {
X
XuanYang-cn 已提交
1225 1226 1227
		return nil, errors.New(describeCollectionResponse.Status.Reason)
	}
	collectionID := describeCollectionResponse.CollectionID
G
godchen 已提交
1228
	showPartitionsResp, err := node.masterService.ShowPartitions(ctx, &milvuspb.ShowPartitionsRequest{
X
XuanYang-cn 已提交
1229
		Base: &commonpb.MsgBase{
1230
			MsgType:   commonpb.MsgType_ShowPartitions,
X
XuanYang-cn 已提交
1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241
			MsgID:     0,
			Timestamp: 0,
			SourceID:  Params.ProxyID,
		},
		DbName:         dbName,
		CollectionName: collectionName,
		CollectionID:   collectionID,
	})
	if err != nil {
		return nil, err
	}
1242
	if showPartitionsResp.Status.ErrorCode != commonpb.ErrorCode_Success {
X
XuanYang-cn 已提交
1243 1244 1245 1246 1247
		return nil, errors.New(showPartitionsResp.Status.Reason)
	}

	ret := make([]UniqueID, 0)
	for _, partitionID := range showPartitionsResp.PartitionIDs {
G
godchen 已提交
1248
		showSegmentResponse, err := node.masterService.ShowSegments(ctx, &milvuspb.ShowSegmentsRequest{
X
XuanYang-cn 已提交
1249
			Base: &commonpb.MsgBase{
1250
				MsgType:   commonpb.MsgType_ShowSegments,
X
XuanYang-cn 已提交
1251 1252 1253 1254 1255 1256 1257 1258 1259 1260
				MsgID:     0,
				Timestamp: 0,
				SourceID:  Params.ProxyID,
			},
			CollectionID: collectionID,
			PartitionID:  partitionID,
		})
		if err != nil {
			return nil, err
		}
1261
		if showSegmentResponse.Status.ErrorCode != commonpb.ErrorCode_Success {
X
XuanYang-cn 已提交
1262 1263 1264 1265 1266 1267
			return nil, errors.New(showSegmentResponse.Status.Reason)
		}
		ret = append(ret, showSegmentResponse.SegmentIDs...)
	}
	return ret, nil
}
1268

G
godchen 已提交
1269 1270
func (node *ProxyNode) RegisterLink(ctx context.Context, req *milvuspb.RegisterLinkRequest) (*milvuspb.RegisterLinkResponse, error) {
	code := node.stateCode.Load().(internalpb.StateCode)
D
dragondriver 已提交
1271 1272 1273 1274
	log.Debug("RegisterLink",
		zap.String("role", Params.RoleName),
		zap.Any("state code of proxynode", code))

G
godchen 已提交
1275
	if code != internalpb.StateCode_Healthy {
1276 1277 1278
		return &milvuspb.RegisterLinkResponse{
			Address: nil,
			Status: &commonpb.Status{
1279
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1280 1281 1282 1283 1284 1285 1286
				Reason:    "proxy node not healthy",
			},
		}, nil
	}
	return &milvuspb.RegisterLinkResponse{
		Address: nil,
		Status: &commonpb.Status{
1287
			ErrorCode: commonpb.ErrorCode_Success,
C
cai.zhang 已提交
1288
			Reason:    os.Getenv("DEPLOY_MODE"),
1289 1290 1291
		},
	}, nil
}