Repository: gocql/gocql
Pull request #1413 (2020-03-02 15:09:41)

Fix data race when accessing hostMetrics
ObservedQuery and ObservedBatch expose a pointer to hostMetrics struct
to the respective observers. However, the fields of the returned
*hostMetrics are not safe to read, since the observer that might access
hostMetrics' fields is called without the lock held.

In order to return a safe value to clients, we need to make a copy of
the hostMetrics struct while the lock is still locked.

This commit also optimizes the count of lock/unlock calls, previously we
did 3 pairs (for updating attempts, for updating latency and for getting
hostMetrics), now we only do one.


Patch:
diff --git a/AUTHORS b/AUTHORS
index 5747e19..4bf368d 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -113,3 +113,4 @@ Thomas Meson <zllak@hycik.org>
 Martin Sucha <martin.sucha@kiwi.com>; <git@mm.ms47.eu>
 Pavel Buchinchik <p.buchinchik@gmail.com>
 Rintaro Okamura <rintaro.okamura@gmail.com>
+Yura Sokolov <y.sokolov@joom.com>; <funny.falcon@gmail.com>
diff --git a/session.go b/session.go
index ca879b0..6c7edff 100644
--- a/session.go
+++ b/session.go
@@ -696,7 +696,11 @@ func (s *Session) MapExecuteBatchCAS(batch *Batch, dest map[string]interface{})
 }
 
 type hostMetrics struct {
+	// Attempts is count of how many times this query has been attempted for this host.
+	// An attempt is either a retry or fetching next page of results.
 	Attempts     int
+
+	// TotalLatency is the sum of attempt latencies for this host in nanoseconds.
 	TotalLatency int64
 }
 
@@ -717,12 +721,15 @@ func preFilledQueryMetrics(m map[string]*hostMetrics) *queryMetrics {
 	return qm
 }
 
-// hostMetricsLocked gets or creates host metrics for given host.
+// hostMetrics returns a snapshot of metrics for given host.
+// If the metrics for host don't exist, they are created.
 func (qm *queryMetrics) hostMetrics(host *HostInfo) *hostMetrics {
 	qm.l.Lock()
 	metrics := qm.hostMetricsLocked(host)
+	copied := new(hostMetrics)
+	*copied = *metrics
 	qm.l.Unlock()
-	return metrics
+	return copied
 }
 
 // hostMetricsLocked gets or creates host metrics for given host.
@@ -746,17 +753,6 @@ func (qm *queryMetrics) attempts() int {
 	return attempts
 }
 
-// addAttempts adds given number of attempts and returns previous total attempts.
-func (qm *queryMetrics) addAttempts(i int, host *HostInfo) int {
-	qm.l.Lock()
-	hostMetric := qm.hostMetricsLocked(host)
-	hostMetric.Attempts += i
-	attempts := qm.totalAttempts
-	qm.totalAttempts += i
-	qm.l.Unlock()
-	return attempts
-}
-
 func (qm *queryMetrics) latency() int64 {
 	qm.l.Lock()
 	var (
@@ -774,11 +770,28 @@ func (qm *queryMetrics) latency() int64 {
 	return 0
 }
 
-func (qm *queryMetrics) addLatency(l int64, host *HostInfo) {
+// attempt adds given number of attempts and latency for given host.
+// It returns previous total attempts.
+// If needsHostMetrics is true, a copy of updated hostMetrics is returned.
+func (qm *queryMetrics) attempt(addAttempts int, addLatency time.Duration,
+	host *HostInfo, needsHostMetrics bool) (int, *hostMetrics) {
 	qm.l.Lock()
-	hostMetric := qm.hostMetricsLocked(host)
-	hostMetric.TotalLatency += l
+
+	totalAttempts := qm.totalAttempts
+	qm.totalAttempts += addAttempts
+
+	updateHostMetrics := qm.hostMetricsLocked(host)
+	updateHostMetrics.Attempts += addAttempts
+	updateHostMetrics.TotalLatency += addLatency.Nanoseconds()
+
+	var hostMetricsCopy *hostMetrics
+	if needsHostMetrics {
+		hostMetricsCopy = new(hostMetrics)
+		*hostMetricsCopy = *updateHostMetrics
+	}
+
 	qm.l.Unlock()
+	return totalAttempts, hostMetricsCopy
 }
 
 // Query represents a CQL statement that can be executed.
@@ -850,7 +863,7 @@ func (q *Query) Attempts() int {
 }
 
 func (q *Query) AddAttempts(i int, host *HostInfo) {
-	q.metrics.addAttempts(i, host)
+	q.metrics.attempt(i, 0, host, false)
 }
 
 //Latency returns the average amount of nanoseconds per attempt of the query.
@@ -859,7 +872,7 @@ func (q *Query) Latency() int64 {
 }
 
 func (q *Query) AddLatency(l int64, host *HostInfo) {
-	q.metrics.addLatency(l, host)
+	q.metrics.attempt(0, time.Duration(l) * time.Nanosecond, host, false)
 }
 
 // Consistency sets the consistency level for this query. If no consistency
@@ -974,8 +987,8 @@ func (q *Query) execute(ctx context.Context, conn *Conn) *Iter {
 }
 
 func (q *Query) attempt(keyspace string, end, start time.Time, iter *Iter, host *HostInfo) {
-	attempt := q.metrics.addAttempts(1, host)
-	q.AddLatency(end.Sub(start).Nanoseconds(), host)
+	latency := end.Sub(start)
+	attempt, metricsForHost := q.metrics.attempt(1, latency, host, q.observer != nil)
 
 	if q.observer != nil {
 		q.observer.ObserveQuery(q.Context(), ObservedQuery{
@@ -985,7 +998,7 @@ func (q *Query) attempt(keyspace string, end, start time.Time, iter *Iter, host
 			End:       end,
 			Rows:      iter.numRows,
 			Host:      host,
-			Metrics:   q.metrics.hostMetrics(host),
+			Metrics:   metricsForHost,
 			Err:       iter.err,
 			Attempt:   attempt,
 		})
@@ -1586,7 +1599,7 @@ func (b *Batch) Attempts() int {
 }
 
 func (b *Batch) AddAttempts(i int, host *HostInfo) {
-	b.metrics.addAttempts(i, host)
+	b.metrics.attempt(i, 0, host, false)
 }
 
 //Latency returns the average number of nanoseconds to execute a single attempt of the batch.
@@ -1595,7 +1608,7 @@ func (b *Batch) Latency() int64 {
 }
 
 func (b *Batch) AddLatency(l int64, host *HostInfo) {
-	b.metrics.addLatency(l, host)
+	b.metrics.attempt(0, time.Duration(l) * time.Nanosecond, host, false)
 }
 
 // GetConsistency returns the currently configured consistency level for the batch
@@ -1719,8 +1732,8 @@ func (b *Batch) WithTimestamp(timestamp int64) *Batch {
 }
 
 func (b *Batch) attempt(keyspace string, end, start time.Time, iter *Iter, host *HostInfo) {
-	b.AddAttempts(1, host)
-	b.AddLatency(end.Sub(start).Nanoseconds(), host)
+	latency := end.Sub(start)
+	_, metricsForHost := b.metrics.attempt(1, latency, host, b.observer != nil)
 
 	if b.observer == nil {
 		return
@@ -1738,7 +1751,7 @@ func (b *Batch) attempt(keyspace string, end, start time.Time, iter *Iter, host
 		End:        end,
 		// Rows not used in batch observations // TODO - might be able to support it when using BatchCAS
 		Host:    host,
-		Metrics: b.metrics.hostMetrics(host),
+		Metrics: metricsForHost,
 		Err:     iter.err,
 	})
 }
@@ -1955,6 +1968,7 @@ type ObservedQuery struct {
 	Err error
 
 	// Attempt is the index of attempt at executing this query.
+	// An attempt might be either retry or fetching next page of a query.
 	// The first attempt is number zero and any retries have non-zero attempt number.
 	Attempt int
 }
diff --git a/topology.go b/topology.go
index 5ff7fd2..3f8a75e 100644
--- a/topology.go
+++ b/topology.go
@@ -171,7 +171,7 @@ func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 	}
 
 	tokens := tokenRing.tokens
-	replicaRing := make(tokenRingReplicas, len(tokens))
+	replicaRing := make(tokenRingReplicas, 0, len(tokens))
 
 	var totalRF int
 	for _, rf := range n.dcs {
@@ -179,6 +179,11 @@ func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 	}
 
 	for i, th := range tokenRing.tokens {
+		if rf := n.dcs[th.host.DataCenter()]; rf == 0 {
+			// skip this token since no replica in this datacenter.
+			continue
+		}
+
 		for k, v := range skipped {
 			skipped[k] = v[:0]
 		}
@@ -202,9 +207,9 @@ func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 			dc := h.DataCenter()
 			rack := h.Rack()
 
-			rf, ok := n.dcs[dc]
-			if !ok {
-				// skip this DC, dont know about it
+			rf := n.dcs[dc]
+			if rf == 0 {
+				// skip this DC, dont know about it or replication factor is zero
 				continue
 			} else if replicasInDC[dc] >= rf {
 				if replicasInDC[dc] > rf {
@@ -261,10 +266,10 @@ func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 			panic(fmt.Sprintf("first replica is not the primary replica for the token: expected %v got %v", replicas[0].ConnectAddress(), th.host.ConnectAddress()))
 		}
 
-		replicaRing[i] = hostTokens{th.token, replicas}
+		replicaRing = append(replicaRing, hostTokens{th.token, replicas})
 	}
 
-	if len(replicaRing) != len(tokens) {
+	if len(n.dcs) == len(dcRacks) && len(replicaRing) != len(tokens) {
 		panic(fmt.Sprintf("token map different size to token ring: got %d expected %d", len(replicaRing), len(tokens)))
 	}
 


Test patch:
diff --git a/topology_test.go b/topology_test.go
index cbc0f83..8c7acb7 100644
--- a/topology_test.go
+++ b/topology_test.go
@@ -90,76 +90,86 @@ func TestPlacementStrategy_NetworkStrategy(t *testing.T) {
 	}
 	sort.Sort(&tokenRing{tokens: tokens})
 
-	strat := &networkTopology{
-		dcs: map[string]int{
-			"dc1": 1,
-			"dc2": 2,
-			"dc3": 3,
+	strats := []*networkTopology{
+		&networkTopology{
+			dcs: map[string]int{
+				"dc1": 1,
+				"dc2": 2,
+				"dc3": 3,
+			},
+		},
+		&networkTopology{
+			dcs: map[string]int{
+				"dc2": 2,
+				"dc3": 3,
+			},
 		},
 	}
 
-	var expReplicas int
-	for _, rf := range strat.dcs {
-		expReplicas += rf
-	}
-
-	tokenReplicas := strat.replicaMap(&tokenRing{hosts: hosts, tokens: tokens})
-	if len(tokenReplicas) != len(tokens) {
-		t.Fatalf("expected replica map to have %d items but has %d", len(tokens), len(tokenReplicas))
-	}
-	if !sort.IsSorted(tokenReplicas) {
-		t.Fatal("replica map was not sorted by token")
-	}
+	for _, strat := range strats {
+		var expReplicas int
+		for _, rf := range strat.dcs {
+			expReplicas += rf
+		}
 
-	for token, replicas := range tokenReplicas {
-		if len(replicas.hosts) != expReplicas {
-			t.Fatalf("expected to have %d replicas got %d for token=%v", expReplicas, len(replicas.hosts), token)
+		tokenReplicas := strat.replicaMap(&tokenRing{hosts: hosts, tokens: tokens})
+		if needTokens := hostsPerDC * len(strat.dcs); len(tokenReplicas) != needTokens {
+			t.Fatalf("expected replica map to have %d items but has %d", needTokens, len(tokenReplicas))
+		}
+		if !sort.IsSorted(tokenReplicas) {
+			t.Fatal("replica map was not sorted by token")
 		}
-	}
 
-	for dc, rf := range strat.dcs {
-		dcTokens := dcRing[dc]
-		for i, th := range dcTokens {
-			token := th.token
-			allReplicas := tokenReplicas.replicasFor(token)
-			if allReplicas.token != token {
-				t.Fatalf("token %v not in replica map", token)
+		for token, replicas := range tokenReplicas {
+			if len(replicas.hosts) != expReplicas {
+				t.Fatalf("expected to have %d replicas got %d for token=%v", expReplicas, len(replicas.hosts), token)
 			}
+		}
 
-			var replicas []*HostInfo
-			for _, replica := range allReplicas.hosts {
-				if replica.dataCenter == dc {
-					replicas = append(replicas, replica)
+		for dc, rf := range strat.dcs {
+			dcTokens := dcRing[dc]
+			for i, th := range dcTokens {
+				token := th.token
+				allReplicas := tokenReplicas.replicasFor(token)
+				if allReplicas.token != token {
+					t.Fatalf("token %v not in replica map", token)
 				}
-			}
 
-			if len(replicas) != rf {
-				t.Fatalf("expected %d replicas in dc %q got %d", rf, dc, len(replicas))
-			}
+				var replicas []*HostInfo
+				for _, replica := range allReplicas.hosts {
+					if replica.dataCenter == dc {
+						replicas = append(replicas, replica)
+					}
+				}
+
+				if len(replicas) != rf {
+					t.Fatalf("expected %d replicas in dc %q got %d", rf, dc, len(replicas))
+				}
 
-			var lastRack string
-			for j, replica := range replicas {
-				// expected is in the next rack
-				var exp *HostInfo
-				if lastRack == "" {
-					// primary, first replica
-					exp = dcTokens[(i+j)%len(dcTokens)].host
-				} else {
-					for k := 0; k < len(dcTokens); k++ {
-						// walk around the ring from i + j to find the next host the
-						// next rack
-						p := (i + j + k) % len(dcTokens)
-						h := dcTokens[p].host
-						if h.rack != lastRack {
-							exp = h
-							break
+				var lastRack string
+				for j, replica := range replicas {
+					// expected is in the next rack
+					var exp *HostInfo
+					if lastRack == "" {
+						// primary, first replica
+						exp = dcTokens[(i+j)%len(dcTokens)].host
+					} else {
+						for k := 0; k < len(dcTokens); k++ {
+							// walk around the ring from i + j to find the next host the
+							// next rack
+							p := (i + j + k) % len(dcTokens)
+							h := dcTokens[p].host
+							if h.rack != lastRack {
+								exp = h
+								break
+							}
+						}
+						if exp.rack == lastRack {
+							panic("no more racks")
 						}
 					}
-					if exp.rack == lastRack {
-						panic("no more racks")
-					}
+					lastRack = replica.rack
 				}
-				lastRack = replica.rack
 			}
 		}
 	}