sql: check exception handler before applying TCO

This commit finishes the tail-call optimization fix begun by the previous
commits, by preventing TCO when it would lose the reference to the calling
routine's exception handler. PL/pgSQL sub-routines always maintain a
reference to their parent's exception handler, so this isn't a problem for
them. However, explicit (user-specified) nested routines do not track the
calling routine's exception handler.

There is no release note because this bug hasn't appeared in any release.

Fixes cockroachdb#120916

Release note: None

pkg/ccl/logictestccl/testdata/logic_test/nested_routines

@@ -25,3 +25,31 @@ query II
 SELECT * FROM f() AS g(x INT, y INT);
 ----
 -4  4
+
+# Case with an exception handler on the parent routine. This prevents TCO,
+# since executing the child routine in the parent's context would lose track
+# of the exception handler.
+statement ok
+DROP FUNCTION f;
+DROP FUNCTION f_nested;
+
+statement ok
+CREATE FUNCTION f_nested() RETURNS INT AS $$
+  BEGIN
+    RETURN 1//0;
+  END
+$$ LANGUAGE PLpgSQL;
+
+statement ok
+CREATE FUNCTION f() RETURNS INT AS $$
+  BEGIN
+    RETURN f_nested();
+  EXCEPTION WHEN division_by_zero THEN
+    RETURN -1;
+  END
+$$ LANGUAGE PLpgSQL;
+
+query I
+SELECT f();
+----
+-1

pkg/sql/routine.go

@@ -114,9 +114,10 @@ func (p *planner) EvalRoutineExpr(
 		return expr.CachedResult, nil
 	}
 
-	if expr.TailCall && !expr.Generator && p.EvalContext().RoutineSender != nil {
+	if tailCallOptimizationEnabled && expr.TailCall && !expr.Generator {
 		// This is a nested routine in tail-call position.
-		if tailCallOptimizationEnabled {
+		sender := p.EvalContext().RoutineSender
+		if sender != nil && sender.CanOptimizeTailCall(expr) {
 			// Tail-call optimizations are enabled. Send the information needed to
 			// evaluate this routine to the parent routine, then return. It is safe to
 			// return NULL here because the parent is guaranteed not to perform any
@@ -484,8 +485,30 @@ var tailCallOptimizationEnabled = util.ConstantWithMetamorphicTestBool(
 	true,
 )
 
-func (g *routineGenerator) SendDeferredRoutine(routine *tree.RoutineExpr, args tree.Datums) {
-	g.deferredRoutine.expr = routine
+func (g *routineGenerator) CanOptimizeTailCall(nestedRoutine *tree.RoutineExpr) bool {
+	// Tail-call optimization is allowed only if the current routine will not
+	// perform any work after its body statements finish executing.
+	//
+	// Note: cursors are opened after the first body statement, and there is
+	// always more than one body statement if a cursor is opened. This is enforced
+	// during exec-building. For this reason, we only have to check for an
+	// exception handler.
+	if g.expr.BlockState != nil {
+		// If the current routine has an exception handler (which is the case when
+		// BlockState is non-nil), the nested routine must either be part of the
+		// same PL/pgSQL block, or a child block. Otherwise, enabling TCO could
+		// cause execution to skip the exception handler.
+		childBlock := nestedRoutine.BlockState
+		if childBlock == nil {
+			return false
+		}
+		return childBlock == g.expr.BlockState || childBlock.Parent == g.expr.BlockState
+	}
+	return true
+}
+
+func (g *routineGenerator) SendDeferredRoutine(nestedRoutine *tree.RoutineExpr, args tree.Datums) {
+	g.deferredRoutine.expr = nestedRoutine
 	g.deferredRoutine.args = args
 }
 

pkg/sql/sem/eval/deps.go

@@ -555,9 +555,13 @@ type ClientNoticeSender interface {
 // for its own evaluation to a parent routine. This is used to defer execution
 // for tail-call optimization. It can only be used during local execution.
 type DeferredRoutineSender interface {
+	// CanOptimizeTailCall determines whether a nested routine in tail-call
+	// position can be executed in its parent's context.
+	CanOptimizeTailCall(nestedRoutine *tree.RoutineExpr) bool
+
 	// SendDeferredRoutine sends a local nested routine and its arguments to its
 	// parent routine.
-	SendDeferredRoutine(expr *tree.RoutineExpr, args tree.Datums)
+	SendDeferredRoutine(nestedRoutine *tree.RoutineExpr, args tree.Datums)
 }
 
 // PrivilegedAccessor gives access to certain queries that would otherwise