From 2aef9c7e1bd7bad452a5a6ec22c09702b7ec1c9b Mon Sep 17 00:00:00 2001 From: Yahor Yuzefovich Date: Thu, 29 Sep 2022 18:37:33 -0700 Subject: [PATCH] execinfra: correctly reset the eval.Context in InternalClose 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 --- pkg/sql/execinfra/processorsbase.go | 6 +++++- pkg/sql/instrumentation.go | 11 ----------- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/pkg/sql/execinfra/processorsbase.go b/pkg/sql/execinfra/processorsbase.go index 1b5b30788c07..f773d38ea420 100644 --- a/pkg/sql/execinfra/processorsbase.go +++ b/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 } diff --git a/pkg/sql/instrumentation.go b/pkg/sql/instrumentation.go index e4eabcdadf76..602d7b088f42 100644 --- a/pkg/sql/instrumentation.go +++ b/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)