sql: fix portals after exhausting rows

Previously, we would erroneously restart the execution from the very
beginning of empty, unclosed portals after they have been fully
exhausted when we should be returning no rows or an error in such
scenarios. This is now fixed by tracking whether a portal is exhausted
or not and intercepting the calls to `execStmt` when the conn executor
state machine is in an open state.

Note that the current solution has known deviations from Postgres:
- when attempting to execute portals of statements that don't return row
sets, on the second and consequent attempt PG returns an error while we
are silently doing nothing (meaning we do not run the statement at all
and return 0 rows)
- we incorrectly populate "command tag" field of pgwire messages of some
rows-returning statements after the portal suspension (for example,
a suspended UPDATE RETURNING in PG will return the total count of "rows
affected" while we will return the count since the last suspension).

These deviations are deemed acceptable since this commit fixes a much
worse problem - re-executing an exhausted portal (which could be
a mutation meaning, previously, we could have executed a mutation
multiple times).

The reasons for why this commit does not address these deviations are:
- Postgres has a concept of "portal strategy"
(see https://github.com/postgres/postgres/blob/2f9661311b83dc481fc19f6e3bda015392010a40/src/include/utils/portal.h#L89).
- Postgres has a concept of "command" type (these are things like
SELECTs, UPDATEs, INSERTs, etc,
see https://github.com/postgres/postgres/blob/1aac32df89eb19949050f6f27c268122833ad036/src/include/nodes/nodes.h#L672).

CRDB doesn't have these concepts, and without them any attempt to
simulate Postgres results in a very error-prone and brittle code.

Release note (bug fix): Previously, CockroachDB would erroneously
restart the execution of empty, unclosed portals after they have been
fully exhausted, and this has been fixed.

pkg/sql/conn_executor.go

@@ -632,12 +632,13 @@ func (s *Server) newConnExecutor(
 	ex.phaseTimes[sessionInit] = timeutil.Now()
 	ex.extraTxnState.prepStmtsNamespace = prepStmtNamespace{
 		prepStmts: make(map[string]*PreparedStatement),
-		portals:   make(map[string]*PreparedPortal),
+		portals:   make(map[string]PreparedPortal),
 	}
 	ex.extraTxnState.prepStmtsNamespaceAtTxnRewindPos = prepStmtNamespace{
 		prepStmts: make(map[string]*PreparedStatement),
-		portals:   make(map[string]*PreparedPortal),
+		portals:   make(map[string]PreparedPortal),
 	}
+	ex.extraTxnState.prepStmtsNamespaceMemAcc = ex.sessionMon.MakeBoundAccount()
 	ex.extraTxnState.descCollection = descs.MakeCollection(s.cfg.LeaseManager,
 		s.cfg.Settings, s.dbCache.getDatabaseCache(), s.dbCache)
 	ex.extraTxnState.txnRewindPos = -1
@@ -854,8 +855,13 @@ func (ex *connExecutor) close(ctx context.Context, closeType closeType) {
 
 	if closeType != panicClose {
 		// Close all statements and prepared portals.
-		ex.extraTxnState.prepStmtsNamespace.resetTo(ctx, prepStmtNamespace{})
-		ex.extraTxnState.prepStmtsNamespaceAtTxnRewindPos.resetTo(ctx, prepStmtNamespace{})
+		ex.extraTxnState.prepStmtsNamespace.resetTo(
+			ctx, prepStmtNamespace{}, &ex.extraTxnState.prepStmtsNamespaceMemAcc,
+		)
+		ex.extraTxnState.prepStmtsNamespaceAtTxnRewindPos.resetTo(
+			ctx, prepStmtNamespace{}, &ex.extraTxnState.prepStmtsNamespaceMemAcc,
+		)
+		ex.extraTxnState.prepStmtsNamespaceMemAcc.Close(ctx)
 	}
 
 	if ex.sessionTracing.Enabled() {
@@ -990,6 +996,12 @@ type connExecutor struct {
 		// collections, but these collections are periodically reconciled.
 		prepStmtsNamespaceAtTxnRewindPos prepStmtNamespace
 
+		// prepStmtsNamespaceMemAcc is the memory account that is shared
+		// between prepStmtsNamespace and prepStmtsNamespaceAtTxnRewindPos. It
+		// tracks the memory usage of portals and should be closed upon
+		// connExecutor's closure.
+		prepStmtsNamespaceMemAcc mon.BoundAccount
+
 		// onTxnFinish (if non-nil) will be called when txn is finished (either
 		// committed or aborted). It is set when txn is started but can remain
 		// unset when txn is executed within another higher-level txn.
@@ -1136,7 +1148,7 @@ type prepStmtNamespace struct {
 	// session.
 	prepStmts map[string]*PreparedStatement
 	// portals contains the portals currently available on the session.
-	portals map[string]*PreparedPortal
+	portals map[string]PreparedPortal
 }
 
 func (ns prepStmtNamespace) String() string {
@@ -1156,13 +1168,15 @@ func (ns prepStmtNamespace) String() string {
 // references are release and all the to's references are duplicated.
 //
 // An empty `to` can be passed in to deallocate everything.
-func (ns *prepStmtNamespace) resetTo(ctx context.Context, to prepStmtNamespace) {
+func (ns *prepStmtNamespace) resetTo(
+	ctx context.Context, to prepStmtNamespace, prepStmtsNamespaceMemAcc *mon.BoundAccount,
+) {
 	for name, p := range ns.prepStmts {
 		p.decRef(ctx)
 		delete(ns.prepStmts, name)
 	}
 	for name, p := range ns.portals {
-		p.decRef(ctx)
+		p.decRef(ctx, prepStmtsNamespaceMemAcc, name)
 		delete(ns.portals, name)
 	}
 
@@ -1193,7 +1207,7 @@ func (ex *connExecutor) resetExtraTxnState(
 
 	// Close all portals.
 	for name, p := range ex.extraTxnState.prepStmtsNamespace.portals {
-		p.decRef(ctx)
+		p.decRef(ctx, &ex.extraTxnState.prepStmtsNamespaceMemAcc, name)
 		delete(ex.extraTxnState.prepStmtsNamespace.portals, name)
 	}
 
@@ -1398,10 +1412,11 @@ func (ex *connExecutor) execCmd(ctx context.Context) error {
 		// ExecPortal is handled like ExecStmt, except that the placeholder info
 		// is taken from the portal.
 
-		portal, ok := ex.extraTxnState.prepStmtsNamespace.portals[tcmd.Name]
+		portalName := tcmd.Name
+		portal, ok := ex.extraTxnState.prepStmtsNamespace.portals[portalName]
 		if !ok {
 			err := pgerror.Newf(
-				pgcode.InvalidCursorName, "unknown portal %q", tcmd.Name)
+				pgcode.InvalidCursorName, "unknown portal %q", portalName)
 			ev = eventNonRetriableErr{IsCommit: fsm.False}
 			payload = eventNonRetriableErrPayload{err: err}
 			res = ex.clientComm.CreateErrorResult(pos)
@@ -1441,18 +1456,11 @@ func (ex *connExecutor) execCmd(ctx context.Context) error {
 			pos, portal.OutFormats,
 			ex.sessionData.DataConversion,
 			tcmd.Limit,
-			tcmd.Name,
+			portalName,
 			ex.implicitTxn(),
 		)
 		res = stmtRes
-		curStmt := Statement{
-			Statement:     portal.Stmt.Statement,
-			Prepared:      portal.Stmt,
-			ExpectedTypes: portal.Stmt.Columns,
-			AnonymizedStr: portal.Stmt.AnonymizedStr,
-		}
-		stmtCtx := withStatement(ctx, ex.curStmt)
-		ev, payload, err = ex.execStmt(stmtCtx, curStmt, stmtRes, pinfo)
+		ev, payload, err = ex.execPortal(ctx, portal, portalName, stmtRes, pinfo)
 		if err != nil {
 			return err
 		}
@@ -1858,14 +1866,16 @@ func (ex *connExecutor) generateID() ClusterWideID {
 // prepStmtsNamespaceAtTxnRewindPos that's not part of prepStmtsNamespace.
 func (ex *connExecutor) commitPrepStmtNamespace(ctx context.Context) {
 	ex.extraTxnState.prepStmtsNamespaceAtTxnRewindPos.resetTo(
-		ctx, ex.extraTxnState.prepStmtsNamespace)
+		ctx, ex.extraTxnState.prepStmtsNamespace, &ex.extraTxnState.prepStmtsNamespaceMemAcc,
+	)
 }
 
 // commitPrepStmtNamespace deallocates everything in prepStmtsNamespace that's
 // not part of prepStmtsNamespaceAtTxnRewindPos.
 func (ex *connExecutor) rewindPrepStmtNamespace(ctx context.Context) {
 	ex.extraTxnState.prepStmtsNamespace.resetTo(
-		ctx, ex.extraTxnState.prepStmtsNamespaceAtTxnRewindPos)
+		ctx, ex.extraTxnState.prepStmtsNamespaceAtTxnRewindPos, &ex.extraTxnState.prepStmtsNamespaceMemAcc,
+	)
 }
 
 // getRewindTxnCapability checks whether rewinding to the position previously
@@ -2580,7 +2590,9 @@ func (ps connExPrepStmtsAccessor) Delete(ctx context.Context, name string) bool
 
 // DeleteAll is part of the preparedStatementsAccessor interface.
 func (ps connExPrepStmtsAccessor) DeleteAll(ctx context.Context) {
-	ps.ex.extraTxnState.prepStmtsNamespace.resetTo(ctx, prepStmtNamespace{})
+	ps.ex.extraTxnState.prepStmtsNamespace.resetTo(
+		ctx, prepStmtNamespace{}, &ps.ex.extraTxnState.prepStmtsNamespaceMemAcc,
+	)
 }
 
 // contextStatementKey is an empty type for the handle associated with the

pkg/sql/conn_executor_exec.go

@@ -137,6 +137,57 @@ func (ex *connExecutor) recordFailure() {
 	ex.metrics.EngineMetrics.FailureCount.Inc(1)
 }
 
+// execPortal executes a prepared statement. It is a "wrapper" around execStmt
+// method that is performing additional work to track portal's state.
+func (ex *connExecutor) execPortal(
+	ctx context.Context,
+	portal PreparedPortal,
+	portalName string,
+	stmtRes CommandResult,
+	pinfo *tree.PlaceholderInfo,
+) (ev fsm.Event, payload fsm.EventPayload, err error) {
+	curStmt := Statement{
+		Statement:     portal.Stmt.Statement,
+		Prepared:      portal.Stmt,
+		ExpectedTypes: portal.Stmt.Columns,
+		AnonymizedStr: portal.Stmt.AnonymizedStr,
+	}
+	stmtCtx := withStatement(ctx, ex.curStmt)
+	switch ex.machine.CurState().(type) {
+	case stateOpen:
+		// We're about to execute the statement in an open state which
+		// could trigger the dispatch to the execution engine. However, it
+		// is possible that we're trying to execute an already exhausted
+		// portal - in such a scenario we should return no rows, but the
+		// execution engine is not aware of that and would run the
+		// statement as if it was running it for the first time. In order
+		// to prevent such behavior, we check whether the portal has been
+		// exhausted and execute the statement only if it hasn't. If it has
+		// been exhausted, then we do not dispatch the query for execution,
+		// but connExecutor will still perform necessary state transitions
+		// which will emit CommandComplete messages and alike (in a sense,
+		// by not calling execStmt we "execute" the portal in such a way
+		// that it returns 0 rows).
+		// Note that here we deviate from Postgres which returns an error
+		// when attempting to execute an exhausted portal which has a
+		// StatementType() different from "Rows".
+		if !portal.exhausted {
+			ev, payload, err = ex.execStmt(stmtCtx, curStmt, stmtRes, pinfo)
+			// Portal suspension is supported via a "side" state machine
+			// (see pgwire.limitedCommandResult for details), so when
+			// execStmt returns, we know for sure that the portal has been
+			// executed to completion, thus, it is exhausted.
+			// Note that the portal is considered exhausted regardless of
+			// the fact whether an error occurred or not - if it did, we
+			// still don't want to re-execute the portal from scratch.
+			ex.exhaustPortal(portalName)
+		}
+	default:
+		ev, payload, err = ex.execStmt(stmtCtx, curStmt, stmtRes, pinfo)
+	}
+	return
+}
+
 // execStmtInOpenState executes one statement in the context of the session's
 // current transaction.
 // It handles statements that affect the transaction state (BEGIN, COMMIT)
@@ -476,10 +527,10 @@ func (ex *connExecutor) execStmtInOpenState(
 	// is re-configured, re-set etc without using NewTxnWithSteppingEnabled().
 	//
 	// Manually hunting them down and calling ConfigureStepping() each
-	// time would be error prone (and increase the change that a future
+	// time would be error prone (and increase the chance that a future
 	// change would forget to add the call).
 	//
-	// TODO(andrei): really the code should be re-architectued to ensure
+	// TODO(andrei): really the code should be rearchitected to ensure
 	// that all uses of SQL execution initialize the client.Txn using a
 	// single/common function. That would be where the stepping mode
 	// gets enabled once for all SQL statements executed "underneath".

pkg/sql/conn_executor_internal_test.go

@@ -38,8 +38,8 @@ import (
 )
 
 // Test portal implicit destruction. Unless destroying a portal is explicitly
-// requested, portals live until the end of the transaction in which
-// they'recreated. If they're created outside of a transaction, they live until
+// requested, portals live until the end of the transaction in which they're
+// created. If they're created outside of a transaction, they live until
 // the next transaction completes (so until the next statement is executed,
 // which statement is expected to be the execution of the portal that was just
 // created).

pkg/sql/conn_executor_prepare.go

@@ -358,9 +358,7 @@ func (ex *connExecutor) execBind(
 	}
 
 	// Create the new PreparedPortal.
-	if err := ex.addPortal(
-		ctx, portalName, bindCmd.PreparedStatementName, ps, qargs, columnFormatCodes,
-	); err != nil {
+	if err := ex.addPortal(ctx, portalName, ps, qargs, columnFormatCodes); err != nil {
 		return retErr(err)
 	}
 
@@ -379,7 +377,6 @@ func (ex *connExecutor) execBind(
 func (ex *connExecutor) addPortal(
 	ctx context.Context,
 	portalName string,
-	psName string,
 	stmt *PreparedStatement,
 	qargs tree.QueryArguments,
 	outFormats []pgwirebase.FormatCode,
@@ -388,7 +385,7 @@ func (ex *connExecutor) addPortal(
 		panic(errors.AssertionFailedf("portal already exists: %q", portalName))
 	}
 
-	portal, err := ex.newPreparedPortal(ctx, portalName, stmt, qargs, outFormats)
+	portal, err := ex.makePreparedPortal(ctx, portalName, stmt, qargs, outFormats)
 	if err != nil {
 		return err
 	}
@@ -397,6 +394,17 @@ func (ex *connExecutor) addPortal(
 	return nil
 }
 
+// exhaustPortal marks a portal with the provided name as "exhausted" and
+// panics if there is no portal with such name.
+func (ex *connExecutor) exhaustPortal(portalName string) {
+	portal, ok := ex.extraTxnState.prepStmtsNamespace.portals[portalName]
+	if !ok {
+		panic(errors.AssertionFailedf("portal %s doesn't exist", portalName))
+	}
+	portal.exhausted = true
+	ex.extraTxnState.prepStmtsNamespace.portals[portalName] = portal
+}
+
 func (ex *connExecutor) deletePreparedStmt(ctx context.Context, name string) {
 	ps, ok := ex.extraTxnState.prepStmtsNamespace.prepStmts[name]
 	if !ok {
@@ -411,7 +419,7 @@ func (ex *connExecutor) deletePortal(ctx context.Context, name string) {
 	if !ok {
 		return
 	}
-	portal.decRef(ctx)
+	portal.decRef(ctx, &ex.extraTxnState.prepStmtsNamespaceMemAcc, name)
 	delete(ex.extraTxnState.prepStmtsNamespace.portals, name)
 }