[pipeline] bug fix: fatal error: concurrent map writes (#1026)

HaozhengAN · web-flow · commit 287d18258a05 · 2023-02-22T18:55:28.000+08:00
diff --git a/pkg/apiserver/controller/pipeline/callback_test.go b/pkg/apiserver/controller/pipeline/callback_test.go
@@ -21,11 +21,14 @@ import (
 	"io/ioutil"
 	"testing"
 
+	"github.com/agiledragon/gomonkey/v2"
+	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/PaddlePaddle/PaddleFlow/pkg/apiserver/common"
 	"github.com/PaddlePaddle/PaddleFlow/pkg/apiserver/models"
 	"github.com/PaddlePaddle/PaddleFlow/pkg/common/logger"
+	"github.com/PaddlePaddle/PaddleFlow/pkg/pipeline"
 	"github.com/PaddlePaddle/PaddleFlow/pkg/storage/driver"
 )
 
@@ -73,3 +76,33 @@ func TestGetJobByRun(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, "job-run-post", jobView.JobID)
 }
+
+func TestUpdateRunByWfEvent(t *testing.T) {
+	driver.InitMockDB()
+	ctx := &logger.RequestContext{UserName: MockRootUser}
+	run := getMockRunWithoutRuntime()
+	runID, err := models.CreateRun(ctx.Logging(), &run)
+	assert.Nil(t, err)
+
+	wfMap.Store(runID, "abc")
+	event := &pipeline.WorkflowEvent{
+		Event: pipeline.WfEventRunUpdate,
+		Extra: map[string]interface{}{
+			common.WfEventKeyRunID:     runID,
+			common.WfEventKeyStatus:    common.StatusRunSucceeded,
+			common.WfEventKeyStartTime: "2022-09-09 10:00:09",
+		},
+		Message: "mesg",
+	}
+
+	patch5 := gomonkey.ApplyFunc(models.GetRunByID, func(logEntry *log.Entry, runID string) (models.Run, error) {
+		return run, nil
+	})
+	defer patch5.Reset()
+
+	_, flag := UpdateRunByWfEvent(runID, event)
+	assert.True(t, flag)
+
+	_, ok := wfMap.Load(runID)
+	assert.False(t, ok)
+}
diff --git a/pkg/apiserver/controller/pipeline/callbacks.go b/pkg/apiserver/controller/pipeline/callbacks.go
@@ -109,7 +109,7 @@ func UpdateRunByWfEvent(id string, event interface{}) (int64, bool) {
 	}
 	if common.IsRunFinalStatus(status) {
 		logging.Debugf("run[%s] has reached final status[%s]", runID, status)
-		delete(wfMap, runID)
+		wfMap.Delete(runID)
 	}
 	startTime, ok := wfEvent.Extra[common.WfEventKeyStartTime].(string)
 	if !ok {
@@ -152,7 +152,7 @@ func UpdateRunByWfEvent(id string, event interface{}) (int64, bool) {
 
 	if common.IsRunFinalStatus(status) {
 		logging.Debugf("run[%s] has reached final status[%s]", runID, status)
-		delete(wfMap, runID)
+		wfMap.Delete(runID)
 
 		// 给scheduler发concurrency channel信号
 		if prevRun.ScheduleID != "" {
diff --git a/pkg/apiserver/controller/pipeline/run.go b/pkg/apiserver/controller/pipeline/run.go
@@ -22,6 +22,7 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"sync"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -43,7 +44,7 @@ import (
 	"github.com/PaddlePaddle/PaddleFlow/pkg/trace_logger"
 )
 
-var wfMap = make(map[string]*pipeline.Workflow, 0)
+var wfMap = sync.Map{}
 
 const (
 	JsonFsOptions   = "fs_options" // 由于在获取BodyMap的FsOptions前已经转为下划线形式，因此这里为fs_options
@@ -913,12 +914,14 @@ func StopRun(logEntry *log.Entry, userName, runID string, request UpdateRunReque
 		return err
 	}
 
-	wf, exist := wfMap[runID]
+	wfInterface, exist := wfMap.Load(runID)
 	if !exist {
 		err := fmt.Errorf("run[%s]'s workflow ptr is lost", runID)
 		logEntry.Errorln(err.Error())
 		return err
 	}
+
+	wf := wfInterface.(*pipeline.Workflow)
 	run.RunOptions.StopForce = request.StopForce
 	runUpdate := models.Run{
 		Status:     common.StatusRunTerminating,
@@ -1121,8 +1124,8 @@ func StartWf(run models.Run, wfPtr *pipeline.Workflow) error {
 		logEntry.Errorf("StartWf failed, error: %s", err.Error())
 		return err
 	}
-	wfMap[run.ID] = wfPtr
 
+	wfMap.Store(run.ID, wfPtr)
 	if err := models.UpdateRunStatus(logEntry, run.ID, common.StatusRunPending); err != nil {
 		return err
 	}
@@ -1216,7 +1219,7 @@ func RestartWf(run models.Run, isResume bool) (string, error) {
 	if isResume {
 		wfPtr.Resume(entryPointDagView, run.PostProcess, run.Status, run.RunOptions.StopForce)
 	} else {
-		wfMap[run.ID] = wfPtr
+		wfMap.Store(run.ID, wfPtr)
 		if err := models.UpdateRunStatus(logEntry, run.ID, common.StatusRunPending); err != nil {
 			return "", err
 		}
@@ -1246,7 +1249,7 @@ func newWorkflowByRun(run models.Run) (*pipeline.Workflow, error) {
 	}
 	// 如果此时没有runID的话，那么在后续有runID之后，需要：1. 填充wfMap 2. 初始化wf.runtime
 	if run.ID != "" {
-		wfMap[run.ID] = wfPtr
+		wfMap.Store(run.ID, wfPtr)
 	}
 	return wfPtr, nil
 }
diff --git a/pkg/apiserver/controller/pipeline/run_test.go b/pkg/apiserver/controller/pipeline/run_test.go
@@ -20,9 +20,11 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"reflect"
 	"testing"
 
 	"github.com/agiledragon/gomonkey/v2"
+	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -289,6 +291,13 @@ func TestNewWorkflowByRun(t *testing.T) {
 	}
 	_, err = newMockWorkflowByRun(run2)
 	assert.Nil(t, err)
+
+	run2.ID = "1445"
+	_, err = newMockWorkflowByRun(run2)
+	assert.Nil(t, err)
+
+	_, ok := wfMap.Load(run2.ID)
+	assert.True(t, ok)
 }
 
 func TestCreateRunByJson(t *testing.T) {
@@ -363,5 +372,118 @@ func TestCreateRun(t *testing.T) {
 	}
 	_, err = CreateRun(ctx, &createRunRequest, map[string]string{})
 	assert.Nil(t, err)
+}
+
+func TestStopRun(t *testing.T) {
+	driver.InitMockDB()
+	run := getMockRunWithoutRuntime()
+	ctx := &logger.RequestContext{UserName: MockRootUser}
+	runID, err := models.CreateRun(ctx.Logging(), &run)
+	assert.Nil(t, err)
+
+	wfMap.Store(runID, &pipeline.Workflow{})
+
+	var wf *pipeline.Workflow
+	patch := gomonkey.ApplyMethod(reflect.TypeOf(wf), "Stop", func(*pipeline.Workflow, bool) {
+		return
+	})
+	defer patch.Reset()
+
+	var r *models.Run
+	patch2 := gomonkey.ApplyMethod(reflect.TypeOf(r), "Encode", func(*models.Run) error {
+		return nil
+	})
+	defer patch2.Reset()
 
+	patch3 := gomonkey.ApplyFunc(GetRunByID, func(logEntry *log.Entry, userName string, runID string) (models.Run, error) {
+		return run, nil
+	})
+	defer patch3.Reset()
+
+	req := UpdateRunRequest{StopForce: false}
+	err = StopRun(ctx.Logging(), "root", runID, req)
+	assert.Nil(t, err)
+}
+
+func TestStartWf(t *testing.T) {
+	run := models.Run{
+		ID: "run=00001",
+	}
+	wfptr := &pipeline.Workflow{}
+	patch := gomonkey.ApplyMethod(reflect.TypeOf(wfptr), "NewWorkflowRuntime", func(*pipeline.Workflow) error {
+		return nil
+	})
+	defer patch.Reset()
+
+	patch2 := gomonkey.ApplyFunc(models.UpdateRunStatus, func(logEntry *log.Entry, runID string, status string) error {
+		return nil
+	})
+	defer patch2.Reset()
+
+	patch3 := gomonkey.ApplyMethod(reflect.TypeOf(wfptr), "Start", func(*pipeline.Workflow) {
+		return
+	})
+	defer patch3.Reset()
+
+	StartWf(run, wfptr)
+	_, ok := wfMap.Load(run.ID)
+	assert.True(t, ok)
+}
+
+func TestRestartWf(t *testing.T) {
+	run := models.Run{
+		ID: "run=00001",
+	}
+	wfptr := &pipeline.Workflow{}
+
+	patch := gomonkey.ApplyFunc(newWorkflowByRun, func(run models.Run) (*pipeline.Workflow, error) {
+		return &pipeline.Workflow{}, nil
+	})
+	defer patch.Reset()
+
+	patch2 := gomonkey.ApplyFunc(models.UpdateRunStatus, func(logEntry *log.Entry, runID string, status string) error {
+		return nil
+	})
+	defer patch2.Reset()
+
+	patch3 := gomonkey.ApplyMethod(reflect.TypeOf(wfptr), "Restart", func(*pipeline.Workflow, *schema.DagView, schema.PostProcessView) {
+		return
+	})
+	defer patch3.Reset()
+
+	patch4 := gomonkey.ApplyFunc(models.GetRunJobsOfRun, func(logEntry *log.Entry, runID string) ([]models.RunJob, error) {
+		return nil, nil
+	})
+	defer patch4.Reset()
+
+	patch5 := gomonkey.ApplyFunc(models.GetRunDagsOfRun, func(logEntry *log.Entry, runID string) ([]models.RunDag, error) {
+		return nil, nil
+	})
+	defer patch5.Reset()
+
+	var r *models.Run
+	patch6 := gomonkey.ApplyMethod(reflect.TypeOf(r), "Encode", func(*models.Run) error {
+		return nil
+	})
+	defer patch6.Reset()
+
+	patch7 := gomonkey.ApplyFunc(models.CreateRun, func(logEntry *log.Entry, run *models.Run) (string, error) {
+		return "", nil
+	})
+	defer patch7.Reset()
+
+	patch8 := gomonkey.ApplyFunc(models.CreateRunDag, func(logEntry *log.Entry, runDag *models.RunDag) (int64, error) {
+		return 234, nil
+	})
+	defer patch8.Reset()
+
+	patch9 := gomonkey.ApplyMethod(reflect.TypeOf(r), "InitRuntime", func(_ *models.Run, jobs []models.RunJob, dags []models.RunDag) error {
+		return nil
+	})
+	defer patch9.Reset()
+
+	id, err := RestartWf(run, false)
+	assert.Nil(t, err)
+	_, ok := wfMap.Load(id)
+	assert.True(t, ok)
 }
diff --git a/pkg/pipeline/dagRuntime.go b/pkg/pipeline/dagRuntime.go
@@ -483,7 +483,7 @@ func (drt *DagRuntime) Resume(dagView *schema.DagView) {
 	for _, name := range sorted {
 		views, ok := dagView.EntryPoints[name]
 		if !ok {
-			// 说面当前节点还没有运行，在此处不进行处理
+			// 说明当前节点还没有运行，在此处不进行处理
 			continue
 		}
 
@@ -734,7 +734,7 @@ func (drt *DagRuntime) scheduleSubComponentAccordingView(dagView *schema.DagView
 			continue
 		}
 
-		// restart 时，所有子节点rumtine都处于终态，可以分成三类：
+		// restart 时，所有子节点rumtime都处于终态，可以分成三类：
 		// succeeded, skipped: 对于这类runtime无需重启，在 subruntime 中记录即可
 		// failed， terminated: 需要重启
 		// cancelled: 分两种情况：
diff --git a/pkg/pipeline/job.go b/pkg/pipeline/job.go
@@ -88,7 +88,7 @@ func NewPaddleFlowJob(name, image, userName string, eventChannel chan<- Workflow
 }
 
 func NewPaddleFlowJobWithJobView(view *schema.JobView, image string, eventChannel chan<- WorkflowEvent,
-	mainFS *schema.FsMount, extraFS []schema.FsMount) *PaddleFlowJob {
+	mainFS *schema.FsMount, extraFS []schema.FsMount, userName string) *PaddleFlowJob {
 	pfj := PaddleFlowJob{
 		BaseJob: BaseJob{
 			ID:         view.JobID,
@@ -106,6 +106,7 @@ func NewPaddleFlowJobWithJobView(view *schema.JobView, image string, eventChanne
 		eventChannel: eventChannel,
 		mainFS:       mainFS,
 		extraFS:      extraFS,
+		userName:     userName,
 	}
 
 	pfj.Status = common.StatusRunRunning
@@ -211,7 +212,7 @@ func (pfj *PaddleFlowJob) Start() (string, error) {
 	}
 
 	if pfj.ID == "" {
-		err = fmt.Errorf("watch paddleflow job[%s] failed, job not started, id is empty!", pfj.Job().Name)
+		err = fmt.Errorf("watch paddleflow job[%s] failed, job not started, id is empty", pfj.Job().Name)
 		return "", err
 	}
 
diff --git a/pkg/pipeline/stepRuntime.go b/pkg/pipeline/stepRuntime.go
@@ -193,10 +193,8 @@ func (srt *StepRuntime) Start() {
 	srt.Execute()
 }
 
-// Restart: 根据 jobView 来重启step
-// 如果 jobView 中的状态为 Succeeded, 则直接返回，无需重启
-// 如果 jobView 中的状态为 Running, 则进入监听即可
-// 否则 创建一个新的job并开始调度执行
+// Restart: 根据 jobView 来重新运行step
+// 只有step的状态为failed，terminated的时候，才会重新运行
 func (srt *StepRuntime) Restart(view *schema.JobView) {
 	srt.logger.Infof("begin to restart step[%s]", srt.name)
 	defer srt.processJobLock.Unlock()
@@ -223,7 +221,7 @@ func (srt *StepRuntime) Resume(view *schema.JobView) {
 	defer srt.catchPanic()
 
 	srt.job = NewPaddleFlowJobWithJobView(view, srt.getWorkFlowStep().DockerEnv,
-		srt.receiveEventChildren, srt.runConfig.mainFS, srt.getWorkFlowStep().ExtraFS)
+		srt.receiveEventChildren, srt.runConfig.mainFS, srt.getWorkFlowStep().ExtraFS, srt.userName)
 
 	srt.pk = view.PK
 	err := srt.updateStatus(view.Status)

Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,7 @@ func NewPaddleFlowJob(name, image, userName string, eventChannel chan<- Workflow`
`88`	`88`	`}`
`89`	`89`
`90`	`90`	`func NewPaddleFlowJobWithJobView(view *schema.JobView, image string, eventChannel chan<- WorkflowEvent,`
`91`		`- mainFS schema.FsMount, extraFS []schema.FsMount) PaddleFlowJob {`
	`91`	`+ mainFS schema.FsMount, extraFS []schema.FsMount, userName string) PaddleFlowJob {`
`92`	`92`	`pfj := PaddleFlowJob{`
`93`	`93`	`BaseJob: BaseJob{`
`94`	`94`	`ID: view.JobID,`
`@@ -106,6 +106,7 @@ func NewPaddleFlowJobWithJobView(view *schema.JobView, image string, eventChanne`
`106`	`106`	`eventChannel: eventChannel,`
`107`	`107`	`mainFS: mainFS,`
`108`	`108`	`extraFS: extraFS,`
	`109`	`+ userName: userName,`
`109`	`110`	`}`
`110`	`111`
`111`	`112`	`pfj.Status = common.StatusRunRunning`
`@@ -211,7 +212,7 @@ func (pfj *PaddleFlowJob) Start() (string, error) {`
`211`	`212`	`}`
`212`	`213`
`213`	`214`	`if pfj.ID == "" {`
`214`		`- err = fmt.Errorf("watch paddleflow job[%s] failed, job not started, id is empty!", pfj.Job().Name)`
	`215`	`+ err = fmt.Errorf("watch paddleflow job[%s] failed, job not started, id is empty", pfj.Job().Name)`
`215`	`216`	`return "", err`
`216`	`217`	`}`
`217`	`218`