*: lock unchanged rows for pessimistic transaction (pingcap#14045)

coocood · Dec 13, 2019 · 79807ba · 79807ba
1 parent be7cb4d
commit 79807ba
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 0 deletions.
diff --git a/executor/adapter.go b/executor/adapter.go
@@ -451,6 +451,7 @@ func (a *ExecStmt) handlePessimisticDML(ctx context.Context, e Executor) error {
 		if err1 != nil {
 			return err1
 		}
+		keys = txnCtx.CollectUnchangedRowKeys(keys)
 		if len(keys) == 0 {
 			return nil
 		}

diff --git a/executor/write.go b/executor/write.go
@@ -23,6 +23,7 @@ import (
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/table"
 	"github.com/pingcap/tidb/table/tables"
+	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/logutil"
 	"go.uber.org/zap"
@@ -116,6 +117,11 @@ func updateRecord(ctx sessionctx.Context, h int64, oldData, newData []types.Datu
 		if ctx.GetSessionVars().ClientCapability&mysql.ClientFoundRows > 0 {
 			sc.AddAffectedRows(1)
 		}
+		unchangedRowKey := tablecodec.EncodeRowKeyWithHandle(t.Meta().ID, h)
+		txnCtx := ctx.GetSessionVars().TxnCtx
+		if txnCtx.IsPessimistic {
+			txnCtx.AddUnchangedRowKey(unchangedRowKey)
+		}
 		return false, false, 0, nil
 	}
 

diff --git a/session/pessimistic_test.go b/session/pessimistic_test.go
@@ -349,6 +349,36 @@ func (s *testPessimisticSuite) TestBankTransfer(c *C) {
 	tk.MustQuery("select sum(c) from accounts").Check(testkit.Rows("300"))
 }
 
+func (s *testPessimisticSuite) TestLockUnchangedRowKey(c *C) {
+	tk := testkit.NewTestKitWithInit(c, s.store)
+	tk2 := testkit.NewTestKitWithInit(c, s.store)
+	tk.MustExec("drop table if exists unchanged")
+	tk.MustExec("create table unchanged (id int primary key, c int)")
+	tk.MustExec("insert unchanged values (1, 1), (2, 2)")
+
+	tk.MustExec("begin pessimistic")
+	tk.MustExec("update unchanged set c = 1 where id < 2")
+
+	tk2.MustExec("begin pessimistic")
+	err := tk2.ExecToErr("select * from unchanged where id = 1 for update nowait")
+	c.Assert(err, NotNil)
+
+	tk.MustExec("rollback")
+
+	tk2.MustQuery("select * from unchanged where id = 1 for update nowait")
+
+	tk.MustExec("begin pessimistic")
+	tk.MustExec("insert unchanged values (2, 2) on duplicate key update c = values(c)")
+
+	err = tk2.ExecToErr("select * from unchanged where id = 2 for update nowait")
+	c.Assert(err, NotNil)
+
+	tk.MustExec("commit")
+
+	tk2.MustQuery("select * from unchanged where id = 1 for update nowait")
+	tk2.MustExec("rollback")
+}
+
 func (s *testPessimisticSuite) TestOptimisticConflicts(c *C) {
 	tk := testkit.NewTestKitWithInit(c, s.store)
 	tk2 := testkit.NewTestKitWithInit(c, s.store)

diff --git a/sessionctx/variable/session.go b/sessionctx/variable/session.go
@@ -112,13 +112,33 @@ type TransactionContext struct {
 	TableDeltaMap map[int64]TableDelta
 	IsPessimistic bool
 
+	// unchangedRowKeys is used to store the unchanged rows that needs to lock for pessimistic transaction.
+	unchangedRowKeys map[string]struct{}
+
 	// CreateTime For metrics.
 	CreateTime     time.Time
 	StatementCount int
 
 	IsBatched bool
 }
 
+// AddUnchangedRowKey adds an unchanged row key in update statement for pessimistic lock.
+func (tc *TransactionContext) AddUnchangedRowKey(key []byte) {
+	if tc.unchangedRowKeys == nil {
+		tc.unchangedRowKeys = map[string]struct{}{}
+	}
+	tc.unchangedRowKeys[string(key)] = struct{}{}
+}
+
+// CollectUnchangedRowKeys collects unchanged row keys for pessimistic lock.
+func (tc *TransactionContext) CollectUnchangedRowKeys(buf []kv.Key) []kv.Key {
+	for key := range tc.unchangedRowKeys {
+		buf = append(buf, kv.Key(key))
+	}
+	tc.unchangedRowKeys = nil
+	return buf
+}
+
 // UpdateDeltaForTable updates the delta info for some table.
 func (tc *TransactionContext) UpdateDeltaForTable(tableID int64, delta int64, count int64, colSize map[int64]int64) {
 	if tc.TableDeltaMap == nil {