From 6e521d2f3f1ec4e8dc7d9a4fb16079fc265da803 Mon Sep 17 00:00:00 2001
From: Gyuho Lee <gyuhox@gmail.com>
Date: Fri, 15 Jun 2018 07:19:11 -0700
Subject: [PATCH] clientv3: add "IsConnCanceled", deprecate
 "grpc.ErrClientConnClosing"

Signed-off-by: Gyuho Lee <gyuhox@gmail.com>
---
 clientv3/client.go                 | 20 ++++++++++++++++++++
 clientv3/doc.go                    |  5 +++++
 clientv3/integration/kv_test.go    |  4 ++--
 clientv3/integration/lease_test.go |  8 +++++---
 4 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/clientv3/client.go b/clientv3/client.go
index 8a0cc1cc4..4d1270e8a 100644
--- a/clientv3/client.go
+++ b/clientv3/client.go
@@ -614,6 +614,26 @@ func canceledByCaller(stopCtx context.Context, err error) bool {
 	return err == context.Canceled || err == context.DeadlineExceeded
 }
 
+// IsConnCanceled returns true, if error is from a closed gRPC connection.
+// ref. https://github.com/grpc/grpc-go/pull/1854
+func IsConnCanceled(err error) bool {
+	if err == nil {
+		return false
+	}
+	// >= gRPC v1.10.x
+	s, ok := status.FromError(err)
+	if ok {
+		// connection is canceled or server has already closed the connection
+		return s.Code() == codes.Canceled || s.Message() == "transport is closing"
+	}
+	// >= gRPC v1.10.x
+	if err == context.Canceled {
+		return true
+	}
+	// <= gRPC v1.7.x returns 'errors.New("grpc: the client connection is closing")'
+	return strings.Contains(err.Error(), "grpc: the client connection is closing")
+}
+
 func getHost(ep string) string {
 	url, uerr := url.Parse(ep)
 	if uerr != nil || !strings.Contains(ep, "://") {
diff --git a/clientv3/doc.go b/clientv3/doc.go
index 717fbe435..0bb686189 100644
--- a/clientv3/doc.go
+++ b/clientv3/doc.go
@@ -87,11 +87,16 @@
 //	go func() { cli.Close() }()
 //	_, err := kvc.Get(ctx, "a")
 //	if err != nil {
+//		// with etcd clientv3 <= v3.3
 //		if err == context.Canceled {
 //			// grpc balancer calls 'Get' with an inflight client.Close
 //		} else if err == grpc.ErrClientConnClosing {
 //			// grpc balancer calls 'Get' after client.Close.
 //		}
+//		// with etcd clientv3 >= v3.4
+//		if clientv3.IsConnCanceled(err) {
+//			// gRPC client connection is closed
+//		}
 //	}
 //
 package clientv3
diff --git a/clientv3/integration/kv_test.go b/clientv3/integration/kv_test.go
index 394a8242b..1c7724acd 100644
--- a/clientv3/integration/kv_test.go
+++ b/clientv3/integration/kv_test.go
@@ -442,7 +442,7 @@ func TestKVGetErrConnClosed(t *testing.T) {
 	go func() {
 		defer close(donec)
 		_, err := cli.Get(context.TODO(), "foo")
-		if err != nil && err != context.Canceled && err != grpc.ErrClientConnClosing {
+		if !clientv3.IsConnCanceled(err) {
 			t.Fatalf("expected %v or %v, got %v", context.Canceled, grpc.ErrClientConnClosing, err)
 		}
 	}()
@@ -474,7 +474,7 @@ func TestKVNewAfterClose(t *testing.T) {
 	donec := make(chan struct{})
 	go func() {
 		_, err := cli.Get(context.TODO(), "foo")
-		if err != context.Canceled && err != grpc.ErrClientConnClosing {
+		if !clientv3.IsConnCanceled(err) {
 			t.Fatalf("expected %v or %v, got %v", context.Canceled, grpc.ErrClientConnClosing, err)
 		}
 		close(donec)
diff --git a/clientv3/integration/lease_test.go b/clientv3/integration/lease_test.go
index 75a0987c5..bd7ee897b 100644
--- a/clientv3/integration/lease_test.go
+++ b/clientv3/integration/lease_test.go
@@ -296,7 +296,7 @@ func TestLeaseGrantErrConnClosed(t *testing.T) {
 	go func() {
 		defer close(donec)
 		_, err := cli.Grant(context.TODO(), 5)
-		if err != nil && err != grpc.ErrClientConnClosing && err != context.Canceled {
+		if !clientv3.IsConnCanceled(err) {
 			// grpc.ErrClientConnClosing if grpc-go balancer calls 'Get' after client.Close.
 			// context.Canceled if grpc-go balancer calls 'Get' with an inflight client.Close.
 			t.Fatalf("expected %v, %v or server unavailable, got %v", err != context.Canceled, grpc.ErrClientConnClosing, err)
@@ -328,7 +328,8 @@ func TestLeaseGrantNewAfterClose(t *testing.T) {
 
 	donec := make(chan struct{})
 	go func() {
-		if _, err := cli.Grant(context.TODO(), 5); err != context.Canceled && err != grpc.ErrClientConnClosing {
+		_, err := cli.Grant(context.TODO(), 5)
+		if !clientv3.IsConnCanceled(err) {
 			t.Fatalf("expected %v, %v or server unavailable, got %v", err != context.Canceled, grpc.ErrClientConnClosing, err)
 		}
 		close(donec)
@@ -360,7 +361,8 @@ func TestLeaseRevokeNewAfterClose(t *testing.T) {
 
 	donec := make(chan struct{})
 	go func() {
-		if _, err := cli.Revoke(context.TODO(), leaseID); err != context.Canceled && err != grpc.ErrClientConnClosing {
+		_, err := cli.Revoke(context.TODO(), leaseID)
+		if !clientv3.IsConnCanceled(err) {
 			t.Fatalf("expected %v, %v or server unavailable, got %v", err != context.Canceled, grpc.ErrClientConnClosing, err)
 		}
 		close(donec)