ddl: Implement basic version of TiFlash replica management

ddl/db_change_test.go

@@ -77,6 +77,7 @@ func (s *testStateChangeSuiteBase) SetUpSuite(c *C) {
 	session.SetSchemaLease(s.lease)
 	s.dom, err = session.BootstrapSession(s.store)
 	c.Assert(err, IsNil)
+	ddl.DisableTiFlashPoll(s.dom.DDL())
 	s.se, err = session.CreateSession4Test(s.store)
 	c.Assert(err, IsNil)
 	_, err = s.se.Execute(context.Background(), "create database test_db_state default charset utf8 default collate utf8_bin")

ddl/db_test.go

@@ -119,6 +119,7 @@ func setUpSuite(s *testDBSuite, c *C) {
 
 	s.dom, err = session.BootstrapSession(s.store)
 	c.Assert(err, IsNil)
+	ddl.DisableTiFlashPoll(s.dom.DDL())
 	s.s, err = session.CreateSession4Test(s.store)
 	c.Assert(err, IsNil)
 	s.ctx = s.s.(sessionctx.Context)

ddl/ddl.go

@@ -193,9 +193,10 @@ type ddl struct {
 	limitJobCh chan *limitJobTask
 
 	*ddlCtx
-	workers     map[workerType]*worker
-	sessPool    *sessionPool
-	delRangeMgr delRangeManager
+	workers           map[workerType]*worker
+	sessPool          *sessionPool
+	delRangeMgr       delRangeManager
+	enableTiFlashPoll bool
 }
 
 // ddlCtx is the context when we use worker to handle DDL jobs.
@@ -230,6 +231,25 @@ func (dc *ddlCtx) isOwner() bool {
 	return isOwner
 }
 
+// EnableTiFlashPoll enables TiFlash poll loop aka pollTiFlashReplicaStatus.
+func EnableTiFlashPoll(d interface{}) {
+	if dd, ok := d.(*ddl); ok {
+		dd.enableTiFlashPoll = true
+	}
+}
+
+// DisableTiFlashPoll disables TiFlash poll loop aka pollTiFlashReplicaStatus.
+func DisableTiFlashPoll(d interface{}) {
+	if dd, ok := d.(*ddl); ok {
+		dd.enableTiFlashPoll = false
+	}
+}
+
+// IsTiFlashPollEnabled reveals enableTiFlashPoll
+func (d *ddl) IsTiFlashPollEnabled() bool {
+	return d.enableTiFlashPoll
+}
+
 // RegisterStatsHandle registers statistics handle and its corresponding even channel for ddl.
 func (d *ddl) RegisterStatsHandle(h *handle.Handle) {
 	d.ddlCtx.statsHandle = h
@@ -312,9 +332,10 @@ func newDDL(ctx context.Context, options ...Option) *ddl {
 	ddlCtx.mu.hook = opt.Hook
 	ddlCtx.mu.interceptor = &BaseInterceptor{}
 	d := &ddl{
-		ctx:        ctx,
-		ddlCtx:     ddlCtx,
-		limitJobCh: make(chan *limitJobTask, batchAddingJobs),
+		ctx:               ctx,
+		ddlCtx:            ddlCtx,
+		limitJobCh:        make(chan *limitJobTask, batchAddingJobs),
+		enableTiFlashPoll: true,
 	}
 
 	return d
@@ -387,6 +408,10 @@ func (d *ddl) Start(ctxPool *pools.ResourcePool) error {
 	variable.RegisterStatistics(d)
 
 	metrics.DDLCounter.WithLabelValues(metrics.CreateDDLInstance).Inc()
+
+	// Start some background routine to manage TiFlash replica.
+	go d.PollTiFlashRoutine()
+
 	return nil
 }
 

ddl/ddl_api.go

@@ -33,6 +33,7 @@ import (
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/ddl/label"
+	"github.com/pingcap/tidb/domain/infosync"
 	"github.com/pingcap/tidb/expression"
 	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/kv"
@@ -4720,6 +4721,27 @@ func (d *ddl) AlterTableSetTiFlashReplica(ctx sessionctx.Context, ident ast.Iden
 		return errors.Trace(err)
 	}
 
+	// We should check this first, in order to avoid creating redundant DDL jobs.
+	tblInfo := tb.Meta()
+	if d.IsTiFlashPollEnabled() {
+		// Must set placement rule before the ActionSetTiFlashReplica job is in queue.
+		if pi := tblInfo.GetPartitionInfo(); pi != nil {
+			logutil.BgLogger().Info("Set TiFlash replica pd rule for partitioned table", zap.Int64("tableID", tblInfo.ID))
+			if e := infosync.ConfigureTiFlashPDForPartitions(false, &pi.Definitions, replicaInfo.Count, &replicaInfo.Labels); e != nil {
+				return errors.Trace(e)
+			}
+			// Partitions that in adding mid-state. They have high priorities, so we should set accordingly pd rules.
+			if e := infosync.ConfigureTiFlashPDForPartitions(true, &pi.AddingDefinitions, replicaInfo.Count, &replicaInfo.Labels); e != nil {
+				return errors.Trace(e)
+			}
+		} else {
+			logutil.BgLogger().Info("Set TiFlash replica pd rule", zap.Int64("tableID", tblInfo.ID))
+			if e := infosync.ConfigureTiFlashPDForTable(tblInfo.ID, replicaInfo.Count, &replicaInfo.Labels); e != nil {
+				return errors.Trace(e)
+			}
+		}
+	}
+
 	job := &model.Job{
 		SchemaID:   schema.ID,
 		TableID:    tb.Meta().ID,
@@ -4812,13 +4834,15 @@ func (d *ddl) AlterTableDropStatistics(ctx sessionctx.Context, ident ast.Ident,
 func (d *ddl) UpdateTableReplicaInfo(ctx sessionctx.Context, physicalID int64, available bool) error {
 	is := d.infoCache.GetLatest()
 	tb, ok := is.TableByID(physicalID)
+
 	if !ok {
 		tb, _, _ = is.FindTableByPartitionID(physicalID)
 		if tb == nil {
 			return infoschema.ErrTableNotExists.GenWithStack("Table which ID = %d does not exist.", physicalID)
 		}
 	}
 	tbInfo := tb.Meta()
+
 	if tbInfo.TiFlashReplica == nil || (tbInfo.ID == physicalID && tbInfo.TiFlashReplica.Available == available) ||
 		(tbInfo.ID != physicalID && available == tbInfo.TiFlashReplica.IsPartitionAvailable(physicalID)) {
 		return nil
@@ -4837,6 +4861,7 @@ func (d *ddl) UpdateTableReplicaInfo(ctx sessionctx.Context, physicalID int64, a
 		BinlogInfo: &model.HistoryInfo{},
 		Args:       []interface{}{available, physicalID},
 	}
+	// Will return when job is finished.
 	err := d.doDDLJob(ctx, job)
 	err = d.callHookOnChanged(err)
 	return errors.Trace(err)