*: save TS into metadata for external engine, and prepare to use it as startTS

DEPS.bzl

@@ -5858,13 +5858,13 @@ def go_deps():
         name = "com_github_pingcap_kvproto",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/pingcap/kvproto",
-        sha256 = "bd7e32754f3923cd66a74d7b10f06aa46af512b81fca73f5c1a22286b4f563a3",
-        strip_prefix = "github.com/pingcap/[email protected]20240206021635-05a3758a1d24",
+        sha256 = "4670b3192a5c733e5c8fe68c28509da8294bd7a3e95cd1da151d45b66900f584",
+        strip_prefix = "github.com/pingcap/[email protected]20240223081609-df7de9808067",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240206021635-05a3758a1d24.zip",
-            "http://ats.apps.svc/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240206021635-05a3758a1d24.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240206021635-05a3758a1d24.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240206021635-05a3758a1d24.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240223081609-df7de9808067.zip",
+            "http://ats.apps.svc/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240223081609-df7de9808067.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240223081609-df7de9808067.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20240223081609-df7de9808067.zip",
         ],
     )
     go_repository(

br/pkg/lightning/backend/backend.go

@@ -118,6 +118,7 @@ type ExternalEngineConfig struct {
 	// TotalKVCount can be an estimated value.
 	TotalKVCount int64
 	CheckHotspot bool
+	TSOfClose    uint64
 }
 
 // CheckCtx contains all parameters used in CheckRequirements

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

@@ -707,8 +707,8 @@ func (m *MemoryIngestData) NewIter(
 	}
 }
 
-// GetTS implements IngestData.GetTS.
-func (m *MemoryIngestData) GetTS() uint64 {
+// GetStartTS implements IngestData.GetStartTS.
+func (m *MemoryIngestData) GetStartTS() uint64 {
 	return m.ts
 }
 

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

@@ -192,7 +192,7 @@ func TestMemoryIngestData(t *testing.T) {
 		ts:         123,
 	}
 
-	require.EqualValues(t, 123, data.GetTS())
+	require.EqualValues(t, 123, data.GetStartTS())
 	testGetFirstAndLastKey(t, data, nil, nil, []byte("key1"), []byte("key5"))
 	testGetFirstAndLastKey(t, data, []byte("key1"), []byte("key6"), []byte("key1"), []byte("key5"))
 	testGetFirstAndLastKey(t, data, []byte("key2"), []byte("key5"), []byte("key2"), []byte("key4"))
@@ -251,7 +251,7 @@ func TestMemoryIngestData(t *testing.T) {
 	data.keys = encodedKeys
 	data.values = encodedValues
 
-	require.EqualValues(t, 234, data.GetTS())
+	require.EqualValues(t, 234, data.GetStartTS())
 	testGetFirstAndLastKey(t, data, nil, nil, []byte("key1"), []byte("key5"))
 	testGetFirstAndLastKey(t, data, []byte("key1"), []byte("key6"), []byte("key1"), []byte("key5"))
 	testGetFirstAndLastKey(t, data, []byte("key2"), []byte("key5"), []byte("key2"), []byte("key4"))

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

@@ -78,7 +78,7 @@ const (
 // engineMeta contains some field that is necessary to continue the engine restore/import process.
 // These field should be written to disk when we update chunk checkpoint
 type engineMeta struct {
-	TS uint64 `json:"ts"`
+	StartTS uint64 `json:"ts"`
 	// Length is the number of KV pairs stored by the engine.
 	Length atomic.Int64 `json:"length"`
 	// TotalSize is the total pre-compressed KV byte size stored by engine.
@@ -1050,9 +1050,9 @@ func (e *Engine) NewIter(
 	)
 }
 
-// GetTS implements IngestData interface.
-func (e *Engine) GetTS() uint64 {
-	return e.TS
+// GetStartTS implements IngestData interface.
+func (e *Engine) GetStartTS() uint64 {
+	return e.StartTS
 }
 
 // IncRef implements IngestData interface.

br/pkg/lightning/backend/local/engine_mgr.go

@@ -25,7 +25,6 @@ import (
 	"github.com/docker/go-units"
 	"github.com/google/uuid"
 	"github.com/pingcap/errors"
-	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/br/pkg/lightning/backend"
 	"github.com/pingcap/tidb/br/pkg/lightning/backend/external"
 	"github.com/pingcap/tidb/br/pkg/lightning/common"
@@ -38,7 +37,6 @@ import (
 	"go.uber.org/atomic"
 	"go.uber.org/zap"
 	"golang.org/x/sync/errgroup"
-	"google.golang.org/grpc/status"
 )
 
 var (
@@ -290,11 +288,6 @@ func (em *engineManager) closeEngine(
 				store.Close()
 			}
 		}()
-		physical, logical, err := em.GetTS(ctx)
-		if err != nil {
-			return err
-		}
-		ts := oracle.ComposeTS(physical, logical)
 		externalEngine := external.NewExternalEngine(
 			store,
 			externalCfg.DataFiles,
@@ -308,7 +301,7 @@ func (em *engineManager) closeEngine(
 			em.duplicateDB,
 			em.DuplicateDetectOpt,
 			em.WorkerConcurrency,
-			ts,
+			externalCfg.TSOfClose,
 			externalCfg.TotalFileSize,
 			externalCfg.TotalKVCount,
 			externalCfg.CheckHotspot,
@@ -339,7 +332,7 @@ func (em *engineManager) closeEngine(
 		engine.db.Store(db)
 		engine.sstIngester = dbSSTIngester{e: engine}
 		if err = engine.loadEngineMeta(); err != nil {
-			return err
+			return errors.Trace(err)
 		}
 		em.engines.Store(engineUUID, engine)
 		return nil
@@ -404,6 +397,7 @@ func (em *engineManager) resetEngine(ctx context.Context, engineUUID uuid.UUID)
 		return nil
 	}
 	defer localEngine.unlock()
+	oldStartTS := localEngine.StartTS
 	if err := localEngine.Close(); err != nil {
 		return err
 	}
@@ -413,37 +407,31 @@ func (em *engineManager) resetEngine(ctx context.Context, engineUUID uuid.UUID)
 	db, err := em.openEngineDB(engineUUID, false)
 	if err == nil {
 		localEngine.db.Store(db)
-		localEngine.engineMeta = engineMeta{}
+		localEngine.engineMeta = engineMeta{StartTS: oldStartTS}
+		if err = localEngine.saveEngineMeta(); err != nil {
+			return errors.Trace(err)
+		}
 		if !common.IsDirExists(localEngine.sstDir) {
 			if err := os.Mkdir(localEngine.sstDir, 0o750); err != nil {
 				return errors.Trace(err)
 			}
 		}
-		if err = em.allocateTSIfNotExists(ctx, localEngine); err != nil {
-			return errors.Trace(err)
-		}
-		failpoint.Inject("mockAllocateTSErr", func() {
-			// mock generate timestamp error when reset engine.
-			localEngine.TS = 0
-			mockGRPCErr, _ := status.FromError(errors.Errorf("mock generate timestamp error"))
-			failpoint.Return(errors.Trace(mockGRPCErr.Err()))
-		})
 	}
 	localEngine.pendingFileSize.Store(0)
 
 	return err
 }
 
 func (em *engineManager) allocateTSIfNotExists(ctx context.Context, engine *Engine) error {
-	if engine.TS > 0 {
+	if engine.StartTS > 0 {
 		return nil
 	}
 	physical, logical, err := em.GetTS(ctx)
 	if err != nil {
 		return err
 	}
 	ts := oracle.ComposeTS(physical, logical)
-	engine.TS = ts
+	engine.StartTS = ts
 	return engine.saveEngineMeta()
 }
 

br/pkg/lightning/backend/local/engine_mgr_test.go

@@ -90,7 +90,6 @@ func TestEngineManager(t *testing.T) {
 
 	// reset non-existent engine should work
 	require.NoError(t, em.resetEngine(ctx, uuid.New()))
-	storeHelper.EXPECT().GetTS(gomock.Any()).Return(int64(0), int64(0), nil)
 	require.NoError(t, em.resetEngine(ctx, engine1ID))
 	require.Equal(t, 1, syncMapLen(&em.engines))
 	_, ok = em.engines.Load(engine1ID)

br/pkg/lightning/backend/local/local.go

@@ -508,9 +508,11 @@ func NewBackend(
 		pdAddrs = pdSvcDiscovery.GetServiceURLs()
 		// TODO(lance6716): if PD client can support creating a client with external
 		// service discovery, we can directly pass pdSvcDiscovery.
-	} else {
+	}
+	if len(pdAddrs) == 0 {
 		pdAddrs = strings.Split(config.PDAddr, ",")
 	}
+
 	pdCli, err := pd.NewClientWithContext(
 		ctx, pdAddrs, tls.ToPDSecurityOption(),
 		pd.WithGRPCDialOptions(maxCallMsgSize...),

br/pkg/lightning/backend/local/local_test.go

@@ -1208,7 +1208,7 @@ func (m mockIngestData) NewIter(_ context.Context, lowerBound, upperBound []byte
 	return &mockIngestIter{data: m, startIdx: i, endIdx: j, curIdx: i}
 }
 
-func (m mockIngestData) GetTS() uint64 { return 0 }
+func (m mockIngestData) GetStartTS() uint64 { return 0 }
 
 func (m mockIngestData) IncRef() {}
 

br/pkg/lightning/backend/local/region_job.go

@@ -258,6 +258,7 @@ func (local *Backend) doWrite(ctx context.Context, j *regionJob) error {
 			End:   lastKey,
 		},
 		ApiVersion: apiVersion,
+		StartTs:    j.ingestData.GetStartTS(),
 	}
 
 	failpoint.Inject("changeEpochVersion", func(val failpoint.Value) {
@@ -318,7 +319,8 @@ func (local *Backend) doWrite(ctx context.Context, j *regionJob) error {
 		clients = append(clients, wstream)
 		allPeers = append(allPeers, peer)
 	}
-	dataCommitTS := j.ingestData.GetTS()
+	dataCommitTS := j.ingestData.GetStartTS()
+	// TODO(lance6716): use local.pdCli.GetTS() to get a new TS as commit TS after deduplication is ready.
 	req.Chunk = &sst.WriteRequest_Batch{
 		Batch: &sst.WriteBatch{
 			CommitTs: dataCommitTS,

br/pkg/lightning/common/ingest_data.go

@@ -34,8 +34,11 @@ type IngestData interface {
 	// these memories will be allocated from given bufPool and be released when the
 	// iterator is closed or ForwardIter.ReleaseBuf is called.
 	NewIter(ctx context.Context, lowerBound, upperBound []byte, bufPool *membuf.Pool) ForwardIter
-	// GetTS will be used as the start/commit TS of the data.
-	GetTS() uint64
+	// GetStartTS return the startTS when ingest. This TS should keep unchanged
+	// during retries of ingest. For local engine, it's set when engine is opened to
+	// avoid change too much code. For external engine, it's set when engine is
+	// closed.
+	GetStartTS() uint64
 	// IncRef should be called every time when IngestData is referred by regionJob.
 	// Multiple regionJob can share one IngestData. Same amount of DecRef should be
 	// called to release the IngestData.

br/pkg/lightning/importer/import.go

@@ -1609,6 +1609,9 @@ func (rc *Controller) importTables(ctx context.Context) (finalErr error) {
 
 		// Disable GC because TiDB enables GC already.
 		urlsWithScheme := rc.pdCli.GetServiceDiscovery().GetServiceURLs()
+		if len(urlsWithScheme) == 0 {
+			urlsWithScheme = strings.Split(rc.cfg.TiDB.PdAddr, ",")
+		}
 		// remove URL scheme
 		urlsWithoutScheme := make([]string, 0, len(urlsWithScheme))
 		for _, u := range urlsWithScheme {
@@ -1840,7 +1843,11 @@ func (rc *Controller) importTables(ctx context.Context) (finalErr error) {
 }
 
 func (rc *Controller) registerTaskToPD(ctx context.Context) (undo func(), _ error) {
-	etcdCli, err := dialEtcdWithCfg(ctx, rc.cfg, rc.pdCli.GetServiceDiscovery().GetServiceURLs())
+	pdAddr := rc.pdCli.GetServiceDiscovery().GetServiceURLs()
+	if len(pdAddr) == 0 {
+		pdAddr = strings.Split(rc.cfg.TiDB.PdAddr, ",")
+	}
+	etcdCli, err := dialEtcdWithCfg(ctx, rc.cfg, pdAddr)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -2142,6 +2149,9 @@ func (rc *Controller) preCheckRequirements(ctx context.Context) error {
 	}
 	if isLocalBackend(rc.cfg) {
 		pdAddrs := rc.pdCli.GetServiceDiscovery().GetServiceURLs()
+		if len(pdAddrs) == 0 {
+			pdAddrs = strings.Split(rc.cfg.TiDB.PdAddr, ",")
+		}
 		pdController, err := pdutil.NewPdController(
 			ctx, pdAddrs, rc.tls.TLSConfig(), rc.tls.ToPDSecurityOption(),
 		)

go.mod

@@ -316,6 +316,8 @@ replace (
 	// fix potential security issue(CVE-2020-26160) introduced by indirect dependency.
 	github.com/dgrijalva/jwt-go => github.com/form3tech-oss/jwt-go v3.2.6-0.20210809144907-32ab6a8243d7+incompatible
 	github.com/go-ldap/ldap/v3 => github.com/YangKeao/ldap/v3 v3.4.5-0.20230421065457-369a3bab1117
+	// remove it after https://github.com/pingcap/kvproto/pull/1224 is merged
+	github.com/pingcap/kvproto v0.0.0-20240206021635-05a3758a1d24 => github.com/pingcap/kvproto v0.0.0-20240223081609-df7de9808067
 	github.com/pingcap/tidb/pkg/parser => ./pkg/parser
 
 	// TODO: `sourcegraph.com/sourcegraph/appdash` has been archived, and the original host has been removed.

go.sum

@@ -713,8 +713,8 @@ github.com/pingcap/fn v1.0.0/go.mod h1:u9WZ1ZiOD1RpNhcI42RucFh/lBuzTu6rw88a+oF2Z
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
-github.com/pingcap/kvproto v0.0.0-20240206021635-05a3758a1d24 h1:pdA3DvkChrIp91JQO89ICT1x/SemOAm7vC848acr5Ik=
-github.com/pingcap/kvproto v0.0.0-20240206021635-05a3758a1d24/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20240223081609-df7de9808067 h1:K0YStutk7w3eMuYljP+rpLOCI/XHhLzRF4P+8CJiBZQ=
+github.com/pingcap/kvproto v0.0.0-20240223081609-df7de9808067/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
 github.com/pingcap/log v1.1.0/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pingcap/log v1.1.1-0.20230317032135-a0d097d16e22 h1:2SOzvGvE8beiC1Y4g9Onkvu6UmuBBOeWRGQEjJaT/JY=

pkg/ddl/backfilling_dist_executor.go

@@ -61,6 +61,7 @@ type BackfillSubTaskMeta struct {
 	RangeSplitKeys [][]byte `json:"range_split_keys,omitempty"`
 	DataFiles      []string `json:"data-files,omitempty"`
 	StatFiles      []string `json:"stat-files,omitempty"`
+	TSOfClose      uint64   `json:"ts_of_close,omitempty"`
 	// Each group of MetaGroups represents a different index kvs meta.
 	MetaGroups []*external.SortedKVMeta `json:"meta_groups,omitempty"`
 	// Only used for adding one single index.

pkg/ddl/backfilling_dist_scheduler.go

@@ -41,6 +41,7 @@ import (
 	"github.com/pingcap/tidb/pkg/table"
 	"github.com/pingcap/tidb/pkg/util/backoff"
 	"github.com/pingcap/tidb/pkg/util/logutil"
+	"github.com/tikv/client-go/v2/oracle"
 	"github.com/tikv/client-go/v2/tikv"
 	"go.uber.org/zap"
 )
@@ -419,6 +420,15 @@ func splitSubtaskMetaForOneKVMetaGroup(
 		// Skip global sort for empty table.
 		return nil, nil
 	}
+	pdCli := store.GetPDClient()
+	physical, logical, err := pdCli.GetTS(ctx)
+	if err != nil {
+		return nil, err
+	}
+	failpoint.Inject("mockPhysicalTSForGlobalSort", func(val failpoint.Value) {
+		i := val.(int)
+		physical = int64(i)
+	})
 	splitter, err := getRangeSplitter(
 		ctx, store, cloudStorageURI, int64(kvMeta.TotalKVSize), instanceCnt, kvMeta.MultipleFilesStats, logger)
 	if err != nil {
@@ -451,6 +461,8 @@ func splitSubtaskMetaForOneKVMetaGroup(
 			return nil, errors.Errorf("invalid range, startKey: %s, endKey: %s",
 				hex.EncodeToString(startKey), hex.EncodeToString(endKey))
 		}
+		ts := oracle.ComposeTS(physical, logical)
+		logical++
 		m := &BackfillSubTaskMeta{
 			MetaGroups: []*external.SortedKVMeta{{
 				StartKey:    startKey,
@@ -460,6 +472,7 @@ func splitSubtaskMetaForOneKVMetaGroup(
 			DataFiles:      dataFiles,
 			StatFiles:      statFiles,
 			RangeSplitKeys: rangeSplitKeys,
+			TSOfClose:      ts,
 		}
 		metaBytes, err := json.Marshal(m)
 		if err != nil {

pkg/ddl/backfilling_import_cloud.go

@@ -93,6 +93,7 @@ func (m *cloudImportExecutor) RunSubtask(ctx context.Context, subtask *proto.Sub
 			TotalFileSize: int64(all.TotalKVSize),
 			TotalKVCount:  0,
 			CheckHotspot:  true,
+			TSOfClose:     sm.TSOfClose,
 		},
 	}, engineUUID)
 	if err != nil {

pkg/disttask/importinto/BUILD.bazel

@@ -30,6 +30,7 @@ go_library(
         "//br/pkg/logutil",
         "//br/pkg/storage",
         "//br/pkg/utils",
+        "//pkg/config",
         "//pkg/disttask/framework/handle",
         "//pkg/disttask/framework/planner",
         "//pkg/disttask/framework/proto",
@@ -66,7 +67,9 @@ go_library(
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_prometheus_client_golang//prometheus",
+        "@com_github_tikv_client_go_v2//oracle",
         "@com_github_tikv_client_go_v2//util",
+        "@com_github_tikv_pd_client//:client",
         "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
         "@org_uber_go_zap//zapcore",