gocron/modules/rpc/client/client.go

88 lines
2.0 KiB
Go

package client
import (
"errors"
"fmt"
"github.com/ouqiang/gocron/modules/logger"
"github.com/ouqiang/gocron/modules/rpc/grpcpool"
pb "github.com/ouqiang/gocron/modules/rpc/proto"
"golang.org/x/net/context"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"time"
"sync"
)
var (
taskMap sync.Map
)
var (
errUnavailable = errors.New("无法连接远程服务器")
)
func generateTaskUniqueKey(ip string, port int, id int64) string {
return fmt.Sprintf("%s:%d:%d", ip, port, id)
}
func Stop(ip string, port int , id int64) {
key := generateTaskUniqueKey(ip, port, id)
cancel, ok := taskMap.Load(key)
if !ok {
return
}
cancel.(context.CancelFunc)()
}
func Exec(ip string, port int, taskReq *pb.TaskRequest) (string, error) {
defer func() {
if err := recover(); err != nil {
logger.Error("panic#rpc/client.go:Exec#", err)
}
}()
addr := fmt.Sprintf("%s:%d", ip, port)
conn, err := grpcpool.Pool.Get(addr)
if err != nil {
return "", err
}
isConnClosed := false
defer func() {
if !isConnClosed {
grpcpool.Pool.Put(addr, conn)
}
}()
c := pb.NewTaskClient(conn)
if taskReq.Timeout <= 0 || taskReq.Timeout > 86400 {
taskReq.Timeout = 86400
}
timeout := time.Duration(taskReq.Timeout) * time.Second
ctx, cancel := context.WithTimeout(context.Background(), timeout)
taskUniqueKey := generateTaskUniqueKey(ip, port, taskReq.Id)
taskMap.Store(taskUniqueKey, cancel)
defer taskMap.Delete(taskUniqueKey)
resp, err := c.Run(ctx, taskReq)
if err != nil {
return parseGRPCError(err, conn, &isConnClosed)
}
if resp.Error == "" {
return resp.Output, nil
}
return resp.Output, errors.New(resp.Error)
}
func parseGRPCError(err error, conn *grpc.ClientConn, connClosed *bool) (string, error) {
switch grpc.Code(err) {
case codes.Unavailable, codes.Internal:
conn.Close()
*connClosed = true
return "", errUnavailable
case codes.DeadlineExceeded:
return "", errors.New("执行超时, 强制结束")
case codes.Canceled:
return "", errors.New("手动停止")
}
return "", err
}