From 3e304d2c08ddcc1eb70d2844b40a25d45f62cf0f Mon Sep 17 00:00:00 2001
From: yedf2 <120050102@qq.com>
Date: Sat, 22 Jan 2022 18:02:36 +0800
Subject: [PATCH] add PrepareAndSubmitBarrier

---
 dtmcli/msg.go | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/dtmcli/msg.go b/dtmcli/msg.go
index 6fba805..6a560b0 100644
--- a/dtmcli/msg.go
+++ b/dtmcli/msg.go
@@ -42,17 +42,25 @@ func (s *Msg) Submit() error {
 
 // PrepareAndSubmit one method for the entire busi->prepare->submit
 func (s *Msg) PrepareAndSubmit(queryPrepared string, db *sql.DB, busiCall BarrierBusiFunc) error {
+	return s.PrepareAndSubmitBarrier(queryPrepared, func(bb *BranchBarrier) error {
+		return bb.CallWithDB(db, busiCall)
+	})
+}
+
+// PrepareAndSubmitBarrier one method for the entire busi->prepare->submit
+func (s *Msg) PrepareAndSubmitBarrier(queryPrepared string, busiCall func(bb *BranchBarrier) error) error {
 	bb, err := BarrierFrom(s.TransType, s.Gid, "00", "msg") // a special barrier for msg QueryPrepared
 	if err == nil {
 		err = s.Prepare(queryPrepared)
 	}
 	if err == nil {
-		defer func() {
-			if err != nil && bb.QueryPrepared(db) == ErrFailure {
-				_ = dtmimp.TransCallDtm(&s.TransBase, s, "abort")
-			}
-		}()
-		err = bb.CallWithDB(db, busiCall)
+		err = busiCall(bb)
+		if err != nil && err != ErrFailure {
+			_, err = dtmimp.TransRequestBranch(&s.TransBase, nil, bb.BranchID, bb.Op, queryPrepared)
+		}
+		if err == ErrFailure {
+			_ = dtmimp.TransCallDtm(&s.TransBase, s, "abort")
+		}
 	}
 	if err == nil {
 		err = s.Submit()