distsql: emit cancellation error in MoveToDraining

Previously if the context used to run a processor was canceled, most of
the time the emitted error would be that context cancellation, however
this depended on the exact implementation of each processor and the code
it called reliably returning the context cancelled error exactly in this
case. In at least some places, this is not reliably done.

This changes DistSQL's base processor to ensure that if the context it
is using to run a processor is canceled, the error emitted in MoveToDraining
is always the context cancellation error, by overriding the error passed
with the context error in MoveToDraining if the former is set.

This ensures callers inspecting errors to determine if they are due to
cancellation can always rely a cancelled flow's emitted error being
a cancellation error.

Fixes #156337.

Release note: none.
Epic: none.

Author: dt

pkg/sql/execinfra/processorsbase.go

@@ -526,7 +526,9 @@ const (
 // at init() time), then we move straight to the StateTrailingMeta.
 //
 // An error can be optionally passed. It will be the first piece of metadata
-// returned by DrainHelper().
+// returned by DrainHelper(), unless the processor's context has already been
+// canceled, in which case the context's error is returned instead (as often an
+// error passed to MoveToDraining() is a consequence of context cancellation).
 //
 // MoveToDraining should only be called from the main goroutine of the
 // processor.
@@ -547,6 +549,18 @@ func (pb *ProcessorBaseNoHelper) MoveToDraining(err error) {
 	}
 
 	if err != nil {
+		// If processor ctx was canceled, reply with that err rather than whatever
+		// error was passed to MoveToDraining by a processor running on top of a
+		// canceled context, which is expected to error. Generally the error passed
+		// in this case will be context.Canceled anyway, but doing this ensures that
+		// distsql can promise that if it cancels a context, the emitted error will
+		// reflect that, making cancellation detectable by callers.
+		if pb.Ctx().Err() != nil {
+			if !errors.Is(err, pb.Ctx().Err()) {
+				log.Dev.Warningf(pb.Ctx(), "overriding non-cancelation emitted after context cancellation: %+v", err)
+			}
+			err = pb.Ctx().Err()
+		}
 		pb.trailingMeta = append(pb.trailingMeta, execinfrapb.ProducerMetadata{Err: err})
 	}
 	if pb.curInputToDrain < len(pb.inputsToDrain) {