Fix the panic when merging empty RegionsInfo

Signed-off-by: JmPotato <[email protected]>
tikv · Mar 22, 2024 · 80bba01 · 80bba01
1 parent 52e8763
commit 80bba01
Show file tree

Hide file tree

Showing 2 changed files with 108 additions and 21 deletions.
diff --git a/client/http/types.go b/client/http/types.go
@@ -133,11 +133,22 @@ type RegionsInfo struct {
 	Regions []RegionInfo `json:"regions"`
 }
 
+func newRegionsInfo(count int64) *RegionsInfo {
+	return &RegionsInfo{
+		Count:   count,
+		Regions: make([]RegionInfo, 0, count),
+	}
+}
+
 // Merge merges two RegionsInfo together and returns a new one.
 func (ri *RegionsInfo) Merge(other *RegionsInfo) *RegionsInfo {
-	newRegionsInfo := &RegionsInfo{
-		Regions: make([]RegionInfo, 0, ri.Count+other.Count),
+	if ri == nil {
+		ri = newRegionsInfo(0)
+	}
+	if other == nil {
+		other = newRegionsInfo(0)
 	}
+	newRegionsInfo := newRegionsInfo(ri.Count + other.Count)
 	m := make(map[int64]RegionInfo, ri.Count+other.Count)
 	for _, region := range ri.Regions {
 		m[region.ID] = region

diff --git a/client/http/types_test.go b/client/http/types_test.go
@@ -23,30 +23,106 @@ import (
 
 func TestMergeRegionsInfo(t *testing.T) {
 	re := require.New(t)
-	regionsInfo1 := &RegionsInfo{
-		Count: 1,
-		Regions: []RegionInfo{
-			{
-				ID:       1,
-				StartKey: "",
-				EndKey:   "a",
+	testCases := []struct {
+		source *RegionsInfo
+		target *RegionsInfo
+	}{
+		{
+			source: &RegionsInfo{
+				Count: 1,
+				Regions: []RegionInfo{
+					{
+						ID:       1,
+						StartKey: "",
+						EndKey:   "a",
+					},
+				},
+			},
+			target: &RegionsInfo{
+				Count: 1,
+				Regions: []RegionInfo{
+					{
+						ID:       2,
+						StartKey: "a",
+						EndKey:   "",
+					},
+				},
 			},
 		},
-	}
-	regionsInfo2 := &RegionsInfo{
-		Count: 1,
-		Regions: []RegionInfo{
-			{
-				ID:       2,
-				StartKey: "a",
-				EndKey:   "",
+		{
+			source: &RegionsInfo{
+				Count: 1,
+				Regions: []RegionInfo{
+					{
+						ID:       1,
+						StartKey: "",
+						EndKey:   "a",
+					},
+				},
+			},
+			target: nil,
+		},
+		{
+			source: nil,
+			target: &RegionsInfo{
+				Count: 1,
+				Regions: []RegionInfo{
+					{
+						ID:       2,
+						StartKey: "a",
+						EndKey:   "",
+					},
+				},
 			},
 		},
+		{
+			source: nil,
+			target: nil,
+		},
+		{
+			source: &RegionsInfo{
+				Count: 1,
+				Regions: []RegionInfo{
+					{
+						ID:       1,
+						StartKey: "",
+						EndKey:   "a",
+					},
+				},
+			},
+			target: newRegionsInfo(0),
+		},
+		{
+			source: newRegionsInfo(0),
+			target: &RegionsInfo{
+				Count: 1,
+				Regions: []RegionInfo{
+					{
+						ID:       2,
+						StartKey: "a",
+						EndKey:   "",
+					},
+				},
+			},
+		},
+		{
+			source: newRegionsInfo(0),
+			target: newRegionsInfo(0),
+		},
+	}
+	for idx, tc := range testCases {
+		regionsInfo := tc.source.Merge(tc.target)
+		if tc.source == nil {
+			tc.source = newRegionsInfo(0)
+		}
+		if tc.target == nil {
+			tc.target = newRegionsInfo(0)
+		}
+		mergedCount := tc.source.Count + tc.target.Count
+		re.Equal(mergedCount, regionsInfo.Count, "case %d", idx)
+		re.Len(regionsInfo.Regions, int(mergedCount), "case %d", idx)
+		re.Subset(regionsInfo.Regions, append(tc.source.Regions, tc.target.Regions...), "case %d", idx)
 	}
-	regionsInfo := regionsInfo1.Merge(regionsInfo2)
-	re.Equal(int64(2), regionsInfo.Count)
-	re.Len(regionsInfo.Regions, 2)
-	re.Subset(regionsInfo.Regions, append(regionsInfo1.Regions, regionsInfo2.Regions...))
 }
 
 func TestRuleStartEndKey(t *testing.T) {