feat(grpc): Implement PickFirst load balancer by nathanielford · Pull Request #2570 · hyperium/tonic

nathanielford · 2026-03-25T15:30:03Z

Motivation

Full implementation of the pick first load balancer, including 'Happy eyeballs' features.

Solution

Load balancing implementation to pick the first available endpoint to connect to, maintaining stickiness across endpoint updates if configured. Handles accepting new LB configuration and subchannel reconstruction.

Prototype is at https://github.com/nathanielford/grpc-rust-testbed/tree/main/pick_first_lib

Notes

Ended up including all happy eyeball features because it wasn't clear where best to slice the line. Considering this a full implementation, and it should be reviewed as such.
This does use tokio::spawn and tokio::time, which may need to be replaced to make things runtime agnostic. Please comment in the PR whether this is the case.

arjan-bal

Leaving initial comments while I review the remaining changes.

…t and stickiness

…ss, and updated sync testing framework

…A61 endpoint handling This commit implements the PickFirst load balancer policy for Tonic gRPC, focusing on: - Efficient subchannel management with backoff preservation. - "Stickiness" support: continuing to use an existing Ready subchannel if it remains in resolver updates. - Compliance with gRFC A61: endpoints are now shuffled before being flattened into an address list, ensuring multiple addresses for a single endpoint (e.g., IPv4/IPv6) stay together. - Clean state reset: subchannels and selected state are now cleared when receiving an empty address list. - Alignment with the updated synchronous testing framework in master. Includes comprehensive test coverage for basic connection, failover, stickiness, exhaustion, deterministic endpoint shuffling, de-duplication, and empty updates.

…active failover This change enhances the PickFirst load balancing policy to better support gRFC A61 (Happy Eyeballs) and improve connection establishment latency. Key changes: - Implement IPv6/IPv4 address interleaving in `compile_address` to ensure subsequent connection attempts alternate between protocol families. - Introduce a `subchannel_states` cache in `PickFirstPolicy` to track the connectivity status of managed subchannels. - Refactor connection logic to use a `frontier_index` and proactively skip subchannels known to be in `TransientFailure` (e.g., during backoff). - Update `advance_frontier` to safely maintain the index within the bounds of the address list, ensuring the policy remains reactive to recovery. - Add deterministic unit tests for shuffling and interleaving logic.

…balls

…ollision unit test

…stLB

…nually written hash impl

…into refactor/AddressHashing

…stLB # Conflicts: # grpc/Cargo.toml

…nformance tests

nathanielford · 2026-05-07T19:47:05Z

Note that this now includes #2631, which will merge first so as to keep a clearer record. However, I needed the fix to get the tests ported from Go to work.

…stLB

arjan-bal · 2026-05-21T11:28:55Z

+
+        // Should NOT have any more events (no Connect, no UpdatePicker),
+        // because it stuck to the original selected subchannel.
+        std::thread::sleep(Duration::from_millis(50));


We should use tokio::time::timeout or tokio::time::sleep to avoid blocking the entire thread. Same comment for other tests.

Good catch. Tokio only runs tests with a single thread, so there's no way for the code under test to fail if the thread is blocked. (https://docs.rs/tokio/latest/tokio/attr.test.html#current-thread-runtime)

I've replaced everything to use helpers to pull from the channel.

This might be less perfect for test cases that know a call into the policy should trigger an event inline, but I think it's OK. Let me know what you think.

This might be less perfect for test cases that know a call into the policy should trigger an event inline.

Looks good to me. I didn't notice that the mpsc channel was not an async channel, so I didn't consider the result of calling recv.

arjan-bal · 2026-05-21T11:45:02Z

+        let mut found = None;
+        for _ in 0..10 {
+            match rx.try_recv() {
+                Ok(event) => {
+                    found = Some(event);
+                    break;
+                }
+                Err(std::sync::mpsc::TryRecvError::Empty) => {
+                    // Yield to runtime to allow timer task to run.
+                    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+                }
+                Err(e) => panic!("error recv: {:?}", e),
+            }
+        }


Instead of looping, we can use a tokio::time::timout of 10s to wrap the rx.recv. We should define 10s as a global const named DEFAULT_TEST_DURATION and re-use it for other tests. For events not expected to happen, we should define a constant DEFAULT_TEST_SHORT_DURATION and set it to 10 millis.

Done. But note that we can't use recv, as that will block the thread. PTAL.

There's an unfortunate difference between go and rust, which is that we don't get the concept of a context with a deadline in it from the start. So if we use the 10s timeout for every operation, it doesn't apply to the whole test. We could follow this convention, but I don't know if it's worth it. For now I'm just making every thing that could timeout use 10s. I don't think we'll have an issue where things will run really slowly and take many tens of seconds and then timeout, so I think this is fine.

But note that we can't use recv, as that will block the thread. PTAL.

I didn't notice that the mpsc channel was not an async channel. I suspect this is because the TestChannelController is using a std::sync::mpsc instead of tokio mpsc channel, making the receiver blocking. In the DNS resolver tests, I've used an unbounded tokio mspc channel for which sends are sync, but receives can be awaited:

tonic/grpc/src/client/name_resolution/test_utils.rs

Lines 62 to 69 in ae8ab11

pub(crate) fn new_pair() -> (Self, mpsc::UnboundedReceiver<ResolverUpdate>) {

let (update_tx, update_rx) = mpsc::unbounded_channel();

let cc = Self {

update_result: Ok(()),

update_tx,

};

(cc, update_rx)

}

Maybe a similar pattern can avoid the need for polling in the pickfirst tests? This can be attempted in a later PR.

I believe the test event channel should just be synchronous since all accesses are done synchronously (from non-async lb policy methods). But the work scheduler should probably be an async channel. We can do that in a follow-up?

arjan-bal · 2026-05-21T12:11:42Z

+
+    // If the timer expires during a connection pass, the LB should advance to
+    // the next subchannel and trigger a connection attempt.
+    #[tokio::test(start_paused = true)]


Is there a reason to use start_paused for this test specifically?

I'm not sure why it would have been like that. Maybe something about the sleep stuff below? I think this is fixed now.

We can revert the Cargo.toml changes now.

arjan-bal

Non-test changes look good. Left some minor comments on the test code.

arjan-bal

LGTM

arjan-bal · 2026-05-21T17:17:34Z

+
+    // If the timer expires during a connection pass, the LB should advance to
+    // the next subchannel and trigger a connection attempt.
+    #[tokio::test(start_paused = true)]


We can revert the Cargo.toml changes now.

arjan-bal · 2026-05-21T17:32:12Z

+
+        // Should NOT have any more events (no Connect, no UpdatePicker),
+        // because it stuck to the original selected subchannel.
+        std::thread::sleep(Duration::from_millis(50));


This might be less perfect for test cases that know a call into the policy should trigger an event inline.

Looks good to me. I didn't notice that the mpsc channel was not an async channel, so I didn't consider the result of calling recv.

arjan-bal · 2026-05-21T17:34:38Z

+        let mut found = None;
+        for _ in 0..10 {
+            match rx.try_recv() {
+                Ok(event) => {
+                    found = Some(event);
+                    break;
+                }
+                Err(std::sync::mpsc::TryRecvError::Empty) => {
+                    // Yield to runtime to allow timer task to run.
+                    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+                }
+                Err(e) => panic!("error recv: {:?}", e),
+            }
+        }


But note that we can't use recv, as that will block the thread. PTAL.

I didn't notice that the mpsc channel was not an async channel. I suspect this is because the TestChannelController is using a std::sync::mpsc instead of tokio mpsc channel, making the receiver blocking. In the DNS resolver tests, I've used an unbounded tokio mspc channel for which sends are sync, but receives can be awaited:

tonic/grpc/src/client/name_resolution/test_utils.rs

Lines 62 to 69 in ae8ab11

pub(crate) fn new_pair() -> (Self, mpsc::UnboundedReceiver<ResolverUpdate>) {

let (update_tx, update_rx) = mpsc::unbounded_channel();

let cc = Self {

update_result: Ok(()),

update_tx,

};

(cc, update_rx)

}

Maybe a similar pattern can avoid the need for polling in the pickfirst tests? This can be attempted in a later PR.

nathanielford requested a review from dfawley March 25, 2026 15:30

nathanielford self-assigned this Mar 25, 2026

nathanielford requested review from arjan-bal and removed request for dfawley March 25, 2026 16:17

nathanielford force-pushed the implement/PickFirstLB branch 3 times, most recently from a442272 to 73397ff Compare March 25, 2026 17:29

nathanielford assigned arjan-bal and unassigned nathanielford and arjan-bal Mar 25, 2026

dfawley assigned arjan-bal Mar 27, 2026

arjan-bal reviewed Mar 30, 2026

View reviewed changes

arjan-bal assigned nathanielford Mar 30, 2026

arjan-bal reviewed Mar 30, 2026

View reviewed changes

arjan-bal removed their assignment Mar 30, 2026

nathanielford force-pushed the implement/PickFirstLB branch from 73397ff to bdf7fc3 Compare April 6, 2026 19:32

nathanielford force-pushed the implement/PickFirstLB branch from 7ea10ad to e1bcaf4 Compare April 27, 2026 20:55

nathanielford requested a review from arjan-bal April 27, 2026 21:38

nathanielford marked this pull request as ready for review April 27, 2026 21:39

arjan-bal assigned arjan-bal and unassigned nathanielford Apr 28, 2026

arjan-bal added the A-grpc-next label Apr 28, 2026

nathanielford added 6 commits April 29, 2026 15:17

Implement PickFirst load balancer with optimized subchannel managemen…

9a3be62

…t and stickiness

WIP: implementation of PickFirstLB

aa39042

Implement PickFirst LB with optimized subchannel management, stickine…

db800a7

…ss, and updated sync testing framework

grpc: implement Racing Selection and Timer Cancellation for Happy Eye…

98b1b41

…balls

nathanielford added 9 commits May 6, 2026 23:56

name_resolution: derive structural equality for Address and add map c…

acd6b2e

…ollision unit test

Merge branch 'master' into refactor/AddressHashing

4d47842

Merge remote-tracking branch 'hyperium/master' into implement/PickFir…

3a32745

…stLB

Update test to remove coverage of derived traits, add coverage for ma…

cc467f5

…nually written hash impl

Merge remote-tracking branch 'nathanielford/refactor/AddressHashing' …

0fa6603

…into refactor/AddressHashing

Merge branch 'master' into refactor/AddressHashing

3f3800f

Merge remote-tracking branch 'hyperium/master' into implement/PickFir…

5ecf935

…stLB # Conflicts: # grpc/Cargo.toml

Merge branch 'refactor/AddressHashing' into implement/PickFirstLB

eb7c378

lb/pick_first: refactor test boilerplate, harden comments, and add co…

c869cb6

…nformance tests

arjan-bal reviewed May 11, 2026

View reviewed changes

nathanielford added 3 commits May 12, 2026 22:30

PR review fixes

c90c5ea

Add Happy Eyeballs regression tests

3e55d3d

Merge branch 'master' into implement/PickFirstLB

ccfd967

arjan-bal reviewed May 13, 2026

View reviewed changes

dfawley added 3 commits May 20, 2026 08:16

Merge remote-tracking branch 'upstream/master' into pf

2433d11

partial progress

3c90553

Merge remote-tracking branch 'upstream/master' into implement/PickFir…

8ed0544

…stLB

dfawley assigned arjan-bal and unassigned nathanielford May 20, 2026

ordering::relaxed is fine for this

9740828

arjan-bal reviewed May 21, 2026

View reviewed changes

arjan-bal assigned dfawley and unassigned arjan-bal May 21, 2026

arjan-bal changed the title ~~Implement PickFirst load balancer~~ feat(grpc): Implement PickFirst load balancer May 21, 2026

test fixes

5036188

arjan-bal approved these changes May 21, 2026

View reviewed changes

dfawley added 2 commits May 21, 2026 12:24

always use try_recv and no polling

24b8f3b

remove unused tokio functionality

a207faf

	pub(crate) fn new_pair() -> (Self, mpsc::UnboundedReceiver<ResolverUpdate>) {
	let (update_tx, update_rx) = mpsc::unbounded_channel();
	let cc = Self {
	update_result: Ok(()),
	update_tx,
	};
	(cc, update_rx)
	}

Conversation

nathanielford commented Mar 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Solution

Notes

Uh oh!

arjan-bal left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nathanielford commented May 7, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

Uh oh!

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!

Uh oh!

Uh oh!

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!

arjan-bal left a comment

Choose a reason for hiding this comment

Uh oh!

arjan-bal left a comment

Choose a reason for hiding this comment

Uh oh!

nathanielford commented Mar 25, 2026 •

edited

Loading

arjan-bal left a comment •

edited

Loading