server: avoid double counting cursor ruv2 result cells

Signed-off-by: disksing <i@disksing.com>

Author: disksing

pkg/server/conn.go

@@ -2526,15 +2526,6 @@ func (cc *clientConn) writeChunksWithFetchSize(ctx context.Context, rs resultset
 	}
 	defer func() {
 		cells := int64(writtenRows) * int64(len(rs.Columns()))
-		if cells > 0 {
-			ruv2Metrics := execdetails.RUV2MetricsFromContext(ctx)
-			if ruv2Metrics == nil {
-				ruv2Metrics = cc.ctx.GetSessionVars().RUV2Metrics
-			}
-			if ruv2Metrics != nil {
-				ruv2Metrics.AddResultChunkCells(cells)
-			}
-		}
 		resultset.ReportCursorRUV2Delta(rs, cells)
 	}()
 	if stmtDetail != nil {

pkg/server/conn_stmt_test.go

@@ -24,6 +24,8 @@ import (
 
 	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
 	"github.com/pingcap/tidb/pkg/config"
+	"github.com/pingcap/tidb/pkg/meta/model"
+	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/planner/core/resolve"
 	"github.com/pingcap/tidb/pkg/server/internal"
@@ -86,6 +88,35 @@ func (*firstNextErrRecordSet) Close() error { return nil }
 
 var _ sqlexec.RecordSet = &firstNextErrRecordSet{}
 
+type singleRowCursorRecordSet struct {
+	returned bool
+}
+
+func (*singleRowCursorRecordSet) Fields() []*resolve.ResultField {
+	return []*resolve.ResultField{{
+		Column:       &model.ColumnInfo{Name: ast.NewCIStr("a"), FieldType: *types.NewFieldType(mysql.TypeLonglong)},
+		ColumnAsName: ast.NewCIStr("a"),
+	}}
+}
+
+func (rs *singleRowCursorRecordSet) Next(_ context.Context, chk *chunk.Chunk) error {
+	chk.Reset()
+	if rs.returned {
+		return nil
+	}
+	chk.AppendInt64(0, 1)
+	rs.returned = true
+	return nil
+}
+
+func (*singleRowCursorRecordSet) NewChunk(chunk.Allocator) *chunk.Chunk {
+	return chunk.NewChunkWithCapacity([]*types.FieldType{types.NewFieldType(mysql.TypeLonglong)}, 1)
+}
+
+func (*singleRowCursorRecordSet) Close() error { return nil }
+
+var _ sqlexec.RecordSet = &singleRowCursorRecordSet{}
+
 func TestCursorExistsFlag(t *testing.T) {
 	store, dom := testkit.CreateMockStoreAndDomain(t)
 	srv := CreateMockServer(t, store)
@@ -225,6 +256,27 @@ func TestCursorWithParams(t *testing.T) {
 		require.Zero(t, execdetails.RUV2MetricsFromContext(ctx).ResultChunkCells())
 	})
 
+	t.Run("write chunks with fetch size records result chunk cells once", func(t *testing.T) {
+		store, dom := testkit.CreateMockStoreAndDomain(t)
+		srv := CreateMockServer(t, store)
+		srv.SetDomain(dom)
+		defer srv.Close()
+
+		c := CreateMockConn(t, srv).(*mockConn)
+		ctx := execdetails.ContextWithInitializedExecDetails(context.Background())
+		ruv2Metrics := execdetails.RUV2MetricsFromContext(ctx)
+		require.NotNil(t, ruv2Metrics)
+
+		rs := resultset.New(&singleRowCursorRecordSet{}, nil)
+		cursorRS := resultset.WrapWithLazyCursor(rs, 1, 1)
+		tracker := resultset.NewCursorRUV2Tracker(nil, "", ruv2Metrics, nil, execdetails.RUV2Weights{})
+		resultset.AttachCursorRUV2Tracker(cursorRS, tracker)
+
+		err := c.writeChunksWithFetchSize(ctx, cursorRS, mysql.ServerStatusAutocommit, 1)
+		require.NoError(t, err)
+		require.Equal(t, int64(1), ruv2Metrics.ResultChunkCells())
+	})
+
 	store, dom := testkit.CreateMockStoreAndDomain(t)
 	srv := CreateMockServer(t, store)
 	srv.SetDomain(dom)