fix(compute): preserve caller context during execution - #1154
fix(compute): preserve caller context during execution#1154fallintoplace wants to merge 3 commits into
Conversation
|
Since this is a draft, i'll hold off on further review until the conflicts are resolved and it's marked ready |
51973d3 to
430755b
Compare
zeroshade
left a comment
There was a problem hiding this comment.
The caller-context propagation is correct in intent, but the new cancellation path exposes result ownership and error-contract problems. In particular, cancellation can return no error or an already-released result, and accumulated scalar outputs can leak. I’ve left details and reproduction cases inline.
This review was drafted by an AI-assisted tool and confirmed by an Apache Arrow Go maintainer. After you’ve addressed the points above and pushed an update, an Apache Arrow Go maintainer — a real person — will take the next look at the PR. If you think a finding is misapplied, please reply on the PR and a maintainer will weigh in.
More on how Apache Arrow Go handles contributions:
CONTRIBUTING.md.
|
|
||
| result = executor.WrapResults(ctx, ch, haveChunkedArray(input.Values)) | ||
| if err == nil { | ||
| if ctx.Err() != nil { |
There was a problem hiding this comment.
Blocking: This cancellation branch drains the channel but does not guarantee a cancellation error or clear a released result. A context-oblivious kernel can finish successfully after cancellation, leaving err == nil. If cancellation occurs before the first output, CallFunction returns (nil, nil); after the first scalar output, WrapResults can return a partial result that line 185 releases but the named return still exposes. I reproduced the latter as a non-nil ArrayDatum with nil data and context.Canceled. Please propagate context.Cause(ctx) when err is nil, and release and set result = nil before returning. A regression test should use a kernel that does not itself return ctx.Err().
| ectx := GetExecCtx(ctx) | ||
|
|
||
| ctx, cancel := context.WithCancel(context.Background()) | ||
| ctx, cancel := context.WithCancel(ctx) |
There was a problem hiding this comment.
Major: Deriving this context from the caller makes the scalar collector’s cancellation branch externally reachable. When scalarExecutor.WrapResults has already moved outputs into its private acc, cancellation returns without releasing those chunks. The drain below cannot reclaim values already consumed from the channel. Please add scalar accumulated-result cleanup equivalent to vectorExecutor.WrapResults’s releaseAccumulated() path.
| return ctx.Ctx.Err() | ||
| }, nil) | ||
| require.NoError(t, fn.AddKernel(kernel)) | ||
| require.True(t, GetFunctionRegistry().AddFunction(fn, false)) |
There was a problem hiding this comment.
Minor: Registering this function in the process-wide registry makes the test non-repeatable: go test ./arrow/compute -run '^TestCallFunctionPreservesCallerCancellation$' -count=2 fails on its second execution. Please install a child registry through ExecCtx for this test.
What
Keep the supplied context attached to compute execution so cancellation reaches result collection. Drain pending results before returning so the executor is not reused while work is still finishing.
Test