From 4571aed22415a0bbd9a6f3f347173533a51f1437 Mon Sep 17 00:00:00 2001
From: crazycs520 <crazycs520@gmail.com>
Date: Tue, 28 Jul 2020 17:19:08 +0800
Subject: [PATCH 1/2] init

Signed-off-by: crazycs520 <crazycs520@gmail.com>
---
 executor/batch_point_get.go      | 17 +++++++++++++++++
 planner/core/integration_test.go | 17 +++++++++++------
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/executor/batch_point_get.go b/executor/batch_point_get.go
index fbae52044c1a8..15f6564346df0 100644
--- a/executor/batch_point_get.go
+++ b/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"
@@ -61,6 +62,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
@@ -97,6 +101,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)
 	}
@@ -105,12 +118,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.DelOption(kv.CollectRuntimeStats)
+	}
 	return nil
 }
 
diff --git a/planner/core/integration_test.go b/planner/core/integration_test.go
index cb5b74eb6b3b0..a1ba1c456c163 100644
--- a/planner/core/integration_test.go
+++ b/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")
 }

From 2dccb9900179997928280dd3b9aabcd6ac0234c4 Mon Sep 17 00:00:00 2001
From: crazycs520 <crazycs520@gmail.com>
Date: Wed, 29 Jul 2020 11:12:31 +0800
Subject: [PATCH 2/2] fix panic

Signed-off-by: crazycs520 <crazycs520@gmail.com>
---
 executor/batch_point_get.go | 2 +-
 executor/point_get.go       | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/executor/batch_point_get.go b/executor/batch_point_get.go
index 2d0f68b73f497..8b0546f1c0907 100644
--- a/executor/batch_point_get.go
+++ b/executor/batch_point_get.go
@@ -126,7 +126,7 @@ func (e *BatchPointGetExec) Open(context.Context) error {
 
 // Close implements the Executor interface.
 func (e *BatchPointGetExec) Close() error {
-	if e.runtimeStats != nil {
+	if e.runtimeStats != nil && e.snapshot != nil {
 		e.snapshot.DelOption(kv.CollectRuntimeStats)
 	}
 	return nil
diff --git a/executor/point_get.go b/executor/point_get.go
index 2d7360a78c0ac..2a744fa84001a 100644
--- a/executor/point_get.go
+++ b/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