Add session vtgate balancer #18552

mhamza15 · 2025-08-17T02:03:58Z

Description

This PR builds on #11959 to create a "sticky" session VTGate balancer. It uses a consistent hash to route connections to the same tablet consistently for its duration, with a preference to local cells. It implements the TabletBalancer interface (with a few changes).

The balancer works by maintaining two hash rings for each target, one for local tablets and one for external tablets. It subscribes to health check events to keep the hash rings updated as tablets go in and out of serving, rather than constructing the hash on-demand.

A new PickOpts is added to the Pick method in the TabletBalancer interface. It currently only contains the current session hash (rather than the UUID so we hash only once on session creation and not on each call to Pick). I thought that as more balancer implementations get added, they might have their own custom parameters, so having an isolated struct so that the signature doesn't get overly long might be desirable.

On call to Pick, the balancer will look up the session hash in PickOpts and use it to find the tablet to route to, first looking at the local tablets. If there are no local tablets, it will look for an external tablet. For a given session UUID/hash, the balancer will route it to the same tablet.

Related Issue(s)

#11971

Checklist

"Backport to:" labels have been added if this change should be back-ported to release branches
If this change is to be back-ported to previous releases, a justification is included in the PR description
Tests were added or are not required
Did the new or modified tests pass consistently locally and on CI?
Documentation was added or is not required

Deployment Notes

vitess-bot · 2025-08-17T02:04:01Z

mhamza15

Some open questions:

mhamza15 · 2025-08-17T02:07:12Z

go/vt/vtgate/balancer/session.go

+// watchHealthCheck watches the health check channel for tablet health changes, and updates hash rings accordingly.
+func (b *SessionBalancer) watchHealthCheck(ctx context.Context, hcChan chan *discovery.TabletHealth) {
+	for {
+		select {
+		case <-ctx.Done():
+			b.hc.Unsubscribe(hcChan)
+			return
+		case tablet := <-hcChan:
+			if tablet == nil {
+				return
+			}
+
+			b.onTabletHealthChange(tablet)
+		}
+	}
+}


I saw two patterns for this type of logic, one using keyspace events and one using this health check. Which one would be recommended in this case?

What is the behavior on a fresh start? How can we "seed" the initial state of the tablets?

I think after subscribing to the health check, you should be able to call GetHealthyTabletStats (before reading from the health check stream returned by Subscribe) and build the initial hashring.

After that, you should be able to consume events from the health check channel.

mhamza15 · 2025-08-17T02:08:45Z

go/vt/vtgate/balancer/session.go

+// NOTE: this currently won't consider any invalid tablets. This means we'll keep returning the same
+// invalid tablet on subsequent tries. We can improve this by maybe returning a random tablet (local
+// cell preferred) when the session hash falls on an invalid tablet.


As the note says, we don't currently have knowledge of what tablets have already been tried so far. We can update the tablet gateway to either pass a list of invalid tablets, or have it double check that the balancer returned a valid tablet. If not, then it'll use a random tablet.

If we pass in a list of invalid tablets, we can instead look for the first tablet (virtual node) that has a hash >= the session hash and isn't invalid.

Yeah, I think passing in a list of tablets to ignore / skip would be the way to go here.

codecov · 2025-08-17T14:14:44Z

Codecov Report

❌ Patch coverage is 86.40227% with 48 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.54%. Comparing base (8d90d26) to head (37b7e62).
⚠️ Report is 22 commits behind head on main.

Files with missing lines	Patch %	Lines
go/vt/vttablet/queryservice/wrapped.go	73.33%	16 Missing ⚠️
go/vt/vtgate/balancer/session.go	87.28%	15 Missing ⚠️
go/vt/vtgate/tabletgateway.go	77.96%	13 Missing ⚠️
go/vt/vtgate/vtgate.go	0.00%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #18552      +/-   ##
==========================================
+ Coverage   67.49%   67.54%   +0.05%     
==========================================
  Files        1607     1609       +2     
  Lines      263104   263608     +504     
==========================================
+ Hits       177569   178042     +473     
- Misses      85535    85566      +31

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

arthurschreiber · 2025-08-18T14:54:41Z

go/vt/proto/vtgate/vtgate.pb.go

-	sizeCache          protoimpl.SizeCache
+	// SessionHash is the xxhash of the Session UUID. Used to route sessions to the same
+	// tablet.
+	SessionHash   uint64 `protobuf:"varint,29,opt,name=SessionHash,proto3" json:"SessionHash,omitempty"`


Will this break the vtgate grpc APIs? What happens if someone does a call with a Session but without the SessionHash set?

Isn't adding new protobuf fields backwards compatible?

Adding new fields is safe.

If you add new fields, any messages serialized by code using your “old” message format can still be parsed by your new generated code. You should keep in mind the default values for these elements so that new code can properly interact with messages generated by old code. Similarly, messages created by your new code can be parsed by your old code: old binaries simply ignore the new field when parsing. See the Unknown Fields section for details.

From https://protobuf.dev/programming-guides/proto3/#wire-safe-changes

Old code should ignore it, and new code handles the missing/default value

Although I just realized that setting it as uint64 means that we can't differentiate between a hash of 0 or the default value, so will need to update that

Updated to *uint64 in fd5c146

arthurschreiber · 2025-08-18T15:03:20Z

go/vt/vtgate/balancer/session.go

+func getOrCreateRing(rings map[discovery.KeyspaceShardTabletType]*hashRing, tablet *discovery.TabletHealth) *hashRing {
+	key := discovery.KeyFromTarget(tablet.Target)
+
+	ring, exists := rings[key]
+	if !exists {
+		ring = newHashRing()
+		rings[key] = ring
+	}
+
+	return ring
+}


Does this mean there will be a new hashring created for every session? Shouldn't we have one hashring per shard known per vtgate?

There should be one hash ring per target (keyspace + shard + tablet type) and shared across all sessions (assuming the tablet gateway that creates the balancer is also shared across all sessions, which I'm pretty sure it is but I didn't actually confirm that).

Ah, right. Thanks for clarifying! ❤️

We could make an optimization here to only consider the keyspaces set by --balancer-keyspaces, so we're not maintaining hash rings for keyspaces that aren't sent through the balancer.

arthurschreiber · 2025-08-21T15:55:29Z

One thing that would be good to add is allowing clients to switch the load balancing mode through a client side flag (in addition to managing the default load balancing behavior through the vtgate cli flag).

proto/vtgate.proto

harshit-gangal · 2025-08-25T16:10:29Z

go/vt/vtgate/balancer/balancer.go

-	Pick(target *querypb.Target, tablets []*discovery.TabletHealth) *discovery.TabletHealth
+	Pick(target *querypb.Target, tablets []*discovery.TabletHealth, invalidTablets map[string]bool, opts *PickOpts) *discovery.TabletHealth



Let's have PickOpts contain all the options needed, we can move invalidTablets in that struct.

We can pass it by value as we do not modify it.

go/vt/vtgate/balancer/session.go

harshit-gangal · 2025-08-25T16:37:55Z

go/vt/vtgate/balancer/session.go

+// tabletTypesToWatch are the tablet types that will be included in the hash rings.
+var tabletTypesToWatch = []topodata.TabletType{topodata.TabletType_PRIMARY, topodata.TabletType_REPLICA, topodata.TabletType_BATCH}
+


Do you need Primary for balancing?

Still most places we refer BATCH as RDONLY, should we use that here?

Removed primaries and left it as RDONLY in 4682bf9

harshit-gangal · 2025-08-25T16:40:28Z

go/vt/vtgate/balancer/session.go

+func (b *SessionBalancer) watchHealthCheck(ctx context.Context, topoServer srvtopo.Server) {
+	// Build initial hash rings
+
+	// Find all the targets we're watching
+	targets, _, err := srvtopo.FindAllTargetsAndKeyspaces(ctx, topoServer, b.localCell, discovery.KeyspacesToWatch, tabletTypesToWatch)
+	if err != nil {
+		log.Errorf("session balancer: failed to find all targets and keyspaces: %q", err)
+		return
+	}
+
+	// Add each tablet to the hash ring
+	for _, target := range targets {
+		tablets := b.hc.GetHealthyTabletStats(target)
+		for _, tablet := range tablets {
+			b.onTabletHealthChange(tablet)
+		}
+	}


looks like this will miss adding new targets to the balancer as it is initialized only once

This is just for the initial hash ring setup, and then any new targets returned by the health check will have a new hash ring created: https://github.com/vitessio/vitess/pull/18552/files#diff-a70b24fd86763104c3092c7da62fa559dce14f6bee289a8c883e26d062b0aafcR172-R181

arthurschreiber · 2025-08-26T08:10:07Z

You changed the order so that first the health check is set up, and then the latest view of tablets is used to setup the hash rings. Is that safe from race conditions? 🤔

mhamza15 · 2025-08-26T13:33:51Z

You changed the order so that first the health check is set up, and then the latest view of tablets is used to setup the hash rings. Is that safe from race conditions? 🤔

Yeah I was worried about the same thing. The reasoning behind why I switched it is that if we build the initial hash rings and then set up the health check after, there might be a little time in between where we lose updates. But this way there might be a case where the health check sees a tablet go out of serving, removes it from the hash ring, then the initial hash ring build overwrites it with a stale serving tablet.

I'll keep thinking about this one.

mhamza15 · 2025-08-26T16:51:59Z

You changed the order so that first the health check is set up, and then the latest view of tablets is used to setup the hash rings. Is that safe from race conditions? 🤔

Changed in 977d2fa:

Set up subscription, but don't start goroutine yet. State changes will be buffered in the meantime
Fetch initial state
Start goroutine and catch up on any state changes that have happened

Add `PickOpts` which allow balancers to accept options specific to their implementation. This allows the `Pick` signature not to get overly long as implementations and their options are added. Signed-off-by: Mohamed Hamza <[email protected]>

Signed-off-by: Mohamed Hamza <[email protected]>

This reverts commit bd6a1b7. Signed-off-by: Mohamed Hamza <[email protected]>

Signed-off-by: Mohamed Hamza <[email protected]>

Update the signature of the wrapper func to accept a new `WrapOpts` struct, which currently contains `ExecuteOptions`, which now contains the session UUID so that it can be passed into `Pick` for the balancer. Signed-off-by: Mohamed Hamza <[email protected]>

Signed-off-by: Mohamed Hamza <[email protected]>

arthurschreiber · 2025-08-28T14:59:58Z

go/vt/vtexplain/vtexplain_vttablet.go

@@ -127,7 +127,7 @@ func (vte *VTExplain) newTablet(ctx context.Context, env *vtenv.Environment, opt

 	tablet.QueryService = queryservice.Wrap(
 		nil,
-		func(ctx context.Context, target *querypb.Target, conn queryservice.QueryService, name string, inTransaction bool, inner func(context.Context, *querypb.Target, queryservice.QueryService) (bool, error)) error {
+		func(ctx context.Context, target *querypb.Target, conn queryservice.QueryService, name string, opts *queryservice.WrapOpts, inner func(context.Context, *querypb.Target, queryservice.QueryService) (bool, error)) error {


Does WrapOpts get modified down the stack? If not, then this should probably not be passed as a reference but by value?

Changed in 2620b76

I want to note I also changed ExecuteOptions inside of WrapOpts to a pointer, as go vet would complain that we are copying a mutex (used inside protobuf's internal structs): 0db53cd

arthurschreiber · 2025-08-28T15:02:02Z

go/vt/vtgate/tabletgateway.go

+			})
+		}
+
+		opts := &balancer.PickOpts{SessionUUID: opts.Options.SessionUUID, InvalidTablets: invalidTablets}


Similar here - I think this might be better passed by value?

Changed in e62c88f

Signed-off-by: Mohamed Hamza <[email protected]>

mhamza15 · 2025-08-29T21:28:39Z

I'm currently working on writing an e2e test.

Signed-off-by: Mohamed Hamza <[email protected]>

mhamza15 changed the title ~~Add consistent hash vtgate balancer~~ [WIP] Add consistent hash vtgate balancer Aug 17, 2025

github-actions bot added this to the v23.0.0 milestone Aug 17, 2025

mhamza15 commented Aug 17, 2025

View reviewed changes

arthurschreiber reviewed Aug 18, 2025

View reviewed changes

arthurschreiber removed NeedsDescriptionUpdate The description is not clear or comprehensive enough, and needs work NeedsIssue A linked issue is missing for this Pull Request labels Aug 21, 2025

harshit-gangal reviewed Aug 25, 2025

View reviewed changes

proto/vtgate.proto Outdated Show resolved Hide resolved

harshit-gangal reviewed Aug 25, 2025

View reviewed changes

go/vt/vtgate/balancer/session.go Outdated Show resolved Hide resolved

harshit-gangal reviewed Aug 25, 2025

View reviewed changes

mhamza15 added 8 commits August 27, 2025 21:16

Initial implementation of the hash ring

f37cdc8

Signed-off-by: Mohamed Hamza <[email protected]>

Set up initial health check subscription

114fece

Signed-off-by: Mohamed Hamza <[email protected]>

Rename to SessionBalancer and create initial tests

4742e8d

Signed-off-by: Mohamed Hamza <[email protected]>

Improvements and add more tests

5be7ac7

Signed-off-by: Mohamed Hamza <[email protected]>

Clarify comment

80eb92f

Signed-off-by: Mohamed Hamza <[email protected]>

Remove brackets

c421c3d

Signed-off-by: Mohamed Hamza <[email protected]>

Clarify comment

34ca9e9

Signed-off-by: Mohamed Hamza <[email protected]>

mhamza15 added 6 commits August 27, 2025 21:16

Remove unused import

f9956e2

Signed-off-by: Mohamed Hamza <[email protected]>

Remove old tablets when a tablet's target changes

d60c593

Signed-off-by: Mohamed Hamza <[email protected]>

Fetch initial tablet state after subscribing to health check

5833a54

Signed-off-by: Mohamed Hamza <[email protected]>

Revert "Remove old tablets when a tablet's target changes"

5bf666d

This reverts commit bd6a1b7. Signed-off-by: Mohamed Hamza <[email protected]>

Remove primary tablets from rings

d8723b3

Signed-off-by: Mohamed Hamza <[email protected]>

undo formatting

479cde4

Signed-off-by: Mohamed Hamza <[email protected]>

mhamza15 force-pushed the mhamza/consistent-hash-balancer branch from 7b7f78e to 479cde4 Compare August 28, 2025 01:16

mhamza15 added 3 commits August 27, 2025 21:28

Fix new types

98a7c8b

Signed-off-by: Mohamed Hamza <[email protected]>

Fix tests

f59a14f

Signed-off-by: Mohamed Hamza <[email protected]>

arthurschreiber reviewed Aug 28, 2025

View reviewed changes

mhamza15 added 4 commits August 28, 2025 19:13

Pass WrapOpts by value

2620b76

Signed-off-by: Mohamed Hamza <[email protected]>

Pass PickOpts by value

e62c88f

Signed-off-by: Mohamed Hamza <[email protected]>

Change ExecuteOptions in WrapOpts as a pointer

0db53cd

Signed-off-by: Mohamed Hamza <[email protected]>

Fix balancer-type help text

cde392c

Signed-off-by: Mohamed Hamza <[email protected]>

mhamza15 marked this pull request as ready for review August 29, 2025 14:32

mhamza15 requested review from mattlord, notfelineit, systay, rohit-nayak-ps, shlomi-noach and frouioui as code owners August 29, 2025 14:32

mhamza15 changed the title ~~[WIP] Add consistent hash vtgate balancer~~ Add consistent hash vtgate balancer Aug 29, 2025

mhamza15 changed the title ~~Add consistent hash vtgate balancer~~ Add session vtgate balancer Aug 29, 2025

mhamza15 added 4 commits August 30, 2025 10:38

Fix some bugs

eadf1ae

Signed-off-by: Mohamed Hamza <[email protected]>

Get cell from tablet alias rather than target

8fae28c

Signed-off-by: Mohamed Hamza <[email protected]>

Initial e2e test

cd30d9d

Signed-off-by: Mohamed Hamza <[email protected]>

Add more e2e tests

37b7e62

Signed-off-by: Mohamed Hamza <[email protected]>

		Pick(target querypb.Target, tablets []discovery.TabletHealth) *discovery.TabletHealth
		Pick(target querypb.Target, tablets []discovery.TabletHealth, invalidTablets map[string]bool, opts PickOpts) discovery.TabletHealth

		// tabletTypesToWatch are the tablet types that will be included in the hash rings.
		var tabletTypesToWatch = []topodata.TabletType{topodata.TabletType_PRIMARY, topodata.TabletType_REPLICA, topodata.TabletType_BATCH}

Add session vtgate balancer #18552

Are you sure you want to change the base?

Add session vtgate balancer #18552

Conversation

mhamza15 commented Aug 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue(s)

Checklist

Deployment Notes

Uh oh!

vitess-bot bot commented Aug 17, 2025

Review Checklist

General

Tests

Documentation

New flags

If a workflow is added or modified:

Backward compatibility

Uh oh!

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

codecov bot commented Aug 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

arthurschreiber Aug 18, 2025 • 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!

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!

arthurschreiber commented Aug 21, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

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!

arthurschreiber commented Aug 26, 2025

Uh oh!

mhamza15 commented Aug 26, 2025

Uh oh!

mhamza15 commented Aug 26, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

mhamza15 commented Aug 17, 2025 •

edited

Loading

codecov bot commented Aug 17, 2025 •

edited

Loading

arthurschreiber Aug 18, 2025 •

edited

Loading