lightning: do not output pre-check info when disable check-requirements

br/pkg/lightning/config/const.go

@@ -21,7 +21,6 @@ import (
 const (
 	// mydumper
 	ReadBlockSize           ByteSize = 64 * units.KiB
-	MinRegionSize           ByteSize = 256 * units.MiB
 	MaxRegionSize           ByteSize = 256 * units.MiB
 	SplitRegionSize         ByteSize = 96 * units.MiB
 	MaxSplitRegionSizeRatio int      = 10

br/pkg/lightning/lightning.go

@@ -61,6 +61,7 @@ type Lightning struct {
 	server     http.Server
 	serverAddr net.Addr
 	serverLock sync.Mutex
+	status     restore.LightningStatus
 
 	cancelLock sync.Mutex
 	curTask    *config.Config
@@ -301,7 +302,8 @@ func (l *Lightning) run(taskCtx context.Context, taskCfg *config.Config, g glue.
 	web.BroadcastInitProgress(dbMetas)
 
 	var procedure *restore.Controller
-	procedure, err = restore.NewRestoreController(ctx, dbMetas, taskCfg, s, g)
+
+	procedure, err = restore.NewRestoreController(ctx, dbMetas, taskCfg, &l.status, s, g)
 	if err != nil {
 		log.L().Error("restore failed", log.ShortError(err))
 		return errors.Trace(err)
@@ -324,6 +326,13 @@ func (l *Lightning) Stop() {
 	l.shutdown()
 }
 
+// Status return the sum size of file which has been imported to TiKV and the total size of source file.
+func (l *Lightning) Status() (int64, int64) {
+	finished := l.status.FinishedFileSize.Load()
+	total := l.status.TotalFileSize.Load()
+	return finished, total
+}
+
 func writeJSONError(w http.ResponseWriter, code int, prefix string, err error) {
 	type errorResponse struct {
 		Error string `json:"error"`

br/pkg/lightning/log/log.go

@@ -116,6 +116,11 @@ func InitLogger(cfg *Config, tidbLoglevel string) error {
 	return nil
 }
 
+// SetAppLogger replaces the default logger in this package to given one
+func SetAppLogger(l *zap.Logger) {
+	appLogger = Logger{l.WithOptions(zap.AddStacktrace(zap.DPanicLevel))}
+}
+
 // L returns the current logger for Lightning.
 func L() Logger {
 	return appLogger

br/pkg/lightning/metric/metric.go

@@ -23,13 +23,10 @@ import (
 
 const (
 	// states used for the TableCounter labels
-	TableStatePending        = "pending"
-	TableStateWritten        = "written"
-	TableStateClosed         = "closed"
-	TableStateImported       = "imported"
-	TableStateAlteredAutoInc = "altered_auto_inc"
-	TableStateChecksum       = "checksum"
-	TableStateCompleted      = "completed"
+	TableStatePending   = "pending"
+	TableStateWritten   = "written"
+	TableStateImported  = "imported"
+	TableStateCompleted = "completed"
 
 	// results used for the TableCounter labels
 	TableResultSuccess = "success"

br/pkg/lightning/restore/check_info.go

@@ -92,8 +92,7 @@ func (rc *Controller) getClusterAvail(ctx context.Context) (uint64, error) {
 	return clusterAvail, nil
 }
 
-// ClusterResource check cluster has enough resource to import data. this test can by skipped.
-func (rc *Controller) ClusterResource(ctx context.Context, localSource int64) error {
+func (rc *Controller) clusterResource(ctx context.Context, localSource int64) error {
 	passed := true
 	message := "Cluster resources are rich for this import task"
 	defer func() {
@@ -314,8 +313,7 @@ func (rc *Controller) checkRegionDistribution(ctx context.Context) error {
 	return nil
 }
 
-// CheckClusterRegion checks cluster if there are too many empty regions or region distribution is unbalanced.
-func (rc *Controller) CheckClusterRegion(ctx context.Context) error {
+func (rc *Controller) checkClusterRegion(ctx context.Context) error {
 	err := rc.taskMgr.CheckTasksExclusively(ctx, func(tasks []taskMeta) ([]taskMeta, error) {
 		restoreStarted := false
 		for _, task := range tasks {
@@ -390,7 +388,7 @@ func (rc *Controller) HasLargeCSV(dbMetas []*mydump.MDDatabaseMeta) error {
 	return nil
 }
 
-func (rc *Controller) EstimateSourceData(ctx context.Context) (int64, error) {
+func (rc *Controller) estimateSourceData(ctx context.Context) (int64, error) {
 	sourceSize := int64(0)
 	originSource := int64(0)
 	bigTableCount := 0
@@ -404,21 +402,32 @@ func (rc *Controller) EstimateSourceData(ctx context.Context) (int64, error) {
 		for _, tbl := range db.Tables {
 			tableInfo, ok := info.Tables[tbl.Name]
 			if ok {
-				if err := rc.SampleDataFromTable(ctx, db.Name, tbl, tableInfo.Core); err != nil {
-					return sourceSize, errors.Trace(err)
-				}
-				sourceSize += int64(float64(tbl.TotalSize) * tbl.IndexRatio)
 				originSource += tbl.TotalSize
-				if tbl.TotalSize > int64(config.DefaultBatchSize)*2 {
-					bigTableCount += 1
-					if !tbl.IsRowOrdered {
-						unSortedTableCount += 1
+				// Do not sample small table because there may a large number of small table and it will take a long
+				// time to sample data for all of them.
+				if tbl.TotalSize < int64(config.SplitRegionSize) {
+					sourceSize += tbl.TotalSize
+					tbl.IndexRatio = 1.0
+					tbl.IsRowOrdered = false
+				} else {
+					if err := rc.sampleDataFromTable(ctx, db.Name, tbl, tableInfo.Core); err != nil {
+						return sourceSize, errors.Trace(err)
+					}
+					sourceSize += int64(float64(tbl.TotalSize) * tbl.IndexRatio)
+					if tbl.TotalSize > int64(config.DefaultBatchSize)*2 {
+						bigTableCount += 1
+						if !tbl.IsRowOrdered {
+							unSortedTableCount += 1
+						}
 					}
 				}
 				tableCount += 1
 			}
 		}
 	}
+	if rc.status != nil {
+		rc.status.TotalFileSize.Store(originSource)
+	}
 
 	// Do not import with too large concurrency because these data may be all unsorted.
 	if bigTableCount > 0 && unSortedTableCount > 0 {
@@ -429,8 +438,7 @@ func (rc *Controller) EstimateSourceData(ctx context.Context) (int64, error) {
 	return sourceSize, nil
 }
 
-// LocalResource checks the local node has enough resources for this import when local backend enabled;
-func (rc *Controller) LocalResource(ctx context.Context, sourceSize int64) error {
+func (rc *Controller) localResource(sourceSize int64) error {
 	if rc.isSourceInLocal() {
 		sourceDir := strings.TrimPrefix(rc.cfg.Mydumper.SourceDir, storage.LocalURIPrefix)
 		same, err := common.SameDisk(sourceDir, rc.cfg.TikvImporter.SortedKVDir)
@@ -449,9 +457,6 @@ func (rc *Controller) LocalResource(ctx context.Context, sourceSize int64) error
 		return errors.Trace(err)
 	}
 	localAvailable := storageSize.Available
-	if err = rc.taskMgr.InitTask(ctx, sourceSize); err != nil {
-		return errors.Trace(err)
-	}
 
 	var message string
 	var passed bool
@@ -732,7 +737,7 @@ func (rc *Controller) SchemaIsValid(ctx context.Context, tableInfo *mydump.MDTab
 	return msgs, nil
 }
 
-func (rc *Controller) SampleDataFromTable(ctx context.Context, dbName string, tableMeta *mydump.MDTableMeta, tableInfo *model.TableInfo) error {
+func (rc *Controller) sampleDataFromTable(ctx context.Context, dbName string, tableMeta *mydump.MDTableMeta, tableInfo *model.TableInfo) error {
 	if len(tableMeta.DataFiles) == 0 {
 		return nil
 	}

br/pkg/lightning/restore/check_template.go

@@ -82,7 +82,7 @@ func (c *SimpleTemplate) Collect(t CheckType, passed bool, msg string) {
 }
 
 func (c *SimpleTemplate) Success() bool {
-	return c.warnFailedCount+c.criticalFailedCount == 0
+	return c.criticalFailedCount == 0
 }
 
 func (c *SimpleTemplate) FailedCount(t CheckType) int {

br/pkg/lightning/restore/restore.go

@@ -257,22 +257,30 @@ type Controller struct {
 	diskQuotaLock  *diskQuotaLock
 	diskQuotaState atomic.Int32
 	compactState   atomic.Int32
+	status         *LightningStatus
+}
+
+type LightningStatus struct {
+	FinishedFileSize atomic.Int64
+	TotalFileSize    atomic.Int64
 }
 
 func NewRestoreController(
 	ctx context.Context,
 	dbMetas []*mydump.MDDatabaseMeta,
 	cfg *config.Config,
+	status *LightningStatus,
 	s storage.ExternalStorage,
 	g glue.Glue,
 ) (*Controller, error) {
-	return NewRestoreControllerWithPauser(ctx, dbMetas, cfg, s, DeliverPauser, g)
+	return NewRestoreControllerWithPauser(ctx, dbMetas, cfg, status, s, DeliverPauser, g)
 }
 
 func NewRestoreControllerWithPauser(
 	ctx context.Context,
 	dbMetas []*mydump.MDDatabaseMeta,
 	cfg *config.Config,
+	status *LightningStatus,
 	s storage.ExternalStorage,
 	pauser *common.Pauser,
 	g glue.Glue,
@@ -379,6 +387,7 @@ func NewRestoreControllerWithPauser(
 		store:          s,
 		metaMgrBuilder: metaBuilder,
 		diskQuotaLock:  newDiskQuotaLock(),
+		status:         status,
 		taskMgr:        nil,
 	}
 
@@ -1716,13 +1725,16 @@ func (rc *Controller) isLocalBackend() bool {
 // 4. Lightning configuration
 // before restore tables start.
 func (rc *Controller) preCheckRequirements(ctx context.Context) error {
-	if err := rc.ClusterIsAvailable(ctx); err != nil {
-		return errors.Trace(err)
-	}
+	if rc.cfg.App.CheckRequirements {
+		if err := rc.ClusterIsAvailable(ctx); err != nil {
+			return errors.Trace(err)
+		}
 
-	if err := rc.StoragePermission(ctx); err != nil {
-		return errors.Trace(err)
+		if err := rc.StoragePermission(ctx); err != nil {
+			return errors.Trace(err)
+		}
 	}
+
 	if err := rc.metaMgrBuilder.Init(ctx); err != nil {
 		return err
 	}
@@ -1741,29 +1753,37 @@ func (rc *Controller) preCheckRequirements(ctx context.Context) error {
 		if err != nil {
 			return errors.Trace(err)
 		}
+
+		// We still need to sample source data even if this task has existed, because we need to judge whether the
+		// source is in order as row key to decide how to sort local data.
+		source, err := rc.estimateSourceData(ctx)
+		if err != nil {
+			return errors.Trace(err)
+		}
 		if !taskExist {
-			source, err := rc.EstimateSourceData(ctx)
-			if err != nil {
+			if err = rc.taskMgr.InitTask(ctx, source); err != nil {
 				return errors.Trace(err)
 			}
-			err = rc.LocalResource(ctx, source)
-			if err != nil {
-				rc.taskMgr.CleanupTask(ctx)
-				return errors.Trace(err)
-			}
-			if err := rc.ClusterResource(ctx, source); err != nil {
-				rc.taskMgr.CleanupTask(ctx)
-				return errors.Trace(err)
-			}
-			if err := rc.CheckClusterRegion(ctx); err != nil {
-				return errors.Trace(err)
+			if rc.cfg.App.CheckRequirements {
+				err = rc.localResource(source)
+				if err != nil {
+					return errors.Trace(err)
+				}
+				if err := rc.clusterResource(ctx, source); err != nil {
+					rc.taskMgr.CleanupTask(ctx)
+					return errors.Trace(err)
+				}
+				if err := rc.checkClusterRegion(ctx); err != nil {
+					return errors.Trace(err)
+				}
 			}
 		}
 	}
-	if rc.tidbGlue.OwnsSQLExecutor() {
+
+	if rc.tidbGlue.OwnsSQLExecutor() && rc.cfg.App.CheckRequirements {
 		// print check info at any time.
 		fmt.Print(rc.checkTemplate.Output())
-		if rc.cfg.App.CheckRequirements && !rc.checkTemplate.Success() {
+		if !rc.checkTemplate.Success() {
 			// if check requirements is true, return error.
 			if !taskExist && rc.taskMgr != nil {
 				rc.taskMgr.CleanupTask(ctx)

br/pkg/lightning/restore/restore_test.go

@@ -1713,7 +1713,7 @@ func (s *tableRestoreSuite) TestCheckClusterResource(c *C) {
 			sourceSize += size
 			return nil
 		})
-		err = rc.ClusterResource(ctx, sourceSize)
+		err = rc.clusterResource(ctx, sourceSize)
 		c.Assert(err, IsNil)
 
 		c.Assert(template.FailedCount(Critical), Equals, ca.expectErrorCount)
@@ -1835,7 +1835,7 @@ func (s *tableRestoreSuite) TestCheckClusterRegion(c *C) {
 		cfg := &config.Config{TiDB: config.DBStore{PdAddr: url}}
 		rc := &Controller{cfg: cfg, tls: tls, taskMgr: mockTaskMetaMgr{}, checkTemplate: template}
 
-		err := rc.CheckClusterRegion(context.Background())
+		err := rc.checkClusterRegion(context.Background())
 		c.Assert(err, IsNil)
 		c.Assert(template.FailedCount(Critical), Equals, ca.expectErrorCnt)
 		c.Assert(template.Success(), Equals, ca.expectResult)

br/pkg/lightning/restore/table_restore.go

@@ -314,11 +314,20 @@ func (tr *TableRestore) restoreEngines(pCtx context.Context, rc *Controller, cp
 						dataWorker := rc.closedEngineLimit.Apply()
 						defer rc.closedEngineLimit.Recycle(dataWorker)
 						err = tr.importEngine(ctx, dataClosedEngine, rc, eid, ecp)
+						if rc.status != nil {
+							for _, chunk := range ecp.Chunks {
+								rc.status.FinishedFileSize.Add(chunk.FileMeta.FileSize)
+							}
+						}
 					}
 					if err != nil {
 						setError(err)
 					}
 				}(restoreWorker, engineID, engine)
+			} else {
+				for _, chunk := range engine.Chunks {
+					rc.status.FinishedFileSize.Add(chunk.FileMeta.FileSize)
+				}
 			}
 		}