Decouple gRPC from the worker protocol and add an in-process transport by stuhood · Pull Request #573 · datafusion-contrib/datafusion-distributed

stuhood · 2026-07-28T19:21:51Z

What

Decouples the datafusion-distributed core logic from the grpc transport and introduces an in-process transport for single-process distributed task execution.

Why

Previously, the worker protocol and the grpc feature were coupled, which meant any use of the distributed execution abstraction required the tonic gRPC stack. Separating the transport layer from the protocol messages allows the framework to be used with alternative transports. The in-process transport provides a mechanism to run and test the abstraction when the grpc feature is disabled.

How

Split worker.proto into messages.proto (data structures) and worker_service.proto (gRPC service definition in the worker_grpc package). This splits the generated code into worker.rs (always compiled) and worker_grpc.rs (compiled only when the grpc feature is enabled).
Added InProcessChannelResolver. It implements WorkerChannel by routing the three protocol methods directly to a co-located Worker instance in the same process, bypassing IPC, networking, and serialization.
Build adjustments:
- Removed the benchmarks dev-dependency to prevent grpc from being enabled in every test build.
- Gated gRPC-coupled test utilities in test_utils behind the grpc feature.
- (Note: The integration feature flag still assumes grpc is enabled because test_utils::InMemoryChannelResolver is used across the integration suite.)

Tests

Added a unit-test-no-grpc CI job that verifies the core library and the in-process transport compile and pass unit tests without the grpc feature.
Added an end-to-end test for InProcessChannelResolver (a distributed GROUP BY across tasks) that runs in the no-gRPC CI job.

The prost message types carry no tonic dependency, so a transport that is not gRPC can speak the same wire shape without pulling in the gRPC stack. Only the tonic client and server stay gated; the generator emits those gates so a regeneration cannot drop them. tonic-prost feeds only the generated client and server, so it moves behind the feature too. Co-authored-by: Stu Hood <stuhood@gmail.com>

The benchmarks crate's dev-dependency on the lib re-unified grpc into every test build, so a genuine no-gRPC test run was impossible; the dataset suites move into the benchmarks crate and the gRPC-coupled test utilities gate behind grpc. A unit-test-no-grpc job then runs the whole lib suite with the feature off. Co-authored-by: Stu Hood <stuhood@gmail.com>

InProcessChannelResolver routes the three protocol methods straight to a co-located Worker, with no gRPC, no IPC, and no serialization round-trip: the reference implementation of the protocol for a co-located worker, and the first transport that exercises the abstraction with grpc off. Its end-to-end test (a distributed GROUP BY across tasks) runs under the no-gRPC CI job. Co-authored-by: Stu Hood <stuhood@gmail.com>

gabotechs

This is starting to look good!

I see there are different pieces of work that can be factored out into independent PRs that we can ship incrementally. I identify the following pieces of work in the following order:

split worker.proto into worker_service.proto and messages.proto, putting some thought into where to place the proto generator script so that we don't copy-paste too much.
Ship the InProcessWorkerChannel, and wire it up to channel_resolver.rs (or somewhere else that makes sense) so that we can start executing queries without grpc. Probably, the interaction of InProcessWorkerChannel and LocalWorkerContext needs to happen here.
Tweak existing tests unit tests that rely on grpc so that they just use in-process comms (mainly the tests about rewritting plans with metrics). That will allow to decouple them for gRPC by using the in-process implementation.
Tweak the remaining test utils, moving some to grpc/ and flagging them appropriately:
- If grpc feature is enabled, then gRPC worker channels are used (still with in-process fallbacks for self urls)
- If grpc is disabled, exclusively in-process worker channels are used. This should be transparent to the tests.

gabotechs · 2026-07-29T05:33:17Z

+        if event::poll(Duration::from_millis(16))?
+            && let Event::Key(key) = event::read()?
+        {
+            input::handle_key_event(app, key);
        }


🤔 this change is unrelated right? if it is, I'd avoid introducing unrelated changes, mainly for better traceability.

gabotechs · 2026-07-29T05:33:29Z

+                    if old_task.status == TaskStatus::Running as i32
+                        && let Some(sk) = &old_task.task_key
+                    {


Same as above

gabotechs · 2026-07-29T05:39:41Z

-    let proto_file = proto_dir.join("worker.proto");
-    let out_dir = repo_root.join("src/protocol/grpc/generated");
+    let protocol_dir = repo_root.join("src/protocol");
+    let messages_proto = protocol_dir.join("messages.proto");


This generator file is under the grpc scope, but it's in charge of generating messages.proto which is not part of grpc.

Probably copy-pasting this generator script one folder above is not worth it. What about moving it to the root of the project under a codegen/ folder or something like that? we can centralize proto generation there.

gabotechs · 2026-07-29T05:41:02Z

+
+#[cfg(feature = "grpc")]
+#[allow(clippy::all)]
+pub mod worker_grpc;


If this is specific to grpc, ideally it should be generated under the grpc/ folder

gabotechs · 2026-07-29T05:42:14Z

+#[cfg(feature = "grpc")]
+mod channel {
+    use super::InMemoryWorkerResolver;


All the test utils that are specific to grpc should probably live under the grpc/ folder, that way they are gated automatically.

gabotechs · 2026-07-29T06:00:27Z

+// `worker_handles` builds `tonic` servers and Flight channels for the benchmark fixtures, which
+// only compile with the gRPC transport.
+#[cfg(all(feature = "grpc", any(test, feature = "integration")))]
 pub(crate) mod test_utils;


This is the pattern we should aim to get rid of in the project: if it's grpc related, it should ideally live under grpc/, this means that we probably need to refactor some tests to rely on the in-memory representation.

gabotechs · 2026-07-29T06:15:35Z

    }

-    #[cfg(test)]
+    #[cfg(all(test, feature = "grpc"))]


Seems like this is needed because this code is used in the task_metrics_rewriter.rs tests, and those tests are still coupled to grpc.

It'd be nice to decouple that test from grpc first so that we can avoid flagging this here.

gabotechs · 2026-07-29T06:16:31Z

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "clickbench", test))]
+#[cfg(all(feature = "clickbench", test))]


These tests should really not be under benchmarks. I think I might be missing the reason for these to be here. Do they really need to be here?

gabotechs · 2026-07-29T06:17:31Z


-    #[cfg(test)]
+    #[cfg(all(test, feature = "grpc"))]
    pub(crate) fn inner(&self) -> &Arc<dyn ExecutionPlan> {


This is only used in task_metrics_rewriter.rs tests. If we manage to decouple those tests from grpc, we can leave this untouched.

gabotechs · 2026-07-29T06:54:45Z

+/// state for every task, keyed by [`crate::TaskKey`], the same way the gRPC worker does when several
+/// tasks of a query land on it.
+#[derive(Clone)]
+pub struct InProcessChannelResolver {


For a follow-up PR: this should somehow interact with channel_resolver.rs and LocalWorkerContext so that, once a request is made to a URL that matches LocalWorkerContext.self_url, the InProcessChannelResolver kicks in automatically.

mdashti and others added 4 commits July 28, 2026 11:30

Split worker.proto into messages.proto and worker_service.proto.

3888470

stuhood marked this pull request as ready for review July 28, 2026 19:29

stuhood force-pushed the stuhood.in-memory-transport branch from 9be56b0 to 3888470 Compare July 28, 2026 19:29

stuhood mentioned this pull request Jul 28, 2026

feat: add an in-process worker transport (no gRPC) #548

Closed

gabotechs reviewed Jul 29, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Decouple gRPC from the worker protocol and add an in-process transport - #573

Decouple gRPC from the worker protocol and add an in-process transport#573
stuhood wants to merge 4 commits into
datafusion-contrib:mainfrom
paradedb:stuhood.in-memory-transport

stuhood commented Jul 28, 2026 •

edited

Loading

Uh oh!

gabotechs left a comment

Uh oh!

gabotechs Jul 29, 2026

Uh oh!

gabotechs Jul 29, 2026

Uh oh!

gabotechs Jul 29, 2026

Uh oh!

gabotechs Jul 29, 2026

Uh oh!

gabotechs Jul 29, 2026

Uh oh!

gabotechs Jul 29, 2026

Uh oh!

gabotechs Jul 29, 2026

Uh oh!

gabotechs Jul 29, 2026

Uh oh!

gabotechs Jul 29, 2026

Uh oh!

gabotechs Jul 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		@@ -1,4 +1,4 @@
		#[cfg(all(feature = "integration", feature = "clickbench", test))]
		#[cfg(all(feature = "clickbench", test))]

Uh oh!

Conversation

stuhood commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Why

How

Tests

Uh oh!

gabotechs left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

stuhood commented Jul 28, 2026 •

edited

Loading