txn: add txn_source field in txn (#617)

Signed-off-by: xiongjiwei <xiongjiwei1996@outlook.com>
tikv · Nov 15, 2022 · e72337b · e72337b
1 parent 9c0835c
commit e72337b
Show file tree

Hide file tree

Showing 8 changed files with 22 additions and 4 deletions.
diff --git a/go.mod b/go.mod
@@ -13,7 +13,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0
 	github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00
 	github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989
-	github.com/pingcap/kvproto v0.0.0-20221026112947-f8d61344b172
+	github.com/pingcap/kvproto v0.0.0-20221114031243-29a30c4ef9c5
 	github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.11.0

diff --git a/go.sum b/go.sum
@@ -154,8 +154,9 @@ github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00 h1:C3N3itkduZXDZ
 github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00/go.mod h1:4qGtCB0QK0wBzKtFEGDhxXnSnbQApw1gc9siScUl8ew=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
-github.com/pingcap/kvproto v0.0.0-20221026112947-f8d61344b172 h1:FYgKV9znRQmzVrrJDZ0gUfMIvKLAMU1tu1UKJib8bEQ=
 github.com/pingcap/kvproto v0.0.0-20221026112947-f8d61344b172/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
+github.com/pingcap/kvproto v0.0.0-20221114031243-29a30c4ef9c5 h1:buJ/WCoxGzznvYge7tY0e/tqSMntiZ7ztCWRnwy9Klc=
+github.com/pingcap/kvproto v0.0.0-20221114031243-29a30c4ef9c5/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
 github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81 h1:URLoJ61DmmY++Sa/yyPEQHG2s/ZBeV1FbIswHEMrdoY=
 github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=

diff --git a/integration_tests/go.mod b/integration_tests/go.mod
@@ -6,7 +6,7 @@ require (
 	github.com/ninedraft/israce v0.0.3
 	github.com/pingcap/errors v0.11.5-0.20220729040631-518f63d66278
 	github.com/pingcap/failpoint v0.0.0-20220423142525-ae43b7f4e5c3
-	github.com/pingcap/kvproto v0.0.0-20221026112947-f8d61344b172
+	github.com/pingcap/kvproto v0.0.0-20221114031243-29a30c4ef9c5
 	github.com/pingcap/tidb v1.1.0-beta.0.20221101102559-97add26c8f84
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.8.0

diff --git a/integration_tests/go.sum b/integration_tests/go.sum
@@ -407,8 +407,9 @@ github.com/pingcap/failpoint v0.0.0-20220423142525-ae43b7f4e5c3/go.mod h1:4qGtCB
 github.com/pingcap/fn v0.0.0-20200306044125-d5540d389059 h1:Pe2LbxRmbTfAoKJ65bZLmhahmvHm7n9DUxGRQT00208=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
-github.com/pingcap/kvproto v0.0.0-20221026112947-f8d61344b172 h1:FYgKV9znRQmzVrrJDZ0gUfMIvKLAMU1tu1UKJib8bEQ=
 github.com/pingcap/kvproto v0.0.0-20221026112947-f8d61344b172/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
+github.com/pingcap/kvproto v0.0.0-20221114031243-29a30c4ef9c5 h1:buJ/WCoxGzznvYge7tY0e/tqSMntiZ7ztCWRnwy9Klc=
+github.com/pingcap/kvproto v0.0.0-20221114031243-29a30c4ef9c5/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
 github.com/pingcap/log v0.0.0-20200511115504-543df19646ad/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81 h1:URLoJ61DmmY++Sa/yyPEQHG2s/ZBeV1FbIswHEMrdoY=
 github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=

diff --git a/txnkv/transaction/2pc.go b/txnkv/transaction/2pc.go
@@ -177,6 +177,9 @@ type twoPhaseCommitter struct {
 	// allowed when tikv disk full happened.
 	diskFullOpt kvrpcpb.DiskFullOpt
 
+	// txnSource is used to record the source of the transaction.
+	txnSource uint8
+
 	// The total number of kv request after batch split.
 	prewriteTotalReqNum int
 
@@ -1060,6 +1063,10 @@ func (c *twoPhaseCommitter) SetDiskFullOpt(level kvrpcpb.DiskFullOpt) {
 	c.diskFullOpt = level
 }
 
+func (c *twoPhaseCommitter) SetTxnSource(txnSource uint8) {
+	c.txnSource = txnSource
+}
+
 type ttlManagerState uint32
 
 const (

diff --git a/txnkv/transaction/commit.go b/txnkv/transaction/commit.go
@@ -75,6 +75,7 @@ func (actionCommit) handleSingleBatch(c *twoPhaseCommitter, bo *retry.Backoffer,
 		SyncLog:                c.syncLog,
 		ResourceGroupTag:       c.resourceGroupTag,
 		DiskFullOpt:            c.diskFullOpt,
+		TxnSource:              uint32(c.txnSource),
 		MaxExecutionDurationMs: uint64(client.MaxWriteExecutionTime.Milliseconds()),
 		RequestSource:          c.txn.GetRequestSource(),
 	})

diff --git a/txnkv/transaction/prewrite.go b/txnkv/transaction/prewrite.go
@@ -182,6 +182,7 @@ func (c *twoPhaseCommitter) buildPrewriteRequest(batch batchMutations, txnSize u
 		SyncLog:                c.syncLog,
 		ResourceGroupTag:       c.resourceGroupTag,
 		DiskFullOpt:            c.diskFullOpt,
+		TxnSource:              uint32(c.txnSource),
 		MaxExecutionDurationMs: uint64(client.MaxWriteExecutionTime.Milliseconds()),
 		RequestSource:          c.txn.GetRequestSource(),
 	})

diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go
@@ -121,6 +121,7 @@ type KVTxn struct {
 	resourceGroupTag        []byte
 	resourceGroupTagger     tikvrpc.ResourceGroupTagger // use this when resourceGroupTag is nil
 	diskFullOpt             kvrpcpb.DiskFullOpt
+	txnSource               uint8
 	commitTSUpperBoundCheck func(uint64) bool
 	// interceptor is used to decorate the RPC request logic related to the txn.
 	interceptor    interceptor.RPCInterceptor
@@ -326,6 +327,11 @@ func (txn *KVTxn) SetDiskFullOpt(level kvrpcpb.DiskFullOpt) {
 	txn.diskFullOpt = level
 }
 
+// SetTxnSource sets the source of the transaction.
+func (txn *KVTxn) SetTxnSource(txnSource uint8) {
+	txn.txnSource = txnSource
+}
+
 // GetDiskFullOpt gets the options of current operation in each TiKV disk usage level.
 func (txn *KVTxn) GetDiskFullOpt() kvrpcpb.DiskFullOpt {
 	return txn.diskFullOpt
@@ -409,6 +415,7 @@ func (txn *KVTxn) Commit(ctx context.Context) error {
 	}
 
 	txn.committer.SetDiskFullOpt(txn.diskFullOpt)
+	txn.committer.SetTxnSource(txn.txnSource)
 
 	defer committer.ttlManager.close()