clientv3/integration: add TestBalancerUnderServerShutdownMutable*

Signed-off-by: Gyu-Ho Lee <[email protected]>

clientv3/integration/server_shutdown_test.go

@@ -17,6 +17,7 @@ package integration
 import (
 	"bytes"
 	"context"
+	"strings"
 	"testing"
 	"time"
 
@@ -116,3 +117,91 @@ func TestBalancerUnderServerShutdownWatch(t *testing.T) {
 		t.Fatal("took too long to receive events")
 	}
 }
+
+func TestBalancerUnderServerShutdownPut(t *testing.T) {
+	testBalancerUnderServerShutdownMutable(t, func(cli *clientv3.Client, ctx context.Context) error {
+		_, err := cli.Put(ctx, "foo", "bar")
+		return err
+	})
+}
+
+func TestBalancerUnderServerShutdownDelete(t *testing.T) {
+	testBalancerUnderServerShutdownMutable(t, func(cli *clientv3.Client, ctx context.Context) error {
+		_, err := cli.Delete(ctx, "foo")
+		return err
+	})
+}
+
+func TestBalancerUnderServerShutdownTxn(t *testing.T) {
+	testBalancerUnderServerShutdownMutable(t, func(cli *clientv3.Client, ctx context.Context) error {
+		_, err := cli.Txn(ctx).
+			If(clientv3.Compare(clientv3.Version("foo"), "=", 0)).
+			Then(clientv3.OpPut("foo", "bar")).
+			Else(clientv3.OpPut("foo", "baz")).Commit()
+		return err
+	})
+}
+
+// testBalancerUnderServerShutdownMutable expects put/delete/txn requests
+// fail when the member of the pinned endpoint is shut down, and
+// following retry succeed with new endpoints.
+func testBalancerUnderServerShutdownMutable(t *testing.T, op func(*clientv3.Client, context.Context) error) {
+	defer testutil.AfterTest(t)
+
+	clus := integration.NewClusterV3(t, &integration.ClusterConfig{
+		Size:               3,
+		SkipCreatingClient: true,
+	})
+	defer clus.Terminate(t)
+
+	eps := []string{clus.Members[0].GRPCAddr(), clus.Members[1].GRPCAddr(), clus.Members[2].GRPCAddr()}
+
+	// pin eps[0]
+	cli, err := clientv3.New(clientv3.Config{Endpoints: []string{eps[0]}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cli.Close()
+
+	// wait for eps[0] to be pinned
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	_, err = cli.Get(ctx, "foo")
+	cancel()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// add all eps to list, so that when the original pined one fails
+	// the client can switch to other available eps
+	cli.SetEndpoints(eps...)
+
+	// shut down eps[0]
+	donec := make(chan struct{})
+	go func() {
+		clus.Members[0].Terminate(t)
+		close(donec)
+	}()
+
+	failed := false
+	for i := 0; i < 30; i++ {
+		cctx, ccancel := context.WithTimeout(context.Background(), time.Second)
+		err = op(cli, cctx)
+		ccancel()
+		if err != nil && (err == context.DeadlineExceeded || strings.Contains(err.Error(), "transport is closing")) {
+			if !failed {
+				failed = true
+				time.Sleep(3 * time.Second) // give more time for endpoint switch in next try
+				continue
+			}
+		}
+		// switched to others when eps[0] was shut down and next retry succeeded
+		if failed {
+			if err == nil {
+				break
+			}
+			t.Fatalf("expected no error in retry from previous failure, got %v", err)
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+	<-donec
+}