From cc6375b8518b938d9a7398c93a5b1f1a68cf4f50 Mon Sep 17 00:00:00 2001
From: smallnest <smallnest@gmail.com>
Date: Tue, 12 Nov 2019 20:26:51 +0800
Subject: [PATCH] don't close the connection if the error is context error

---
 client/xclient.go      | 45 +++++++++++++++++++++++++++++++++---------
 client/xclient_test.go | 23 ++++++++++++++++++++-
 2 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/client/xclient.go b/client/xclient.go
index 53b71d0..95b4d4c 100644
--- a/client/xclient.go
+++ b/client/xclient.go
@@ -425,7 +425,9 @@ func (c *xClient) Call(ctx context.Context, serviceMethod string, args interface
 				}
 			}
 
-			c.removeClient(k, client)
+			if uncoverError(err) {
+				c.removeClient(k, client)
+			}
 			client, e = c.getCachedClient(k)
 		}
 		if err == nil {
@@ -447,7 +449,9 @@ func (c *xClient) Call(ctx context.Context, serviceMethod string, args interface
 				}
 			}
 
-			c.removeClient(k, client)
+			if uncoverError(err) {
+				c.removeClient(k, client)
+			}
 			//select another server
 			k, client, e = c.selectClient(ctx, c.servicePath, serviceMethod, args)
 		}
@@ -487,7 +491,7 @@ func (c *xClient) Call(ctx context.Context, serviceMethod string, args interface
 		}
 		_, err2 := c.Go(ctx, serviceMethod, args, reply2, call2)
 		if err2 != nil {
-			if _, ok := err.(ServiceError); !ok {
+			if uncoverError(err2) {
 				c.removeClient(k, client)
 			}
 			err = err1
@@ -513,7 +517,7 @@ func (c *xClient) Call(ctx context.Context, serviceMethod string, args interface
 	default: //Failfast
 		err = c.wrapCall(ctx, client, serviceMethod, args, reply)
 		if err != nil {
-			if _, ok := err.(ServiceError); !ok {
+			if uncoverError(err) {
 				c.removeClient(k, client)
 			}
 		}
@@ -522,6 +526,21 @@ func (c *xClient) Call(ctx context.Context, serviceMethod string, args interface
 	}
 }
 
+func uncoverError(err error) bool {
+	if _, ok := err.(ServiceError); ok {
+		return false
+	}
+
+	if err == context.DeadlineExceeded {
+		return false
+	}
+
+	if err == context.Canceled {
+		return false
+	}
+
+	return true
+}
 func (c *xClient) SendRaw(ctx context.Context, r *protocol.Message) (map[string]string, []byte, error) {
 	if c.isShutdown {
 		return nil, nil, ErrXClientShutdown
@@ -566,7 +585,9 @@ func (c *xClient) SendRaw(ctx context.Context, r *protocol.Message) (map[string]
 				}
 			}
 
-			c.removeClient(k, client)
+			if uncoverError(err) {
+				c.removeClient(k, client)
+			}
 			client, e = c.getCachedClient(k)
 		}
 
@@ -588,7 +609,9 @@ func (c *xClient) SendRaw(ctx context.Context, r *protocol.Message) (map[string]
 				}
 			}
 
-			c.removeClient(k, client)
+			if uncoverError(err) {
+				c.removeClient(k, client)
+			}
 			//select another server
 			k, client, e = c.selectClient(ctx, r.ServicePath, r.ServiceMethod, r.Payload)
 		}
@@ -602,7 +625,7 @@ func (c *xClient) SendRaw(ctx context.Context, r *protocol.Message) (map[string]
 		m, payload, err := client.SendRaw(ctx, r)
 
 		if err != nil {
-			if _, ok := err.(ServiceError); !ok {
+			if uncoverError(err) {
 				c.removeClient(k, client)
 			}
 		}
@@ -666,7 +689,9 @@ func (c *xClient) Broadcast(ctx context.Context, serviceMethod string, args inte
 			e := c.wrapCall(ctx, client, serviceMethod, args, reply)
 			done <- (e == nil)
 			if e != nil {
-				c.removeClient(k, client)
+				if uncoverError(err) {
+					c.removeClient(k, client)
+				}
 				err.Append(e)
 			}
 		}()
@@ -743,7 +768,9 @@ func (c *xClient) Fork(ctx context.Context, serviceMethod string, args interface
 			}
 			done <- (e == nil)
 			if e != nil {
-				c.removeClient(k, client)
+				if uncoverError(err) {
+					c.removeClient(k, client)
+				}
 				err.Append(e)
 			}
 
diff --git a/client/xclient_test.go b/client/xclient_test.go
index e8c363a..74c80bd 100644
--- a/client/xclient_test.go
+++ b/client/xclient_test.go
@@ -2,12 +2,13 @@ package client
 
 import (
 	"context"
+	"errors"
 	"testing"
 	"time"
 
 	"fmt"
 
-	"github.com/smallnest/rpcx/_testutils"
+	testutils "github.com/smallnest/rpcx/_testutils"
 	"github.com/smallnest/rpcx/protocol"
 	"github.com/smallnest/rpcx/server"
 	"github.com/smallnest/rpcx/share"
@@ -99,3 +100,23 @@ func TestXClient_filterByStateAndGroup(t *testing.T) {
 		t.Error("node must be removed")
 	}
 }
+
+func TestUncoverError(t *testing.T) {
+	var e error = ServiceError("error")
+	if uncoverError(e) {
+		t.Fatalf("expect false but get true")
+	}
+
+	if uncoverError(context.DeadlineExceeded) {
+		t.Fatalf("expect false but get true")
+	}
+
+	if uncoverError(context.Canceled) {
+		t.Fatalf("expect false but get true")
+	}
+
+	e = errors.New("error")
+	if !uncoverError(e) {
+		t.Fatalf("expect true but get false")
+	}
+}