From 139fd0cf6ed3ea711d6f81bd4096d7463debc6e0 Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Wed, 25 Aug 2021 19:15:31 -0700
Subject: [PATCH] sql: fix tracking of query level stats for sub- and
 post-queries

Previously, we were creating a new `topLevelQueryStats` object for each
sub- and post-query, and as a result we would accumulate the top level
query stats only for the main query. I think this is undesirable and is
now fixed.

Release note: None

Release justification: low-risk bug fix.
---
 pkg/sql/conn_executor_exec.go | 6 +++---
 pkg/sql/distsql_running.go    | 8 +++-----
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/pkg/sql/conn_executor_exec.go b/pkg/sql/conn_executor_exec.go
index 1ef68c627455..dcdca099884c 100644
--- a/pkg/sql/conn_executor_exec.go
+++ b/pkg/sql/conn_executor_exec.go
@@ -1272,7 +1272,7 @@ func (ex *connExecutor) execWithDistSQLEngine(
 		if !ex.server.cfg.DistSQLPlanner.PlanAndRunSubqueries(
 			ctx, planner, evalCtxFactory, planner.curPlan.subqueryPlans, recv, &subqueryResultMemAcc,
 		) {
-			return recv.stats, recv.commErr
+			return *recv.stats, recv.commErr
 		}
 	}
 	recv.discardRows = planner.instrumentation.ShouldDiscardRows()
@@ -1285,14 +1285,14 @@ func (ex *connExecutor) execWithDistSQLEngine(
 	// need to have access to the main query tree.
 	defer cleanup()
 	if recv.commErr != nil || res.Err() != nil {
-		return recv.stats, recv.commErr
+		return *recv.stats, recv.commErr
 	}
 
 	ex.server.cfg.DistSQLPlanner.PlanAndRunCascadesAndChecks(
 		ctx, planner, evalCtxFactory, &planner.curPlan.planComponents, recv,
 	)
 
-	return recv.stats, recv.commErr
+	return *recv.stats, recv.commErr
 }
 
 // beginTransactionTimestampsAndReadMode computes the timestamps and
diff --git a/pkg/sql/distsql_running.go b/pkg/sql/distsql_running.go
index a8193874d744..094a0c2477cc 100644
--- a/pkg/sql/distsql_running.go
+++ b/pkg/sql/distsql_running.go
@@ -608,11 +608,7 @@ type DistSQLReceiver struct {
 	// this node's clock.
 	clockUpdater clockUpdater
 
-	// TODO(yuzefovich): I believe these stats currently only include the
-	// metrics from the main query, and not from any sub- or post-queries
-	// because we use DistSQLReceiver.clone() for those. Think through whether
-	// this is expected or not.
-	stats topLevelQueryStats
+	stats *topLevelQueryStats
 
 	expectedRowsRead int64
 	progressAtomic   *uint64
@@ -757,6 +753,7 @@ func MakeDistSQLReceiver(
 		rangeCache:         rangeCache,
 		txn:                txn,
 		clockUpdater:       clockUpdater,
+		stats:              &topLevelQueryStats{},
 		stmtType:           stmtType,
 		tracing:            tracing,
 		contentionRegistry: contentionRegistry,
@@ -782,6 +779,7 @@ func (r *DistSQLReceiver) clone() *DistSQLReceiver {
 		rangeCache:         r.rangeCache,
 		txn:                r.txn,
 		clockUpdater:       r.clockUpdater,
+		stats:              r.stats,
 		stmtType:           tree.Rows,
 		tracing:            r.tracing,
 		contentionRegistry: r.contentionRegistry,