Add session algorithm to tablet balancer

changelog/24.0/24.0.0/summary.md

@@ -9,6 +9,7 @@
 - **[Minor Changes](#minor-changes)**
     - **[VTGate](#minor-changes-vtgate)**
         - [New default for `--legacy-replication-lag-algorithm` flag](#vtgate-new-default-legacy-replication-lag-algorithm)
+        - [New "session" mode for `--vtgate-balancer-mode` flag](#vtgate-session-balancer-mode)
     - **[Query Serving](#minor-changes-query-serving)**
         - [JSON_EXTRACT now supports dynamic path arguments](#query-serving-json-extract-dynamic-args)
     - **[VTTablet](#minor-changes-vttablet)**
@@ -42,6 +43,16 @@ Instead, a simpler algorithm purely based on low lag, high lag and minimum numbe
 
 In v25 this flag will become deprecated and in the following release it will be removed. In the meantime, the legacy behaviour can be used by setting `--legacy-replication-lag-algorithm=true`. This deprecation is tracked in https://github.com/vitessio/vitess/issues/18914.
 
+#### <a id="vtgate-session-balancer-mode"/>New "session" mode for `--vtgate-balancer-mode` flag</a>
+
+The VTGate flag `--vtgate-balancer-mode` now supports a new "session" mode in addition to the existing "cell", "prefer-cell", and "random" modes. Session mode routes each session consistently to the same tablet for the session's duration.
+
+To enable session mode, set the flag when starting VTGate:
+
+```
+--vtgate-balancer-mode=session
+```
+
 ### <a id="minor-changes-query-serving"/>Query Serving</a>
 
 #### <a id="query-serving-json-extract-dynamic-args"/>JSON_EXTRACT now supports dynamic path arguments</a>

go/flags/endtoend/vtgate.txt

@@ -246,7 +246,7 @@ Flags:
   -v, --version                                                          print binary version
       --vmodule vModuleFlag                                              comma-separated list of pattern=N settings for file-filtered logging
       --vschema-ddl-authorized-users string                              List of users authorized to execute vschema ddl operations, or '%' to allow all users.
-      --vtgate-balancer-mode string                                      Tablet balancer mode (options: cell, prefer-cell, random). Defaults to 'cell' which shuffles tablets in the local cell.
+      --vtgate-balancer-mode string                                      Tablet balancer mode (options: cell, prefer-cell, random, session). Defaults to 'cell' which shuffles tablets in the local cell.
       --vtgate-config-terse-errors                                       prevent bind vars from escaping in returned errors
       --warming-reads-concurrency int                                    Number of concurrent warming reads allowed (default 500)
       --warming-reads-percent int                                        Percentage of reads on the primary to forward to replicas. Useful for keeping buffer pools warm

go/test/endtoend/cluster/cluster_process.go

@@ -45,12 +45,14 @@ import (
 	"vitess.io/vitess/go/vt/grpcclient"
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/vterrors"
+	"vitess.io/vitess/go/vt/vtgate/executorcontext"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
 	"vitess.io/vitess/go/vt/vtgate/vtgateconn"
 	"vitess.io/vitess/go/vt/vttablet/tabletconn"
 
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+	vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 
 	// Ensure dialers are registered (needed by ExecOnTablet and ExecOnVTGate).
@@ -807,7 +809,7 @@ func NewBareCluster(cell string, hostname string) *LocalProcessCluster {
 		// path/to/whatever exists
 		cluster.ReusingVTDATAROOT = true
 	} else {
-		err = createDirectory(cluster.CurrentVTDATAROOT, 0700)
+		err = createDirectory(cluster.CurrentVTDATAROOT, 0o700)
 		if err != nil {
 			log.Fatal(err)
 		}
@@ -948,11 +950,12 @@ func (cluster *LocalProcessCluster) ExecOnTablet(ctx context.Context, vttablet *
 
 	txID, reservedID := 0, 0
 
-	return conn.Execute(ctx, &querypb.Target{
+	session := executorcontext.NewSafeSession(&vtgatepb.Session{Options: opts})
+	return conn.Execute(ctx, session, &querypb.Target{
 		Keyspace:   tablet.Keyspace,
 		Shard:      tablet.Shard,
 		TabletType: tablet.Type,
-	}, sql, bindvars, int64(txID), int64(reservedID), opts)
+	}, sql, bindvars, int64(txID), int64(reservedID))
 }
 
 // ExecOnVTGate executes a query on a local cluster VTGate with the provided
@@ -1165,7 +1168,8 @@ func (cluster *LocalProcessCluster) waitForMySQLProcessToExit(mysqlctlProcessLis
 
 // StartVtbackup starts a vtbackup
 func (cluster *LocalProcessCluster) StartVtbackup(newInitDBFile string, initialBackup bool,
-	keyspace string, shard string, cell string, extraArgs ...string) error {
+	keyspace string, shard string, cell string, extraArgs ...string,
+) error {
 	log.Info("Starting vtbackup")
 	cluster.VtbackupProcess = *VtbackupProcessInstance(
 		cluster.GetAndReserveTabletUID(),
@@ -1195,7 +1199,6 @@ func (cluster *LocalProcessCluster) GetAndReservePort() int {
 		cluster.nextPortForProcess = cluster.nextPortForProcess + 1
 		log.Infof("Attempting to reserve port: %v", cluster.nextPortForProcess)
 		ln, err := net.Listen("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(cluster.nextPortForProcess)))
-
 		if err != nil {
 			log.Errorf("Can't listen on port %v: %s, trying next port", cluster.nextPortForProcess, err)
 			continue
@@ -1218,7 +1221,7 @@ const portFileTimeout = 1 * time.Hour
 // If yes, then return that port, and save port + 200 in the same file
 // here, assumptions is 200 ports might be consumed for all tests in a package
 func getPort() int {
-	portFile, err := os.OpenFile(path.Join(os.TempDir(), "endtoend.port"), os.O_CREATE|os.O_RDWR, 0644)
+	portFile, err := os.OpenFile(path.Join(os.TempDir(), "endtoend.port"), os.O_CREATE|os.O_RDWR, 0o644)
 	if err != nil {
 		panic(err)
 	}

go/test/endtoend/vtgate/tabletbalancer/session_test.go

@@ -0,0 +1,223 @@
+/*
+Copyright 2025 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package tabletbalancer
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/test/endtoend/cluster"
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
+)
+
+// TestSessionModeBalancer tests the "session" mode routes each session consistently to the same tablet.
+func TestSessionModeBalancer(t *testing.T) {
+	vtgateProcess, vtParams, _, _ := setupCluster(t)
+	defer vtgateProcess.TearDown()
+
+	// Create 2 session connections that route to different tablets
+	conns := createSessionConnections(t, &vtParams, 2)
+	for conn := range conns {
+		defer conn.Close()
+	}
+
+	verifyStickiness(t, conns, 20)
+}
+
+// TestSessionModeRemoveTablet tests that when a tablet is killed, connections switch to remaining tablets
+func TestSessionModeRemoveTablet(t *testing.T) {
+	vtgateProcess, vtParams, replicaTablets, aliases := setupCluster(t)
+	defer vtgateProcess.TearDown()
+
+	// Create 2 connections to different tablets
+	conns := createSessionConnections(t, &vtParams, 2)
+	for conn := range conns {
+		defer conn.Close()
+	}
+
+	// Find the first replica tablet that one of our connections is using
+	var tabletToKill *cluster.Vttablet
+	var affectedConn *mysql.Conn
+	var killedServerID int64
+
+	for _, tablet := range replicaTablets {
+		tabletServerID := aliases[tablet.Alias]
+
+		// Check if any connection is using this tablet
+		for conn, connServerID := range conns {
+			if connServerID != tabletServerID {
+				continue
+			}
+
+			// We found a connection that's using this tablet, let's kill this tablet
+			tabletToKill = tablet
+			affectedConn = conn
+			killedServerID = tabletServerID
+			break
+		}
+
+		// We found a tablet, no need to check other tablets
+		if tabletToKill != nil {
+			break
+		}
+	}
+
+	require.NotNil(t, tabletToKill, "Should find a tablet to kill")
+
+	// Kill the tablet immediately
+	err := tabletToKill.VttabletProcess.Kill()
+	require.Error(t, err)
+
+	// Wait for the connection to switch to a new tablet and update the map
+	require.Eventually(t, func() bool {
+		newServerID := getServerID(t, affectedConn)
+		if newServerID != killedServerID {
+			conns[affectedConn] = newServerID
+			return true
+		}
+
+		return false
+	}, 10*time.Millisecond, 1*time.Millisecond, "Connection should switch to a different tablet")
+
+	verifyStickiness(t, conns, 20)
+}
+
+// setupCluster sets up a cluster with a vtgate using the session balancer.
+func setupCluster(t *testing.T) (*cluster.VtgateProcess, mysql.ConnParams, []*cluster.Vttablet, map[string]int64) {
+	t.Helper()
+
+	// Start vtgate in cell1 with session mode
+	vtgateProcess := cluster.VtgateProcessInstance(
+		clusterInstance.GetAndReservePort(),
+		clusterInstance.GetAndReservePort(),
+		clusterInstance.GetAndReservePort(),
+		cell1,
+		fmt.Sprintf("%s,%s", cell1, cell2),
+		clusterInstance.Hostname,
+		replicaStr,
+		clusterInstance.TopoProcess.Port,
+		clusterInstance.TmpDirectory,
+		[]string{
+			"--vtgate-balancer-mode", "session",
+		},
+		plancontext.PlannerVersion(0),
+	)
+	require.NoError(t, vtgateProcess.Setup())
+	require.True(t, vtgateProcess.WaitForStatus())
+
+	vtParams := mysql.ConnParams{
+		Host: clusterInstance.Hostname,
+		Port: vtgateProcess.MySQLServerPort,
+	}
+
+	allTablets := clusterInstance.Keyspaces[0].Shards[0].Vttablets
+	shardName := clusterInstance.Keyspaces[0].Shards[0].Name
+	replicaTablets := replicaTablets(allTablets)
+
+	conn, err := mysql.Connect(context.Background(), &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	// Wait for tablets to be discovered
+	err = vtgateProcess.WaitForStatusOfTabletInShard(fmt.Sprintf("%s.%s.primary", keyspaceName, shardName), 1, 30*time.Second)
+	require.NoError(t, err)
+
+	err = vtgateProcess.WaitForStatusOfTabletInShard(fmt.Sprintf("%s.%s.replica", keyspaceName, shardName), len(replicaTablets), 30*time.Second)
+	require.NoError(t, err)
+
+	aliases := mapTabletAliasToMySQLServerID(t, allTablets)
+
+	// Insert test data
+	testValue := fmt.Sprintf("session_test_%d", time.Now().UnixNano())
+	_, err = conn.ExecuteFetch(fmt.Sprintf("INSERT INTO balancer_test (value) VALUES ('%s')", testValue), 1, false)
+	require.NoError(t, err)
+	waitForReplication(t, replicaTablets, testValue)
+
+	return vtgateProcess, vtParams, replicaTablets, aliases
+}
+
+// getServerID returns the server ID that the connection is currently routing to.
+func getServerID(t *testing.T, conn *mysql.Conn) int64 {
+	t.Helper()
+
+	res, err := conn.ExecuteFetch("SELECT @@server_id", 1, false)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(res.Rows), "expected one row from server_id query")
+
+	serverID, err := res.Rows[0][0].ToInt64()
+	require.NoError(t, err)
+
+	return serverID
+}
+
+// createSessionConnections creates `n` connections that route to different tablets.
+// Returns a map of mysql.Conn -> serverID.
+func createSessionConnections(t *testing.T, vtParams *mysql.ConnParams, numConnections int) map[*mysql.Conn]int64 {
+	t.Helper()
+
+	conns := make(map[*mysql.Conn]int64)
+	seenServerIDs := make(map[int64]bool)
+
+	// Try up to 50 times to get numConnections with different server IDs
+	for range 50 {
+		conn, err := mysql.Connect(context.Background(), vtParams)
+		require.NoError(t, err)
+
+		_, err = conn.ExecuteFetch("USE @replica", 1, false)
+		require.NoError(t, err)
+
+		// Get the server ID this connection routes to
+		serverID := getServerID(t, conn)
+
+		// If this is a new tablet, keep the connection
+		if !seenServerIDs[serverID] {
+			seenServerIDs[serverID] = true
+			conns[conn] = serverID
+
+			// If we have enough connections, return
+			if len(conns) == numConnections {
+				return conns
+			}
+
+			continue
+		}
+
+		// Already seen this tablet, close and try again
+		conn.Close()
+	}
+
+	t.Fatalf("could not create %d connections with different tablets after 50 attempts, only got %d", numConnections, len(conns))
+	return nil
+}
+
+// verifyStickiness validates whether the given connections remain connected to the same
+// server `n` times in a row.
+func verifyStickiness(t *testing.T, conns map[*mysql.Conn]int64, n uint) {
+	t.Helper()
+
+	for conn, expectedServerID := range conns {
+		for range n {
+			currentServerID := getServerID(t, conn)
+			require.Equal(t, expectedServerID, currentServerID, "Connection should stick to tablet %d, got %d", expectedServerID, currentServerID)
+		}
+	}
+}