Allow workers to reject work for a query (before side effects)

[lesam]

Right now, once the coordinator selects a worker for a query task, that worker is committed to the query and any failure in the worker will fail the entire query, across many workers. This commitment happens before any communication with the worker.

For load balancing, rolling restarts, etc - it is valuable for a worker to be able to locally decide to refuse some new work to maintain throughput.

I see two options to achieve this:

1. The worker returns a yes/no to the coordinator about accepting the work (this can happen as part of establishing the initial gRPC stream), and the coordinator retries on a different worker only on an explicit 'no' which is guaranteed to be sent before any side effects.

2. After establishing the gRPC stream, the worker waits for an explicit 'go-ahead' from the coordinator before executing anything side-effectfull (e.g., pulling data from lower stages in the plan). This means that timeouts and other ambiguous failures can also be retried on other workers from the coordinator, as long as the coordinator has not sent a 'go-ahead'.

My preference is for option 2, as this provides more ability to work around an unhealthy worker with a cost of ~half a coordinator-worker round-trip latency.

[gabotechs]

After establishing the gRPC stream, the worker waits for an explicit 'go-ahead' from the coordinator before executing anything side-effectfull

Whenever a gRPC stream is stablished, that pretty much means already a 'go-ahead' implicitly, side-effectfull things are expected to happen immediately in every case.

The worker returns a yes/no to the coordinator about accepting the work (this can happen as part of establishing the initial gRPC stream), and the coordinator retries on a different worker only on an explicit 'no' which is guaranteed to be sent before any side effects.

This sounds reasonable. The gRPC protocol already allows to say no by returning an error, so a flow that sounds reasonable would be something like:

1. The coordinator tries to establish a CoordinatorChannel stream with the worker.
2. Before the stream is stablished, the server side CoordinatorChannel handle returns an error.
3. The coordinator reads the error, flags the worker as "unreliable" or something, and chooses a different one.

We have good support for sending DataFusionError messages over the wire, we can leverage that code for sending informative errors that are not necessarily coupled to gRPC.

[lesam]

Sounds like you prefer option 1. Just a note that if we go that direction, we cannot retry timeouts establishing the stream, because we don't know which of the following cases we're in:

• the worker failed, so we're ok to retry
• the worker finished establishing its end of the stream and moved on to side-effectful work, but the coordinator never received the initial headers that the worker sent - in this case it would be an error to retry.

To allow retrying timeouts, the worker has to wait to do side-effectful things until the worker is sure that the coordinator actually received the success that the worker sent (which is what the extra message would be for in option 2).

[lesam]

Also with option 1 we'll have to be extremely careful about which errors we retry - we have to classify all errors as either 'definitely sent by the worker with the intention of stopping and not performing any side effect' or 'potentially side effects happened after this error was sent'. Generic 'connection errors' are hard to classify because it's possible that the worker views a success but the coordinator views a failure, and these must not be retried by the coordinator.

[lesam]

I was leaning toward option 2 especially because of the potential case of "worker is so overloaded that it can't even successfully send an error response in a reasonable time", which under option 2 the coordinator would be able to retry somewhere else, but under option 1 the whole query has to fail.

[gabotechs]

I don't see how option 2 solves any of that. The only difference I see with option 2 is that there's an extra 'go-ahead' message sent by the coordinator, but from there on the problems look applicable to both cases.

The thing is, there's no need for an extra 'go-ahead' message, establishing the connection with the worker is already the 'go-ahead' message we need.

[lesam]

Currently, SamplerExec::kick_off_first_sampler(Arc::clone(&plan), Arc::clone(&task_ctx))?; happens in the worker before we even complete the handshake with the coordinator (before Worker::coordinator_channel returns).

We could almost fix this by waiting in the worker until after the worker completes the handshake to kick off the sampler, but it is not a fix because there is still a racy bug: the worker can complete the handshake successfully but the coordinator views a failure.

Concretely, without the go-ahead message and with sampling post-handshake, the failure case for retrying a timeout is:

1. the coordinator sends the first message to a stage-2 worker to try to establish the stream
2. the worker initializes the stream successfully on the worker end
3. the stage-2 worker pulls some of the stream from stage-1 (this can happen as part of sampling for AQE, which kicks off immediately)
4. the successful initialization of the worker is not yet communicated to the coordinator, and on the coordinator side the stream initialization times out
5. the coordinator retries the initialization on a different worker, which also tries to pull from the same stage-1 💥

If the worker instead has to wait for a go-ahead from the coordinator, step 4 cannot happen after step 3, and so the coordinator can retry timeouts. Otherwise, a timeout (for AQE at least, where the worker could have started a sampler) is ambiguous and the coordinator cannot assume a timeout from the worker means no side effects happened.

I think static plans work without a go-ahead today, but that also seems brittle - we'd have to make sure that nothing ever proactively pulls data from lower stages before the higher stage is 100% successful in the coordinator (note the 'sampler' case for dynamic plans is exactly this kind of proactive data pull).

I've put up a test at https://github.com/lesam/datafusion-distributed/pull/1/changes to prove that the worker and the coordinator can disagree about the success of the handshake.

[gabotechs]

Indeed, you are right, the current placement of SamplerExec::kick_off_first_sampler can be problematic here. We do might need an explicit sampler kick-off message in the coordinator->worker stream.

For static planning, there do are no side effects happening during the coordinator->worker channel creation, and the expectation is that there are never some, so I think we should be fine there.

I see two separate pieces of work here:

• How to react in the coordinator to workers saying "no" for more work.
• Introducing the extra message in the coordinator->worker stream for kicking off samplers in aqe.

I get the impression that neither of those should be very difficult.

[lesam]

I have to backburner this in favour of other stuff right now.

Notes for that later:

Claude noticed that SessionBuilder and especially PhysicalExtensionCodec need to be idempotent/retryable, so we need to document that (assuming that we let them happen during the initial handshake).

Additionally, with no retries we fire-and-forget all the SetPlanRequest calls, but with retries we have to serialize them from bottom to top as we don't know the correct worker URL for stage 1 until a worker accepts, and we need to send that URL in the stage 2 plan. This means that even with no additional messages and static plans we need to wait for N RTT's where N is the depth of the plan.

This makes me want to move the other way - when retries are enabled, have a new optional 'StartStream' message, which carries minimal info (and doesn't carry the full plan/lower stage URLs), which the worker fails if it can't serve, and wait to send SetPlanRequest until we get the successful stream init from StartStream. That way, under retries, we know that failing/retrying StartStream is always safe (the worker doesn't know anything yet so it can't do bad side effects) and also we can StartStream everything in parallel, then fire-and-forget all the SetPlanRequest calls with the URLs filled in, which moves us from N RTT's to 2 RTT's for initialization latency. In the future we could potentially include some notion of 'task heaviness'/'task priority' in StartStream to give the worker a chance to reject work based on relevant properties, but for now I think we can have StartStream carry literally no data.

[gabotechs]

I have added a new test that I expect to be capable of replicating this kind of errors:

• Add chaos test #667

That way we can qualify further improvements with an actual integration test enforced in the CI.