seaweedfs/weed/pb/grpc_client_server.go
2020-03-04 00:39:47 -08:00

183 lines
5.1 KiB
Go

package pb
import (
"context"
"fmt"
"net/http"
"strconv"
"strings"
"sync"
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/keepalive"
"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
)
var (
// cache grpc connections
grpcClients = make(map[string]*grpc.ClientConn)
grpcClientsLock sync.Mutex
)
func init() {
http.DefaultTransport.(*http.Transport).MaxIdleConnsPerHost = 1024
}
func NewGrpcServer(opts ...grpc.ServerOption) *grpc.Server {
var options []grpc.ServerOption
options = append(options, grpc.KeepaliveParams(keepalive.ServerParameters{
Time: 10 * time.Second, // wait time before ping if no activity
Timeout: 20 * time.Second, // ping timeout
}), grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
MinTime: 60 * time.Second, // min time a client should wait before sending a ping
PermitWithoutStream: true,
}))
for _, opt := range opts {
if opt != nil {
options = append(options, opt)
}
}
return grpc.NewServer(options...)
}
func GrpcDial(ctx context.Context, address string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
// opts = append(opts, grpc.WithBlock())
// opts = append(opts, grpc.WithTimeout(time.Duration(5*time.Second)))
var options []grpc.DialOption
options = append(options,
// grpc.WithInsecure(),
grpc.WithKeepaliveParams(keepalive.ClientParameters{
Time: 30 * time.Second, // client ping server if no activity for this long
Timeout: 20 * time.Second,
PermitWithoutStream: true,
}))
for _, opt := range opts {
if opt != nil {
options = append(options, opt)
}
}
return grpc.DialContext(ctx, address, options...)
}
func WithCachedGrpcClient(fn func(*grpc.ClientConn) error, address string, opts ...grpc.DialOption) error {
grpcClientsLock.Lock()
existingConnection, found := grpcClients[address]
if found {
grpcClientsLock.Unlock()
err := fn(existingConnection)
if err != nil {
grpcClientsLock.Lock()
delete(grpcClients, address)
grpcClientsLock.Unlock()
existingConnection.Close()
}
return err
}
grpcConnection, err := GrpcDial(context.Background(), address, opts...)
if err != nil {
grpcClientsLock.Unlock()
return fmt.Errorf("fail to dial %s: %v", address, err)
}
grpcClients[address] = grpcConnection
grpcClientsLock.Unlock()
err = fn(grpcConnection)
if err != nil {
grpcClientsLock.Lock()
delete(grpcClients, address)
grpcClientsLock.Unlock()
grpcConnection.Close()
}
return err
}
func ParseServerToGrpcAddress(server string) (serverGrpcAddress string, err error) {
colonIndex := strings.LastIndex(server, ":")
if colonIndex < 0 {
return "", fmt.Errorf("server should have hostname:port format: %v", server)
}
port, parseErr := strconv.ParseUint(server[colonIndex+1:], 10, 64)
if parseErr != nil {
return "", fmt.Errorf("server port parse error: %v", parseErr)
}
grpcPort := int(port) + 10000
return fmt.Sprintf("%s:%d", server[:colonIndex], grpcPort), nil
}
func ServerToGrpcAddress(server string) (serverGrpcAddress string) {
hostnameAndPort := strings.Split(server, ":")
if len(hostnameAndPort) != 2 {
return fmt.Sprintf("unexpected server address: %s", server)
}
port, parseErr := strconv.ParseUint(hostnameAndPort[1], 10, 64)
if parseErr != nil {
return fmt.Sprintf("failed to parse port for %s:%s", hostnameAndPort[0], hostnameAndPort[1])
}
grpcPort := int(port) + 10000
return fmt.Sprintf("%s:%d", hostnameAndPort[0], grpcPort)
}
func WithMasterClient(master string, grpcDialOption grpc.DialOption, fn func(client master_pb.SeaweedClient) error) error {
masterGrpcAddress, parseErr := ParseServerToGrpcAddress(master)
if parseErr != nil {
return fmt.Errorf("failed to parse master grpc %v: %v", master, parseErr)
}
return WithCachedGrpcClient(func(grpcConnection *grpc.ClientConn) error {
client := master_pb.NewSeaweedClient(grpcConnection)
return fn(client)
}, masterGrpcAddress, grpcDialOption)
}
func WithFilerClient(filer string, grpcDialOption grpc.DialOption, fn func(client filer_pb.SeaweedFilerClient) error) error {
filerGrpcAddress, parseErr := ParseServerToGrpcAddress(filer)
if parseErr != nil {
return fmt.Errorf("failed to parse filer grpc %v: %v", filer, parseErr)
}
return WithGrpcFilerClient(filerGrpcAddress, grpcDialOption, fn)
}
func WithGrpcFilerClient(filerGrpcAddress string, grpcDialOption grpc.DialOption, fn func(client filer_pb.SeaweedFilerClient) error) error {
return WithCachedGrpcClient(func(grpcConnection *grpc.ClientConn) error {
client := filer_pb.NewSeaweedFilerClient(grpcConnection)
return fn(client)
}, filerGrpcAddress, grpcDialOption)
}
func ParseFilerGrpcAddress(filer string) (filerGrpcAddress string, err error) {
hostnameAndPort := strings.Split(filer, ":")
if len(hostnameAndPort) != 2 {
return "", fmt.Errorf("filer should have hostname:port format: %v", hostnameAndPort)
}
filerPort, parseErr := strconv.ParseUint(hostnameAndPort[1], 10, 64)
if parseErr != nil {
return "", fmt.Errorf("filer port parse error: %v", parseErr)
}
filerGrpcPort := int(filerPort) + 10000
return fmt.Sprintf("%s:%d", hostnameAndPort[0], filerGrpcPort), nil
}