From fab2f630b43fb444d673742713cc97f9210b87b5 Mon Sep 17 00:00:00 2001
From: Ben Ash <32777270+benashz@users.noreply.github.com>
Date: Tue, 21 Dec 2021 16:14:39 -0500
Subject: [PATCH] Fix properly initialize replicateStateStore from
 SetReadYourWrites() (#13486)

Fixes an issue where the `replicateStateStore` was being set to `nil`
upon consecutive calls to `client.SetReadYourWrites(true)`.
---
 api/client.go       | 12 +++++----
 api/client_test.go  | 66 +++++++++++++++++++++++++++++++++++++++++++++
 changelog/13486.txt |  3 +++
 3 files changed, 76 insertions(+), 5 deletions(-)
 create mode 100644 changelog/13486.txt

diff --git a/api/client.go b/api/client.go
index 2d1c3b683..b6c497d0b 100644
--- a/api/client.go
+++ b/api/client.go
@@ -20,9 +20,9 @@ import (
 	"unicode"
 
 	"github.com/hashicorp/errwrap"
-	cleanhttp "github.com/hashicorp/go-cleanhttp"
-	retryablehttp "github.com/hashicorp/go-retryablehttp"
-	rootcerts "github.com/hashicorp/go-rootcerts"
+	"github.com/hashicorp/go-cleanhttp"
+	"github.com/hashicorp/go-retryablehttp"
+	"github.com/hashicorp/go-rootcerts"
 	"github.com/hashicorp/go-secure-stdlib/parseutil"
 	"golang.org/x/net/http2"
 	"golang.org/x/time/rate"
@@ -880,8 +880,10 @@ func (c *Client) SetReadYourWrites(preventStaleReads bool) {
 	c.config.modifyLock.Lock()
 	defer c.config.modifyLock.Unlock()
 
-	if preventStaleReads && c.replicationStateStore == nil {
-		c.replicationStateStore = &replicationStateStore{}
+	if preventStaleReads {
+		if c.replicationStateStore == nil {
+			c.replicationStateStore = &replicationStateStore{}
+		}
 	} else {
 		c.replicationStateStore = nil
 	}
diff --git a/api/client_test.go b/api/client_test.go
index f335a765b..89d238c3e 100644
--- a/api/client_test.go
+++ b/api/client_test.go
@@ -986,3 +986,69 @@ func TestClient_ReadYourWrites(t *testing.T) {
 		})
 	}
 }
+
+func TestClient_SetReadYourWrites(t *testing.T) {
+	tests := []struct {
+		name   string
+		config *Config
+		calls  []bool
+	}{
+		{
+			name:   "false",
+			config: &Config{},
+			calls:  []bool{false},
+		},
+		{
+			name:   "true",
+			config: &Config{},
+			calls:  []bool{true},
+		},
+		{
+			name:   "multi-false",
+			config: &Config{},
+			calls:  []bool{false, false},
+		},
+		{
+			name:   "multi-true",
+			config: &Config{},
+			calls:  []bool{true, true},
+		},
+		{
+			name:   "multi-mix",
+			config: &Config{},
+			calls:  []bool{false, true, false, true},
+		},
+	}
+
+	assertSetReadYourRights := func(t *testing.T, c *Client, v bool, s *replicationStateStore) {
+		t.Helper()
+		c.SetReadYourWrites(v)
+		if c.config.ReadYourWrites != v {
+			t.Fatalf("expected config.ReadYourWrites %#v, actual %#v", v, c.config.ReadYourWrites)
+		}
+		if !reflect.DeepEqual(s, c.replicationStateStore) {
+			t.Fatalf("expected replicationStateStore %#v, actual %#v", s, c.replicationStateStore)
+		}
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := &Client{
+				config: tt.config,
+			}
+			for i, v := range tt.calls {
+				var expectStateStore *replicationStateStore
+				if v {
+					if c.replicationStateStore == nil {
+						c.replicationStateStore = &replicationStateStore{
+							store: []string{},
+						}
+					}
+					c.replicationStateStore.store = append(c.replicationStateStore.store,
+						fmt.Sprintf("%s-%d", tt.name, i))
+					expectStateStore = c.replicationStateStore
+				}
+				assertSetReadYourRights(t, c, v, expectStateStore)
+			}
+		})
+	}
+}
diff --git a/changelog/13486.txt b/changelog/13486.txt
new file mode 100644
index 000000000..493d36342
--- /dev/null
+++ b/changelog/13486.txt
@@ -0,0 +1,3 @@
+```release-note:bug
+api/client: Fixes an issue where the `replicateStateStore` was being set to `nil` upon consecutive calls to `client.SetReadYourWrites(true)`.
+```