feat: XA connect manager (#433)

* xa connect manager

go.mod

@@ -42,6 +42,7 @@ require (
 	github.com/benbjohnson/clock v1.1.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/bits-and-blooms/bitset v1.2.0 // indirect
+	github.com/bluele/gcache v0.0.2 // indirect
 	github.com/buger/jsonparser v1.1.1 // indirect
 	github.com/census-instrumentation/opencensus-proto v0.2.1 // indirect
 	github.com/cespare/xxhash/v2 v2.1.2 // indirect

go.sum

@@ -122,6 +122,8 @@ github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kB
 github.com/bits-and-blooms/bitset v1.2.0 h1:Kn4yilvwNtMACtf1eYDlG8H77R07mZSPbMjLyS07ChA=
 github.com/bits-and-blooms/bitset v1.2.0/go.mod h1:gIdJ4wp64HaoK2YrL1Q5/N7Y16edYb8uY+O0FJTyyDA=
 github.com/bketelsen/crypt v0.0.3-0.20200106085610-5cbc8cc4026c/go.mod h1:MKsuJmJgSg28kpZDP6UIiPt0e0Oz0kqKNGyRaWEPv84=
+github.com/bluele/gcache v0.0.2 h1:WcbfdXICg7G/DGBh1PFfcirkWOQV+v077yF1pSy3DGw=
+github.com/bluele/gcache v0.0.2/go.mod h1:m15KV+ECjptwSPxKhOhQoAFQVtUFjTVkc3H8o0t/fp0=
 github.com/buger/jsonparser v0.0.0-20181115193947-bf1c66bbce23/go.mod h1:bbYlZJ7hK1yFx9hf58LP0zeX7UjIGs20ufpu3evjr+s=
 github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs=
 github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=

pkg/datasource/sql/conn.go

@@ -237,6 +237,10 @@ func (c *Conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, e
 	)
 }
 
+func (c *Conn) GetAutoCommit() bool {
+	return c.autoCommit
+}
+
 // Close invalidates and potentially stops any current
 // prepared statements and transactions, marking this
 // connection as no longer in use.

pkg/datasource/sql/datasource_resource.go

@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package sql
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/bluele/gcache"
+
+	"github.com/seata/seata-go/pkg/datasource/sql/types"
+	"github.com/seata/seata-go/pkg/protocol/branch"
+)
+
+type Holdable interface {
+	SetHeld(held bool)
+	IsHeld() bool
+	ShouldBeHeld() bool
+}
+
+type BaseDataSourceResource struct {
+	db           *DBResource
+	shouldBeHeld bool
+	keeper       map[string]interface{}
+	Cache        map[string]branch.BranchStatus
+}
+
+var BranchStatusCache = gcache.New(1024).LRU().Expiration(time.Minute * 10).Build()
+
+func (b *BaseDataSourceResource) init() error {
+	return nil
+}
+
+func (b *BaseDataSourceResource) GetDB() *DBResource {
+	return b.db
+}
+
+func (b *BaseDataSourceResource) SetDB(db *DBResource) {
+	b.db = db
+}
+
+func (b *BaseDataSourceResource) IsShouldBeHeld() bool {
+	return b.shouldBeHeld
+}
+
+func (b *BaseDataSourceResource) SetShouldBeHeld(shouldBeHeld bool) {
+	b.shouldBeHeld = shouldBeHeld
+}
+
+func (b *BaseDataSourceResource) GetKeeper() map[string]interface{} {
+	return b.keeper
+}
+
+func (b *BaseDataSourceResource) SetKeeper(keeper map[string]interface{}) {
+	b.keeper = keeper
+}
+
+func (b *BaseDataSourceResource) GetCache() map[string]branch.BranchStatus {
+	return b.Cache
+}
+
+func (b *BaseDataSourceResource) SetCache(cache map[string]branch.BranchStatus) {
+	b.Cache = cache
+}
+
+func (b *BaseDataSourceResource) GetResourceId() string {
+	return b.db.GetResourceId()
+}
+
+func (b *BaseDataSourceResource) Hold(key string, value Holdable) (interface{}, error) {
+	if value.IsHeld() {
+		var x = b.keeper[key]
+		if x != value {
+			return nil, fmt.Errorf("something wrong with keeper, keeping[%v] but[%v] is also kept with the same key[%v]", x, value, key)
+		}
+		return value, nil
+	}
+	var x = b.keeper[key]
+	b.keeper[key] = value
+	value.SetHeld(true)
+	return x, nil
+}
+
+func (b *BaseDataSourceResource) Release(key string, value Holdable) (interface{}, error) {
+	if value.IsHeld() {
+		var x = b.keeper[key]
+		if x != value {
+			return nil, fmt.Errorf("something wrong with keeper, keeping[%v] but[%v] is also kept with the same key[%v]", x, value, key)
+		}
+		return value, nil
+	}
+	var x = b.keeper[key]
+	b.keeper[key] = value
+	value.SetHeld(true)
+	return x, nil
+}
+
+func (b *BaseDataSourceResource) GetBranchStatus(xaBranchXid string) (interface{}, error) {
+	branchStatus, err := BranchStatusCache.GetIFPresent(xaBranchXid)
+	return branchStatus, err
+}
+
+func (b *BaseDataSourceResource) GetDbType() string {
+	return b.db.dbType.String()
+}
+
+func (b *BaseDataSourceResource) SetDbType(dbType types.DBType) {
+	b.db.dbType = dbType
+}

pkg/datasource/sql/root_context.go

@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package sql
+
+import (
+	"github.com/seata/seata-go/pkg/protocol/branch"
+)
+
+type RootContext interface {
+	RootContext()
+	SetDefaultBranchType(branchType branch.BranchType)
+	GetXID() string
+	Bind(xid string)
+	GetTimeout() (int, bool)
+	SetTimeout(timeout int)
+}

pkg/datasource/sql/xa/xa_connection.go

@@ -0,0 +1,26 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package xa
+
+import (
+	"github.com/seata/seata-go/pkg/datasource/sql/exec"
+)
+
+type XAConnection interface {
+	getXAResource() (exec.XAResource, error)
+}

pkg/datasource/sql/xa/xa_connection_proxy.go

@@ -0,0 +1,325 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package xa
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/seata/seata-go/pkg/datasource/sql"
+	"github.com/seata/seata-go/pkg/datasource/sql/datasource"
+	"github.com/seata/seata-go/pkg/datasource/sql/exec"
+	"github.com/seata/seata-go/pkg/protocol/branch"
+	"github.com/seata/seata-go/pkg/protocol/message"
+	"github.com/seata/seata-go/pkg/rm"
+)
+
+type ConnectionProxyXA struct {
+	xaBranchXid             *XABranchXid
+	currentAutoCommitStatus bool  `default:"true"`
+	xaActive                bool  `default:"false"`
+	kept                    bool  `default:"false"`
+	rollBacked              bool  `default:"false"`
+	branchRegisterTime      int64 `default:"0"`
+	prepareTime             int64 `default:"0"`
+	timeout                 int   `default:"0"`
+	proxyShouldBeHeld       bool  `default:"false"`
+	originalConnection      sql.Conn
+	xaConnection            XAConnection
+	xaResource              exec.XAResource
+	resource                sql.BaseDataSourceResource
+	xid                     string
+}
+
+const timeout int = 60000
+
+func NewConnectionProxyXA(originalConnection sql.Conn, xaConnection XAConnection, resource sql.BaseDataSourceResource, xid string) (*ConnectionProxyXA, error) {
+	connectionProxyXA := &ConnectionProxyXA{}
+
+	connectionProxyXA.originalConnection = originalConnection
+	connectionProxyXA.xaConnection = xaConnection
+	connectionProxyXA.resource = resource
+	connectionProxyXA.xid = xid
+
+	connectionProxyXA.proxyShouldBeHeld = connectionProxyXA.resource.IsShouldBeHeld()
+
+	xaResource, err := xaConnection.getXAResource()
+	if err != nil {
+		return nil, fmt.Errorf("get xa resource failed")
+	} else {
+		connectionProxyXA.xaResource = xaResource
+	}
+	var rootContext sql.RootContext
+	transactionTimeout, ok := rootContext.GetTimeout()
+	if !ok {
+		transactionTimeout = timeout
+	}
+	if transactionTimeout < timeout {
+		transactionTimeout = timeout
+	}
+	connectionProxyXA.timeout = transactionTimeout
+	connectionProxyXA.currentAutoCommitStatus = connectionProxyXA.originalConnection.GetAutoCommit()
+	if !connectionProxyXA.currentAutoCommitStatus {
+		return nil, fmt.Errorf("connection[autocommit=false] as default is NOT supported")
+	}
+
+	return connectionProxyXA, nil
+}
+
+func (c *ConnectionProxyXA) keepIfNecessary() {
+	if c.ShouldBeHeld() {
+		c.resource.Hold(c.xaBranchXid.String(), c)
+	}
+}
+
+func (c *ConnectionProxyXA) releaseIfNecessary() {
+	if c.ShouldBeHeld() {
+		if c.xaBranchXid == nil {
+			if c.IsHeld() {
+				c.resource.Release(c.xaBranchXid.String(), c)
+			}
+		}
+	}
+}
+
+func (c *ConnectionProxyXA) XaCommit(xid string, branchId int64) error {
+	xaXid := Build(xid, branchId)
+	err := c.xaResource.Commit(xaXid.String(), false)
+	c.releaseIfNecessary()
+	return err
+}
+
+func (c *ConnectionProxyXA) XaRollbackByBranchId(xid string, branchId int64) {
+	xaXid := Build(xid, branchId)
+	c.XaRollback(xaXid)
+}
+
+func (c *ConnectionProxyXA) XaRollback(xaXid XAXid) error {
+	err := c.xaResource.Rollback(xaXid.GetGlobalXid())
+	c.releaseIfNecessary()
+	return err
+}
+
+func (c *ConnectionProxyXA) SetAutoCommit(autoCommit bool) error {
+	if c.currentAutoCommitStatus == autoCommit {
+		return nil
+	}
+	if autoCommit {
+		if c.xaActive {
+			_ = c.Commit()
+		}
+	} else {
+		if c.xaActive {
+			return fmt.Errorf("should NEVER happen: setAutoCommit from true to false while xa branch is active")
+		}
+
+		c.branchRegisterTime = time.Now().UnixMilli()
+		var branchRegisterParam rm.BranchRegisterParam
+		branchRegisterParam.BranchType = branch.BranchTypeXA
+		branchRegisterParam.ResourceId = c.resource.GetResourceId()
+		branchRegisterParam.Xid = c.xid
+		branchId, err := datasource.GetDataSourceManager(branch.BranchTypeXA).BranchRegister(context.TODO(), branchRegisterParam)
+		if err != nil {
+			c.cleanXABranchContext()
+			return fmt.Errorf("failed to register xa branch [%v]", c.xid)
+		}
+		c.xaBranchXid = Build(c.xid, branchId)
+		c.keepIfNecessary()
+		err = c.start()
+		if err != nil {
+			c.cleanXABranchContext()
+			return fmt.Errorf("failed to start xa branch [%v]", c.xid)
+		}
+		c.xaActive = true
+	}
+	c.currentAutoCommitStatus = autoCommit
+	return nil
+}
+
+func (c *ConnectionProxyXA) GetAutoCommit() bool {
+	return c.currentAutoCommitStatus
+}
+
+func (c *ConnectionProxyXA) Commit() error {
+	if c.currentAutoCommitStatus {
+		return nil
+	}
+	if !c.xaActive || c.xaBranchXid == nil {
+		return fmt.Errorf("should NOT commit on an inactive session")
+	}
+	now := time.Now().UnixMilli()
+	if c.end(exec.TMSUCCESS) != nil {
+		return c.commitErrorHandle()
+	}
+	if c.checkTimeout(now) != nil {
+		return c.commitErrorHandle()
+	}
+	if c.xaResource.XAPrepare(c.xaBranchXid.String()) != nil {
+		return c.commitErrorHandle()
+	}
+	return nil
+}
+
+func (c *ConnectionProxyXA) commitErrorHandle() error {
+	req := message.BranchReportRequest{
+		BranchType:      branch.BranchTypeXA,
+		Xid:             c.xid,
+		BranchId:        c.xaBranchXid.GetBranchId(),
+		Status:          branch.BranchStatusPhaseoneFailed,
+		ApplicationData: nil,
+		ResourceId:      c.resource.GetResourceId(),
+	}
+	if datasource.NewBasicSourceManager().BranchReport(context.TODO(), req) != nil {
+		c.cleanXABranchContext()
+		return fmt.Errorf("Failed to report XA branch commit-failure on [%v] - [%v]", c.xid, c.xaBranchXid.GetBranchId())
+	}
+	c.cleanXABranchContext()
+	return fmt.Errorf("Failed to end(TMSUCCESS)/prepare xa branch on [%v] - [%v]", c.xid, c.xaBranchXid.GetBranchId())
+}
+
+func (c *ConnectionProxyXA) Rollback() error {
+	if c.currentAutoCommitStatus {
+		return nil
+	}
+	if !c.xaActive || c.xaBranchXid == nil {
+		return fmt.Errorf("should NOT rollback on an inactive session")
+	}
+	if !c.rollBacked {
+		if c.xaResource.End(c.xaBranchXid.String(), exec.TMFAIL) != nil {
+			return c.rollbackErrorHandle()
+		}
+		if c.XaRollback(c.xaBranchXid) != nil {
+			c.cleanXABranchContext()
+			return c.rollbackErrorHandle()
+		}
+		req := message.BranchReportRequest{
+			BranchType:      branch.BranchTypeXA,
+			Xid:             c.xid,
+			BranchId:        c.xaBranchXid.GetBranchId(),
+			Status:          branch.BranchStatusPhaseoneFailed,
+			ApplicationData: nil,
+			ResourceId:      c.resource.GetResourceId(),
+		}
+		if datasource.NewBasicSourceManager().BranchReport(context.TODO(), req) != nil {
+			c.cleanXABranchContext()
+			return fmt.Errorf("failed to report XA branch commit-failure on [%v] - [%v]", c.xid, c.xaBranchXid.GetBranchId())
+		}
+	}
+	c.cleanXABranchContext()
+	return nil
+}
+
+func (c *ConnectionProxyXA) rollbackErrorHandle() error {
+	return fmt.Errorf("failed to end(TMFAIL) xa branch on [%v] - [%v]", c.xid, c.xaBranchXid.GetBranchId())
+}
+
+func (c *ConnectionProxyXA) start() error {
+	err := c.xaResource.Start(c.xaBranchXid.String(), exec.TMNOFLAGS)
+	if err := c.termination(c.xaBranchXid.String()); err != nil {
+		c.xaResource.End(c.xaBranchXid.String(), exec.TMFAIL)
+		c.XaRollback(c.xaBranchXid)
+		return err
+	}
+	return err
+}
+
+func (c *ConnectionProxyXA) end(flags int) error {
+	err := c.termination(c.xaBranchXid.String())
+	if err != nil {
+		return err
+	}
+	err = c.xaResource.End(c.xaBranchXid.String(), flags)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (c *ConnectionProxyXA) cleanXABranchContext() {
+	c.branchRegisterTime = 0
+	c.prepareTime = 0
+	c.timeout = 0
+	c.xaActive = false
+	if !c.IsHeld() {
+		c.xaBranchXid = nil
+	}
+}
+
+func (c *ConnectionProxyXA) checkTimeout(now int64) error {
+	if now-c.branchRegisterTime > int64(c.timeout) {
+		c.XaRollback(c.xaBranchXid)
+		return fmt.Errorf("XA branch timeout error")
+	}
+	return nil
+}
+
+func (c *ConnectionProxyXA) Close() error {
+	c.rollBacked = false
+	if c.IsHeld() && c.ShouldBeHeld() {
+		return nil
+	}
+	c.cleanXABranchContext()
+	if err := c.originalConnection.Close(); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (c *ConnectionProxyXA) CloseForce() error {
+	physicalConn := c.originalConnection
+	if err := physicalConn.Close(); err != nil {
+		return err
+	}
+	c.rollBacked = false
+	c.cleanXABranchContext()
+	if err := c.originalConnection.Close(); err != nil {
+		return err
+	}
+	c.releaseIfNecessary()
+	return nil
+}
+
+func (c *ConnectionProxyXA) SetHeld(kept bool) {
+	c.kept = kept
+}
+
+func (c *ConnectionProxyXA) IsHeld() bool {
+	return c.kept
+}
+
+func (c *ConnectionProxyXA) ShouldBeHeld() bool {
+	return c.proxyShouldBeHeld || c.resource.GetDB() != nil
+}
+
+func (c *ConnectionProxyXA) GetPrepareTime() int64 {
+	return c.prepareTime
+}
+
+func (c *ConnectionProxyXA) setPrepareTime(prepareTime int64) {
+	c.prepareTime = prepareTime
+}
+
+func (c *ConnectionProxyXA) termination(xaBranchXid string) error {
+	branchStatus, err := c.resource.GetBranchStatus(xaBranchXid)
+	if err != nil {
+		c.releaseIfNecessary()
+		return fmt.Errorf("failed xa branch [%v] the global transaction has finish, branch status: [%v]", c.xid, branchStatus)
+	}
+	return nil
+}