From e37ea6b66176190a45fad88f2faac207f36dcdcc Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Thu, 7 Jul 2022 18:32:07 -0700
Subject: [PATCH] sql: make sure to close mem acc of prepared stmt in case of
 an error

Previously, it was possible that we would not close the memory account
created for a prepared statement when an error is encountered. This was
the case because we would not include the prepared stmt into the prep
stmts namespace, so it would just get lost. However, up until recently
this was not an issue since we mistakenly cleared that memory account
when creating the prepared statement.

Release note: None
---
 pkg/sql/conn_executor_prepare.go | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/pkg/sql/conn_executor_prepare.go b/pkg/sql/conn_executor_prepare.go
index 6dac885c538b..6a3d56b03752 100644
--- a/pkg/sql/conn_executor_prepare.go
+++ b/pkg/sql/conn_executor_prepare.go
@@ -105,12 +105,14 @@ func (ex *connExecutor) addPreparedStmt(
 	}
 
 	if len(prepared.TypeHints) > pgwirebase.MaxPreparedStatementArgs {
+		prepared.memAcc.Close(ctx)
 		return nil, pgwirebase.NewProtocolViolationErrorf(
 			"more than %d arguments to prepared statement: %d",
 			pgwirebase.MaxPreparedStatementArgs, len(prepared.TypeHints))
 	}
 
 	if err := prepared.memAcc.Grow(ctx, int64(len(name))); err != nil {
+		prepared.memAcc.Close(ctx)
 		return nil, err
 	}
 	ex.extraTxnState.prepStmtsNamespace.prepStmts[name] = prepared
@@ -135,16 +137,13 @@ func (ex *connExecutor) addPreparedStmt(
 //
 // placeholderHints may contain partial type information for placeholders.
 // prepare will populate the missing types. It can be nil.
-//
-// The PreparedStatement is returned (or nil if there are no results). The
-// returned PreparedStatement needs to be close()d once its no longer in use.
 func (ex *connExecutor) prepare(
 	ctx context.Context,
 	stmt Statement,
 	placeholderHints tree.PlaceholderTypes,
 	rawTypeHints []oid.Oid,
 	origin PreparedStatementOrigin,
-) (*PreparedStatement, error) {
+) (_ *PreparedStatement, retErr error) {
 
 	prepared := &PreparedStatement{
 		memAcc:   ex.sessionMon.MakeBoundAccount(),
@@ -153,6 +152,12 @@ func (ex *connExecutor) prepare(
 		createdAt: timeutil.Now(),
 		origin:    origin,
 	}
+	defer func() {
+		// Make sure to close the memory account if an error is returned.
+		if retErr != nil {
+			prepared.memAcc.Close(ctx)
+		}
+	}()
 
 	if stmt.AST == nil {
 		return prepared, nil