From d7e454a7a0e108209e6abdf8e3972b7b253db171 Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Thu, 21 Jul 2022 08:16:59 -0700
Subject: [PATCH] sql: fix recent leak of a context

Release note: None
---
 pkg/sql/conn_executor_exec.go | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/pkg/sql/conn_executor_exec.go b/pkg/sql/conn_executor_exec.go
index 9e6e4b02a1b2..1680e3bd0287 100644
--- a/pkg/sql/conn_executor_exec.go
+++ b/pkg/sql/conn_executor_exec.go
@@ -279,9 +279,6 @@ func (ex *connExecutor) execStmtInOpenState(
 	ast := parserStmt.AST
 	ctx = withStatement(ctx, ast)
 
-	var cancelQuery context.CancelFunc
-	ctx, cancelQuery = contextutil.WithCancel(ctx)
-
 	makeErrEvent := func(err error) (fsm.Event, fsm.EventPayload, error) {
 		ev, payload := ex.makeErrEvent(err, ast)
 		return ev, payload, nil
@@ -327,10 +324,9 @@ func (ex *connExecutor) execStmtInOpenState(
 		ex.planner.EvalContext().Placeholders = pinfo
 	}
 
+	var cancelQuery context.CancelFunc
+	ctx, cancelQuery = contextutil.WithCancel(ctx)
 	ex.addActiveQuery(ast, formatWithPlaceholders(ast, ex.planner.EvalContext()), queryID, cancelQuery)
-	if ex.executorType != executorTypeInternal {
-		ex.metrics.EngineMetrics.SQLActiveStatements.Inc(1)
-	}
 
 	// Make sure that we always unregister the query. It also deals with
 	// overwriting res.Error to a more user-friendly message in case of query
@@ -344,6 +340,7 @@ func (ex *connExecutor) execStmtInOpenState(
 			}
 		}
 		ex.removeActiveQuery(queryID, ast)
+		cancelQuery()
 		if ex.executorType != executorTypeInternal {
 			ex.metrics.EngineMetrics.SQLActiveStatements.Dec(1)
 		}
@@ -386,6 +383,10 @@ func (ex *connExecutor) execStmtInOpenState(
 		}
 	}(ctx, res)
 
+	if ex.executorType != executorTypeInternal {
+		ex.metrics.EngineMetrics.SQLActiveStatements.Inc(1)
+	}
+
 	p := &ex.planner
 	stmtTS := ex.server.cfg.Clock.PhysicalTime()
 	ex.statsCollector.Reset(ex.applicationStats, ex.phaseTimes)
@@ -505,7 +506,7 @@ func (ex *connExecutor) execStmtInOpenState(
 		timeoutTicker = time.AfterFunc(
 			timerDuration,
 			func() {
-				ex.cancelQuery(queryID)
+				cancelQuery()
 				queryTimedOut = true
 				doneAfterFunc <- struct{}{}
 			})