From 8e6b212b1dd4df64e51341eb82467a4d05fcb84a Mon Sep 17 00:00:00 2001
From: ouqiang <qingqianludao@gmail.com>
Date: Mon, 29 May 2017 10:19:58 +0800
Subject: [PATCH] =?UTF-8?q?gRPC=E5=AE=A2=E6=88=B7=E7=AB=AF=E4=B8=8D?=
 =?UTF-8?q?=E5=8F=AF=E7=94=A8=E8=BF=9E=E6=8E=A5,=20=E4=B8=8D=E6=94=BE?=
 =?UTF-8?q?=E5=9B=9E=E8=BF=9E=E6=8E=A5=E6=B1=A0,=20=E7=9B=B4=E6=8E=A5?=
 =?UTF-8?q?=E5=85=B3=E9=97=AD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 modules/rpc/client/client.go | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/modules/rpc/client/client.go b/modules/rpc/client/client.go
index a9c2b6e..86ca5ec 100644
--- a/modules/rpc/client/client.go
+++ b/modules/rpc/client/client.go
@@ -7,16 +7,21 @@ import (
     "time"
     "errors"
     "github.com/ouqiang/gocron/modules/rpc/grpcpool"
+    "google.golang.org/grpc/codes"
+    "google.golang.org/grpc"
 )
 
 func Exec(ip string, port int, taskReq *pb.TaskRequest) (string, error)  {
-    addr := fmt.Sprintf("%s:%d", ip, port);
+    addr := fmt.Sprintf("%s:%d", ip, port)
     conn, err := grpcpool.Pool.Get(addr)
     if err != nil {
         return "", err
     }
+    isConnClosed := false
     defer func() {
-        grpcpool.Pool.Put(addr, conn)
+        if !isConnClosed {
+            grpcpool.Pool.Put(addr, conn)
+        }
     }()
     c := pb.NewTaskClient(conn)
     if taskReq.Timeout <= 0 || taskReq.Timeout > 86400 {
@@ -27,7 +32,7 @@ func Exec(ip string, port int, taskReq *pb.TaskRequest) (string, error)  {
     defer cancel()
     resp, err := c.Run(ctx, taskReq)
     if err != nil {
-        return "", err
+       return parseGRPCError(err, conn, &isConnClosed)
     }
 
     if resp.Error == "" {
@@ -35,4 +40,16 @@ func Exec(ip string, port int, taskReq *pb.TaskRequest) (string, error)  {
     }
 
     return resp.Output, errors.New(resp.Error)
-}
\ No newline at end of file
+}
+
+func parseGRPCError(err error, conn *grpc.ClientConn, connClosed *bool) (string, error) {
+    switch grpc.Code(err) {
+        case codes.Unavailable:
+            conn.Close()
+            *connClosed = true
+            return "", errors.New("无法连接远程服务器")
+        case codes.DeadlineExceeded:
+            return "", errors.New("执行超时, 强制结束")
+    }
+    return "", err
+}