// 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 // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // 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. package types import ( "context" "io" "github.com/tikv/client-go/v2/txnkv" clientv3 "go.etcd.io/etcd/client/v3" "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/indexpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/proxypb" "github.com/milvus-io/milvus/internal/proto/querypb" "github.com/milvus-io/milvus/internal/proto/rootcoordpb" ) // Limiter defines the interface to perform request rate limiting. // If Limit function return true, the request will be rejected. // Otherwise, the request will pass. Limit also returns limit of limiter. type Limiter interface { Check(collectionIDs []int64, rt internalpb.RateType, n int) error } // Component is the interface all services implement type Component interface { Init() error Start() error Stop() error Register() error } // DataNodeClient is the client interface for datanode server type DataNodeClient interface { io.Closer datapb.DataNodeClient } // DataNode is the interface `datanode` package implements type DataNode interface { Component datapb.DataNodeServer } // DataNodeComponent is used by grpc server of DataNode // //go:generate mockery --name=DataNodeComponent --structname=MockDataNode --output=../mocks --filename=mock_datanode.go --with-expecter 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) // SetRootCoordClient set SetRootCoordClient for DataNode // `rootCoord` is a client of root coordinator. // // Return a generic error in status: // If the rootCoord is nil or the rootCoord has been set before. // Return nil in status: // The rootCoord is not nil. SetRootCoordClient(rootCoord RootCoordClient) error // SetDataCoordClient set DataCoord for DataNode // `dataCoord` is a client of data coordinator. // // Return a generic error in status: // If the dataCoord is nil or the dataCoord has been set before. // Return nil in status: // The dataCoord is not nil. SetDataCoordClient(dataCoord DataCoordClient) error } // DataCoordClient is the client interface for datacoord server type DataCoordClient interface { io.Closer datapb.DataCoordClient } // DataCoord is the interface `datacoord` package implements type DataCoord interface { Component datapb.DataCoordServer } // DataCoordComponent defines the interface of DataCoord component. // //go:generate mockery --name=DataCoordComponent --structname=MockDataCoord --output=../mocks --filename=mock_datacoord.go --with-expecter 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) SetRootCoordClient(rootCoord RootCoordClient) // SetDataNodeCreator set DataNode client creator func for DataCoord SetDataNodeCreator(func(context.Context, string, int64) (DataNodeClient, error)) // SetIndexNodeCreator set Index client creator func for DataCoord SetIndexNodeCreator(func(context.Context, string, int64) (IndexNodeClient, error)) } // IndexNodeClient is the client interface for indexnode server type IndexNodeClient interface { io.Closer indexpb.IndexNodeClient } // IndexNode is the interface `indexnode` package implements type IndexNode interface { Component indexpb.IndexNodeServer } // IndexNodeComponent is used by grpc server of IndexNode type IndexNodeComponent interface { IndexNode SetAddress(address string) GetAddress() string // SetEtcdClient set etcd client for IndexNodeComponent SetEtcdClient(etcdClient *clientv3.Client) // UpdateStateCode updates state code for IndexNodeComponent // `stateCode` is current statement of this QueryCoord, indicating whether it's healthy. UpdateStateCode(stateCode commonpb.StateCode) } // RootCoordClient is the client interface for rootcoord server type RootCoordClient interface { io.Closer rootcoordpb.RootCoordClient } // RootCoord is the interface `rootcoord` package implements // //go:generate mockery --name=RootCoord --output=../mocks --filename=mock_rootcoord.go --with-expecter type RootCoord interface { Component rootcoordpb.RootCoordServer } // RootCoordComponent is used by grpc server of RootCoord 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) // SetDataCoordClient set SetDataCoordClient for RootCoord // `dataCoord` is a client of data coordinator. // // Always return nil. SetDataCoordClient(dataCoord DataCoordClient) error // SetQueryCoord set QueryCoord for RootCoord // `queryCoord` is a client of query coordinator. // // Always return nil. SetQueryCoordClient(queryCoord QueryCoordClient) 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) } // ProxyClient is the client interface for proxy server type ProxyClient interface { io.Closer proxypb.ProxyClient } // Proxy is the interface `proxy` package implements type Proxy interface { Component proxypb.ProxyServer milvuspb.MilvusServiceServer } // ProxyComponent defines the interface of proxy component. // //go:generate mockery --name=ProxyComponent --structname=MockProxy --output=../mocks --filename=mock_proxy.go --with-expecter type ProxyComponent interface { Proxy SetAddress(address string) GetAddress() string // SetEtcdClient set EtcdClient for Proxy // `etcdClient` is a client of etcd SetEtcdClient(etcdClient *clientv3.Client) // SetRootCoordClient set RootCoord for Proxy // `rootCoord` is a client of root coordinator. SetRootCoordClient(rootCoord RootCoordClient) // SetDataCoordClient set DataCoord for Proxy // `dataCoord` is a client of data coordinator. SetDataCoordClient(dataCoord DataCoordClient) // SetIndexCoordClient set IndexCoord for Proxy // `indexCoord` is a client of index coordinator. // SetIndexCoordClient(indexCoord IndexCoord) // SetQueryCoordClient set QueryCoord for Proxy // `queryCoord` is a client of query coordinator. SetQueryCoordClient(queryCoord QueryCoordClient) // 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 QueryNodeClient interface { io.Closer querypb.QueryNodeClient } // QueryNode is the interface `querynode` package implements type QueryNode interface { Component querypb.QueryNodeServer } // QueryNodeComponent is used by grpc server of QueryNode // //go:generate mockery --name=QueryNodeComponent --structname=MockQueryNode --output=../mocks --filename=mock_querynode.go --with-expecter 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) } // QueryCoordClient is the client interface for querycoord server type QueryCoordClient interface { io.Closer querypb.QueryCoordClient } // QueryCoord is the interface `querycoord` package implements type QueryCoord interface { Component querypb.QueryCoordServer } // QueryCoordComponent is used by grpc server of QueryCoord // //go:generate mockery --name=QueryCoordComponent --structname=MockQueryCoord --output=../mocks --filename=mock_querycoord.go --with-expecter 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) // SetDataCoordClient set SetDataCoordClient for QueryCoord // `dataCoord` is a client of data coordinator. // // Return a generic error in status: // If the dataCoord is nil. // Return nil in status: // The dataCoord is not nil. SetDataCoordClient(dataCoord DataCoordClient) error // SetRootCoordClient set SetRootCoordClient for QueryCoord // `rootCoord` is a client of root coordinator. // // Return a generic error in status: // If the rootCoord is nil. // Return nil in status: // The rootCoord is not nil. SetRootCoordClient(rootCoord RootCoordClient) error // SetQueryNodeCreator set QueryNode client creator func for QueryCoord SetQueryNodeCreator(func(ctx context.Context, addr string, nodeID int64) (QueryNodeClient, error)) }