planner: do not push down lock to pointGet/bacthPointGet when selection exists

planner/core/optimizer.go

@@ -238,6 +238,7 @@ func eliminateUnionScanAndLock(sctx sessionctx.Context, p PhysicalPlan) Physical
 	var batchPointGet *BatchPointGetPlan
 	var physLock *PhysicalLock
 	var unionScan *PhysicalUnionScan
+	var selection *PhysicalSelection
 	iteratePhysicalPlan(p, func(p PhysicalPlan) bool {
 		if len(p.Children()) > 1 {
 			return false
@@ -251,6 +252,8 @@ func eliminateUnionScanAndLock(sctx sessionctx.Context, p PhysicalPlan) Physical
 			physLock = x
 		case *PhysicalUnionScan:
 			unionScan = x
+		case *PhysicalSelection:
+			selection = x
 		}
 		return true
 	})
@@ -262,7 +265,7 @@ func eliminateUnionScanAndLock(sctx sessionctx.Context, p PhysicalPlan) Physical
 	}
 	if physLock != nil {
 		lock, waitTime := getLockWaitTime(sctx, physLock.Lock)
-		if !lock {
+		if !lock || selection != nil {
 			return p
 		}
 		if pointGet != nil {

planner/core/point_get_plan_test.go

@@ -515,15 +515,17 @@ func (s *testPointGetSuite) TestIssue20692(c *C) {
 	tk3.MustExec("begin pessimistic;")
 	tk3.MustExec("use test")
 	tk1.MustExec("delete from t where id = 1 and v = 1 and vv = 1;")
-	stop1, stop2 := make(chan struct{}), make(chan struct{})
+	stop1, stop2 := make(chan error), make(chan error)
 	go func() {
-		tk2.MustExec("insert into t values(1, 2, 3, 4);")
-		stop1 <- struct{}{}
-		tk3.MustExec("update t set id = 10, v = 20, vv = 30, vvv = 40 where id = 1 and v = 2 and vv = 3;")
-		stop2 <- struct{}{}
+		stop1 <- tk2.ExecToErr("insert into t values(1, 2, 3, 4);")
 	}()
+	time.Sleep(50 * time.Millisecond)
+	go func() {
+		stop2 <- tk3.ExecToErr("update t set id = 10, v = 20, vv = 30, vvv = 40 where id = 1 and v = 2 and vv = 3;")
+	}()
+	time.Sleep(50 * time.Millisecond)
 	tk1.MustExec("commit;")
-	<-stop1
+	c.Check(<-stop1, IsNil)
 
 	// wait 50ms to ensure tk3 is blocked by tk2
 	select {
@@ -533,7 +535,7 @@ func (s *testPointGetSuite) TestIssue20692(c *C) {
 	}
 
 	tk2.MustExec("commit;")
-	<-stop2
+	c.Check(<-stop2, IsNil)
 	tk3.MustExec("commit;")
 	tk3.MustQuery("select * from t;").Check(testkit.Rows("10 20 30 40"))
 }

session/pessimistic_test.go

@@ -1091,13 +1091,19 @@ func (s *testPessimisticSuite) TestPessimisticLockNonExistsKey(c *C) {
 	tk.MustQuery("select c + 1 from t where k = 2 and c = 2 for update").Check(testkit.Rows())
 	explainStr := tk.MustQuery("explain select c + 1 from t where k = 2 and c = 2 for update").Rows()[0][0].(string)
 	c.Assert(strings.Contains(explainStr, "UnionScan"), IsFalse)
-	tk.MustQuery("select * from t where k in (4, 5, 7) for update").Check(testkit.Rows("5 5"))
+	tk.MustQuery("select * from t where k = 3 and c = 3 for update").Check(testkit.Rows("3 3"))
+	tk.MustQuery("select c + 1 from t where k = 5 for update").Check(testkit.Rows("6"))
+	tk.MustQuery("select * from t where k in (4, 7) for update").Check(testkit.Rows())
 
 	tk1.MustExec("begin pessimistic")
 	err := tk1.ExecToErr("select * from t where k = 2 for update nowait")
+	c.Check(err, IsNil)
+	err = tk1.ExecToErr("select * from t where k = 3 for update nowait")
 	c.Check(tikv.ErrLockAcquireFailAndNoWaitSet.Equal(err), IsTrue)
 	err = tk1.ExecToErr("select * from t where k = 4 for update nowait")
 	c.Check(tikv.ErrLockAcquireFailAndNoWaitSet.Equal(err), IsTrue)
+	err = tk1.ExecToErr("select * from t where k = 5 for update nowait")
+	c.Check(tikv.ErrLockAcquireFailAndNoWaitSet.Equal(err), IsTrue)
 	err = tk1.ExecToErr("select * from t where k = 7 for update nowait")
 	c.Check(tikv.ErrLockAcquireFailAndNoWaitSet.Equal(err), IsTrue)
 	tk.MustExec("rollback")
@@ -1111,9 +1117,11 @@ func (s *testPessimisticSuite) TestPessimisticLockNonExistsKey(c *C) {
 
 	tk1.MustExec("begin pessimistic")
 	err = tk1.ExecToErr("select * from t where k = 2 for update nowait")
+	c.Check(err, IsNil)
+	err = tk1.ExecToErr("select * from t where k = 5 for update nowait")
 	c.Check(tikv.ErrLockAcquireFailAndNoWaitSet.Equal(err), IsTrue)
 	err = tk1.ExecToErr("select * from t where k = 6 for update nowait")
-	c.Check(tikv.ErrLockAcquireFailAndNoWaitSet.Equal(err), IsTrue)
+	c.Check(err, IsNil)
 	tk.MustExec("rollback")
 	tk1.MustExec("rollback")
 }
@@ -2221,3 +2229,41 @@ func (s *testPessimisticSuite) TestAmendForUniqueIndex(c *C) {
 	tk2.MustExec("admin check table t")
 	*/
 }
+
+func (s *testPessimisticSuite) TestIssue21688(c *C) {
+	tk := testkit.NewTestKitWithInit(c, s.store)
+	tk2 := testkit.NewTestKitWithInit(c, s.store)
+	tk.MustExec("use test")
+	tk2.MustExec("use test")
+
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t (k1 int, k2 int, v int, unique key (k1))")
+	tk.MustExec("insert into t values (1, 1, null), (2, 2, 2), (3, 3, 3), (4, 4, 4), (5, 5, 5)")
+
+	tk.MustExec("begin pessimistic")
+	tk2.MustExec("begin pessimistic")
+	tk.MustExec("select * from t where (k1, v) in ((1, null)) for update")
+	tk2.MustExec("select * from t where (k1, v) in ((1, null)) for update nowait")
+	tk.MustExec("select * from t where (k1, v) in ((2, null), (6, null)) for update")
+	tk2.MustExec("select * from t where (k1, v) in ((2, null), (6, null)) for update nowait")
+	tk.MustExec("rollback")
+	tk2.MustExec("rollback")
+
+	tk.MustExec("begin pessimistic")
+	tk2.MustExec("begin pessimistic")
+	tk.MustExec("select * from t where (k2, v) in ((1, null)) for update")
+	tk2.MustExec("select * from t where (k2, v) in ((1, null)) for update nowait")
+	tk.MustExec("select * from t where (k2, v) in ((2, null), (6, null)) for update")
+	tk2.MustExec("select * from t where (k2, v) in ((2, null), (6, null)) for update nowait")
+	tk.MustExec("rollback")
+	tk2.MustExec("rollback")
+
+	tk.MustExec("begin pessimistic")
+	tk2.MustExec("begin pessimistic")
+	tk.MustExec("select * from t where (k1, v) in ((1, 2)) for update")
+	tk2.MustExec("select * from t where (k1, v) in ((1, 3)) for update nowait")
+	tk.MustExec("select * from t where (k1, v) in ((2, 3), (4, 4)) for update")
+	tk2.MustExec("select * from t where (k1, v) in ((2, 3), (5, 5)) for update nowait")
+	tk.MustExec("rollback")
+	tk2.MustExec("rollback")
+}