networkTopology: optimise replicaMap slightly (#1355)

```
name                   old time/op    new time/op    delta
Topology_ReplicaMap-8    83.5ms ± 1%    54.0ms ± 1%  -35.29%  (p=0.000 n=8+9)

name                   old alloc/op   new alloc/op   delta
Topology_ReplicaMap-8    5.73MB ± 0%    0.83MB ± 0%  -85.44%  (p=0.000 n=9+8)

name                   old allocs/op  new allocs/op  delta
Topology_ReplicaMap-8     89.2k ± 0%      6.9k ± 0%  -92.22%  (p=0.000 n=10+10)
```

host_source.go

@@ -228,8 +228,9 @@ func (h *HostInfo) PreferredIP() net.IP {
 
 func (h *HostInfo) DataCenter() string {
 	h.mu.RLock()
-	defer h.mu.RUnlock()
-	return h.dataCenter
+	dc := h.dataCenter
+	h.mu.RUnlock()
+	return dc
 }
 
 func (h *HostInfo) setDataCenter(dataCenter string) *HostInfo {
@@ -241,8 +242,9 @@ func (h *HostInfo) setDataCenter(dataCenter string) *HostInfo {
 
 func (h *HostInfo) Rack() string {
 	h.mu.RLock()
-	defer h.mu.RUnlock()
-	return h.rack
+	rack := h.rack
+	h.mu.RUnlock()
+	return rack
 }
 
 func (h *HostInfo) setRack(rack string) *HostInfo {

topology.go

@@ -150,6 +150,12 @@ func (n *networkTopology) haveRF(replicaCounts map[string]int) bool {
 
 func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 	dcRacks := make(map[string]map[string]struct{}, len(n.dcs))
+	// skipped hosts in a dc
+	skipped := make(map[string][]*HostInfo, len(n.dcs))
+	// number of replicas per dc
+	replicasInDC := make(map[string]int, len(n.dcs))
+	// dc -> racks
+	seenDCRacks := make(map[string]map[string]struct{}, len(n.dcs))
 
 	for _, h := range tokenRing.hosts {
 		dc := h.DataCenter()
@@ -163,6 +169,11 @@ func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 		racks[rack] = struct{}{}
 	}
 
+	for dc, racks := range dcRacks {
+		replicasInDC[dc] = 0
+		seenDCRacks[dc] = make(map[string]struct{}, len(racks))
+	}
+
 	tokens := tokenRing.tokens
 	replicaRing := make(tokenRingReplicas, len(tokens))
 
@@ -172,18 +183,25 @@ func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 	}
 
 	for i, th := range tokenRing.tokens {
-		// number of replicas per dc
-		// TODO: recycle these
-		replicasInDC := make(map[string]int, len(n.dcs))
-		// dc -> racks
-		seenDCRacks := make(map[string]map[string]struct{}, len(n.dcs))
-		// skipped hosts in a dc
-		skipped := make(map[string][]*HostInfo, len(n.dcs))
+		for k, v := range skipped {
+			skipped[k] = v[:0]
+		}
+
+		for dc := range n.dcs {
+			replicasInDC[dc] = 0
+			for rack := range seenDCRacks[dc] {
+				delete(seenDCRacks[dc], rack)
+			}
+		}
 
 		replicas := make([]*HostInfo, 0, totalRF)
-		for j := 0; j < len(tokens) && !n.haveRF(replicasInDC); j++ {
+		for j := 0; j < len(tokens) && (len(replicas) < totalRF && !n.haveRF(replicasInDC)); j++ {
 			// TODO: ensure we dont add the same host twice
-			h := tokens[(i+j)%len(tokens)].host
+			p := i + j
+			if p >= len(tokens) {
+				p -= len(tokens)
+			}
+			h := tokens[p].host
 
 			dc := h.DataCenter()
 			rack := h.Rack()
@@ -202,13 +220,6 @@ func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 			} else if _, ok := dcRacks[dc][rack]; !ok {
 				// dont know about this rack
 				continue
-			} else if len(replicas) >= totalRF {
-				if replicasInDC[dc] > rf {
-					panic(fmt.Sprintf("replica overflow. total rf=%d have=%d", totalRF, len(replicas)))
-				}
-
-				// we now have enough replicas
-				break
 			}
 
 			racks := seenDCRacks[dc]
@@ -225,7 +236,7 @@ func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 				// new rack
 				racks[rack] = struct{}{}
 				replicas = append(replicas, h)
-				replicasInDC[dc]++
+				r := replicasInDC[dc] + 1
 
 				if len(racks) == len(dcRacks[dc]) {
 					// if we have been through all the racks, drain the rest of the skipped
@@ -233,13 +244,14 @@ func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 					// above
 					skippedHosts := skipped[dc]
 					var k int
-					for ; k < len(skippedHosts) && replicasInDC[dc] < rf; k++ {
+					for ; k < len(skippedHosts) && r+k < rf; k++ {
 						sh := skippedHosts[k]
 						replicas = append(replicas, sh)
-						replicasInDC[dc]++
 					}
+					r += k
 					skipped[dc] = skippedHosts[k:]
 				}
+				replicasInDC[dc] = r
 			} else {
 				// already seen this rack, keep hold of this host incase
 				// we dont get enough for rf
@@ -260,7 +272,5 @@ func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 		panic(fmt.Sprintf("token map different size to token ring: got %d expected %d", len(replicaRing), len(tokens)))
 	}
 
-	sort.Sort(replicaRing)
-
 	return replicaRing
 }

topology_test.go

@@ -107,6 +107,9 @@ func TestPlacementStrategy_NetworkStrategy(t *testing.T) {
 	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 token, replicas := range tokenReplicas {
 		if len(replicas.hosts) != expReplicas {