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

Code before:
// See the License for the specific language governing permissions and
// limitations under the License.

package plancache

import (
	"testing"

	plannercore "github.com/pingcap/tidb/pkg/planner/core"
	"github.com/pingcap/tidb/pkg/sessionctx"
	"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)
	}

Review comment (suggestion):
we still can use `require.NoError(b, err)` if you want

Code after:
// See the License for the specific language governing permissions and
// limitations under the License.

package plancache

import (
	"testing"

	plannercore "github.com/pingcap/tidb/pkg/planner/core"
	"github.com/pingcap/tidb/pkg/sessionctx"
	"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)