external: set the period of checking merge iter hotspot according to the key size

br/pkg/lightning/backend/external/iter.go

@@ -81,15 +81,16 @@ func (h *mergeHeap[T]) Pop() interface{} {
 }
 
 type mergeIter[T heapElem, R sortedReader[T]] struct {
-	h               mergeHeap[T]
-	readers         []*R
-	curr            T
-	lastReaderIdx   int
-	err             error
-	hotspotMap      map[int]int
-	checkHotspotCnt int
-	lastHotspotIdx  int
-	elemFromHotspot *T
+	h                  mergeHeap[T]
+	readers            []*R
+	curr               T
+	lastReaderIdx      int
+	err                error
+	hotspotMap         map[int]int
+	checkHotspotCnt    int
+	checkHotspotPeriod int
+	lastHotspotIdx     int
+	elemFromHotspot    *T
 
 	logger *zap.Logger
 }
@@ -151,6 +152,8 @@ func newMergeIter[
 		hotspotMap:    make(map[int]int),
 		logger:        logger,
 	}
+	sampleKeySize := 0
+	sampleKeyCnt := 0
 	for j := range i.readers {
 		if i.readers[j] == nil {
 			continue
@@ -179,6 +182,15 @@ func newMergeIter[
 			elem:      e,
 			readerIdx: j,
 		})
+		sampleKeySize += len(e.sortKey())
+		sampleKeyCnt++
+	}
+	// We check the hotspot when the elements size is almost the same as the concurrent reader buffer size.
+	// So that we don't drop too many bytes if the hotspot shifts to other files.
+	if sampleKeySize == 0 || sampleKeySize/sampleKeyCnt == 0 {
+		i.checkHotspotPeriod = 10000
+	} else {
+		i.checkHotspotPeriod = max(1000, ConcurrentReaderBufferSizePerConc*ConcurrentReaderConcurrency/(sampleKeySize/sampleKeyCnt))
 	}
 	heap.Init(&i.h)
 	return i, nil
@@ -210,8 +222,6 @@ func (i *mergeIter[T, R]) currElem() T {
 	return i.curr
 }
 
-var checkHotspotPeriod = 1000
-
 // next forwards the iterator to the next element. It returns false if there is
 // no available element.
 func (i *mergeIter[T, R]) next() bool {
@@ -222,12 +232,12 @@ func (i *mergeIter[T, R]) next() bool {
 		i.checkHotspotCnt++
 
 		// check hotspot every checkPeriod times
-		if i.checkHotspotCnt == checkHotspotPeriod {
+		if i.checkHotspotCnt == i.checkHotspotPeriod {
 			oldHotspotIdx := i.lastHotspotIdx
 			i.lastHotspotIdx = -1
 			for idx, cnt := range i.hotspotMap {
 				// currently only one reader will become hotspot
-				if cnt > (checkHotspotPeriod / 2) {
+				if cnt > (i.checkHotspotPeriod / 2) {
 					i.lastHotspotIdx = idx
 					break
 				}

br/pkg/lightning/backend/external/iter_test.go

@@ -327,12 +327,6 @@ func testMergeIterSwitchMode(t *testing.T, f func([]byte, int) []byte) {
 }
 
 func TestHotspot(t *testing.T) {
-	backup := checkHotspotPeriod
-	checkHotspotPeriod = 2
-	t.Cleanup(func() {
-		checkHotspotPeriod = backup
-	})
-
 	ctx := context.Background()
 	store := storage.NewMemStorage()
 
@@ -364,6 +358,7 @@ func TestHotspot(t *testing.T) {
 	// readerBufSize = 8+5+8+5, every KV will cause reload
 	iter, err := NewMergeKVIter(ctx, filenames, make([]uint64, len(filenames)), store, 26)
 	require.NoError(t, err)
+	iter.iter.checkHotspotPeriod = 2
 	// after read key00 and key01 from reader_0, it becomes hotspot
 	require.True(t, iter.Next())
 	require.Equal(t, "key00", string(iter.Key()))
@@ -413,15 +408,10 @@ func TestHotspot(t *testing.T) {
 }
 
 func TestMemoryUsageWhenHotspotChange(t *testing.T) {
-	backup := checkHotspotPeriod
-	checkHotspotPeriod = 10
-	t.Cleanup(func() {
-		checkHotspotPeriod = backup
-	})
-	backup2 := ConcurrentReaderBufferSizePerConc
+	backup := ConcurrentReaderBufferSizePerConc
 	ConcurrentReaderBufferSizePerConc = 100 * 1024 * 1024 // 100MB, make memory leak more obvious
 	t.Cleanup(func() {
-		ConcurrentReaderBufferSizePerConc = backup2
+		ConcurrentReaderBufferSizePerConc = backup
 	})
 
 	getMemoryInUse := func() uint64 {
@@ -452,7 +442,7 @@ func TestMemoryUsageWhenHotspotChange(t *testing.T) {
 		rc.reset()
 		kvStore, err := NewKeyValueStore(ctx, writer, rc)
 		require.NoError(t, err)
-		for j := 0; j < checkHotspotPeriod; j++ {
+		for j := 0; j < 1000; j++ {
 			key := fmt.Sprintf("key%06d", cur)
 			val := fmt.Sprintf("value%06d", cur)
 			err = kvStore.addEncodedData(getEncodedData([]byte(key), []byte(val)))
@@ -472,6 +462,7 @@ func TestMemoryUsageWhenHotspotChange(t *testing.T) {
 
 	iter, err := NewMergeKVIter(ctx, filenames, make([]uint64, len(filenames)), store, 1024)
 	require.NoError(t, err)
+	iter.iter.checkHotspotPeriod = 10
 	i := 0
 	for cur > 0 {
 		cur--