ddl: manage job lifecycle in worker (#53747)

ref #53246

pkg/ddl/BUILD.bazel

@@ -37,6 +37,7 @@ go_library(
         "delete_range.go",
         "delete_range_util.go",
         "dist_owner.go",
+        "doc.go",
         "foreign_key.go",
         "generated_column.go",
         "index.go",
@@ -76,6 +77,7 @@ go_library(
         "//pkg/ddl/placement",
         "//pkg/ddl/resourcegroup",
         "//pkg/ddl/syncer",
+        "//pkg/ddl/systable",
         "//pkg/ddl/util",
         "//pkg/distsql",
         "//pkg/distsql/context",

pkg/ddl/cluster.go

@@ -73,16 +73,16 @@ const (
 	keyRangesOffset
 )
 
-func closePDSchedule() error {
+func closePDSchedule(ctx context.Context) error {
 	closeMap := make(map[string]any)
 	for _, key := range pdScheduleKey {
 		closeMap[key] = 0
 	}
-	return infosync.SetPDScheduleConfig(context.Background(), closeMap)
+	return infosync.SetPDScheduleConfig(ctx, closeMap)
 }
 
-func savePDSchedule(job *model.Job) error {
-	retValue, err := infosync.GetPDScheduleConfig(context.Background())
+func savePDSchedule(ctx context.Context, job *model.Job) error {
+	retValue, err := infosync.GetPDScheduleConfig(ctx)
 	if err != nil {
 		return err
 	}
@@ -94,11 +94,11 @@ func savePDSchedule(job *model.Job) error {
 	return nil
 }
 
-func recoverPDSchedule(pdScheduleParam map[string]any) error {
+func recoverPDSchedule(ctx context.Context, pdScheduleParam map[string]any) error {
 	if pdScheduleParam == nil {
 		return nil
 	}
-	return infosync.SetPDScheduleConfig(context.Background(), pdScheduleParam)
+	return infosync.SetPDScheduleConfig(ctx, pdScheduleParam)
 }
 
 func getStoreGlobalMinSafeTS(s kv.Storage) time.Time {
@@ -228,7 +228,7 @@ func checkAndSetFlashbackClusterInfo(ctx context.Context, se sessionctx.Context,
 	if err = gcutil.DisableGC(se); err != nil {
 		return err
 	}
-	if err = closePDSchedule(); err != nil {
+	if err = closePDSchedule(ctx); err != nil {
 		return err
 	}
 	if err = setTiDBEnableAutoAnalyze(ctx, se, variable.Off); err != nil {
@@ -664,7 +664,7 @@ func (w *worker) onFlashbackCluster(d *ddlCtx, t *meta.Meta, job *model.Job) (ve
 	switch job.SchemaState {
 	// Stage 1, check and set FlashbackClusterJobID, and update job args.
 	case model.StateNone:
-		if err = savePDSchedule(job); err != nil {
+		if err = savePDSchedule(w.ctx, job); err != nil {
 			job.State = model.JobStateCancelled
 			return ver, errors.Trace(err)
 		}
@@ -801,7 +801,7 @@ func finishFlashbackCluster(w *worker, job *model.Job) error {
 	defer w.sessPool.Put(sess)
 
 	err = kv.RunInNewTxn(w.ctx, w.store, true, func(context.Context, kv.Transaction) error {
-		if err = recoverPDSchedule(pdScheduleValue); err != nil {
+		if err = recoverPDSchedule(w.ctx, pdScheduleValue); err != nil {
 			return err
 		}
 		if gcEnabled {

pkg/ddl/constraint.go

@@ -15,7 +15,6 @@
 package ddl
 
 import (
-	"context"
 	"fmt"
 	"strings"
 
@@ -374,7 +373,7 @@ func (w *worker) verifyRemainRecordsForCheckConstraint(dbInfo *model.DBInfo, tab
 	// can let the check expression restored string as the filter in where clause directly.
 	// Prepare internal SQL to fetch data from physical table under this filter.
 	sql := fmt.Sprintf("select 1 from `%s`.`%s` where not %s limit 1", dbInfo.Name.L, tableInfo.Name.L, constr.ExprString)
-	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnDDL)
+	ctx := kv.WithInternalSourceType(w.ctx, kv.InternalTxnDDL)
 	rows, _, err := sctx.GetRestrictedSQLExecutor().ExecRestrictedSQL(ctx, nil, sql)
 	if err != nil {
 		return errors.Trace(err)

pkg/ddl/ddl.go

@@ -838,7 +838,7 @@ func (d *ddl) Start(ctxPool *pools.ResourcePool) error {
 	d.wg.Run(func() {
 		d.limitDDLJobs(d.limitJobChV2, d.addBatchLocalDDLJobs)
 	})
-	d.sessPool = sess.NewSessionPool(ctxPool, d.store)
+	d.sessPool = sess.NewSessionPool(ctxPool)
 
 	d.delRangeMgr = d.newDeleteRangeManager(ctxPool == nil)
 
@@ -1597,7 +1597,7 @@ type Info struct {
 // GetDDLInfoWithNewTxn returns DDL information using a new txn.
 func GetDDLInfoWithNewTxn(s sessionctx.Context) (*Info, error) {
 	se := sess.NewSession(s)
-	err := se.Begin()
+	err := se.Begin(context.Background())
 	if err != nil {
 		return nil, err
 	}
@@ -1757,7 +1757,7 @@ func processJobs(process func(*sess.Session, *model.Job, model.AdminCommandOpera
 			idsStr = append(idsStr, strconv.FormatInt(id, 10))
 		}
 
-		err = ns.Begin()
+		err = ns.Begin(context.Background())
 		if err != nil {
 			return nil, err
 		}
@@ -1797,7 +1797,7 @@ func processJobs(process func(*sess.Session, *model.Job, model.AdminCommandOpera
 		})
 
 		// There may be some conflict during the update, try it again
-		if err = ns.Commit(); err != nil {
+		if err = ns.Commit(context.Background()); err != nil {
 			continue
 		}
 
@@ -1848,7 +1848,7 @@ func processAllJobs(process func(*sess.Session, *model.Job, model.AdminCommandOp
 	var jobErrs = make(map[int64]error)
 
 	ns := sess.NewSession(se)
-	err = ns.Begin()
+	err = ns.Begin(context.Background())
 	if err != nil {
 		return nil, err
 	}
@@ -1894,7 +1894,7 @@ func processAllJobs(process func(*sess.Session, *model.Job, model.AdminCommandOp
 		jobID = jobIDMax + 1
 	}
 
-	err = ns.Commit()
+	err = ns.Commit(context.Background())
 	if err != nil {
 		return nil, err
 	}

pkg/ddl/ddl_api_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/model"
 	sessiontypes "github.com/pingcap/tidb/pkg/session/types"
 	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
 	"github.com/pingcap/tidb/pkg/util/chunk"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/sync/errgroup"
@@ -163,14 +164,14 @@ func TestCreateViewConcurrently(t *testing.T) {
 		counterErr error
 		counter    int
 	)
-	failpoint.EnableCall("github.com/pingcap/tidb/pkg/ddl/onDDLCreateView", func(job *model.Job) {
+	testfailpoint.EnableCall(t, "github.com/pingcap/tidb/pkg/ddl/onDDLCreateView", func(job *model.Job) {
 		counter++
 		if counter > 1 {
 			counterErr = fmt.Errorf("create view job should not run concurrently")
 			return
 		}
 	})
-	failpoint.EnableCall("github.com/pingcap/tidb/pkg/ddl/afterDelivery2Worker", func(job *model.Job) {
+	testfailpoint.EnableCall(t, "github.com/pingcap/tidb/pkg/ddl/afterDelivery2Worker", func(job *model.Job) {
 		if job.Type == model.ActionCreateView {
 			counter--
 		}

pkg/ddl/ddl_history.go

@@ -40,8 +40,8 @@ const (
 )
 
 // AddHistoryDDLJob record the history job.
-func AddHistoryDDLJob(sess *sess.Session, t *meta.Meta, job *model.Job, updateRawArgs bool) error {
-	err := addHistoryDDLJob2Table(sess, job, updateRawArgs)
+func AddHistoryDDLJob(ctx context.Context, sess *sess.Session, t *meta.Meta, job *model.Job, updateRawArgs bool) error {
+	err := addHistoryDDLJob2Table(ctx, sess, job, updateRawArgs)
 	if err != nil {
 		logutil.DDLLogger().Info("failed to add DDL job to history table", zap.Error(err))
 	}
@@ -50,12 +50,12 @@ func AddHistoryDDLJob(sess *sess.Session, t *meta.Meta, job *model.Job, updateRa
 }
 
 // addHistoryDDLJob2Table adds DDL job to history table.
-func addHistoryDDLJob2Table(sess *sess.Session, job *model.Job, updateRawArgs bool) error {
+func addHistoryDDLJob2Table(ctx context.Context, sess *sess.Session, job *model.Job, updateRawArgs bool) error {
 	b, err := job.Encode(updateRawArgs)
 	if err != nil {
 		return err
 	}
-	_, err = sess.Execute(context.Background(),
+	_, err = sess.Execute(ctx,
 		fmt.Sprintf("insert ignore into mysql.tidb_ddl_history(job_id, job_meta, db_name, table_name, schema_ids, table_ids, create_time) values (%d, %s, %s, %s, %s, %s, %v)",
 			job.ID, util.WrapKey2String(b), strconv.Quote(job.SchemaName), strconv.Quote(job.TableName),
 			strconv.Quote(strconv.FormatInt(job.SchemaID, 10)),

pkg/ddl/ddl_history_test.go

@@ -38,15 +38,15 @@ func TestDDLHistoryBasic(t *testing.T) {
 		newTk := testkit.NewTestKit(t, store)
 		return newTk.Session(), nil
 	}, 8, 8, 0)
-	sessPool := session.NewSessionPool(rs, store)
+	sessPool := session.NewSessionPool(rs)
 	sessCtx, err := sessPool.Get()
 	require.NoError(t, err)
 	sess := session.NewSession(sessCtx)
 
 	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnLightning)
 	err = kv.RunInNewTxn(ctx, store, false, func(ctx context.Context, txn kv.Transaction) error {
 		t := meta.NewMeta(txn)
-		return ddl.AddHistoryDDLJob(sess, t, &model.Job{
+		return ddl.AddHistoryDDLJob(context.Background(), sess, t, &model.Job{
 			ID: 1,
 		}, false)
 	})
@@ -55,7 +55,7 @@ func TestDDLHistoryBasic(t *testing.T) {
 
 	err = kv.RunInNewTxn(ctx, store, false, func(ctx context.Context, txn kv.Transaction) error {
 		t := meta.NewMeta(txn)
-		return ddl.AddHistoryDDLJob(sess, t, &model.Job{
+		return ddl.AddHistoryDDLJob(context.Background(), sess, t, &model.Job{
 			ID: 2,
 		}, false)
 	})

pkg/ddl/ddl_running_jobs.go

@@ -53,40 +53,38 @@ func (j *runningJobs) clear() {
 	j.unfinishedSchema = make(map[string]map[string]struct{})
 }
 
-func (j *runningJobs) add(job *model.Job) {
+func (j *runningJobs) add(jobID int64, involvedSchemaInfos []model.InvolvingSchemaInfo) {
 	j.Lock()
 	defer j.Unlock()
-	j.processingIDs[job.ID] = struct{}{}
+	j.processingIDs[jobID] = struct{}{}
 	j.updateInternalRunningJobIDs()
 
-	if _, ok := j.unfinishedIDs[job.ID]; ok {
+	if _, ok := j.unfinishedIDs[jobID]; ok {
 		// Already exists, no need to add it again.
 		return
 	}
-	j.unfinishedIDs[job.ID] = struct{}{}
-	for _, info := range job.GetInvolvingSchemaInfo() {
+	j.unfinishedIDs[jobID] = struct{}{}
+	for _, info := range involvedSchemaInfos {
 		if _, ok := j.unfinishedSchema[info.Database]; !ok {
 			j.unfinishedSchema[info.Database] = make(map[string]struct{})
 		}
 		j.unfinishedSchema[info.Database][info.Table] = struct{}{}
 	}
 }
 
-func (j *runningJobs) remove(job *model.Job) {
+func (j *runningJobs) remove(jobID int64, involvedSchemaInfos []model.InvolvingSchemaInfo) {
 	j.Lock()
 	defer j.Unlock()
-	delete(j.processingIDs, job.ID)
+	delete(j.processingIDs, jobID)
 	j.updateInternalRunningJobIDs()
 
-	if job.IsFinished() || job.IsSynced() {
-		delete(j.unfinishedIDs, job.ID)
-		for _, info := range job.GetInvolvingSchemaInfo() {
-			if db, ok := j.unfinishedSchema[info.Database]; ok {
-				delete(db, info.Table)
-			}
-			if len(j.unfinishedSchema[info.Database]) == 0 {
-				delete(j.unfinishedSchema, info.Database)
-			}
+	delete(j.unfinishedIDs, jobID)
+	for _, info := range involvedSchemaInfos {
+		if db, ok := j.unfinishedSchema[info.Database]; ok {
+			delete(db, info.Table)
+		}
+		if len(j.unfinishedSchema[info.Database]) == 0 {
+			delete(j.unfinishedSchema, info.Database)
 		}
 	}
 }

pkg/ddl/ddl_running_jobs_test.go

@@ -66,8 +66,8 @@ func TestRunningJobs(t *testing.T) {
 	require.True(t, runnable)
 	job1 := mkJob(1, "db1.t1", "db1.t2")
 	job2 := mkJob(2, "db2.t3")
-	j.add(job1)
-	j.add(job2)
+	j.add(job1.ID, job1.GetInvolvingSchemaInfo())
+	j.add(job2.ID, job2.GetInvolvingSchemaInfo())
 	require.Equal(t, "1,2", orderedAllIDs(j.allIDs()))
 	runnable = j.checkRunnable(mkJob(0, "db1.t1"))
 	require.False(t, runnable)
@@ -79,37 +79,37 @@ func TestRunningJobs(t *testing.T) {
 	require.True(t, runnable)
 
 	job3 := mkJob(3, "db1.*")
-	j.add(job3)
+	j.add(job3.ID, job3.GetInvolvingSchemaInfo())
 	require.Equal(t, "1,2,3", orderedAllIDs(j.allIDs()))
 	runnable = j.checkRunnable(mkJob(0, "db1.t100"))
 	require.False(t, runnable)
 
 	job4 := mkJob(4, "db4.")
-	j.add(job4)
+	j.add(job4.ID, job4.GetInvolvingSchemaInfo())
 	require.Equal(t, "1,2,3,4", orderedAllIDs(j.allIDs()))
 	runnable = j.checkRunnable(mkJob(0, "db4.t100"))
 	require.True(t, runnable)
 
 	job5 := mkJob(5, "*.*")
-	j.add(job5)
+	j.add(job5.ID, job5.GetInvolvingSchemaInfo())
 	require.Equal(t, "1,2,3,4,5", orderedAllIDs(j.allIDs()))
 	runnable = j.checkRunnable(mkJob(0, "db100.t100"))
 	require.False(t, runnable)
 
 	job5.State = model.JobStateDone
-	j.remove(job5)
+	j.remove(job5.ID, job5.GetInvolvingSchemaInfo())
 	require.Equal(t, "1,2,3,4", orderedAllIDs(j.allIDs()))
 	runnable = j.checkRunnable(mkJob(0, "db100.t100"))
 	require.True(t, runnable)
 
 	job3.State = model.JobStateDone
-	j.remove(job3)
+	j.remove(job3.ID, job3.GetInvolvingSchemaInfo())
 	require.Equal(t, "1,2,4", orderedAllIDs(j.allIDs()))
 	runnable = j.checkRunnable(mkJob(0, "db1.t100"))
 	require.True(t, runnable)
 
 	job1.State = model.JobStateDone
-	j.remove(job1)
+	j.remove(job1.ID, job1.GetInvolvingSchemaInfo())
 	require.Equal(t, "2,4", orderedAllIDs(j.allIDs()))
 	runnable = j.checkRunnable(mkJob(0, "db1.t1"))
 	require.True(t, runnable)
@@ -119,15 +119,15 @@ func TestOwnerRetireThenToBeOwner(t *testing.T) {
 	j := newRunningJobs()
 	require.Equal(t, "", j.allIDs())
 	job := mkJob(1, "test.t1")
-	j.add(job)
+	j.add(job.ID, job.GetInvolvingSchemaInfo())
 	require.False(t, j.checkRunnable(job))
 	// retire
 	j.clear()
 	// to be owner, try to start a new job.
 	require.False(t, j.checkRunnable(job))
 	// previous job removed.
-	j.remove(job)
+	j.remove(job.ID, job.GetInvolvingSchemaInfo())
 	require.True(t, j.checkRunnable(job))
-	j.add(job)
+	j.add(job.ID, job.GetInvolvingSchemaInfo())
 	require.False(t, j.checkRunnable(job))
 }