ddl: expend interface of lightning for add index acceleration

br/pkg/lightning/backend/backend.go

@@ -215,6 +215,9 @@ type AbstractBackend interface {
 	// ResolveDuplicateRows resolves duplicated rows by deleting/inserting data
 	// according to the required algorithm.
 	ResolveDuplicateRows(ctx context.Context, tbl table.Table, tableName string, algorithm config.DuplicateResolutionAlgorithm) error
+
+	// TotalMemoryConsume counts total memory usage. This is only used for local backend
+	TotalMemoryConsume() int64
 }
 
 // Backend is the delivery target for Lightning
@@ -286,6 +289,10 @@ func (be Backend) FlushAll(ctx context.Context) error {
 	return be.abstract.FlushAllEngines(ctx)
 }
 
+func (be Backend) TotalMemoryConsume() int64 {
+	return be.abstract.TotalMemoryConsume()
+}
+
 // CheckDiskQuota verifies if the total engine file size is below the given
 // quota. If the quota is exceeded, this method returns an array of engines,
 // which after importing can decrease the total size below quota.
@@ -411,6 +418,10 @@ func (engine *OpenedEngine) LocalWriter(ctx context.Context, cfg *LocalWriterCon
 	return &LocalEngineWriter{writer: w, tableName: engine.tableName}, nil
 }
 
+func (engine *OpenedEngine) TotalMemoryConsume() int64 {
+	return engine.engine.backend.TotalMemoryConsume()
+}
+
 // WriteRows writes a collection of encoded rows into the engine.
 func (w *LocalEngineWriter) WriteRows(ctx context.Context, columnNames []string, rows kv.Rows) error {
 	return w.writer.AppendRows(ctx, w.tableName, columnNames, rows)
@@ -501,3 +512,7 @@ type EngineWriter interface {
 	IsSynced() bool
 	Close(ctx context.Context) (ChunkFlushStatus, error)
 }
+
+func (engine *OpenedEngine) GetEngineUuid() uuid.UUID {
+	return engine.uuid
+}

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

@@ -235,6 +235,20 @@ func (e *Engine) unlock() {
 	e.mutex.Unlock()
 }
 
+func (e *Engine) TotalMemorySize() int64 {
+	var memSize int64 = 0
+	e.localWriters.Range(func(k, v interface{}) bool {
+		w := k.(*Writer)
+		if w.kvBuffer != nil {
+			w.Lock()
+			memSize += w.kvBuffer.TotalSize()
+			w.Unlock()
+		}
+		return true
+	})
+	return memSize
+}
+
 type rangeOffsets struct {
 	Size uint64
 	Keys uint64

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

@@ -495,6 +495,18 @@ func NewLocalBackend(
 	return backend.MakeBackend(local), nil
 }
 
+func (local *local) TotalMemoryConsume() int64 {
+	var memConsume int64 = 0
+	local.engines.Range(func(k, v interface{}) bool {
+		e := v.(*Engine)
+		if e != nil {
+			memConsume += e.TotalMemorySize()
+		}
+		return true
+	})
+	return memConsume
+}
+
 func (local *local) checkMultiIngestSupport(ctx context.Context) error {
 	stores, err := local.pdCtl.GetPDClient().GetAllStores(ctx, pd.WithExcludeTombstone())
 	if err != nil {

br/pkg/lightning/backend/noop/noop.go

@@ -155,6 +155,10 @@ func (b noopBackend) ResolveDuplicateRows(ctx context.Context, tbl table.Table,
 	return nil
 }
 
+func (b noopBackend) TotalMemoryConsume() int64 {
+	return 0
+}
+
 type noopEncoder struct{}
 
 // Close the encoder.

br/pkg/lightning/backend/tidb/tidb.go

@@ -617,6 +617,10 @@ rowLoop:
 	return nil
 }
 
+func (be *tidbBackend) TotalMemoryConsume() int64 {
+	return 0
+}
+
 type stmtTask struct {
 	rows tidbRows
 	stmt string