Fix race condition with tokenAwareHostPolicy SetPartitioner (apache#988)

- it wasn't using the lock and causing data races
jameshartig · Oct 16, 2017 · 04cbe7a · 04cbe7a
1 parent 2416cf3
commit 04cbe7a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/AUTHORS b/AUTHORS
@@ -96,3 +96,4 @@ Dharmendra Parsaila <d4dharmu@gmail.com>
 Nayef Ghattas <nayef.ghattas@datadoghq.com>
 Michał Matczuk <mmatczuk@gmail.com>
 Ben Krebsbach <ben.krebsbach@gmail.com>
+Vivian Mathews <vivian.mathews.3@gmail.com>
diff --git a/policies.go b/policies.go
@@ -297,6 +297,9 @@ type tokenAwareHostPolicy struct {
 }
 
 func (t *tokenAwareHostPolicy) SetPartitioner(partitioner string) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
 	if t.partitioner != partitioner {
 		t.fallback.SetPartitioner(partitioner)
 		t.partitioner = partitioner
@@ -306,13 +309,19 @@ func (t *tokenAwareHostPolicy) SetPartitioner(partitioner string) {
 }
 
 func (t *tokenAwareHostPolicy) AddHost(host *HostInfo) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
 	t.hosts.add(host)
 	t.fallback.AddHost(host)
 
 	t.resetTokenRing()
 }
 
 func (t *tokenAwareHostPolicy) RemoveHost(host *HostInfo) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
 	t.hosts.remove(host.ConnectAddress())
 	t.fallback.RemoveHost(host)
 
@@ -328,9 +337,6 @@ func (t *tokenAwareHostPolicy) HostDown(host *HostInfo) {
 }
 
 func (t *tokenAwareHostPolicy) resetTokenRing() {
-	t.mu.Lock()
-	defer t.mu.Unlock()
-
 	if t.partitioner == "" {
 		// partitioner not yet set
 		return