external: make sure read after switchConcurrentMode will see EOF (#48232

) close #47587, close #48223
pingcap · Nov 2, 2023 · c416ca1 · c416ca1
1 parent 9defab4
commit c416ca1
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 2 deletions.
diff --git a/br/pkg/lightning/backend/external/BUILD.bazel b/br/pkg/lightning/backend/external/BUILD.bazel
@@ -64,13 +64,15 @@ go_test(
     ],
     embed = [":external"],
     flaky = True,
-    shard_count = 42,
+    shard_count = 43,
     deps = [
         "//br/pkg/lightning/backend/kv",
         "//br/pkg/lightning/common",
+        "//br/pkg/membuf",
         "//br/pkg/storage",
         "//pkg/kv",
         "//pkg/util/codec",
+        "//pkg/util/logutil",
         "//pkg/util/size",
         "@com_github_aws_aws_sdk_go//aws",
         "@com_github_aws_aws_sdk_go//aws/credentials",

diff --git a/br/pkg/lightning/backend/external/byte_reader.go b/br/pkg/lightning/backend/external/byte_reader.go
@@ -278,6 +278,8 @@ func (r *byteReader) reload() error {
 	if err != nil {
 		switch err {
 		case io.EOF:
+			// move curBufOffset so following read will also find EOF
+			r.curBufOffset = len(r.curBuf)
 			return err
 		case io.ErrUnexpectedEOF:
 			// The last batch.

diff --git a/br/pkg/lightning/backend/external/iter_test.go b/br/pkg/lightning/backend/external/iter_test.go
@@ -24,7 +24,9 @@ import (
 	"time"
 
 	"github.com/pingcap/tidb/br/pkg/lightning/common"
+	"github.com/pingcap/tidb/br/pkg/membuf"
 	"github.com/pingcap/tidb/br/pkg/storage"
+	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
 	"golang.org/x/exp/rand"
@@ -326,6 +328,43 @@ func testMergeIterSwitchMode(t *testing.T, f func([]byte, int) []byte) {
 	require.NoError(t, err)
 }
 
+type eofReader struct {
+	storage.ExternalFileReader
+}
+
+func (r eofReader) Seek(_ int64, _ int) (int64, error) {
+	return 0, nil
+}
+
+func (r eofReader) Read(_ []byte) (int, error) {
+	return 0, io.EOF
+}
+
+func TestReadAfterCloseConnReader(t *testing.T) {
+	ctx := context.Background()
+
+	reader := &byteReader{
+		ctx:           ctx,
+		storageReader: eofReader{},
+		smallBuf:      []byte{0, 255, 255, 255, 255, 255, 255, 255},
+		curBufOffset:  8,
+		logger:        logutil.Logger(ctx),
+	}
+	reader.curBuf = reader.smallBuf
+	pool := membuf.NewPool()
+	reader.concurrentReader.largeBufferPool = pool.NewBuffer()
+	reader.concurrentReader.store = storage.NewMemStorage()
+
+	// set current reader to concurrent reader, and then close it
+	reader.concurrentReader.now = true
+	err := reader.switchConcurrentMode(false)
+	require.NoError(t, err)
+
+	wrapKVReader := &kvReader{reader}
+	_, _, err = wrapKVReader.nextKV()
+	require.ErrorIs(t, err, io.EOF)
+}
+
 func TestHotspot(t *testing.T) {
 	ctx := context.Background()
 	store := storage.NewMemStorage()

diff --git a/pkg/ddl/backfilling_merge_sort.go b/pkg/ddl/backfilling_merge_sort.go
@@ -78,9 +78,9 @@ func (m *mergeSortExecutor) RunSubtask(ctx context.Context, subtask *proto.Subta
 		return err
 	}
 
-	m.mu.Lock()
 	m.subtaskSortedKVMeta = &external.SortedKVMeta{}
 	onClose := func(summary *external.WriterSummary) {
+		m.mu.Lock()
 		m.subtaskSortedKVMeta.MergeSummary(summary)
 		m.mu.Unlock()
 	}