GH-37796: [C++][Acero] Fix race condition caused by straggling input in the as-of-join node

[JerAguilon]

Rationale for this change

What changes are included in this PR?

While asofjoining some large parquet datasets with many row groups, I ran into a deadlock that I described here: #37796. Copy pasting below for convenience:

1. The left hand side of the asofjoin completes and is matched with the right hand tables, so InputFinished proceeds as expected. So far so good
2. The right hand table(s) of the join are a huge dataset scan. They're still streaming and can legally still call AsofJoinNode::InputReceived all they want (doc ref)
3. Each input batch is blindly pushed to the InputStates, which in turn defer to BackpressureHandlers to decide whether to pause inputs. (code pointer)
4. If enough batches come in right after EndFromProcessThread is called, then we might exceed the high_threshold and tell the input node to pause via the BackpressureController
5. At this point, the process thread has stopped for the asofjoiner, so the right hand table(s) won't be dequeue'd, meaning BackpressureController::Resume() will never be called. This causes a deadlock

TLDR this is caused by a straggling input node being paused due to backpressure after the process thread has ended. And since every PauseInput needs a corresponding ResumeInput to exit gracefully, we deadlock.

Turns out this is fairly easy to reproduce with small tables, if you make a slow input node composed of 1-row record batches with a synthetic delay.

My solution is to:

1. Create a ForceShutdown hook that puts the input nodes in a resumed state, and for good measure we call StopProducing
2. Also for good measure, if nodes come after the process thread exits, we short circuit and return OK. This is because InputReceived can be called an arbitrary number of times after StopProducing, so it makes sense to not enqueue useless batches.

Are these changes tested?

Yes, I added a delay to the batches of one of the already-existing asofjoin backpressure tests. Checkout out main, we get a timeout failure. With my changes, it passes.

I considered a more deterministic test, but I struggled to create callbacks in a way that wasn't invasive to the Asof implementation. The idea of using delays was inspired by things I saw in source_node_test.cc

[JerAguilon]

Perhaps one thing to clarify is whether ResumeInput behaves idempotently? I.e., is it OK to always call resume, even though only some inputs hit this PauseInput race condition?

My perusal of source_node.cc tells me this is OK, but LMK if this is a poor assumption to make.

[JerAguilon]

Oops... Will remove

[westonpace]

This seems like a good idea to me.

@icexelloss @rtpsw do either of you want to take a look?

[westonpace]

I assume this new option triggers the deadlock on the unfixed code?

[JerAguilon]

Correct

[westonpace]

So if I understand correctly this means we will call StopProducing on all right hand side nodes once:

• The left hand side has finished
• The right hand side has caught up

If so, then I agree this is a valid thing to do.

[JerAguilon]

Yep.

As an aside, I feel like a more invasive change could fix this issue in the general case. If a node (in this example asof join) has:

• Called output->InputFinished() AND
• Called output_->InputReceived for however many record batches it advertised on InputFinished

We should be able to shut down execution, even if the node's inputs:

• are paused or
• not done streaming
• haven't called InputFinished

But I think this is a more invasive change to exec_plan.h and might have some hairy issues that I'm not thinking of.

[icexelloss]

This looks reasonable to me. Free feel to merge.

[bkietz]

Suggested change

	// InputReceived may be called after execution was finished. Pushing it to the
	// InputState may cause the BackPressureController to pause the input, causing a
	// deadlock
	// InputReceived may be called after execution was finished. Pushing it to the
	// InputState is unnecessary since we're done (and anyway may cause the
	// BackPressureController to pause the input, causing a deadlock), so drop it.

Do we still deadlock with this short circuit but without ForceShutdown etc?

[JerAguilon]

Yes, the forceShutdown is still necessary. there's nothing stopping this order of events:

1. We receive enough data to finish the as of join.
2. Right before we finish processing and shut down the worker thread, lots of unneeded batches come in from input A. Input A pauses
3. We shut down the thread, and input A can't be unpaused

Put another way, forceShutdown keeps us from deadlocking when we ingest unneeded data before the worker thread exits. And this block keeps us from deadlocking when we ingest unneeded data after the worker thread exits.

But your comment change suggestions sound good to me

[JerAguilon]

Clarifying comment for @bkietz added. Ready for more thoughts

[JerAguilon]

Forgive the ignorance - first time making a PR on arrow. There's no further action needed from me to merge, correct?

[bkietz]

Could you rebase to pick up the fix #37867 ? I think CI should be green after that

[JerAguilon]

Sadly there are some seemingly unrelated failures: TestS3FS.GetFileInfoGeneratorStress and arrow-threading-utility-test

[bkietz]

CI failures seem unrelated. I'll merge. Thanks for working on this!

[conbench-apache-arrow]

After merging your PR, Conbench analyzed the 6 benchmarking runs that have been run so far on merge-commit e3d6b9b.

There were no benchmark performance regressions. 🎉

The full Conbench report has more details. It also includes information about 1 possible false positive for unstable benchmarks that are known to sometimes produce them.

[raulcd]

Closes: #37796