From 0975ae5f32ca58710cc61d10bd810fccb59170fc Mon Sep 17 00:00:00 2001
From: yedf2 <120050102@qq.com>
Date: Sat, 9 Jul 2022 17:43:32 +0800
Subject: [PATCH] change implementation of CompensateErrorBranch

---
 dtmgrpc/workflow/imp.go      | 10 ++++++----
 dtmgrpc/workflow/workflow.go |  7 +++----
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/dtmgrpc/workflow/imp.go b/dtmgrpc/workflow/imp.go
index 26b62a3..6693735 100644
--- a/dtmgrpc/workflow/imp.go
+++ b/dtmgrpc/workflow/imp.go
@@ -18,7 +18,6 @@ type workflowImp struct {
 	currentActionAdded   bool                   //nolint
 	currentCommitAdded   bool                   //nolint
 	currentRollbackAdded bool                   //nolint
-	currentRollbackItem  *workflowPhase2Item    // nolint
 	progresses           map[string]*stepResult //nolint
 	currentOp            string
 	succeededOps         []workflowPhase2Item
@@ -161,10 +160,13 @@ func (wf *Workflow) callPhase2(branchID string, fn WfPhase2Func) error {
 
 func (wf *Workflow) recordedDo(fn func(bb *dtmcli.BranchBarrier) *stepResult) *stepResult {
 	sr := wf.recordedDoInner(fn)
-	if wf.currentRollbackItem != nil && (sr.Status == dtmcli.StatusSucceed || sr.Status == dtmcli.StatusFailed && wf.Options.CompensateErrorBranch) {
-		wf.failedOps = append(wf.failedOps, *wf.currentRollbackItem)
+	// donot compensate the failed branch if !CompensateErrorBranch
+	if !wf.Options.CompensateErrorBranch && sr.Status == dtmcli.StatusFailed {
+		lastFailed := len(wf.failedOps) - 1
+		if lastFailed >= 0 && wf.failedOps[lastFailed].branchID == wf.currentBranch {
+			wf.failedOps = wf.failedOps[:lastFailed]
+		}
 	}
-	wf.currentRollbackItem = nil
 	return sr
 }
 
diff --git a/dtmgrpc/workflow/workflow.go b/dtmgrpc/workflow/workflow.go
index e9a66aa..0e13723 100644
--- a/dtmgrpc/workflow/workflow.go
+++ b/dtmgrpc/workflow/workflow.go
@@ -122,12 +122,11 @@ func (wf *Workflow) OnRollback(compensate WfPhase2Func) *Workflow {
 	branchID := wf.currentBranch
 	dtmimp.PanicIf(wf.currentRollbackAdded, fmt.Errorf("one branch can only add one rollback callback"))
 	wf.currentRollbackAdded = true
-	item := workflowPhase2Item{
+	wf.failedOps = append(wf.failedOps, workflowPhase2Item{
 		branchID: branchID,
-		op:       dtmimp.OpRollback,
+		op:       dtmimp.OpCommit,
 		fn:       compensate,
-	}
-	wf.currentRollbackItem = &item
+	})
 	return wf
 }