Repository: gocql/gocql
Pull request #1552 (2021-04-13 16:01:56)

Fix panic in replicaMap when rf=0
In case a network topology strategy with zero replication factor
for at least one DC is encountered, the driver panicked with
token map different size to token ring.

This is because len(dcRacks) includes the datacenters where replication
factor is zero, but these datacenters should not have any replicas.


Patch:
diff --git a/topology.go b/topology.go
index 885e35d..a46c665 100644
--- a/topology.go
+++ b/topology.go
@@ -279,7 +279,14 @@ func (n *networkTopology) replicaMap(tokenRing *tokenRing) tokenRingReplicas {
 		replicaRing = append(replicaRing, hostTokens{th.token, replicas})
 	}
 
-	if len(n.dcs) == len(dcRacks) && len(replicaRing) != len(tokens) {
+	dcsWithReplicas := 0
+	for _, dc := range n.dcs {
+		if dc > 0 {
+			dcsWithReplicas++
+		}
+	}
+
+	if dcsWithReplicas == 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 8c7acb7..724f1b4 100644
--- a/topology_test.go
+++ b/topology_test.go
@@ -49,128 +49,157 @@ func TestPlacementStrategy_SimpleStrategy(t *testing.T) {
 }
 
 func TestPlacementStrategy_NetworkStrategy(t *testing.T) {
-	var (
-		hosts  []*HostInfo
-		tokens []hostToken
-	)
-
 	const (
 		totalDCs   = 3
 		racksPerDC = 3
 		hostsPerDC = 5
 	)
 
-	dcRing := make(map[string][]hostToken, totalDCs)
-	for i := 0; i < totalDCs; i++ {
-		var dcTokens []hostToken
-		dc := fmt.Sprintf("dc%d", i+1)
-
-		for j := 0; j < hostsPerDC; j++ {
-			rack := fmt.Sprintf("rack%d", (j%racksPerDC)+1)
-
-			h := &HostInfo{hostId: fmt.Sprintf("%s:%s:%d", dc, rack, j), dataCenter: dc, rack: rack}
-
-			token := hostToken{
-				token: orderedToken([]byte(h.hostId)),
-				host:  h,
-			}
-
-			tokens = append(tokens, token)
-			dcTokens = append(dcTokens, token)
-
-			hosts = append(hosts, h)
-		}
-
-		sort.Sort(&tokenRing{tokens: dcTokens})
-		dcRing[dc] = dcTokens
-	}
-
-	if len(tokens) != hostsPerDC*totalDCs {
-		t.Fatalf("expected %d tokens in the ring got %d", hostsPerDC*totalDCs, len(tokens))
-	}
-	sort.Sort(&tokenRing{tokens: tokens})
-
-	strats := []*networkTopology{
-		&networkTopology{
-			dcs: map[string]int{
-				"dc1": 1,
-				"dc2": 2,
-				"dc3": 3,
+	tests := []struct{
+		name                   string
+		strat                  *networkTopology
+		expectedReplicaMapSize int
+	}{
+		{
+			name: "full",
+			strat: &networkTopology{
+				dcs: map[string]int{
+					"dc1": 1,
+					"dc2": 2,
+					"dc3": 3,
+				},
+			},
+			expectedReplicaMapSize: hostsPerDC*totalDCs,
+		},
+		{
+			name: "missing",
+			strat: &networkTopology{
+				dcs: map[string]int{
+					"dc2": 2,
+					"dc3": 3,
+				},
 			},
+			expectedReplicaMapSize: hostsPerDC*2,
 		},
-		&networkTopology{
-			dcs: map[string]int{
-				"dc2": 2,
-				"dc3": 3,
+		{
+			name: "zero",
+			strat: &networkTopology{
+				dcs: map[string]int{
+					"dc1": 0,
+					"dc2": 2,
+					"dc3": 3,
+				},
 			},
+			expectedReplicaMapSize: hostsPerDC*2,
 		},
 	}
 
-	for _, strat := range strats {
-		var expReplicas int
-		for _, rf := range strat.dcs {
-			expReplicas += rf
-		}
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			var (
+				hosts  []*HostInfo
+				tokens []hostToken
+			)
+			dcRing := make(map[string][]hostToken, totalDCs)
+			for i := 0; i < totalDCs; i++ {
+				var dcTokens []hostToken
+				dc := fmt.Sprintf("dc%d", i+1)
+
+				for j := 0; j < hostsPerDC; j++ {
+					rack := fmt.Sprintf("rack%d", (j%racksPerDC)+1)
+
+					h := &HostInfo{hostId: fmt.Sprintf("%s:%s:%d", dc, rack, j), dataCenter: dc, rack: rack}
+
+					token := hostToken{
+						token: orderedToken([]byte(h.hostId)),
+						host:  h,
+					}
 
-		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")
-		}
+					tokens = append(tokens, token)
+					dcTokens = append(dcTokens, token)
+
+					hosts = append(hosts, h)
+				}
 
-		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)
+				sort.Sort(&tokenRing{tokens: dcTokens})
+				dcRing[dc] = dcTokens
 			}
-		}
 
-		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(tokens) != hostsPerDC*totalDCs {
+				t.Fatalf("expected %d tokens in the ring got %d", hostsPerDC*totalDCs, len(tokens))
+			}
+			sort.Sort(&tokenRing{tokens: tokens})
 
-				var replicas []*HostInfo
-				for _, replica := range allReplicas.hosts {
-					if replica.dataCenter == dc {
-						replicas = append(replicas, replica)
-					}
+			var expReplicas int
+			for _, rf := range test.strat.dcs {
+				expReplicas += rf
+			}
+
+			tokenReplicas := test.strat.replicaMap(&tokenRing{hosts: hosts, tokens: tokens})
+			if len(tokenReplicas) != test.expectedReplicaMapSize {
+				t.Fatalf("expected replica map to have %d items but has %d", test.expectedReplicaMapSize,
+					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 {
+					t.Fatalf("expected to have %d replicas got %d for token=%v", expReplicas, len(replicas.hosts), token)
 				}
+			}
 
-				if len(replicas) != rf {
-					t.Fatalf("expected %d replicas in dc %q got %d", rf, dc, len(replicas))
+			for dc, rf := range test.strat.dcs {
+				if rf == 0 {
+					continue
 				}
+				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)
+					}
 
-				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 replicas []*HostInfo
+					for _, replica := range allReplicas.hosts {
+						if replica.dataCenter == dc {
+							replicas = append(replicas, replica)
 						}
-						if exp.rack == lastRack {
-							panic("no more racks")
+					}
+
+					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
+								}
+							}
+							if exp.rack == lastRack {
+								panic("no more racks")
+							}
 						}
+						lastRack = replica.rack
 					}
-					lastRack = replica.rack
 				}
 			}
-		}
+		})
 	}
 }