milvus/internal/master/grpc_service.go

package master

import (
	"context"
	"time"

	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
	"github.com/zilliztech/milvus-distributed/internal/proto/internalpb"
	"github.com/zilliztech/milvus-distributed/internal/proto/servicepb"
)

const slowThreshold = 5 * time.Millisecond

func (s *Master) CreateCollection(ctx context.Context, in *internalpb.CreateCollectionRequest) (*commonpb.Status, error) {
	var t task = &createCollectionTask{
		req: in,
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
	}

	response := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
	}

	var err = s.scheduler.Enqueue(t)
	if err != nil {
		response.Reason = "Enqueue failed: " + err.Error()
		return response, nil
	}

	err = t.WaitToFinish(ctx)
	if err != nil {
		response.Reason = "Create collection failed: " + err.Error()
		return response, nil
	}

	response.ErrorCode = commonpb.ErrorCode_SUCCESS
	return response, nil
}

func (s *Master) DropCollection(ctx context.Context, in *internalpb.DropCollectionRequest) (*commonpb.Status, error) {
	var t task = &dropCollectionTask{
		req: in,
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
	}

	response := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
	}

	var err = s.scheduler.Enqueue(t)
	if err != nil {
		response.Reason = "Enqueue failed: " + err.Error()
		return response, nil
	}

	err = t.WaitToFinish(ctx)
	if err != nil {
		response.Reason = "Drop collection failed: " + err.Error()
		return response, nil
	}

	response.ErrorCode = commonpb.ErrorCode_SUCCESS
	return response, nil
}

func (s *Master) HasCollection(ctx context.Context, in *internalpb.HasCollectionRequest) (*servicepb.BoolResponse, error) {
	var t task = &hasCollectionTask{
		req: in,
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
		hasCollection: false,
	}

	st := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
	}

	response := &servicepb.BoolResponse{
		Status: st,
		Value:  false,
	}

	var err = s.scheduler.Enqueue(t)
	if err != nil {
		st.Reason = "Enqueue failed: " + err.Error()
		return response, nil
	}

	err = t.WaitToFinish(ctx)
	if err != nil {
		st.Reason = "Has collection failed: " + err.Error()
		return response, nil
	}

	st.ErrorCode = commonpb.ErrorCode_SUCCESS
	response.Value = t.(*hasCollectionTask).hasCollection
	return response, nil
}

func (s *Master) DescribeCollection(ctx context.Context, in *internalpb.DescribeCollectionRequest) (*servicepb.CollectionDescription, error) {
	var t task = &describeCollectionTask{
		req: in,
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
		description: nil,
	}

	response := &servicepb.CollectionDescription{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
		},
		Schema: nil,
	}

	t.(*describeCollectionTask).description = response

	var err = s.scheduler.Enqueue(t)
	if err != nil {
		response.Status.Reason = "Enqueue failed: " + err.Error()
		return response, nil
	}

	err = t.WaitToFinish(ctx)
	if err != nil {
		response.Status.Reason = "Describe collection failed: " + err.Error()
		return response, nil
	}

	response.Status.ErrorCode = commonpb.ErrorCode_SUCCESS
	return response, nil
}

func (s *Master) ShowCollections(ctx context.Context, in *internalpb.ShowCollectionRequest) (*servicepb.StringListResponse, error) {
	var t task = &showCollectionsTask{
		req: in,
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
		stringListResponse: nil,
	}

	response := &servicepb.StringListResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
			Reason:    "",
		},
		Values: nil,
	}

	t.(*showCollectionsTask).stringListResponse = response

	var err = s.scheduler.Enqueue(t)
	if err != nil {
		response.Status.Reason = "Enqueue filed: " + err.Error()
		return response, nil
	}

	err = t.WaitToFinish(ctx)
	if err != nil {
		response.Status.Reason = "Show Collections failed: " + err.Error()
		return response, nil
	}

	response.Status.ErrorCode = commonpb.ErrorCode_SUCCESS
	return response, nil
}

//////////////////////////////////////////////////////////////////////////
func (s *Master) CreatePartition(ctx context.Context, in *internalpb.CreatePartitionRequest) (*commonpb.Status, error) {
	var t task = &createPartitionTask{
		req: in,
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
	}

	var err = s.scheduler.Enqueue(t)
	if err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
			Reason:    "Enqueue failed",
		}, nil
	}

	err = t.WaitToFinish(ctx)
	if err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
			Reason:    "WaitToFinish failed",
		}, nil
	}

	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_SUCCESS,
	}, nil
}

func (s *Master) DropPartition(ctx context.Context, in *internalpb.DropPartitionRequest) (*commonpb.Status, error) {
	var t task = &dropPartitionTask{
		req: in,
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
	}

	var err = s.scheduler.Enqueue(t)
	if err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
			Reason:    "Enqueue failed",
		}, nil
	}

	err = t.WaitToFinish(ctx)
	if err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
			Reason:    "WaitToFinish failed",
		}, nil
	}

	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_SUCCESS,
	}, nil
}

func (s *Master) HasPartition(ctx context.Context, in *internalpb.HasPartitionRequest) (*servicepb.BoolResponse, error) {
	var t task = &hasPartitionTask{
		req: in,
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
		hasPartition: false,
	}

	var err = s.scheduler.Enqueue(t)
	if err != nil {
		return &servicepb.BoolResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
				Reason:    "Enqueue failed",
			},
			Value: t.(*hasPartitionTask).hasPartition,
		}, nil
	}

	err = t.WaitToFinish(ctx)
	if err != nil {
		return &servicepb.BoolResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
				Reason:    err.Error(),
			},
			Value: t.(*hasPartitionTask).hasPartition,
		}, nil
	}

	return &servicepb.BoolResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_SUCCESS,
		},
		Value: t.(*hasPartitionTask).hasPartition,
	}, nil
}

func (s *Master) DescribePartition(ctx context.Context, in *internalpb.DescribePartitionRequest) (*servicepb.PartitionDescription, error) {
	var t task = &describePartitionTask{
		req: in,
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
		description: nil,
	}

	var err = s.scheduler.Enqueue(t)
	if err != nil {
		return &servicepb.PartitionDescription{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
				Reason:    "Enqueue failed",
			},
			Name:       in.PartitionName,
			Statistics: nil,
		}, nil
	}

	err = t.WaitToFinish(ctx)
	if err != nil {
		return &servicepb.PartitionDescription{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
				Reason:    "WaitToFinish failed",
			},
			Name:       in.PartitionName,
			Statistics: nil,
		}, nil
	}

	return t.(*describePartitionTask).description, nil
}

func (s *Master) ShowPartitions(ctx context.Context, in *internalpb.ShowPartitionRequest) (*servicepb.StringListResponse, error) {
	var t task = &showPartitionTask{
		req: in,
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
		stringListResponse: nil,
	}

	var err = s.scheduler.Enqueue(t)
	if err != nil {
		return &servicepb.StringListResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
				Reason:    "Enqueue failed",
			},
			Values: nil,
		}, nil
	}

	err = t.WaitToFinish(ctx)
	if err != nil {
		return &servicepb.StringListResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
				Reason:    "WaitToFinish failed",
			},
			Values: nil,
		}, nil
	}

	return t.(*showPartitionTask).stringListResponse, nil
}

func (s *Master) GetSysConfigs(ctx context.Context, in *internalpb.SysConfigRequest) (*servicepb.SysConfigResponse, error) {
	var t task = &getSysConfigsTask{
		req:      in,
		configkv: s.kvBase,
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
		keys:   []string{},
		values: []string{},
	}

	response := &servicepb.SysConfigResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
		},
	}

	var err = s.scheduler.Enqueue(t)
	if err != nil {
		response.Status.Reason = "Enqueue failed: " + err.Error()
		return response, nil
	}

	err = t.WaitToFinish(ctx)
	if err != nil {
		response.Status.Reason = "Get System Config failed: " + err.Error()
		return response, nil
	}

	response.Keys = t.(*getSysConfigsTask).keys
	response.Values = t.(*getSysConfigsTask).values
	response.Status.ErrorCode = commonpb.ErrorCode_SUCCESS
	return response, nil
}

//----------------------------------------Internal GRPC Service--------------------------------

func (s *Master) AllocTimestamp(ctx context.Context, request *internalpb.TsoRequest) (*internalpb.TsoResponse, error) {
	count := request.GetCount()
	ts, err := s.tsoAllocator.Alloc(count)

	if err != nil {
		return &internalpb.TsoResponse{
			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR},
		}, nil
	}

	response := &internalpb.TsoResponse{
		Status:    &commonpb.Status{ErrorCode: commonpb.ErrorCode_SUCCESS},
		Timestamp: ts,
		Count:     count,
	}

	return response, nil
}

func (s *Master) AllocID(ctx context.Context, request *internalpb.IDRequest) (*internalpb.IDResponse, error) {
	count := request.GetCount()
	ts, err := s.idAllocator.AllocOne()

	if err != nil {
		return &internalpb.IDResponse{
			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR},
		}, nil
	}

	response := &internalpb.IDResponse{
		Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_SUCCESS},
		ID:     ts,
		Count:  count,
	}

	return response, nil
}

func (s *Master) AssignSegmentID(ctx context.Context, request *internalpb.AssignSegIDRequest) (*internalpb.AssignSegIDResponse, error) {
	segInfos, _ := s.segmentManager.AssignSegment(request.GetPerChannelReq())
	return &internalpb.AssignSegIDResponse{
		PerChannelAssignment: segInfos,
	}, nil
}

func (s *Master) CreateIndex(ctx context.Context, req *internalpb.CreateIndexRequest) (*commonpb.Status, error) {
	task := &createIndexTask{
		baseTask: baseTask{
			sch: s.scheduler,
			mt:  s.metaTable,
			cv:  make(chan error),
		},
		req:                 req,
		indexBuildScheduler: s.indexBuildSch,
		indexLoadScheduler:  s.indexLoadSch,
		segManager:          s.segmentManager,
	}

	err := s.scheduler.Enqueue(task)
	if err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
			Reason:    "Enqueue failed: " + err.Error(),
		}, nil
	}

	err = task.WaitToFinish(ctx)
	if err != nil {
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR,
			Reason:    "Create Index error: " + err.Error(),
		}, nil
	}

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

func (s *Master) DescribeIndex(context.Context, *internalpb.DescribeIndexRequest) (*servicepb.DescribeIndexResponse, error) {
	return nil, nil
}

func (s *Master) DescribeIndexProgress(context.Context, *internalpb.DescribeIndexProgressRequest) (*servicepb.BoolResponse, error) {
	return nil, nil
}