Initial commit (code only without large binaries)
This commit is contained in:
robin
282
EdgeDNS/internal/rpc/rpc_client.go
Normal file
282
EdgeDNS/internal/rpc/rpc_client.go
Normal file
@@ -0,0 +1,282 @@
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
|
||||
"github.com/TeaOSLab/EdgeDNS/internal/configs"
|
||||
teaconst "github.com/TeaOSLab/EdgeDNS/internal/const"
|
||||
"github.com/TeaOSLab/EdgeDNS/internal/encrypt"
|
||||
"github.com/TeaOSLab/EdgeDNS/internal/utils"
|
||||
"github.com/iwind/TeaGo/maps"
|
||||
"github.com/iwind/TeaGo/rands"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/connectivity"
|
||||
"google.golang.org/grpc/credentials"
|
||||
"google.golang.org/grpc/credentials/insecure"
|
||||
"google.golang.org/grpc/encoding/gzip"
|
||||
"google.golang.org/grpc/keepalive"
|
||||
"google.golang.org/grpc/metadata"
|
||||
"net/url"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// RPCClient RPC客户端
|
||||
type RPCClient struct {
|
||||
apiConfig *configs.APIConfig
|
||||
conns []*grpc.ClientConn
|
||||
|
||||
locker sync.RWMutex
|
||||
|
||||
APINodeRPC pb.APINodeServiceClient
|
||||
NodeLogRPC pb.NodeLogServiceClient
|
||||
NSNodeRPC pb.NSNodeServiceClient
|
||||
NSRecordRPC pb.NSRecordServiceClient
|
||||
NSDomainRPC pb.NSDomainServiceClient
|
||||
NSRouteRPC pb.NSRouteServiceClient
|
||||
NSKeyRPC pb.NSKeyServiceClient
|
||||
NSAccessLogRPC pb.NSAccessLogServiceClient
|
||||
NSRecordHourlyStatRPC pb.NSRecordHourlyStatServiceClient
|
||||
NSQuestionOptionRPC pb.NSQuestionOptionServiceClient
|
||||
NodeValueRPC pb.NodeValueServiceClient
|
||||
NodeTaskRPC pb.NodeTaskServiceClient
|
||||
RegionCountryRPC pb.RegionCountryServiceClient
|
||||
RegionProvinceRPC pb.RegionProvinceServiceClient
|
||||
RegionCityRPC pb.RegionCityServiceClient
|
||||
RegionProviderRPC pb.RegionProviderServiceClient
|
||||
IPLibraryArtifactRPC pb.IPLibraryArtifactServiceClient
|
||||
FileChunkRPC pb.FileChunkServiceClient
|
||||
ClientAgentIPRPC pb.ClientAgentIPServiceClient
|
||||
}
|
||||
|
||||
// NewRPCClient 构造新的RPC客户端
|
||||
func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) {
|
||||
if apiConfig == nil {
|
||||
return nil, errors.New("api config should not be nil")
|
||||
}
|
||||
|
||||
var client = &RPCClient{
|
||||
apiConfig: apiConfig,
|
||||
}
|
||||
|
||||
// 初始化RPC实例
|
||||
client.APINodeRPC = pb.NewAPINodeServiceClient(client)
|
||||
client.NodeLogRPC = pb.NewNodeLogServiceClient(client)
|
||||
client.NSNodeRPC = pb.NewNSNodeServiceClient(client)
|
||||
client.NSRecordRPC = pb.NewNSRecordServiceClient(client)
|
||||
client.NSDomainRPC = pb.NewNSDomainServiceClient(client)
|
||||
client.NSRouteRPC = pb.NewNSRouteServiceClient(client)
|
||||
client.NSKeyRPC = pb.NewNSKeyServiceClient(client)
|
||||
client.NSAccessLogRPC = pb.NewNSAccessLogServiceClient(client)
|
||||
client.NSRecordHourlyStatRPC = pb.NewNSRecordHourlyStatServiceClient(client)
|
||||
client.NSQuestionOptionRPC = pb.NewNSQuestionOptionServiceClient(client)
|
||||
client.NodeValueRPC = pb.NewNodeValueServiceClient(client)
|
||||
client.NodeTaskRPC = pb.NewNodeTaskServiceClient(client)
|
||||
client.RegionCountryRPC = pb.NewRegionCountryServiceClient(client)
|
||||
client.RegionProvinceRPC = pb.NewRegionProvinceServiceClient(client)
|
||||
client.RegionCityRPC = pb.NewRegionCityServiceClient(client)
|
||||
client.RegionProviderRPC = pb.NewRegionProviderServiceClient(client)
|
||||
client.IPLibraryArtifactRPC = pb.NewIPLibraryArtifactServiceClient(client)
|
||||
client.FileChunkRPC = pb.NewFileChunkServiceClient(client)
|
||||
client.ClientAgentIPRPC = pb.NewClientAgentIPServiceClient(client)
|
||||
|
||||
err := client.init()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return client, nil
|
||||
}
|
||||
|
||||
// Context 构造用户上下文
|
||||
func (this *RPCClient) Context() context.Context {
|
||||
ctx := context.Background()
|
||||
m := maps.Map{
|
||||
"timestamp": time.Now().Unix(),
|
||||
"type": "dns",
|
||||
"userId": 0,
|
||||
}
|
||||
method, err := encrypt.NewMethodInstance(teaconst.EncryptMethod, this.apiConfig.Secret, this.apiConfig.NodeId)
|
||||
if err != nil {
|
||||
utils.PrintError(err)
|
||||
return context.Background()
|
||||
}
|
||||
data, err := method.Encrypt(m.AsJSON())
|
||||
if err != nil {
|
||||
utils.PrintError(err)
|
||||
return context.Background()
|
||||
}
|
||||
token := base64.StdEncoding.EncodeToString(data)
|
||||
ctx = metadata.AppendToOutgoingContext(ctx, "nodeId", this.apiConfig.NodeId, "token", token)
|
||||
return ctx
|
||||
}
|
||||
|
||||
// UpdateConfig 修改配置
|
||||
func (this *RPCClient) UpdateConfig(config *configs.APIConfig) error {
|
||||
this.apiConfig = config
|
||||
this.locker.Lock()
|
||||
err := this.init()
|
||||
this.locker.Unlock()
|
||||
return err
|
||||
}
|
||||
|
||||
// TestEndpoints 测试Endpoints是否可用
|
||||
func (this *RPCClient) TestEndpoints(endpoints []string) bool {
|
||||
if len(endpoints) == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
var wg = sync.WaitGroup{}
|
||||
wg.Add(len(endpoints))
|
||||
|
||||
var ok = false
|
||||
|
||||
for _, endpoint := range endpoints {
|
||||
go func(endpoint string) {
|
||||
defer wg.Done()
|
||||
|
||||
u, err := url.Parse(endpoint)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancelFunc := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer func() {
|
||||
cancelFunc()
|
||||
}()
|
||||
var conn *grpc.ClientConn
|
||||
if u.Scheme == "http" {
|
||||
conn, err = grpc.DialContext(ctx, u.Host, grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithBlock())
|
||||
} else if u.Scheme == "https" {
|
||||
conn, err = grpc.DialContext(ctx, u.Host, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{
|
||||
InsecureSkipVerify: true,
|
||||
})), grpc.WithBlock())
|
||||
} else {
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if conn == nil {
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
_ = conn.Close()
|
||||
}()
|
||||
|
||||
var pingService = pb.NewPingServiceClient(conn)
|
||||
_, err = pingService.Ping(this.Context(), &pb.PingRequest{})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
ok = true
|
||||
}(endpoint)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
return ok
|
||||
}
|
||||
|
||||
// 初始化
|
||||
func (this *RPCClient) init() error {
|
||||
// 重新连接
|
||||
conns := []*grpc.ClientConn{}
|
||||
for _, endpoint := range this.apiConfig.RPCEndpoints {
|
||||
u, err := url.Parse(endpoint)
|
||||
if err != nil {
|
||||
return fmt.Errorf("parse endpoint failed: %w", err)
|
||||
}
|
||||
var conn *grpc.ClientConn
|
||||
var callOptions = grpc.WithDefaultCallOptions(
|
||||
grpc.MaxCallRecvMsgSize(128<<20),
|
||||
grpc.MaxCallSendMsgSize(128<<20),
|
||||
grpc.UseCompressor(gzip.Name))
|
||||
var keepaliveParams = grpc.WithKeepaliveParams(keepalive.ClientParameters{
|
||||
Time: 30 * time.Second,
|
||||
})
|
||||
if u.Scheme == "http" {
|
||||
conn, err = grpc.Dial(u.Host, grpc.WithTransportCredentials(insecure.NewCredentials()), callOptions, keepaliveParams)
|
||||
} else if u.Scheme == "https" {
|
||||
conn, err = grpc.Dial(u.Host, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{
|
||||
InsecureSkipVerify: true,
|
||||
})), callOptions, keepaliveParams)
|
||||
} else {
|
||||
return errors.New("parse endpoint failed: invalid scheme '" + u.Scheme + "'")
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
conns = append(conns, conn)
|
||||
}
|
||||
if len(conns) == 0 {
|
||||
return errors.New("[RPC]no available endpoints")
|
||||
}
|
||||
|
||||
// 这里不需要加锁，因为会和pickConn冲突
|
||||
this.conns = conns
|
||||
return nil
|
||||
}
|
||||
|
||||
// 随机选择一个连接
|
||||
func (this *RPCClient) pickConn() *grpc.ClientConn {
|
||||
this.locker.Lock()
|
||||
defer this.locker.Unlock()
|
||||
|
||||
// 检查连接状态
|
||||
var countConns = len(this.conns)
|
||||
if countConns > 0 {
|
||||
if countConns == 1 {
|
||||
return this.conns[0]
|
||||
}
|
||||
for _, state := range []connectivity.State{
|
||||
connectivity.Ready,
|
||||
connectivity.Idle,
|
||||
connectivity.Connecting,
|
||||
connectivity.TransientFailure,
|
||||
} {
|
||||
var availableConns = []*grpc.ClientConn{}
|
||||
for _, conn := range this.conns {
|
||||
if conn.GetState() == state {
|
||||
availableConns = append(availableConns, conn)
|
||||
}
|
||||
}
|
||||
if len(availableConns) > 0 {
|
||||
return this.randConn(availableConns)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return this.randConn(this.conns)
|
||||
}
|
||||
|
||||
func (this *RPCClient) Invoke(ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption) error {
|
||||
var conn = this.pickConn()
|
||||
if conn == nil {
|
||||
return errors.New("can not get available grpc connection")
|
||||
}
|
||||
return conn.Invoke(ctx, method, args, reply, opts...)
|
||||
}
|
||||
|
||||
func (this *RPCClient) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
|
||||
var conn = this.pickConn()
|
||||
if conn == nil {
|
||||
return nil, errors.New("can not get available grpc connection")
|
||||
}
|
||||
return conn.NewStream(ctx, desc, method, opts...)
|
||||
}
|
||||
|
||||
func (this *RPCClient) randConn(conns []*grpc.ClientConn) *grpc.ClientConn {
|
||||
var l = len(conns)
|
||||
if l == 0 {
|
||||
return nil
|
||||
}
|
||||
if l == 1 {
|
||||
return conns[0]
|
||||
}
|
||||
return conns[rands.Int(0, l-1)]
|
||||
}
|
||||
Reference in New Issue
Block a user