*: merge statement buffer when BatchGetValues (#9374)

ddl/index.go

@@ -691,7 +691,7 @@ func (w *addIndexWorker) batchCheckUniqueKey(txn kv.Transaction, idxRecords []*i
  w.distinctCheckFlags = append(w.distinctCheckFlags, distinct)
  }
 
- batchVals, err := kv.BatchGetValues(txn, w.batchCheckKeys)
+ batchVals, err := txn.BatchGet(w.batchCheckKeys)
  if err != nil {
  return errors.Trace(err)
  }

executor/admin.go

@@ -376,7 +376,7 @@ func (e *RecoverIndexExec) batchMarkDup(txn kv.Transaction, rows []recoverRows)
  distinctFlags[i] = distinct
  }
 
- values, err := kv.BatchGetValues(txn, e.batchKeys)
+ values, err := txn.BatchGet(e.batchKeys)
  if err != nil {
  return errors.Trace(err)
  }
@@ -500,7 +500,7 @@ func (e *CleanupIndexExec) batchGetRecord(txn kv.Transaction) (map[string][]byte
  for handle := range e.idxValues {
  e.batchKeys = append(e.batchKeys, e.table.RecordKey(handle))
  }
- values, err := kv.BatchGetValues(txn, e.batchKeys)
+ values, err := txn.BatchGet(e.batchKeys)
  if err != nil {
  return nil, errors.Trace(err)
  }

executor/batch_checker.go

@@ -56,7 +56,7 @@ func (b *batchChecker) batchGetOldValues(ctx sessionctx.Context, batchKeys []kv.
  if err != nil {
  return errors.Trace(err)
  }
- values, err := kv.BatchGetValues(txn, batchKeys)
+ values, err := txn.BatchGet(batchKeys)
  if err != nil {
  return errors.Trace(err)
  }
@@ -213,7 +213,7 @@ func (b *batchChecker) batchGetInsertKeys(ctx sessionctx.Context, t table.Table,
  if err != nil {
  return errors.Trace(err)
  }
- b.dupKVs, err = kv.BatchGetValues(txn, batchKeys)
+ b.dupKVs, err = txn.BatchGet(batchKeys)
  return errors.Trace(err)
 }
 

executor/write_test.go

@@ -285,6 +285,17 @@ func (s *testSuite) TestInsert(c *C) {
  tk.MustQuery("select * from t").Check(testkit.Rows("1 1"))
 }
 
+func (s *testSuite) TestMultiBatch(c *C) {
+ tk := testkit.NewTestKit(c, s.store)
+ tk.MustExec("use test")
+ tk.MustExec("create table t0 (i int)")
+ tk.MustExec("insert into t0 values (1), (1)")
+ tk.MustExec("create table t (i int unique key)")
+ tk.MustExec("set @@tidb_dml_batch_size = 1")
+ tk.MustExec("insert ignore into t select * from t0")
+ tk.MustExec("admin check table t")
+}
+
 func (s *testSuite) TestInsertAutoInc(c *C) {
  tk := testkit.NewTestKit(c, s.store)
  tk.MustExec("use test")

kv/fault_injection.go

@@ -98,6 +98,16 @@ func (t *InjectedTransaction) Get(k Key) ([]byte, error) {
  return t.Transaction.Get(k)
 }
 
+// BatchGet returns an error if cfg.getError is set.
+func (t *InjectedTransaction) BatchGet(keys []Key) (map[string][]byte, error) {
+ t.cfg.RLock()
+ defer t.cfg.RUnlock()
+ if t.cfg.getError != nil {
+ return nil, t.cfg.getError
+ }
+ return t.Transaction.BatchGet(keys)
+}
+
 // Commit returns an error if cfg.commitError is set.
 func (t *InjectedTransaction) Commit(ctx context.Context) error {
  t.cfg.RLock()
@@ -108,14 +118,6 @@ func (t *InjectedTransaction) Commit(ctx context.Context) error {
  return t.Transaction.Commit(ctx)
 }
 
-// GetSnapshot implements Transaction GetSnapshot method.
-func (t *InjectedTransaction) GetSnapshot() Snapshot {
- return &InjectedSnapshot{
- Snapshot: t.Transaction.GetSnapshot(),
- cfg: t.cfg,
- }
-}
-
 // InjectedSnapshot wraps a Snapshot with injections.
 type InjectedSnapshot struct {
  Snapshot

kv/kv.go

@@ -150,10 +150,10 @@ type Transaction interface {
  Valid() bool
  // GetMemBuffer return the MemBuffer binding to this transaction.
  GetMemBuffer() MemBuffer
- // GetSnapshot returns the snapshot of this transaction.
- GetSnapshot() Snapshot
  // SetVars sets variables to the transaction.
  SetVars(vars *Variables)
+ // BatchGet gets kv from the memory buffer of statement and transaction, and the kv storage.
+ BatchGet(keys []Key) (map[string][]byte, error)
 }
 
 // Client is used to send request to KV layer.

kv/mock.go

@@ -68,6 +68,10 @@ func (t *mockTxn) Get(k Key) ([]byte, error) {
  return nil, nil
 }
 
+func (t *mockTxn) BatchGet(keys []Key) (map[string][]byte, error) {
+ return nil, nil
+}
+
 func (t *mockTxn) Iter(k Key, upperBound Key) (Iterator, error) {
  return nil, nil
 }
@@ -99,12 +103,6 @@ func (t *mockTxn) GetMemBuffer() MemBuffer {
  return nil
 }
 
-func (t *mockTxn) GetSnapshot() Snapshot {
- return &mockSnapshot{
- store: NewMemDbBuffer(DefaultTxnMembufCap),
- }
-}
-
 func (t *mockTxn) SetCap(cap int) {
 
 }

kv/txn.go

@@ -89,37 +89,3 @@ func BackOff(attempts uint) int {
  time.Sleep(sleep)
  return int(sleep)
 }
-
-// BatchGetValues gets values in batch.
-// The values from buffer in transaction and the values from the storage node are merged together.
-func BatchGetValues(txn Transaction, keys []Key) (map[string][]byte, error) {
- if txn.IsReadOnly() {
- return txn.GetSnapshot().BatchGet(keys)
- }
- bufferValues := make([][]byte, len(keys))
- shrinkKeys := make([]Key, 0, len(keys))
- for i, key := range keys {
- val, err := txn.GetMemBuffer().Get(key)
- if IsErrNotFound(err) {
- shrinkKeys = append(shrinkKeys, key)
- continue
- }
- if err != nil {
- return nil, errors.Trace(err)
- }
- if len(val) != 0 {
- bufferValues[i] = val
- }
- }
- storageValues, err := txn.GetSnapshot().BatchGet(shrinkKeys)
- if err != nil {
- return nil, errors.Trace(err)
- }
- for i, key := range keys {
- if bufferValues[i] == nil {
- continue
- }
- storageValues[string(key)] = bufferValues[i]
- }
- return storageValues, nil
-}

session/txn.go

@@ -212,6 +212,36 @@ func (st *TxnState) Get(k kv.Key) ([]byte, error) {
  return val, nil
 }
 
+// BatchGet overrides the Transaction interface.
+func (st *TxnState) BatchGet(keys []kv.Key) (map[string][]byte, error) {
+ bufferValues := make([][]byte, len(keys))
+ shrinkKeys := make([]kv.Key, 0, len(keys))
+ for i, key := range keys {
+ val, err := st.buf.Get(key)
+ if kv.IsErrNotFound(err) {
+ shrinkKeys = append(shrinkKeys, key)
+ continue
+ }
+ if err != nil {
+ return nil, errors.Trace(err)
+ }
+ if len(val) != 0 {
+ bufferValues[i] = val
+ }
+ }
+ storageValues, err := st.Transaction.BatchGet(shrinkKeys)
+ if err != nil {
+ return nil, errors.Trace(err)
+ }
+ for i, key := range keys {
+ if bufferValues[i] == nil {
+ continue
+ }
+ storageValues[string(key)] = bufferValues[i]
+ }
+ return storageValues, nil
+}
+
 // Set overrides the Transaction interface.
 func (st *TxnState) Set(k kv.Key, v []byte) error {
  return st.buf.Set(k, v)

store/tikv/txn.go

@@ -108,6 +108,38 @@ func (txn *tikvTxn) Get(k kv.Key) ([]byte, error) {
  return ret, nil
 }
 
+func (txn *tikvTxn) BatchGet(keys []kv.Key) (map[string][]byte, error) {
+ if txn.IsReadOnly() {
+ return txn.snapshot.BatchGet(keys)
+ }
+ bufferValues := make([][]byte, len(keys))
+ shrinkKeys := make([]kv.Key, 0, len(keys))
+ for i, key := range keys {
+ val, err := txn.GetMemBuffer().Get(key)
+ if kv.IsErrNotFound(err) {
+ shrinkKeys = append(shrinkKeys, key)
+ continue
+ }
+ if err != nil {
+ return nil, errors.Trace(err)
+ }
+ if len(val) != 0 {
+ bufferValues[i] = val
+ }
+ }
+ storageValues, err := txn.snapshot.BatchGet(shrinkKeys)
+ if err != nil {
+ return nil, errors.Trace(err)
+ }
+ for i, key := range keys {
+ if bufferValues[i] == nil {
+ continue
+ }
+ storageValues[string(key)] = bufferValues[i]
+ }
+ return storageValues, nil
+}
+
 func (txn *tikvTxn) Set(k kv.Key, v []byte) error {
  txn.setCnt++
 
@@ -287,7 +319,3 @@ func (txn *tikvTxn) Size() int {
 func (txn *tikvTxn) GetMemBuffer() kv.MemBuffer {
  return txn.us.GetMemBuffer()
 }
-
-func (txn *tikvTxn) GetSnapshot() kv.Snapshot {
- return txn.snapshot
-}