From 67af366093a45e480d45933c7f3b3a751361b829 Mon Sep 17 00:00:00 2001
From: Raphael 'kena' Poss <knz@cockroachlabs.com>
Date: Wed, 23 Jan 2019 13:11:18 +0100
Subject: [PATCH] sql: avoid calling planNode.Close() during execution

We want that a planNode does not get closed during
execution. This was already mostly true, except for
`unionNode`. This patch modifies `unionNode` to make it compliant
with the new protocol, and clarifies the constraint in the
interface documentation for `planNode.Close()`.

Release note: None
---
 pkg/sql/plan.go     |  4 ++++
 pkg/sql/subquery.go |  3 ---
 pkg/sql/union.go    | 29 ++++++++++++-----------------
 3 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/pkg/sql/plan.go b/pkg/sql/plan.go
index 34f30915ad10..1ff19ada855c 100644
--- a/pkg/sql/plan.go
+++ b/pkg/sql/plan.go
@@ -142,6 +142,10 @@ type planNode interface {
 	// This method should be called if the node has been used in any way (any
 	// methods on it have been called) after it was constructed. Note that this
 	// doesn't imply that startPlan() has been necessarily called.
+	//
+	// This method must not be called during execution - the planNode
+	// tree must remain "live" and readable via walk() even after
+	// execution completes.
 	Close(ctx context.Context)
 }
 
diff --git a/pkg/sql/subquery.go b/pkg/sql/subquery.go
index 86b1f820798b..f182cf5e7797 100644
--- a/pkg/sql/subquery.go
+++ b/pkg/sql/subquery.go
@@ -86,9 +86,6 @@ func (p *planTop) evalSubqueries(params runParams) error {
 }
 
 func (s *subquery) doEval(params runParams) (result tree.Datum, err error) {
-	// After evaluation, there is no plan remaining.
-	defer func() { s.plan.Close(params.ctx); s.plan = nil }()
-
 	switch s.execMode {
 	case distsqlrun.SubqueryExecModeExists:
 		// For EXISTS expressions, all we want to know is if there is at least one
diff --git a/pkg/sql/union.go b/pkg/sql/union.go
index 735d774eeef6..4d722872329f 100644
--- a/pkg/sql/union.go
+++ b/pkg/sql/union.go
@@ -65,7 +65,9 @@ import (
 type unionNode struct {
 	// right and left are the data source operands.
 	// right is read first, to populate the `emit` field.
-	right, left planNode
+	right, left             planNode
+	rightClosed, leftClosed bool
+
 	// columns contains the metadata for the results of this node.
 	columns sqlbase.ResultColumns
 	// inverted, when true, indicates that the right plan corresponds to
@@ -198,34 +200,28 @@ func (n *unionNode) Next(params runParams) (bool, error) {
 	if err := params.p.cancelChecker.Check(); err != nil {
 		return false, err
 	}
-	if n.right != nil {
+	if !n.rightClosed {
 		return n.readRight(params)
 	}
-	if n.left != nil {
+	if !n.leftClosed {
 		return n.readLeft(params)
 	}
 	return false, nil
 }
 
 func (n *unionNode) Values() tree.Datums {
-	if n.right != nil {
+	if !n.rightClosed {
 		return n.right.Values()
 	}
-	if n.left != nil {
+	if !n.leftClosed {
 		return n.left.Values()
 	}
 	return nil
 }
 
 func (n *unionNode) Close(ctx context.Context) {
-	if n.right != nil {
-		n.right.Close(ctx)
-		n.right = nil
-	}
-	if n.left != nil {
-		n.left.Close(ctx)
-		n.left = nil
-	}
+	n.right.Close(ctx)
+	n.left.Close(ctx)
 }
 
 func (n *unionNode) readRight(params runParams) (bool, error) {
@@ -250,8 +246,7 @@ func (n *unionNode) readRight(params runParams) (bool, error) {
 		return false, err
 	}
 
-	n.right.Close(params.ctx)
-	n.right = nil
+	n.rightClosed = true
 	return n.readLeft(params)
 }
 
@@ -273,8 +268,8 @@ func (n *unionNode) readLeft(params runParams) (bool, error) {
 	if err != nil {
 		return false, err
 	}
-	n.left.Close(params.ctx)
-	n.left = nil
+
+	n.leftClosed = true
 	return false, nil
 }