Merge #89093

89093: execinfra: correctly reset the eval.Context in InternalClose r=yuzefovich a=yuzefovich

This commit fixes the way we restore the `eval.Context.Context` in
`ProcessorBase.InternalClose`. Whenever a processor is started with
`StartInternal`, we derive a tracing span, and in order for all things
related to that processor (including the evaluation of the builtin
functions) to be correctly attributed to the tracing span of the processor,
we update the eval context with the new context. Most processors get
a fresh copy of the eval context via `FlowCtx.NewEvalCtx`, but a couple
(that are deemed to be "safe") do not do that (since they don't actually
evaluate any expressions) and, instead, reuse the "main" eval context.
Then, when the processor is closed, we want to reset the eval context to
its original form, and we had a bug here - we would use the context that
was passed into `StartInternal` rather than the "original" one, so we
could leak the context with an already finished tracing span outside.
This is now fixed.

This bug could occur for example when we have a `FlowCoordinator` (which
doesn't create a copy of the eval context) at the root of the main query
and then there are some postqueries - the eval context used by the
postqueries would refer to the context from the main query with
a finished tracing span. However, the bug was mitigated by explicitly
setting the context before running each postquery (due to an unrelated
old issue).

One could argue that since the processors modify the eval context, they
should always use a copy of the eval context. That would be unfortunate
since it would lead to increased allocations with the only modification
being this context-mangering business. The fix in this commit seems more
appropriate (both philosophically speaking and performance-wise).

Release note: None

Co-authored-by: Yahor Yuzefovich <yahor@cockroachlabs.com>

pkg/sql/execinfra/processorsbase.go

@@ -359,6 +359,9 @@ type ProcessorBaseNoHelper struct {
 	// origCtx is the context from which ctx was derived. InternalClose() resets
 	// ctx to this.
 	origCtx context.Context
+	// evalOrigCtx is the original context that was stored in the eval.Context.
+	// InternalClose() uses it to correctly reset the eval.Context.
+	evalOrigCtx context.Context
 
 	State procState
 
@@ -865,6 +868,7 @@ func (pb *ProcessorBaseNoHelper) StartInternal(ctx context.Context, name string)
 			pb.span.SetTag(execinfrapb.ProcessorIDTagKey, attribute.IntValue(int(pb.ProcessorID)))
 		}
 	}
+	pb.evalOrigCtx = pb.EvalCtx.Context
 	pb.EvalCtx.Context = pb.Ctx
 	return pb.Ctx
 }
@@ -896,7 +900,7 @@ func (pb *ProcessorBaseNoHelper) InternalClose() bool {
 	// Reset the context so that any incidental uses after this point do not
 	// access the finished span.
 	pb.Ctx = pb.origCtx
-	pb.EvalCtx.Context = pb.origCtx
+	pb.EvalCtx.Context = pb.evalOrigCtx
 	return true
 }
 

pkg/sql/instrumentation.go

@@ -712,18 +712,7 @@ func (ih *instrumentationHelper) SetIndexRecommendations(
 		isInternal,
 	) {
 		f := opc.optimizer.Factory()
-		// EvalContext() has the context with the already closed span, so we
-		// need to update with the current context.
-		// The replacement of the context here isn't ideal, but the current
-		// implementation of contexts would need to change
-		// significantly to accommodate this case.
 		evalCtx := opc.p.EvalContext()
-		oldCtx := evalCtx.Context
-		evalCtx.Context = ctx
-		defer func() {
-			evalCtx.Context = oldCtx
-		}()
-
 		f.Init(evalCtx, &opc.catalog)
 		f.FoldingControl().AllowStableFolds()
 		bld := optbuilder.New(ctx, &opc.p.semaCtx, evalCtx, &opc.catalog, f, opc.p.stmt.AST)