ddl: make reorgCtx safer to use (#41421)

close #41418
pingcap · Feb 15, 2023 · a56fa8e · a56fa8e
1 parent 44e4381
commit a56fa8e
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 24 deletions.
diff --git a/ddl/ddl.go b/ddl/ddl.go
@@ -526,6 +526,13 @@ func (dc *ddlCtx) getReorgCtx(jobID int64) *reorgCtx {
 }
 
 func (dc *ddlCtx) newReorgCtx(jobID int64, startKey []byte, currElement *meta.Element, rowCount int64) *reorgCtx {
+	dc.reorgCtx.Lock()
+	defer dc.reorgCtx.Unlock()
+	existedRC, ok := dc.reorgCtx.reorgCtxMap[jobID]
+	if ok {
+		existedRC.references.Add(1)
+		return existedRC
+	}
 	rc := &reorgCtx{}
 	rc.doneCh = make(chan error, 1)
 	// initial reorgCtx
@@ -535,22 +542,10 @@ func (dc *ddlCtx) newReorgCtx(jobID int64, startKey []byte, currElement *meta.El
 	rc.mu.warnings = make(map[errors.ErrorID]*terror.Error)
 	rc.mu.warningsCount = make(map[errors.ErrorID]int64)
 	rc.references.Add(1)
-	dc.reorgCtx.Lock()
-	defer dc.reorgCtx.Unlock()
 	dc.reorgCtx.reorgCtxMap[jobID] = rc
 	return rc
 }
 
-func (dc *ddlCtx) setReorgCtxForBackfill(bfJob *BackfillJob) {
-	rc := dc.getReorgCtx(bfJob.JobID)
-	if rc == nil {
-		ele := &meta.Element{ID: bfJob.EleID, TypeKey: bfJob.EleKey}
-		dc.newReorgCtx(bfJob.JobID, bfJob.Meta.StartKey, ele, bfJob.Meta.RowCount)
-	} else {
-		rc.references.Add(1)
-	}
-}
-
 func (dc *ddlCtx) removeReorgCtx(jobID int64) {
 	dc.reorgCtx.Lock()
 	defer dc.reorgCtx.Unlock()

diff --git a/ddl/ddl_test.go b/ddl/ddl_test.go
@@ -44,7 +44,6 @@ type DDLForTest interface {
 	// SetInterceptor sets the interceptor.
 	SetInterceptor(h Interceptor)
 	NewReorgCtx(jobID int64, startKey []byte, currElement *meta.Element, rowCount int64) *reorgCtx
-	SetReorgCtxForBackfill(bfJob *BackfillJob)
 	GetReorgCtx(jobID int64) *reorgCtx
 	RemoveReorgCtx(id int64)
 }
@@ -67,11 +66,6 @@ func (d *ddl) NewReorgCtx(jobID int64, startKey []byte, currElement *meta.Elemen
 	return d.newReorgCtx(jobID, startKey, currElement, rowCount)
 }
 
-// SetReorgCtxForBackfill exports for testing.
-func (d *ddl) SetReorgCtxForBackfill(bfJob *BackfillJob) {
-	d.setReorgCtxForBackfill(bfJob)
-}
-
 // GetReorgCtx exports for testing.
 func (d *ddl) GetReorgCtx(jobID int64) *reorgCtx {
 	return d.getReorgCtx(jobID)

diff --git a/ddl/ddl_worker_test.go b/ddl/ddl_worker_test.go
@@ -317,10 +317,10 @@ func TestUsingReorgCtx(t *testing.T) {
 	wg := util.WaitGroupWrapper{}
 	wg.Run(func() {
 		jobID := int64(1)
-		m := &model.BackfillMeta{StartKey: []byte("skey"), RowCount: 1}
-		bfJob := &ddl.BackfillJob{JobID: jobID, EleID: 1, EleKey: nil, Meta: m}
-		for i := 0; i < 100; i++ {
-			d.(ddl.DDLForTest).SetReorgCtxForBackfill(bfJob)
+		startKey := []byte("skey")
+		ele := &meta.Element{ID: 1, TypeKey: nil}
+		for i := 0; i < 500; i++ {
+			d.(ddl.DDLForTest).NewReorgCtx(jobID, startKey, ele, 0)
 			d.(ddl.DDLForTest).GetReorgCtx(jobID).IsReorgCanceled()
 			d.(ddl.DDLForTest).RemoveReorgCtx(jobID)
 		}
@@ -329,7 +329,7 @@ func TestUsingReorgCtx(t *testing.T) {
 		jobID := int64(1)
 		startKey := []byte("skey")
 		ele := &meta.Element{ID: 1, TypeKey: nil}
-		for i := 0; i < 100; i++ {
+		for i := 0; i < 500; i++ {
 			d.(ddl.DDLForTest).NewReorgCtx(jobID, startKey, ele, 0)
 			d.(ddl.DDLForTest).GetReorgCtx(jobID).IsReorgCanceled()
 			d.(ddl.DDLForTest).RemoveReorgCtx(jobID)

diff --git a/ddl/job_table.go b/ddl/job_table.go
@@ -402,7 +402,7 @@ func (d *ddl) loadBackfillJobAndRun() {
 		return
 	}
 	// TODO: Adjust how the non-owner uses ReorgCtx.
-	d.setReorgCtxForBackfill(bfJob)
+	d.newReorgCtx(bfJob.JobID, bfJob.Meta.StartKey, &meta.Element{ID: bfJob.EleID, TypeKey: bfJob.EleKey}, bfJob.Meta.RowCount)
 	d.wg.Run(func() {
 		defer func() {
 			tidbutil.Recover(metrics.LabelDistReorg, "runBackfillJobs", nil, false)