From c39a49fda9ee5a4873133a8f49a741d7d71f3b1e Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Thu, 9 Sep 2021 15:02:51 +0000
Subject: [PATCH] colexec: adjust the eager cancellation in parallel unordered
 sync a bit

This commit adjusts the logic for swallowing errors in the parallel
unordered synchronizer because of the eager cancellation when the
synchronizer transitions into the draining state to swallow all errors
coming from an input if the input's context has been canceled by the
synchronizer.

Release note: None

Release justification: update to the new functionality.
---
 pkg/sql/colexec/BUILD.bazel                   |  1 -
 .../parallel_unordered_synchronizer.go        | 36 ++++++-------------
 2 files changed, 10 insertions(+), 27 deletions(-)

diff --git a/pkg/sql/colexec/BUILD.bazel b/pkg/sql/colexec/BUILD.bazel
index 74c128b876a0..32632d8151f4 100644
--- a/pkg/sql/colexec/BUILD.bazel
+++ b/pkg/sql/colexec/BUILD.bazel
@@ -71,7 +71,6 @@ go_library(
         "//pkg/sql/sqltelemetry",  # keep
         "//pkg/sql/types",
         "//pkg/util",
-        "//pkg/util/cancelchecker",
         "//pkg/util/duration",  # keep
         "//pkg/util/encoding",  # keep
         "//pkg/util/humanizeutil",
diff --git a/pkg/sql/colexec/parallel_unordered_synchronizer.go b/pkg/sql/colexec/parallel_unordered_synchronizer.go
index 8cfb01c1d7b7..59c375b10754 100644
--- a/pkg/sql/colexec/parallel_unordered_synchronizer.go
+++ b/pkg/sql/colexec/parallel_unordered_synchronizer.go
@@ -22,8 +22,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/colexecop"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
-	"github.com/cockroachdb/cockroach/pkg/util"
-	"github.com/cockroachdb/cockroach/pkg/util/cancelchecker"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/errors"
 )
@@ -249,30 +247,16 @@ func (s *ParallelUnorderedSynchronizer) init() {
 				switch state {
 				case parallelUnorderedSynchronizerStateRunning:
 					if err := colexecerror.CatchVectorizedRuntimeError(s.nextBatch[inputIdx]); err != nil {
-						if s.Ctx.Err() == nil && s.cancelLocalInput[inputIdx] != nil {
-							if errors.Is(err, context.Canceled) || errors.Is(err, cancelchecker.QueryCanceledError) {
-								// The input context has been canceled, yet the
-								// main context of the synchronizer has not.
-								// This indicates that the synchronizer has
-								// transitioned into draining state and wanted
-								// this goroutine to stop whatever it was doing.
-								// Therefore, we swallow the error and proceed
-								// to draining.
-								//
-								// Note that we need the second part of the
-								// conditional in case the context cancellation
-								// was observed by the CancelChecker and was
-								// propagated as a query canceled error.
-								if util.CrdbTestBuild {
-									if s.getState() != parallelUnorderedSynchronizerStateDraining {
-										colexecerror.InternalError(errors.AssertionFailedf(
-											"unexpectedly the input context is canceled, the main " +
-												"context is not, and not in the draining state",
-										))
-									}
-								}
-								continue
-							}
+						if s.getState() == parallelUnorderedSynchronizerStateDraining && s.Ctx.Err() == nil && s.cancelLocalInput[inputIdx] != nil {
+							// The synchronizer has just transitioned into the
+							// draining state and eagerly canceled work of this
+							// input. That cancellation is likely to manifest
+							// itself as the context.Canceled error, but it
+							// could be another error too; in any case, we will
+							// swallow the error because the user of the
+							// synchronizer is only interested in the metadata
+							// at this point.
+							continue
 						}
 						sendErr(err)
 						return