executor: add runtime information for the batch-point-get executor

executor/batch_point_get.go

@@ -23,6 +23,7 @@ import (
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
+	"github.com/pingcap/tidb/store/tikv"
 	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/chunk"
@@ -62,6 +63,9 @@ type BatchPointGetExec struct {
 
 	// virtualColumnRetFieldTypes records the RetFieldTypes of virtual columns.
 	virtualColumnRetFieldTypes []*types.FieldType
+
+	snapshot kv.Snapshot
+	stats    *pointGetRuntimeStats
 }
 
 // buildVirtualColumnInfo saves virtual column indices and sort them in definition order
@@ -98,6 +102,15 @@ func (e *BatchPointGetExec) Open(context.Context) error {
 			return err
 		}
 	}
+	if e.runtimeStats != nil {
+		snapshotStats := &tikv.SnapshotRuntimeStats{}
+		e.stats = &pointGetRuntimeStats{
+			BasicRuntimeStats:    e.runtimeStats,
+			SnapshotRuntimeStats: snapshotStats,
+		}
+		snapshot.SetOption(kv.CollectRuntimeStats, snapshotStats)
+		e.ctx.GetSessionVars().StmtCtx.RuntimeStatsColl.RegisterStats(e.id.String(), e.stats)
+	}
 	if e.ctx.GetSessionVars().GetReplicaRead().IsFollowerRead() {
 		snapshot.SetOption(kv.ReplicaRead, kv.ReplicaReadFollower)
 	}
@@ -106,12 +119,16 @@ func (e *BatchPointGetExec) Open(context.Context) error {
 	if txn.Valid() {
 		batchGetter = kv.NewBufferBatchGetter(txn.GetMemBuffer(), &PessimisticLockCacheGetter{txnCtx: txnCtx}, snapshot)
 	}
+	e.snapshot = snapshot
 	e.batchGetter = batchGetter
 	return nil
 }
 
 // Close implements the Executor interface.
 func (e *BatchPointGetExec) Close() error {
+	if e.runtimeStats != nil && e.snapshot != nil {
+		e.snapshot.DelOption(kv.CollectRuntimeStats)
+	}
 	return nil
 }
 

executor/point_get.go

@@ -155,7 +155,7 @@ func (e *PointGetExecutor) Open(context.Context) error {
 
 // Close implements the Executor interface.
 func (e *PointGetExecutor) Close() error {
-	if e.runtimeStats != nil {
+	if e.runtimeStats != nil && e.snapshot != nil {
 		e.snapshot.DelOption(kv.CollectRuntimeStats)
 	}
 	return nil

planner/core/integration_test.go

@@ -1431,11 +1431,16 @@ func (s *testIntegrationSerialSuite) TestExplainAnalyzePointGet(c *C) {
 	tk.MustExec("insert into t values (1,1)")
 
 	res := tk.MustQuery("explain analyze select * from t where a=1;")
-	resBuff := bytes.NewBufferString("")
-	for _, row := range res.Rows() {
-		fmt.Fprintf(resBuff, "%s\n", row)
+	checkExplain := func(rpc string) {
+		resBuff := bytes.NewBufferString("")
+		for _, row := range res.Rows() {
+			fmt.Fprintf(resBuff, "%s\n", row)
+		}
+		explain := resBuff.String()
+		c.Assert(strings.Contains(explain, rpc+":{num_rpc:"), IsTrue, Commentf("%s", explain))
+		c.Assert(strings.Contains(explain, "total_time:"), IsTrue, Commentf("%s", explain))
 	}
-	explain := resBuff.String()
-	c.Assert(strings.Contains(explain, "Get:{num_rpc:"), IsTrue, Commentf("%s", explain))
-	c.Assert(strings.Contains(explain, "total_time:"), IsTrue, Commentf("%s", explain))
+	checkExplain("Get")
+	res = tk.MustQuery("explain analyze select * from t where a in (1,2,3);")
+	checkExplain("BatchGet")
 }