lightning: reuse dupDetector to implement in-memory ingest data (#46270)

ref #45719

br/pkg/lightning/backend/external/BUILD.bazel

@@ -19,14 +19,17 @@ go_library(
  "//br/pkg/lightning/backend",
  "//br/pkg/lightning/backend/encode",
  "//br/pkg/lightning/backend/kv",
+ "//br/pkg/lightning/backend/local",
  "//br/pkg/lightning/common",
+ "//br/pkg/lightning/log",
  "//br/pkg/membuf",
  "//br/pkg/storage",
  "//kv",
  "//util/hack",
  "//util/logutil",
  "//util/mathutil",
  "//util/size",
+ "@com_github_cockroachdb_pebble//:pebble",
  "@com_github_pingcap_errors//:errors",
  "@org_golang_x_sync//errgroup",
  "@org_uber_go_zap//:zap",
@@ -47,11 +50,14 @@ go_test(
  ],
  embed = [":external"],
  flaky = True,
- shard_count = 19,
+ shard_count = 20,
  deps = [
  "//br/pkg/lightning/backend/kv",
+ "//br/pkg/lightning/backend/local",
  "//br/pkg/lightning/common",
  "//br/pkg/storage",
+ "//util/codec",
+ "@com_github_cockroachdb_pebble//:pebble",
  "@com_github_pingcap_errors//:errors",
  "@com_github_stretchr_testify//require",
  "@org_golang_x_exp//rand",

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

@@ -15,10 +15,15 @@
 package external
 
 import (
+ "bytes"
  "context"
  "encoding/hex"
+ "sort"
 
+ "github.com/cockroachdb/pebble"
  "github.com/pingcap/errors"
+ "github.com/pingcap/tidb/br/pkg/lightning/backend/local"
+ "github.com/pingcap/tidb/br/pkg/lightning/log"
  "github.com/pingcap/tidb/br/pkg/storage"
  "github.com/pingcap/tidb/kv"
  "github.com/pingcap/tidb/util/logutil"
@@ -58,3 +63,200 @@ func (e *Engine) createMergeIter(ctx context.Context, start kv.Key) (*MergeKVIte
  }
  return iter, nil
 }
+
+// MemoryIngestData is the in-memory implementation of IngestData.
+type MemoryIngestData struct {
+ keyAdapter local.KeyAdapter
+ duplicateDetection bool
+ duplicateDB *pebble.DB
+ dupDetectOpt local.DupDetectOpt
+
+ keys [][]byte
+ values [][]byte
+ ts uint64
+}
+
+var _ local.IngestData = (*MemoryIngestData)(nil)
+
+func (m *MemoryIngestData) firstAndLastKeyIndex(lowerBound, upperBound []byte) (int, int) {
+ firstKeyIdx := 0
+ if len(lowerBound) > 0 {
+ lowerBound = m.keyAdapter.Encode(nil, lowerBound, local.MinRowID)
+ firstKeyIdx = sort.Search(len(m.keys), func(i int) bool {
+ return bytes.Compare(lowerBound, m.keys[i]) <= 0
+ })
+ if firstKeyIdx == len(m.keys) {
+ return -1, -1
+ }
+ }
+
+ lastKeyIdx := len(m.keys) - 1
+ if len(upperBound) > 0 {
+ upperBound = m.keyAdapter.Encode(nil, upperBound, local.MinRowID)
+ i := sort.Search(len(m.keys), func(i int) bool {
+ reverseIdx := len(m.keys) - 1 - i
+ return bytes.Compare(upperBound, m.keys[reverseIdx]) > 0
+ })
+ if i == len(m.keys) {
+ // should not happen
+ return -1, -1
+ }
+ lastKeyIdx = len(m.keys) - 1 - i
+ }
+ return firstKeyIdx, lastKeyIdx
+}
+
+// GetFirstAndLastKey implements IngestData.GetFirstAndLastKey.
+func (m *MemoryIngestData) GetFirstAndLastKey(lowerBound, upperBound []byte) ([]byte, []byte, error) {
+ firstKeyIdx, lastKeyIdx := m.firstAndLastKeyIndex(lowerBound, upperBound)
+ if firstKeyIdx < 0 || firstKeyIdx > lastKeyIdx {
+ return nil, nil, nil
+ }
+ firstKey, err := m.keyAdapter.Decode(nil, m.keys[firstKeyIdx])
+ if err != nil {
+ return nil, nil, err
+ }
+ lastKey, err := m.keyAdapter.Decode(nil, m.keys[lastKeyIdx])
+ if err != nil {
+ return nil, nil, err
+ }
+ return firstKey, lastKey, nil
+}
+
+type memoryDataIter struct {
+ keys [][]byte
+ values [][]byte
+
+ firstKeyIdx int
+ lastKeyIdx int
+ curIdx int
+}
+
+// First implements ForwardIter.
+func (m *memoryDataIter) First() bool {
+ if m.firstKeyIdx < 0 {
+ return false
+ }
+ m.curIdx = m.firstKeyIdx
+ return true
+}
+
+// Valid implements ForwardIter.
+func (m *memoryDataIter) Valid() bool {
+ return m.firstKeyIdx <= m.curIdx && m.curIdx <= m.lastKeyIdx
+}
+
+// Next implements ForwardIter.
+func (m *memoryDataIter) Next() bool {
+ m.curIdx++
+ return m.Valid()
+}
+
+// Key implements ForwardIter.
+func (m *memoryDataIter) Key() []byte {
+ return m.keys[m.curIdx]
+}
+
+// Value implements ForwardIter.
+func (m *memoryDataIter) Value() []byte {
+ return m.values[m.curIdx]
+}
+
+// Close implements ForwardIter.
+func (m *memoryDataIter) Close() error {
+ return nil
+}
+
+// Error implements ForwardIter.
+func (m *memoryDataIter) Error() error {
+ return nil
+}
+
+type memoryDataDupDetectIter struct {
+ iter *memoryDataIter
+ dupDetector *local.DupDetector
+ err error
+ curKey, curVal []byte
+}
+
+// First implements ForwardIter.
+func (m *memoryDataDupDetectIter) First() bool {
+ if m.err != nil || !m.iter.First() {
+ return false
+ }
+ m.curKey, m.curVal, m.err = m.dupDetector.Init(m.iter)
+ return m.Valid()
+}
+
+// Valid implements ForwardIter.
+func (m *memoryDataDupDetectIter) Valid() bool {
+ return m.err == nil && m.iter.Valid()
+}
+
+// Next implements ForwardIter.
+func (m *memoryDataDupDetectIter) Next() bool {
+ if m.err != nil {
+ return false
+ }
+ key, val, ok, err := m.dupDetector.Next(m.iter)
+ if err != nil {
+ m.err = err
+ return false
+ }
+ if !ok {
+ return false
+ }
+ m.curKey, m.curVal = key, val
+ return true
+}
+
+// Key implements ForwardIter.
+func (m *memoryDataDupDetectIter) Key() []byte {
+ return m.curKey
+}
+
+// Value implements ForwardIter.
+func (m *memoryDataDupDetectIter) Value() []byte {
+ return m.curVal
+}
+
+// Close implements ForwardIter.
+func (m *memoryDataDupDetectIter) Close() error {
+ return m.dupDetector.Close()
+}
+
+// Error implements ForwardIter.
+func (m *memoryDataDupDetectIter) Error() error {
+ return m.err
+}
+
+// NewIter implements IngestData.NewIter.
+func (m *MemoryIngestData) NewIter(ctx context.Context, lowerBound, upperBound []byte) local.ForwardIter {
+ firstKeyIdx, lastKeyIdx := m.firstAndLastKeyIndex(lowerBound, upperBound)
+ iter := &memoryDataIter{
+ keys: m.keys,
+ values: m.values,
+ firstKeyIdx: firstKeyIdx,
+ lastKeyIdx: lastKeyIdx,
+ }
+ if !m.duplicateDetection {
+ return iter
+ }
+ logger := log.FromContext(ctx)
+ detector := local.NewDupDetector(m.keyAdapter, m.duplicateDB.NewBatch(), logger, m.dupDetectOpt)
+ return &memoryDataDupDetectIter{
+ iter: iter,
+ dupDetector: detector,
+ }
+}
+
+// GetTS implements IngestData.GetTS.
+func (m *MemoryIngestData) GetTS() uint64 {
+ return m.ts
+}
+
+// Finish implements IngestData.Finish.
+func (m *MemoryIngestData) Finish(totalBytes, totalCount int64) {
+ //TODO implement me
+ panic("implement me")
+}