diff --git a/executor/insert_common.go b/executor/insert_common.go index dbd4a5ae264cd..d321f02a2f4fd 100644 --- a/executor/insert_common.go +++ b/executor/insert_common.go @@ -1008,7 +1008,7 @@ func (e *InsertValues) allocAutoRandomID(ctx context.Context, fieldType *types.F if err != nil { return 0, err } - currentShard := e.ctx.GetSessionVars().TxnCtx.GetCurrentShard(1) + currentShard := e.ctx.GetSessionVars().GetCurrentShard(1) return shardFmt.Compose(currentShard, autoRandomID), nil } diff --git a/session/bench_test.go b/session/bench_test.go index 742f63aa00f01..ece43c39cdc77 100644 --- a/session/bench_test.go +++ b/session/bench_test.go @@ -1855,6 +1855,22 @@ func BenchmarkCompileStmt(b *testing.B) { b.StopTimer() } +func BenchmarkAutoIncrement(b *testing.B) { + se, do, st := prepareBenchSession() + defer func() { + se.Close() + do.Close() + st.Close() + }() + mustExecute(se, "create table auto_inc (id int unsigned key nonclustered auto_increment) shard_row_id_bits=4 auto_id_cache 1;") + mustExecute(se, "set @@tidb_enable_mutation_checker = false") + b.ResetTimer() + for i := 0; i < b.N; i++ { + mustExecute(se, "insert into auto_inc values ()") + } + b.StopTimer() +} + // TestBenchDaily collects the daily benchmark test result and generates a json output file. // The format of the json output is described by the BenchOutput. // Used by this command in the Makefile @@ -1887,5 +1903,6 @@ func TestBenchDaily(t *testing.T) { BenchmarkHashPartitionPruningMultiSelect, BenchmarkInsertIntoSelect, BenchmarkCompileStmt, + BenchmarkAutoIncrement, ) } diff --git a/sessionctx/variable/session.go b/sessionctx/variable/session.go index d95eb071c7b8b..8c977aa832579 100644 --- a/sessionctx/variable/session.go +++ b/sessionctx/variable/session.go @@ -196,7 +196,6 @@ type TxnCtxNoNeedToRestore struct { ShardStep int shardRemain int currentShard int64 - shardRand *rand.Rand // unchangedRowKeys is used to store the unchanged rows that needs to lock for pessimistic transaction. unchangedRowKeys map[string]struct{} @@ -246,21 +245,22 @@ type SavepointRecord struct { } // GetCurrentShard returns the shard for the next `count` IDs. -func (tc *TransactionContext) GetCurrentShard(count int) int64 { - if tc.shardRand == nil { - tc.shardRand = rand.New(rand.NewSource(int64(tc.StartTS))) // #nosec G404 +func (s *SessionVars) GetCurrentShard(count int) int64 { + tc := s.TxnCtx + if s.shardRand == nil { + s.shardRand = rand.New(rand.NewSource(int64(tc.StartTS))) // #nosec G404 } if tc.shardRemain <= 0 { - tc.updateShard() + tc.updateShard(s.shardRand) tc.shardRemain = tc.ShardStep } tc.shardRemain -= count return tc.currentShard } -func (tc *TransactionContext) updateShard() { +func (tc *TransactionContext) updateShard(shardRand *rand.Rand) { var buf [8]byte - binary.LittleEndian.PutUint64(buf[:], tc.shardRand.Uint64()) + binary.LittleEndian.PutUint64(buf[:], shardRand.Uint64()) tc.currentShard = int64(murmur3.Sum32(buf[:])) } @@ -1323,6 +1323,9 @@ type SessionVars struct { // StoreBatchSize indicates the batch size limit of store batch, set this field to 0 to disable store batch. StoreBatchSize int + // shardRand is used by TxnCtx, for the GetCurrentShard() method. + shardRand *rand.Rand + // Resource group name ResourceGroupName string } diff --git a/table/tables/tables.go b/table/tables/tables.go index da20b1647fbd8..46d037c3197ac 100644 --- a/table/tables/tables.go +++ b/table/tables/tables.go @@ -1521,8 +1521,7 @@ func allocHandleIDs(ctx context.Context, sctx sessionctx.Context, t table.Table, // shard = 0010000000000000000000000000000000000000000000000000000000000000 return 0, 0, autoid.ErrAutoincReadFailed } - txnCtx := sctx.GetSessionVars().TxnCtx - shard := txnCtx.GetCurrentShard(int(n)) + shard := sctx.GetSessionVars().GetCurrentShard(int(n)) base = shardFmt.Compose(shard, base) maxID = shardFmt.Compose(shard, maxID) }