execinfra: correctly reset the eval.Context in InternalClose

pkg/sql/execinfra/processorsbase.go

@@ -358,6 +358,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
 
@@ -862,6 +865,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
 }
@@ -893,7 +897,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

@@ -710,18 +710,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)