types

type Component interface {
	Init() error
	Start() error
	Stop() error
	Register() error
}

type DataCoord interface {
	Component
	datapb.DataCoordServer
}

type DataCoordClient interface {
	io.Closer
	datapb.DataCoordClient
	indexpb.IndexCoordClient
}

type DataCoordComponent interface {
	DataCoord

	SetAddress(address string)
	// SetEtcdClient set EtcdClient for DataCoord
	// `etcdClient` is a client of etcd
	SetEtcdClient(etcdClient *clientv3.Client)

	// SetTiKVClient set TiKV client for QueryNode
	SetTiKVClient(client *txnkv.Client)

	SetMixCoord(mixCoord MixCoord)

	// SetDataNodeCreator set DataNode client creator func for DataCoord
	SetDataNodeCreator(func(context.Context, string, int64) (DataNodeClient, error))
}

type DataNode interface {
	Component
	datapb.DataNodeServer
	workerpb.IndexNodeServer
}

type DataNodeClient interface {
	io.Closer
	datapb.DataNodeClient
	workerpb.IndexNodeClient
}

type DataNodeComponent interface {
	DataNode

	// UpdateStateCode updates state code for DataNode
	//  `stateCode` is current statement of this data node, indicating whether it's healthy.
	UpdateStateCode(stateCode commonpb.StateCode)

	// GetStateCode return state code of this data node
	GetStateCode() commonpb.StateCode

	SetAddress(address string)
	GetAddress() string
	GetNodeID() int64

	// SetEtcdClient set etcd client for DataNode
	SetEtcdClient(etcdClient *clientv3.Client)
}

type Limiter interface {
	Check(dbID int64, collectionIDToPartIDs map[int64][]int64, rt internalpb.RateType, n int) error
	Alloc(ctx context.Context, dbID int64, collectionIDToPartIDs map[int64][]int64, rt internalpb.RateType, n int) error
}

type MixCoord interface {
	Component
	rootcoordpb.RootCoordServer
	querypb.QueryCoordServer
	datapb.DataCoordServer

	// GetMetrics notifies MixCoordComponent to collect metrics for specified component
	GetDcMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)

	// GetMetrics notifies MixCoordComponent to collect metrics for specified component
	GetQcMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)
	// GetMetrics notifies MixCoordComponent to collect metrics for specified component
	NotifyDropPartition(ctx context.Context, channel string, partitionIDs []int64) error
}

type MixCoordClient interface {
	io.Closer
	rootcoordpb.RootCoordClient
	querypb.QueryCoordClient
	datapb.DataCoordClient
	indexpb.IndexCoordClient
}

type MixCoordComponent interface {
	MixCoord

	SetAddress(address string)
	// SetEtcdClient set EtcdClient for MixCoord
	// `etcdClient` is a client of etcd
	SetEtcdClient(etcdClient *clientv3.Client)

	// SetTiKVClient set TiKV client for MixCoord
	SetTiKVClient(client *txnkv.Client)

	// UpdateStateCode updates state code for MixCoord
	// State includes: Initializing, Healthy and Abnormal
	UpdateStateCode(commonpb.StateCode)

	// GetMetrics notifies MixCoordComponent to collect metrics for specified component
	GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)

	RegisterStreamingCoordGRPCService(server *grpc.Server)

	GracefulStop()

	SetMixCoordClient(client MixCoordClient)
}

type Proxy interface {
	Component
	proxypb.ProxyServer
	milvuspb.MilvusServiceServer

	ImportV2(context.Context, *internalpb.ImportRequest) (*internalpb.ImportResponse, error)
	GetImportProgress(context.Context, *internalpb.GetImportProgressRequest) (*internalpb.GetImportProgressResponse, error)
	ListImports(context.Context, *internalpb.ListImportsRequest) (*internalpb.ListImportsResponse, error)
}

type ProxyClient interface {
	io.Closer
	proxypb.ProxyClient
}

type ProxyComponent interface {
	Proxy

	SetAddress(address string)
	GetAddress() string

	// SetMixCoordClient set MixCoord for Proxy
	// `mixCoord` is a client of mix coordinator.
	SetMixCoordClient(rootCoord MixCoordClient)

	// SetQueryNodeCreator set QueryNode client creator func for Proxy
	SetQueryNodeCreator(func(ctx context.Context, addr string, nodeID int64) (QueryNodeClient, error))

	// GetRateLimiter returns the rateLimiter in Proxy
	GetRateLimiter() (Limiter, error)

	// UpdateStateCode updates state code for Proxy
	//  `stateCode` is current statement of this proxy node, indicating whether it's healthy.
	UpdateStateCode(stateCode commonpb.StateCode)
}

type QueryCoord interface {
	Component
	querypb.QueryCoordServer
}

type QueryCoordClient interface {
	io.Closer
	querypb.QueryCoordClient
}

type QueryCoordComponent interface {
	QueryCoord

	SetAddress(address string)

	// SetEtcdClient set etcd client for QueryCoord
	SetEtcdClient(etcdClient *clientv3.Client)

	// SetTiKVClient set TiKV client for QueryCoord
	SetTiKVClient(client *txnkv.Client)

	// UpdateStateCode updates state code for QueryCoord
	//  `stateCode` is current statement of this QueryCoord, indicating whether it's healthy.
	UpdateStateCode(stateCode commonpb.StateCode)

	// SetQueryNodeCreator set QueryNode client creator func for QueryCoord
	SetQueryNodeCreator(func(ctx context.Context, addr string, nodeID int64) (QueryNodeClient, error))

	SetMixCoord(mixCoord MixCoord)
}

type QueryNode interface {
	Component
	querypb.QueryNodeServer
}

type QueryNodeClient interface {
	io.Closer
	querypb.QueryNodeClient
}

type QueryNodeComponent interface {
	QueryNode

	// UpdateStateCode updates state code for QueryNode
	//  `stateCode` is current statement of this query node, indicating whether it's healthy.
	UpdateStateCode(stateCode commonpb.StateCode)

	SetAddress(address string)
	GetAddress() string
	GetNodeID() int64

	// SetEtcdClient set etcd client for QueryNode
	SetEtcdClient(etcdClient *clientv3.Client)
}

type RootCoord interface {
	Component
	rootcoordpb.RootCoordServer
}

type RootCoordClient interface {
	io.Closer
	rootcoordpb.RootCoordClient
}

type RootCoordComponent interface {
	RootCoord

	SetAddress(address string)
	// SetEtcdClient set EtcdClient for RootCoord
	// `etcdClient` is a client of etcd
	SetEtcdClient(etcdClient *clientv3.Client)

	// SetTiKVClient set TiKV client for RootCoord
	SetTiKVClient(client *txnkv.Client)

	// UpdateStateCode updates state code for RootCoord
	// State includes: Initializing, Healthy and Abnormal
	UpdateStateCode(commonpb.StateCode)

	// SetMixCoord set SetMixCoord for RootCoord
	// `dataCoord` is a client of data coordinator.
	//
	// Always return nil.
	SetMixCoord(mixCoord MixCoord) error

	// SetProxyCreator set Proxy client creator func for RootCoord
	SetProxyCreator(func(ctx context.Context, addr string, nodeID int64) (ProxyClient, error))

	// GetMetrics notifies RootCoordComponent to collect metrics for specified component
	GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)
}