clientv3: WithRequireLeader

clientv3/client.go

@@ -24,9 +24,12 @@ import (
 	"sync"
 	"time"
 
+	"github.com/coreos/etcd/etcdserver/api/v3rpc/rpctypes"
+
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/metadata"
 )
 
 var (
@@ -160,6 +163,13 @@ func (c *Client) Dial(endpoint string) (*grpc.ClientConn, error) {
 	return conn, nil
 }
 
+// WithRequireLeader requires client requests to only succeed
+// when the cluster has a leader.
+func WithRequireLeader(ctx context.Context) context.Context {
+	md := metadata.Pairs(rpctypes.MetadataRequireLeaderKey, rpctypes.MetadataHasLeader)
+	return metadata.NewContext(ctx, md)
+}
+
 func newClient(cfg *Config) (*Client, error) {
 	if cfg == nil {
 		cfg = &Config{RetryDialer: dialEndpointList}

clientv3/integration/kv_test.go

@@ -110,6 +110,32 @@ func TestKVPut(t *testing.T) {
 	}
 }
 
+func TestKVPutWithRequireLeader(t *testing.T) {
+	// this test might block for a few seconds, make it parallel to speed up the test.
+	t.Parallel()
+
+	defer testutil.AfterTest(t)
+
+	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 3})
+	defer clus.Terminate(t)
+
+	clus.Members[1].Stop(t)
+	clus.Members[2].Stop(t)
+
+	// wait for election timeout, then member[0] will not have a leader.
+	var (
+		electionTicks = 10
+		tickDuration  = 10 * time.Millisecond
+	)
+	time.Sleep(time.Duration(3*electionTicks) * tickDuration)
+
+	kv := clientv3.NewKV(clus.Client(0))
+	_, err := kv.Put(clientv3.WithRequireLeader(context.Background()), "foo", "bar")
+	if err != rpctypes.ErrNoLeader {
+		t.Fatal(err)
+	}
+}
+
 func TestKVRange(t *testing.T) {
 	defer testutil.AfterTest(t)