fix(scheduler): set TCP keepalive on scheduler gRPC connections by ump45nose · Pull Request #2816 · apache/datafusion-ballista

ump45nose · 2026-11-08T14:22:31Z

Which issue does this PR close?

Closes #2808.

Rationale for this change

The scheduler serves its gRPC listener with axum::serve, which sets no TCP keepalive on accepted connections. A scheduler connection whose peer disappears without closing it (a killed executor host, a dropped NAT or load-balancer entry) is therefore never reaped by OS keepalive probes, while executor and result-service listeners get a 3600 s keepalive from create_grpc_server_incoming.

What changes are included in this PR?

start_grpc_service_with_listener: the existing tap_io hook now goes through a small configure_connection helper that sets TCP_NODELAY (unchanged) and TCP keepalive via socket2::SockRef::set_tcp_keepalive. Failures are logged as warnings, like the TCP_NODELAY case.
New scheduler option grpc_server_tcp_keepalive_seconds (--grpc-server-tcp-keepalive-seconds, default 3600 to match GrpcServerConfig, 0 disables) with a SchedulerConfig::with_grpc_server_tcp_keepalive_seconds builder.
Standalone scheduler: its TcpIncoming now also gets with_keepalive from the default GrpcServerConfig, since tonic ignores the builder's tcp_keepalive on a pre-bound listener.
socket2 (already in the lockfile at 0.6.5 through tokio/hyper) becomes a direct dependency of ballista-scheduler.
Unit test configures_nodelay_and_keepalive_on_accepted_connections checks TCP_NODELAY, SO_KEEPALIVE and the keepalive idle time on an accepted socket, and that 0 leaves keepalive off.

Not covered: the HTTP/2 keepalive ping interval/timeout that create_grpc_server sets for tonic servers. axum::serve does not expose those settings, so matching them would mean serving with hyper directly; I left that out to keep this change small.

Are there any user-facing changes?

One new scheduler CLI flag / SchedulerConfig field, grpc_server_tcp_keepalive_seconds (default 3600). Scheduler connections now get TCP keepalive by default. No API removals.

Testing

cargo test -p ballista-scheduler --lib (611 passed, 2 ignored)
cargo clippy -p ballista-scheduler --all-targets -- -D warnings
cargo check -p ballista-scheduler --all-features --bins
cargo fmt --all

This PR was made with AI assistance. I reviewed the diff and ran the tests above.

The scheduler serves its gRPC listener with axum::serve, which sets no TCP keepalive, so a connection whose peer disappears without closing it is never reaped. Executors and the result service get a 3600 s keepalive from create_grpc_server_incoming. Set keepalive in the same tap_io hook that sets TCP_NODELAY, from a new grpc_server_tcp_keepalive_seconds scheduler option (default 3600, 0 disables). The standalone scheduler now applies GrpcServerConfig's keepalive to its TcpIncoming as well.

andygrove

Thanks for picking this up! The configure_connection helper reads really nicely, and it's good to have the scheduler match the executor and result service keepalive. The one thing holding this up is the red python Cargo.lock up to date check, details inline. Happy to approve once that's green.

andygrove · 2026-11-08T14:49:57Z

 rand = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true, optional = true }
+socket2 = { version = "0.6", features = ["all"] }


The python Cargo.lock up to date check is failing because python/Cargo.toml pulls in ballista-scheduler by path, so python/Cargo.lock needs the new socket2 entry too. Running cd python && cargo update -p ballista-scheduler and committing the lockfile should fix it.

andygrove · 2026-11-08T14:49:57Z

 rand = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true, optional = true }
+socket2 = { version = "0.6", features = ["all"] }


Small one. I think features = ["all"] is only needed for tcp_keepalive_time() in the test. set_tcp_keepalive and TcpKeepalive::with_time work without it. Could the main dependency drop the feature, with socket2 = { version = "0.6", features = ["all"] } under [dev-dependencies] instead?

andygrove · 2026-11-08T14:49:57Z

+    /// Connections accepted by the scheduler get the same TCP keepalive as
+    /// executor and result-service connections, and none when it is disabled.
+    #[tokio::test]
+    async fn configures_nodelay_and_keepalive_on_accepted_connections() {


Nit: SockRef::tcp_keepalive_time() isn't available on Windows in socket2 0.6, so this test won't compile there. We don't run Windows tests in CI, but a #[cfg(not(windows))] would keep it portable.

python/Cargo.toml depends on ballista-scheduler by path, so the new socket2 dependency must appear in python/Cargo.lock. Keep features = ["all"] on the dev-dependency only (needed for tcp_keepalive_time in the test) and skip that test on Windows, where socket2 0.6 does not expose the method.

ump45nose · 2026-11-09T07:32:28Z

Thanks @andygrove — addressed the review notes:

python/Cargo.lock now includes socket2 on ballista-scheduler (cargo update -p ballista-scheduler).
features = ["all"] moved to a [dev-dependencies] socket2 entry; the main dependency no longer enables it.
Gated configures_nodelay_and_keepalive_on_accepted_connections with #[cfg(not(windows))] because tcp_keepalive_time() is unavailable there in socket2 0.6.

Signed-off-by: Vincent <ump45nose@users.noreply.github.com>

ump45nose · 2026-11-09T13:19:16Z

Follow-up on CI for ee1cb9b:

Check Cargo.toml formatting failed because taplo wants the new socket2 dev-dependency in sorted order. Fixed in a8156f5; taplo format --check (0.10.0) passes locally on all 15 manifests.
Test Python Release fails in test_typed_setters[with_shuffle_compression_codec-...-zstd] (assert 'zstd' != 'zstd'). That looks unrelated to this PR: perf: default the shuffle compression codec to zstd #2805 made zstd the default shuffle codec, and python/python/tests/test_session_config.py still uses zstd as the non-default value. I left it alone here so this PR stays focused; happy to change the test value here or in a separate PR if you prefer.
Build Web TUI WASM32 application hit a 504 while downloading a release archive, and Trigger (labeler) currently fails on other PRs too.

(AI-assisted reply.)

andygrove

Thanks for picking this up. The tap_io change is small and the test checks the options on a real accepted socket, which is nice. A couple of things inline.

The Python Release failures (test_typed_setters[with_shuffle_compression_codec-...-zstd]) aren't from this PR. That test was fixed on main in #2829, so a rebase should turn CI green.

andygrove · 2026-11-10T12:59:03Z

        "Ballista Scheduler v{BALLISTA_VERSION} (DataFusion v{DATAFUSION_VERSION}) listening on {addr:?}"
    );
+    let grpc_server_config = GrpcServerConfig::default();
+    let incoming = TcpIncoming::from(listener)


Good catch that tonic ignores tcp_keepalive on a pre-bound listener here too. This now duplicates the with_nodelay/with_keepalive chain in create_grpc_server_incoming in ballista/core/src/utils.rs. Could you pull that into a small helper in ballista_core::utils that takes a TcpListener and a &GrpcServerConfig? Then create_grpc_server_incoming and this call site can both use it and won't drift apart.

Done in 64f7275: tcp_incoming_from_listener in ballista_core::utils now owns the with_nodelay / with_keepalive chain. create_grpc_server_incoming and the standalone scheduler both use it, so those settings cannot drift.

(AI-assisted reply.)

andygrove · 2026-11-10T12:59:03Z


+/// Converts a keepalive setting in seconds into the idle time passed to
+/// [`configure_connection`], where 0 disables TCP keepalive.
+fn tcp_keepalive(seconds: u64) -> Option<Duration> {


Here 0 disables keepalive, but GrpcServerConfig::tcp_keepalive_seconds = 0 still passes Some(Duration::from_secs(0)) to with_keepalive in create_grpc_server_incoming. Could you make the two agree, so 0 means disabled everywhere? Reusing this tcp_keepalive conversion in core (alongside the helper from the other comment) would cover it.

Done in 64f7275: tcp_keepalive now lives in ballista_core::utils (0 means disabled) and is used by tcp_incoming_from_listener, create_grpc_server_incoming, create_grpc_server, and the scheduler's accepted-connection path. GrpcServerConfig::tcp_keepalive_seconds = 0 no longer passes Some(Duration::from_secs(0)).

(AI-assisted reply.)

…keepalive

…sabled Review follow-up for apache#2816: move the TcpIncoming nodelay/keepalive setup into ballista_core::utils so the standalone scheduler and create_grpc_server_incoming cannot drift, and use the same conversion so tcp_keepalive_seconds = 0 disables keepalive on both paths (and on create_grpc_server). Signed-off-by: Vincent <ump45nose@users.noreply.github.com>

ump45nose · 2026-11-10T15:41:30Z

Thanks @andygrove — addressed in 64f7275, and I merged current main (including the Python Release fix from #2829) so the unrelated test_typed_setters failures should clear after CI reruns.

Shared helper: tcp_incoming_from_listener in ballista_core::utils is used by both create_grpc_server_incoming and the standalone scheduler.
0 disables keepalive everywhere those server paths set it (tcp_incoming_from_listener, create_grpc_server, and the scheduler tap_io path), via one tcp_keepalive conversion.

Locally, with protoc available:

cargo test -p ballista-core --lib test_tcp_keepalive_zero_disables test_create_grpc_server_incoming_binds_eagerly test_create_grpc_server_incoming_port_in_use (3 passed)
cargo test -p ballista-scheduler --lib configures_nodelay_and_keepalive_on_accepted_connections (1 passed)
rustfmt on the three changed files

(AI-assisted reply.)

andygrove

Thanks for this, it's a nice tidy fix. Moving the socket setup into configure_connection reads well, and the test checks the real socket options on an accepted connection, including the disabled case.

One thing outside the diff. ballista/executor/src/standalone.rs still builds its TcpIncoming by hand, around line 188, with Some(Duration::from_secs(..)). Would you be open to switching it to tcp_incoming_from_listener(grpc_listener, &process_config.grpc_server_config) too? Then every server path handles 0 the same way.

A couple of small inline notes below. None of them block merging.

andygrove · 2026-11-10T22:35:24Z


+/// Converts a TCP keepalive interval in seconds to the idle time passed to
+/// tonic and `socket2`. A value of 0 disables keepalive.
+pub fn tcp_keepalive(seconds: u64) -> Option<Duration> {


Treating 0 as disabled makes sense here. The client builders in this file (create_grpc_client_connection and create_grpc_client_endpoint) still pass Some(Duration::from_secs(config.tcp_keepalive_seconds)), though. So ballista.client.grpc_tcp_keepalive_seconds = 0 still behaves the old way. Would it make sense to use tcp_keepalive there as well? A follow-up issue is fine too if you'd rather keep this PR small.

andygrove · 2026-11-10T22:35:24Z

-            warn!("Failed to set TCP_NODELAY on a scheduler connection: {e}");
-        }
-    });
+    let tcp_keepalive = tcp_keepalive(config.grpc_server_tcp_keepalive_seconds);


Small nit. This local shadows the tcp_keepalive function it calls. It works fine, but something like keepalive might read a little clearer.

andygrove reviewed Nov 8, 2026

View reviewed changes

github-actions Bot added the python label Nov 9, 2026

chore(scheduler): sort socket2 dev-dependency for taplo

a8156f5

Signed-off-by: Vincent <ump45nose@users.noreply.github.com>

andygrove reviewed Nov 10, 2026

View reviewed changes

ump45nose added 2 commits November 10, 2026 23:30

Merge remote-tracking branch 'upstream/main' into fix-scheduler-grpc-…

d2eccd2

…keepalive

andygrove approved these changes Nov 10, 2026

View reviewed changes

Conversation

ump45nose commented Nov 8, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Testing

Uh oh!

andygrove 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!

ump45nose commented Nov 9, 2026

Uh oh!

ump45nose commented Nov 9, 2026

Uh oh!

andygrove 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!

ump45nose commented Nov 10, 2026

Uh oh!

andygrove 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!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants