kratos-utils/grpc/client.go

package grpc

import (
	"context"
	"sync"
	"time"

	"git.ikuban.com/server/kratos-nacos/registry"
	"git.ikuban.com/server/kratos-utils/http/middleware/logging"
	"github.com/go-kratos/kratos/v2/log"
	"github.com/go-kratos/kratos/v2/middleware"
	"github.com/go-kratos/kratos/v2/transport/grpc"

	grpc2 "google.golang.org/grpc"
)

var (
	ConnMap  map[string]*grpc2.ClientConn
	ConnLock sync.Mutex
)

type DialOption struct {
	Middlewares []middleware.Middleware
	Timeout     time.Duration
	GrpcOptions []grpc2.DialOption
}

func GetDialInsecure(
	r *registry.Registry,
	clientName string,
	endpointNameKey string, option *DialOption) (*grpc2.ClientConn, error) {
	option = checkOption(option)
	endpoint := grpc.WithEndpoint(endpointNameKey)

	if r == nil {
		return nil, nil
	}
	ConnLock.Lock()
	conn, ok := ConnMap[clientName]
	ConnLock.Unlock()
	if ok {
		return conn, nil
	}
	var err error
	conn, err = grpc.DialInsecure(context.Background(),
		endpoint,
		grpc.WithDiscovery(r),
		grpc.WithMiddleware(option.Middlewares...),
		grpc.WithTimeout(option.Timeout),
		grpc.WithOptions(option.GrpcOptions...),
	)
	if err != nil {
		return nil, err
	}
	ConnLock.Lock()
	ConnMap[clientName] = conn
	ConnLock.Unlock()
	return conn, nil
}

func GetDialSimple(clientName string, endpoint string, option *DialOption) (*grpc2.ClientConn, error) {
	option = checkOption(option)
	ConnLock.Lock()
	conn, ok := ConnMap[clientName]
	ConnLock.Unlock()
	if ok {
		return conn, nil
	}
	var err error
	conn, err = grpc.DialInsecure(context.Background(),
		grpc.WithEndpoint(endpoint),
		grpc.WithMiddleware(option.Middlewares...),
		grpc.WithTimeout(option.Timeout),
		grpc.WithOptions(option.GrpcOptions...),
	)
	if err != nil {
		return nil, err
	}
	ConnLock.Lock()
	ConnMap[clientName] = conn
	ConnLock.Unlock()
	return conn, nil
}

func NewGrpcClientOption(logger log.Logger) *DialOption {
	option := &DialOption{
		Middlewares: make([]middleware.Middleware, 0),
		Timeout:     time.Second * 60,
	}

	if len(option.GrpcOptions) == 0 {
		option.GrpcOptions = append(option.GrpcOptions, grpc2.WithDefaultCallOptions(grpc2.MaxCallRecvMsgSize(ReceiveMsgSize)))
	}
	if logger != nil {
		option.Middlewares = append(option.Middlewares, logging.Client(logger))
	}
	return option
}

func checkOption(option *DialOption) *DialOption {
	if option == nil {
		option = &DialOption{
			Middlewares: make([]middleware.Middleware, 0),
			Timeout:     time.Second * 60,
		}
	}
	if option.Timeout.Milliseconds() == 0 {
		option.Timeout = time.Second * 60
	}
	if len(option.GrpcOptions) == 0 {
		option.GrpcOptions = append(option.GrpcOptions, grpc2.WithDefaultCallOptions(grpc2.MaxCallRecvMsgSize(ReceiveMsgSize)))
	}
	return option
}