refine code, add more tests.

pingcap · Nov 22, 2023 · 9b74653 · 9b74653
1 parent 624d538
commit 9b74653
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 16 deletions.
diff --git a/br/pkg/lightning/backend/external/bench_test.go b/br/pkg/lightning/backend/external/bench_test.go
@@ -233,18 +233,18 @@ func TestCompareWriter(t *testing.T) {
 		afterWriterClose:   afterClose,
 	}
 
-	// writePlainFile(suite)
-	// baseSpeed := float64(source.outputSize()) / elapsed.Seconds() / 1024 / 1024
-	// t.Logf("base speed for %d bytes: %.2f MB/s", source.outputSize(), baseSpeed)
+	writePlainFile(suite)
+	baseSpeed := float64(source.outputSize()) / elapsed.Seconds() / 1024 / 1024
+	t.Logf("base speed for %d bytes: %.2f MB/s", source.outputSize(), baseSpeed)
 
-	// suite.source = newAscendingKeySource(20, 100, 10000000)
-	// writeExternalFile(suite)
-	// writerSpeed := float64(source.outputSize()) / elapsed.Seconds() / 1024 / 1024
-	// t.Logf("writer speed for %d bytes: %.2f MB/s", source.outputSize(), writerSpeed)
+	suite.source = newAscendingKeySource(20, 100, 10000000)
+	writeExternalFile(suite)
+	writerSpeed := float64(source.outputSize()) / elapsed.Seconds() / 1024 / 1024
+	t.Logf("writer speed for %d bytes: %.2f MB/s", source.outputSize(), writerSpeed)
 
-	suite.source = newAscendingKeySource(20, 100, 10000)
+	suite.source = newAscendingKeySource(20, 100, 10000000)
 	writeExternalOneFile(suite)
-	writerSpeed := float64(source.outputSize()) / elapsed.Seconds() / 1024 / 1024
+	writerSpeed = float64(source.outputSize()) / elapsed.Seconds() / 1024 / 1024
 	t.Logf("one file writer speed for %d bytes: %.2f MB/s", source.outputSize(), writerSpeed)
 }
 

diff --git a/br/pkg/lightning/backend/external/merge.go b/br/pkg/lightning/backend/external/merge.go
@@ -202,10 +202,7 @@ func MergeOverlappingFilesV2(
 		return err
 	}
 	maxKey = tidbkv.Key(iter.Key()).Clone()
-	err = writer.Close(ctx)
-	if err != nil {
-		return err
-	}
+
 	var stat MultipleFilesStat
 	stat.Filenames = append(stat.Filenames,
 		[2]string{writer.dataFile, writer.statFile})
@@ -220,5 +217,10 @@ func MergeOverlappingFilesV2(
 			MultipleFilesStats: []MultipleFilesStat{stat},
 		})
 	}
+
+	err = writer.Close(ctx)
+	if err != nil {
+		return err
+	}
 	return nil
 }
diff --git a/br/pkg/lightning/backend/external/onefile_writer.go b/br/pkg/lightning/backend/external/onefile_writer.go
@@ -90,7 +90,6 @@ func (w *OneFileWriter) WriteRow(ctx context.Context, idxKey, idxVal []byte) err
 	_, buf, _, allocated := w.kvBuffer.Alloc(length)
 	if !allocated {
 		w.kvBuffer.reset()
-		logutil.BgLogger().Info("ywq test here...")
 		_, buf, _, allocated = w.kvBuffer.Alloc(length)
 		// we now don't support KV larger than blockSize
 		if !allocated {
@@ -110,7 +109,7 @@ func (w *OneFileWriter) WriteRow(ctx context.Context, idxKey, idxVal []byte) err
 	binary.BigEndian.AppendUint64(buf[lengthBytes+keyLen:lengthBytes+keyLen], uint64(len(idxVal)))
 	copy(buf[lengthBytes*2+keyLen:], idxVal)
 	w.kvStore.addEncodedData(buf[:length])
-	w.totalSize = uint64(keyLen + len(idxVal))
+	w.totalSize += uint64(keyLen + len(idxVal))
 	return nil
 }
 

diff --git a/br/pkg/lightning/backend/external/onefile_writer_test.go b/br/pkg/lightning/backend/external/onefile_writer_test.go
@@ -45,7 +45,6 @@ func TestOnefileWriterBasic(t *testing.T) {
 	// 1. write into one file.
 	// 2. read kv file and check result.
 	// 3. read stat file and check result.
-
 	writer := NewWriterBuilder().
 		SetPropSizeDistance(100).
 		SetPropKeysDistance(2).
@@ -281,16 +280,20 @@ func TestOnefileWriterManyRows(t *testing.T) {
 	require.NoError(t, err)
 
 	kvCnt := 100000
+	expectedTotalSize := 0
 	kvs := make([]common.KvPair, kvCnt)
 	for i := 0; i < kvCnt; i++ {
 		randLen := rand.Intn(10) + 1
 		kvs[i].Key = make([]byte, randLen)
 		_, err := rand.Read(kvs[i].Key)
+		expectedTotalSize += randLen
+
 		require.NoError(t, err)
 		randLen = rand.Intn(10) + 1
 		kvs[i].Val = make([]byte, randLen)
 		_, err = rand.Read(kvs[i].Val)
 		require.NoError(t, err)
+		expectedTotalSize += randLen
 	}
 
 	slices.SortFunc(kvs, func(i, j common.KvPair) int {
@@ -352,4 +355,5 @@ func TestOnefileWriterManyRows(t *testing.T) {
 	require.EqualValues(t, expected.MaxKey, resSummary.Max)
 	require.Equal(t, expected.Filenames, resSummary.MultipleFilesStats[0].Filenames)
 	require.Equal(t, expected.MaxOverlappingNum, resSummary.MultipleFilesStats[0].MaxOverlappingNum)
+	require.EqualValues(t, expectedTotalSize, resSummary.TotalSize)
 }
diff --git a/pkg/ddl/backfilling_merge_sort.go b/pkg/ddl/backfilling_merge_sort.go
@@ -81,6 +81,9 @@ func (m *mergeSortExecutor) RunSubtask(ctx context.Context, subtask *proto.Subta
 	m.subtaskSortedKVMeta = &external.SortedKVMeta{}
 	onClose := func(summary *external.WriterSummary) {
 		m.mu.Lock()
+		logutil.BgLogger().Error("ywq test summary",
+			zap.String("category", "ddl"),
+			zap.Any("summary", summary))
 		m.subtaskSortedKVMeta.MergeSummary(summary)
 		m.mu.Unlock()
 	}