rework tabletbalancer interface to avoid unnecessary sort

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

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
@@ -268,6 +268,22 @@ func TestBalancedShuffle(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}
@@ -293,13 +309,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,37 +350,37 @@ 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)
 
 		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 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)
 		}
 	}

go/vt/vtgate/tabletgateway.go

@@ -75,6 +75,17 @@ func init() {
 	})
 }
 
+// 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
+		}
+	}
+
+	return false
+}
+
 // TabletGateway implements the Gateway interface.
 // This implementation uses the new healthcheck module.
 type TabletGateway struct {
@@ -356,35 +367,38 @@ 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
+		var th *discovery.TabletHealth
+
+		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))
+				for _, t := range tablets {
+					if _, ok := invalidTablets[topoproto.TabletAliasString(t.Tablet.Alias)]; !ok {
+						validTablets = append(validTablets, th)
 					}
 				}
-			} else {
-				useBalancer = true
+				tablets = validTablets
 			}
-		}
 
-		if useBalancer {
-			gw.balancer.ShuffleTablets(target, tablets)
+			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 {