*: implement renew write lock lease for cached table

go.mod

@@ -65,7 +65,7 @@ require (
  github.com/spf13/pflag v1.0.5
  github.com/stretchr/testify v1.7.0
  github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2
- github.com/tikv/client-go/v2 v2.0.0-rc.0.20211213075151-b147ced35a14
+ github.com/tikv/client-go/v2 v2.0.0-rc.0.20211214093715-605f49d3ba50
  github.com/tikv/pd v1.1.0-beta.0.20211118054146-02848d2660ee
  github.com/twmb/murmur3 v1.1.3
  github.com/uber/jaeger-client-go v2.22.1+incompatible

go.sum

@@ -712,8 +712,8 @@ github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2/go.mod h1:2PfK
 github.com/tidwall/gjson v1.3.5/go.mod h1:P256ACg0Mn+j1RXIDXoss50DeIABTYK1PULOJHhxOls=
 github.com/tidwall/match v1.0.1/go.mod h1:LujAq0jyVjBy028G1WhWfIzbpQfMO8bBZ6Tyb0+pL9E=
 github.com/tidwall/pretty v1.0.0/go.mod h1:XNkn88O1ChpSDQmQeStsy+sBenx6DDtFZJxhVysOjyk=
-github.com/tikv/client-go/v2 v2.0.0-rc.0.20211213075151-b147ced35a14 h1:l2T+gfgYpwmLRY5geDq1zM4Lz4X2mi1ruO18/bDGo70=
-github.com/tikv/client-go/v2 v2.0.0-rc.0.20211213075151-b147ced35a14/go.mod h1:wRuh+W35daKTiYBld0oBlT6PSkzEVr+pB/vChzJZk+8=
+github.com/tikv/client-go/v2 v2.0.0-rc.0.20211214093715-605f49d3ba50 h1:B+cAIm2P1/SNsVV1vL9/mRaGUVl/vdgV8MU03O0vY28=
+github.com/tikv/client-go/v2 v2.0.0-rc.0.20211214093715-605f49d3ba50/go.mod h1:wRuh+W35daKTiYBld0oBlT6PSkzEVr+pB/vChzJZk+8=
 github.com/tikv/pd v1.1.0-beta.0.20211029083450-e65f0c55b6ae/go.mod h1:varH0IE0jJ9E9WN2Ei/N6pajMlPkcXdDEf7f5mmsUVQ=
 github.com/tikv/pd v1.1.0-beta.0.20211118054146-02848d2660ee h1:rAAdvQ8Hh36syHr92g0VmZEpkH+40RGQBpFL2121xMs=
 github.com/tikv/pd v1.1.0-beta.0.20211118054146-02848d2660ee/go.mod h1:lRbwxBAhnTQR5vqbTzeI/Bj62bD2OvYYuFezo2vrmeI=

kv/option.go

@@ -71,6 +71,9 @@ const (
 
  // SnapInterceptor is used for setting the interceptor for snapshot
  SnapInterceptor
+ // CommitTSUpperBoundChec is used by cached table
+ // The commitTS must be greater than all the write lock lease of the visited cached table.
+ CommitTSUpperBoundCheck
 )
 
 // ReplicaReadType is the type of replica to read data from

session/session.go

@@ -44,6 +44,7 @@ import (
  "github.com/pingcap/tidb/parser/model"
  "github.com/pingcap/tidb/parser/mysql"
  "github.com/pingcap/tidb/parser/terror"
+ "github.com/pingcap/tidb/table/tables"
  "github.com/pingcap/tidb/table/temptable"
  "github.com/pingcap/tidb/util/topsql"
  "github.com/pingcap/tipb/go-binlog"
@@ -89,6 +90,7 @@ import (
  "github.com/pingcap/tidb/util/tableutil"
  "github.com/pingcap/tidb/util/timeutil"
  tikvstore "github.com/tikv/client-go/v2/kv"
+ "github.com/tikv/client-go/v2/oracle"
  "github.com/tikv/client-go/v2/tikv"
  tikvutil "github.com/tikv/client-go/v2/util"
 )
@@ -558,10 +560,103 @@ func (s *session) doCommit(ctx context.Context) error {
  if tables := sessVars.TxnCtx.TemporaryTables; len(tables) > 0 {
  s.txn.SetOption(kv.KVFilter, temporaryTableKVFilter(tables))
  }
+ if tables := sessVars.TxnCtx.CachedTables; len(tables) > 0 {
+ c := cachedTableRenewLease{tables: tables}
+ now := time.Now()
+ err := c.start(ctx)
+ defer c.stop(ctx)
+ sessVars.StmtCtx.WaitLockLeaseTime += time.Since(now)
+ if err != nil {
+ return errors.Trace(err)
+ }
+ s.txn.SetOption(kv.CommitTSUpperBoundCheck, c.commitTSCheck)
+ }
 
  return s.commitTxnWithTemporaryData(tikvutil.SetSessionID(ctx, sessVars.ConnectionID), &s.txn)
 }
 
+type cachedTableRenewLease struct {
+ tables map[int64]interface{}
+ lease []uint64 // Lease for each visited cached tables.
+ exit chan struct{}
+}
+
+func (c *cachedTableRenewLease) start(ctx context.Context) error {
+ c.exit = make(chan struct{})
+ c.lease = make([]uint64, len(c.tables))
+ wg := make(chan error)
+ ith := 0
+ for tid, raw := range c.tables {
+ go c.keepAlive(ctx, wg, raw.(tables.StateRemote), tid, &c.lease[ith])
+ ith++
+ }
+
+ // Wait for all LockForWrite() return, this function can return.
+ var err error
+ for ; ith > 0; ith-- {
+ tmp := <-wg
+ if tmp != nil {
+ err = tmp
+ }
+ }
+ return err
+}
+
+const cacheTableWriteLease = 5 * time.Second
+
+func (c *cachedTableRenewLease) keepAlive(ctx context.Context, wg chan error, handle tables.StateRemote, tid int64, leasePtr *uint64) {
+ writeLockLease, err := handle.LockForWrite(ctx, tid)
+ atomic.StoreUint64(leasePtr, writeLockLease)
+ wg <- err
+ if err != nil {
+ logutil.Logger(ctx).Warn("[cached table] lock for write lock fail", zap.Error(err))
+ return
+ }
+
+ t := time.NewTicker(cacheTableWriteLease)
+ defer t.Stop()
+ for {
+ select {
+ case <-t.C:
+ if err := c.renew(ctx, handle, tid, leasePtr); err != nil {
+ logutil.Logger(ctx).Warn("[cached table] renew write lock lease fail", zap.Error(err))
+ return
+ }
+ case <-c.exit:
+ return
+ }
+ }
+}
+
+func (c *cachedTableRenewLease) renew(ctx context.Context, handle tables.StateRemote, tid int64, leasePtr *uint64) error {
+ oldLease := atomic.LoadUint64(leasePtr)
+ physicalTime := oracle.GetTimeFromTS(oldLease)
+ newLease := oracle.GoTimeToTS(physicalTime.Add(cacheTableWriteLease))
+
+ succ, err := handle.RenewLease(ctx, tid, newLease, tables.RenewWriteLease)
+ if err != nil {
+ return errors.Trace(err)
+ }
+ if succ {
+ atomic.StoreUint64(leasePtr, newLease)
+ }
+ return nil
+}
+
+func (c *cachedTableRenewLease) stop(ctx context.Context) {
+ close(c.exit)
+}
+
+func (c *cachedTableRenewLease) commitTSCheck(commitTS uint64) bool {
+ for i := 0; i < len(c.lease); i++ {
+ lease := atomic.LoadUint64(&c.lease[i])
+ if commitTS >= lease {
+ return false
+ }
+ }
+ return true
+}
+
 func (s *session) commitTxnWithTemporaryData(ctx context.Context, txn kv.Transaction) error {
  sessVars := s.sessionVars
  txnTempTables := sessVars.TxnCtx.TemporaryTables

session/session_test.go

@@ -5883,3 +5883,35 @@ func (s *testSessionSuite) TestSameNameObjectWithLocalTemporaryTable(c *C) {
  " `cs1` int(11) DEFAULT NULL\n" +
  ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin"))
 }
+
+func (s *testSessionSuite) TestWriteOnMultipleCachedTable(c *C) {
+ tk := testkit.NewTestKit(c, s.store)
+ tk.MustExec("use test")
+ tk.MustExec("drop table if exists ct1, ct2")
+ tk.MustExec("create table ct1 (id int, c int)")
+ tk.MustExec("create table ct2 (id int, c int)")
+ tk.MustExec("alter table ct1 cache")
+ tk.MustExec("alter table ct2 cache")
+ tk.MustQuery("select * from ct1").Check(testkit.Rows())
+ tk.MustQuery("select * from ct2").Check(testkit.Rows())
+
+ cached := false
+ for i := 0; i < 50; i++ {
+ if tk.HasPlan("select * from ct1", "Union") {
+ if tk.HasPlan("select * from ct2", "Union") {
+ cached = true
+ break
+ }
+ }
+ time.Sleep(100 * time.Millisecond)
+ }
+ c.Assert(cached, IsTrue)
+
+ tk.MustExec("begin")
+ tk.MustExec("insert into ct1 values (3, 4)")
+ tk.MustExec("insert into ct2 values (5, 6)")
+ tk.MustExec("commit")
+
+ tk.MustQuery("select * from ct1").Check(testkit.Rows("3 4"))
+ tk.MustQuery("select * from ct2").Check(testkit.Rows("5 6"))
+}

sessionctx/variable/session.go

@@ -180,6 +180,9 @@ type TransactionContext struct {
  // TemporaryTables is used to store transaction-specific information for global temporary tables.
  // It can also be stored in sessionCtx with local temporary tables, but it's easier to clean this data after transaction ends.
  TemporaryTables map[int64]tableutil.TempTable
+
+ // CachedTables is not nil if the transaction write on cached table.
+ CachedTables map[int64]interface{}
 }
 
 // GetShard returns the shard prefix for the next `count` rowids.

store/driver/txn/txn_driver.go

@@ -230,6 +230,8 @@ func (txn *tikvTxn) SetOption(opt int, val interface{}) {
  txn.KVTxn.SetKVFilter(val.(tikv.KVFilter))
  case kv.SnapInterceptor:
  txn.snapshotInterceptor = val.(kv.SnapshotInterceptor)
+ case kv.CommitTSUpperBoundCheck:
+ txn.KVTxn.SetCommitTSUpperBoundCheck(val.(func(commitTS uint64) bool))
  }
 }
 

table/tables/cache.go

@@ -183,51 +183,31 @@ func (c *cachedTable) UpdateLockForRead(ctx context.Context, store kv.Storage, t
 }
 
 // AddRecord implements the AddRecord method for the table.Table interface.
-func (c *cachedTable) AddRecord(ctx sessionctx.Context, r []types.Datum, opts ...table.AddRecordOption) (recordID kv.Handle, err error) {
- txn, err := ctx.Txn(true)
- if err != nil {
- return nil, err
+func (c *cachedTable) AddRecord(sctx sessionctx.Context, r []types.Datum, opts ...table.AddRecordOption) (recordID kv.Handle, err error) {
+ txnCtxAddCachedTable(sctx, c.Meta().ID, c.handle)
+ return c.TableCommon.AddRecord(sctx, r, opts...)
+}
+
+func txnCtxAddCachedTable(sctx sessionctx.Context, tid int64, handle StateRemote) {
+ txnCtx := sctx.GetSessionVars().TxnCtx
+ if txnCtx.CachedTables == nil {
+ txnCtx.CachedTables = make(map[int64]interface{})
  }
- now := txn.StartTS()
- start := time.Now()
- err = c.handle.LockForWrite(context.Background(), c.Meta().ID, leaseFromTS(now))
- if err != nil {
- return nil, errors.Trace(err)
+ if _, ok := txnCtx.CachedTables[tid]; !ok {
+ txnCtx.CachedTables[tid] = handle
  }
- ctx.GetSessionVars().StmtCtx.WaitLockLeaseTime += time.Since(start)
- return c.TableCommon.AddRecord(ctx, r, opts...)
 }
 
 // UpdateRecord implements table.Table
 func (c *cachedTable) UpdateRecord(ctx context.Context, sctx sessionctx.Context, h kv.Handle, oldData, newData []types.Datum, touched []bool) error {
- txn, err := sctx.Txn(true)
- if err != nil {
- return err
- }
- now := txn.StartTS()
- start := time.Now()
- err = c.handle.LockForWrite(ctx, c.Meta().ID, leaseFromTS(now))
- if err != nil {
- return errors.Trace(err)
- }
- sctx.GetSessionVars().StmtCtx.WaitLockLeaseTime += time.Since(start)
+ txnCtxAddCachedTable(sctx, c.Meta().ID, c.handle)
  return c.TableCommon.UpdateRecord(ctx, sctx, h, oldData, newData, touched)
 }
 
 // RemoveRecord implements table.Table RemoveRecord interface.
-func (c *cachedTable) RemoveRecord(ctx sessionctx.Context, h kv.Handle, r []types.Datum) error {
- txn, err := ctx.Txn(true)
- if err != nil {
- return err
- }
- now := txn.StartTS()
- start := time.Now()
- err = c.handle.LockForWrite(context.Background(), c.Meta().ID, leaseFromTS(now))
- if err != nil {
- return errors.Trace(err)
- }
- ctx.GetSessionVars().StmtCtx.WaitLockLeaseTime += time.Since(start)
- return c.TableCommon.RemoveRecord(ctx, h, r)
+func (c *cachedTable) RemoveRecord(sctx sessionctx.Context, h kv.Handle, r []types.Datum) error {
+ txnCtxAddCachedTable(sctx, c.Meta().ID, c.handle)
+ return c.TableCommon.RemoveRecord(sctx, h, r)
 }
 
 func (c *cachedTable) renewLease(ts uint64, op RenewLeaseType, data *cacheData) func() {