From f5522fca287b576c79dfd9db65b40a2c2b1ea409 Mon Sep 17 00:00:00 2001 From: Michael Demmer Date: Mon, 8 Jul 2024 12:01:29 -0700 Subject: [PATCH 1/4] rework tabletbalancer interface to avoid unnecessary sort There's actually no need in the tabletgateway to ever sort the full list of potential tablet candidates, because each query attempt only goes to a single tablet, and then if that fails, the withRetry loop records which tablet failed and then re-enters the loop to try again. Therefore, refactor the balancer to replace ShuffleTablets with a simple Pick interface that just returns the best tablet to route to. --- go/vt/vtgate/balancer/balancer.go | 30 ++++++++++++------------- go/vt/vtgate/balancer/balancer_test.go | 18 +++++++-------- go/vt/vtgate/tabletgateway.go | 31 ++++++++++++++++++-------- 3 files changed, 45 insertions(+), 34 deletions(-) diff --git a/go/vt/vtgate/balancer/balancer.go b/go/vt/vtgate/balancer/balancer.go index 462ccfda901..30cf862f890 100644 --- a/go/vt/vtgate/balancer/balancer.go +++ b/go/vt/vtgate/balancer/balancer.go @@ -88,8 +88,9 @@ converge on the desired balanced query load. */ type TabletBalancer interface { - // Randomly shuffle the tablets into an order for routing queries - ShuffleTablets(target *querypb.Target, tablets []*discovery.TabletHealth) + // Pick is the main entry point to the balancer. Returns the best tablet out of the list + // for a given query to maintain the desired balanced allocation over multiple executions. + Pick(target *querypb.Target, tablets []*discovery.TabletHealth) *discovery.TabletHealth // Balancer debug page request DebugHandler(w http.ResponseWriter, r *http.Request) @@ -161,33 +162,30 @@ func (b *tabletBalancer) DebugHandler(w http.ResponseWriter, _ *http.Request) { fmt.Fprintf(w, "Allocations: %v\r\n", string(allocations)) } -// ShuffleTablets is the main entry point to the balancer. +// Pick is the main entry point to the balancer. // -// It shuffles the tablets into a preference list for routing a given query. -// However, since all tablets should be healthy, the query will almost always go -// to the first tablet in the list, so the balancer ranking algoritm randomly -// shuffles the list to break ties, then chooses a weighted random selection -// based on the balance alorithm to promote to the first in the set. -func (b *tabletBalancer) ShuffleTablets(target *querypb.Target, tablets []*discovery.TabletHealth) { +// Given the total allocation for the set of tablets, choose the best target +// by a weighted random sample so that over time the system will achieve the +// desired balanced allocation. +func (b *tabletBalancer) Pick(target *querypb.Target, tablets []*discovery.TabletHealth) *discovery.TabletHealth { numTablets := len(tablets) + if numTablets == 0 { + return nil + } allocationMap, totalAllocation := b.getAllocation(target, tablets) - rand.Shuffle(numTablets, func(i, j int) { tablets[i], tablets[j] = tablets[j], tablets[i] }) - - // Do another O(n) seek through the list to effect the weighted sample by picking - // a random point in the allocation space and seeking forward in the list of (randomized) - // tablets to that point, promoting the tablet to the head of the list. r := rand.Intn(totalAllocation) for i := 0; i < numTablets; i++ { flow := allocationMap[tablets[i].Tablet.Alias.Uid] if r < flow { - tablets[0], tablets[i] = tablets[i], tablets[0] - break + return tablets[i] } r -= flow } + + return tablets[0] } // To stick with integer arithmetic, use 1,000,000 as the full load diff --git a/go/vt/vtgate/balancer/balancer_test.go b/go/vt/vtgate/balancer/balancer_test.go index 1eb9e69fadf..b8a1095465e 100644 --- a/go/vt/vtgate/balancer/balancer_test.go +++ b/go/vt/vtgate/balancer/balancer_test.go @@ -223,7 +223,7 @@ func TestAllocateFlows(t *testing.T) { } } -func TestBalancedShuffle(t *testing.T) { +func TestBalancedPick(t *testing.T) { cases := []struct { test string tablets []*discovery.TabletHealth @@ -293,13 +293,13 @@ func TestBalancedShuffle(t *testing.T) { b := NewTabletBalancer(localCell, vtGateCells).(*tabletBalancer) for i := 0; i < N/len(vtGateCells); i++ { - b.ShuffleTablets(target, tablets) + th := b.Pick(target, tablets) if i == 0 { t.Logf("Target Flows %v, Balancer: %s\n", expectedPerCell, b.print()) t.Logf(b.print()) } - routed[tablets[0].Tablet.Alias.Uid]++ + routed[th.Tablet.Alias.Uid]++ } } @@ -334,25 +334,25 @@ func TestTopologyChanged(t *testing.T) { tablets = tablets[0:2] for i := 0; i < N; i++ { - b.ShuffleTablets(target, tablets) + th := b.Pick(target, tablets) allocation, totalAllocation := b.getAllocation(target, tablets) if totalAllocation != ALLOCATION/2 { t.Errorf("totalAllocation mismatch %s", b.print()) } - if allocation[allTablets[0].Tablet.Alias.Uid] != ALLOCATION/4 { + if allocation[th.Tablet.Alias.Uid] != ALLOCATION/4 { t.Errorf("allocation mismatch %s, cell %s", b.print(), allTablets[0].Tablet.Alias.Cell) } - if tablets[0].Tablet.Alias.Cell != "a" { + if th.Tablet.Alias.Cell != "a" { t.Errorf("shuffle promoted wrong tablet from cell %s", tablets[0].Tablet.Alias.Cell) } } // Run again with the full topology. Now traffic should go to cell b for i := 0; i < N; i++ { - b.ShuffleTablets(target, allTablets) + th := b.Pick(target, allTablets) allocation, totalAllocation := b.getAllocation(target, allTablets) @@ -360,11 +360,11 @@ func TestTopologyChanged(t *testing.T) { t.Errorf("totalAllocation mismatch %s", b.print()) } - if allocation[allTablets[0].Tablet.Alias.Uid] != ALLOCATION/4 { + if allocation[th.Tablet.Alias.Uid] != ALLOCATION/4 { t.Errorf("allocation mismatch %s, cell %s", b.print(), allTablets[0].Tablet.Alias.Cell) } - if allTablets[0].Tablet.Alias.Cell != "b" { + if th.Tablet.Alias.Cell != "b" { t.Errorf("shuffle promoted wrong tablet from cell %s", allTablets[0].Tablet.Alias.Cell) } } diff --git a/go/vt/vtgate/tabletgateway.go b/go/vt/vtgate/tabletgateway.go index eff7cdd02f4..558d1e852e0 100644 --- a/go/vt/vtgate/tabletgateway.go +++ b/go/vt/vtgate/tabletgateway.go @@ -371,20 +371,33 @@ func (gw *TabletGateway) withRetry(ctx context.Context, target *querypb.Target, } } + var th *discovery.TabletHealth if useBalancer { - gw.balancer.ShuffleTablets(target, tablets) + // filter out the tablets that we've tried before (if any), then pick the best one + if len(invalidTablets) > 0 { + validTablets := make([]*discovery.TabletHealth, len(tablets)) + for _, t := range tablets { + if _, ok := invalidTablets[topoproto.TabletAliasString(t.Tablet.Alias)]; !ok { + validTablets = append(validTablets, th) + } + } + tablets = validTablets + } + + th = gw.balancer.Pick(target, tablets) + } else { + // shuffle sort the set of tablets with AZ-affinity, then pick the first one in the + // result set that we haven't tried before gw.shuffleTablets(gw.localCell, tablets) - } - - var th *discovery.TabletHealth - // skip tablets we tried before - for _, t := range tablets { - if _, ok := invalidTablets[topoproto.TabletAliasString(t.Tablet.Alias)]; !ok { - th = t - break + for _, t := range tablets { + if _, ok := invalidTablets[topoproto.TabletAliasString(t.Tablet.Alias)]; !ok { + th = t + break + } } } + if th == nil { // do not override error from last attempt. if err == nil { From 308d7c8ce56a2bbb2e33eb1299ad9966be2799fa Mon Sep 17 00:00:00 2001 From: Michael Demmer Date: Mon, 8 Jul 2024 12:14:57 -0700 Subject: [PATCH 2/4] add a couple tests --- go/vt/vtgate/balancer/balancer_test.go | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/go/vt/vtgate/balancer/balancer_test.go b/go/vt/vtgate/balancer/balancer_test.go index b8a1095465e..0318fda1391 100644 --- a/go/vt/vtgate/balancer/balancer_test.go +++ b/go/vt/vtgate/balancer/balancer_test.go @@ -268,6 +268,22 @@ func TestBalancedPick(t *testing.T) { []string{"a", "b", "c", "d"}, }, + { + "one target same cell", + []*discovery.TabletHealth{ + createTestTablet("a"), + }, + + []string{"a"}, + }, + { + "one target other cell", + []*discovery.TabletHealth{ + createTestTablet("a"), + }, + + []string{"b", "c", "d"}, + }, } target := &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA} From 157b40fa4ffe0c19f06e9d1b663bc25f385a5e63 Mon Sep 17 00:00:00 2001 From: Michael Demmer Date: Mon, 8 Jul 2024 13:03:49 -0700 Subject: [PATCH 3/4] simplify the useBalancer logic --- go/vt/vtgate/tabletgateway.go | 33 +++++++++++++++++---------------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/go/vt/vtgate/tabletgateway.go b/go/vt/vtgate/tabletgateway.go index 558d1e852e0..0f1ac0dbfe4 100644 --- a/go/vt/vtgate/tabletgateway.go +++ b/go/vt/vtgate/tabletgateway.go @@ -75,6 +75,22 @@ func init() { }) } +func useBalancer(keyspace string) bool { + if balancerEnabled { + if len(balancerKeyspaces) == 0 { + return true + } + + for _, k := range balancerKeyspaces { + if keyspace == k { + return true + } + } + } + + return false +} + // TabletGateway implements the Gateway interface. // This implementation uses the new healthcheck module. type TabletGateway struct { @@ -356,23 +372,8 @@ func (gw *TabletGateway) withRetry(ctx context.Context, target *querypb.Target, break } - // Determine whether or not to use the balancer or the standard affinity-based shuffle - useBalancer := false - if balancerEnabled { - if len(balancerKeyspaces) != 0 { - for _, keyspace := range balancerKeyspaces { - if keyspace == target.Keyspace { - useBalancer = true - break - } - } - } else { - useBalancer = true - } - } - var th *discovery.TabletHealth - if useBalancer { + if useBalancer(target.Keyspace) { // filter out the tablets that we've tried before (if any), then pick the best one if len(invalidTablets) > 0 { validTablets := make([]*discovery.TabletHealth, len(tablets)) From fe2520db276a40ddf05667be355efb1b0f0654ea Mon Sep 17 00:00:00 2001 From: Michael Demmer Date: Mon, 8 Jul 2024 14:08:05 -0700 Subject: [PATCH 4/4] one more refactor to align more with upstream --- go/vt/vtgate/tabletgateway.go | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/go/vt/vtgate/tabletgateway.go b/go/vt/vtgate/tabletgateway.go index 0f1ac0dbfe4..b2674e96a8d 100644 --- a/go/vt/vtgate/tabletgateway.go +++ b/go/vt/vtgate/tabletgateway.go @@ -75,17 +75,12 @@ func init() { }) } -func useBalancer(keyspace string) bool { - if balancerEnabled { - if len(balancerKeyspaces) == 0 { +// this utility can be replaced with slices.Contains in a future iteration +func isBalancerKeyspaceEnabled(keyspace string) bool { + for _, k := range balancerKeyspaces { + if keyspace == k { return true } - - for _, k := range balancerKeyspaces { - if keyspace == k { - return true - } - } } return false @@ -373,7 +368,12 @@ func (gw *TabletGateway) withRetry(ctx context.Context, target *querypb.Target, } var th *discovery.TabletHealth - if useBalancer(target.Keyspace) { + + useBalancer := balancerEnabled + if balancerEnabled && len(balancerKeyspaces) > 0 { + useBalancer = isBalancerKeyspaceEnabled(target.Keyspace) + } + if useBalancer { // filter out the tablets that we've tried before (if any), then pick the best one if len(invalidTablets) > 0 { validTablets := make([]*discovery.TabletHealth, len(tablets))