Repository: pingcap/tidb
File: pkg/planner/core/casetest/plancache/plan_cache_test.go (Go)

Code before:
	"github.com/pingcap/tidb/pkg/testkit"
	"github.com/pingcap/tidb/pkg/util/benchdaily"
)

func BenchmarkNewPlanCacheKey(b *testing.B) {
	store := testkit.CreateMockStore(b)
	tk := testkit.NewTestKit(b, store)
	tk.MustExec("use test")
	tk.MustExec("create table t (a int, b int, c varchar(100), key(a), key(b, c))")
	tk.MustExec("prepare st from 'select * from t where a = ? and b = ? and c = ?'")
	tk.MustExec("set @a=1, @b=2, @c='test'")
	tk.MustExec("execute st using @a, @b, @c") // first execution to prepare the statement

	sctx := tk.Session().(sessionctx.Context)
	stmtID, _, _, err := tk.Session().PrepareStmt("select * from t where a = ? and b = ? and c = ?")
	if err != nil {
		b.Fatal(err)
	}
	prepStmt, err := sctx.GetSessionVars().GetPreparedStmtByID(stmtID)
	if err != nil {
		b.Fatal(err)
	}
	stmt := prepStmt.(*plannercore.PlanCacheStmt)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, _, _, _, _ = plannercore.NewPlanCacheKey(sctx, stmt)
	}
}

func BenchmarkNewPlanCacheKeyInTxn(b *testing.B) {
	store := testkit.CreateMockStore(b)
	tk := testkit.NewTestKit(b, store)
	tk.MustExec("use test")
	tk.MustExec("create table t (a int, b int, key(a))")
	tk.MustExec("insert into t values (1, 1), (2, 2)")

	stmtID, _, _, err := tk.Session().PrepareStmt("select * from t where a = ?")
	if err != nil {
		b.Fatal(err)
	}
	sctx := tk.Session().(sessionctx.Context)
	prepStmt, err := sctx.GetSessionVars().GetPreparedStmtByID(stmtID)
	if err != nil {
		b.Fatal(err)
	}
	stmt := prepStmt.(*plannercore.PlanCacheStmt)

	// Start a transaction and make the table dirty
	tk.MustExec("begin")
	tk.MustExec("insert into t values (3, 3)")

Review comment (suggestion):
after we have upgraded to go 1.25, we can use b.Loop() now https://pkg.go.dev/testing#B.Loop https://go.dev/blog/testing-b-loop 

Code after:
	"github.com/pingcap/tidb/pkg/testkit"
	"github.com/pingcap/tidb/pkg/util/benchdaily"
	"github.com/stretchr/testify/require"
)

func BenchmarkNewPlanCacheKey(b *testing.B) {
	store := testkit.CreateMockStore(b)
	tk := testkit.NewTestKit(b, store)
	tk.MustExec("use test")
	tk.MustExec("create table t (a int, b int, c varchar(100), key(a), key(b, c))")
	tk.MustExec("prepare st from 'select * from t where a = ? and b = ? and c = ?'")
	tk.MustExec("set @a=1, @b=2, @c='test'")
	tk.MustExec("execute st using @a, @b, @c") // first execution to prepare the statement

	sctx := tk.Session().(sessionctx.Context)
	stmtID, _, _, err := tk.Session().PrepareStmt("select * from t where a = ? and b = ? and c = ?")
	require.NoError(b, err)
	prepStmt, err := sctx.GetSessionVars().GetPreparedStmtByID(stmtID)
	require.NoError(b, err)
	stmt := prepStmt.(*plannercore.PlanCacheStmt)

	b.ResetTimer()
	for b.Loop() {
		_, _, _, _, _ = plannercore.NewPlanCacheKey(sctx, stmt)
	}
}

func BenchmarkNewPlanCacheKeyInTxn(b *testing.B) {
	store := testkit.CreateMockStore(b)
	tk := testkit.NewTestKit(b, store)
	tk.MustExec("use test")
	tk.MustExec("create table t (a int, b int, key(a))")
	tk.MustExec("insert into t values (1, 1), (2, 2)")

	stmtID, _, _, err := tk.Session().PrepareStmt("select * from t where a = ?")
	require.NoError(b, err)
	sctx := tk.Session().(sessionctx.Context)
	prepStmt, err := sctx.GetSessionVars().GetPreparedStmtByID(stmtID)
	require.NoError(b, err)
	stmt := prepStmt.(*plannercore.PlanCacheStmt)

	// Start a transaction and make the table dirty
	tk.MustExec("begin")
	tk.MustExec("insert into t values (3, 3)")

	b.ResetTimer()
	for b.Loop() {
		_, _, _, _, _ = plannercore.NewPlanCacheKey(sctx, stmt)
	}
	b.StopTimer()
	tk.MustExec("rollback")