Gate vttablet conn pools with Snake load shedder (#864)

* Gate OLTP read pool with Snake load shedder

Wire the CoDel-based Snake gate into QueryExecutor.getConn() so that
OLTP read requests are admitted or shed before consuming a MySQL
connection. Controlled by --snake-enabled flag (off by default).

ContentionID is extracted from the unique_id SQL margin comment that
webapp injects, with a random 16-char hex fallback for requests
without one.

AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Brett Wines <bwines@slack-corp.com>

* Use proto UniqueId as Snake contentionID, bypass valve when unset

Replace the extractUniqueID regex (which targeted a comment format
webapp never sends) with the proto-based approach: read
qre.options.GetUniqueId() and only enter the Snake gate when it's
non-empty. Requests without a unique_id bypass the gate entirely
rather than sharing a single valve slot or using random IDs.

Add `string unique_id = 22` to ExecuteOptions in proto/query.proto.
This is an optional field — callers that don't set it get zero-cost
bypass. The sqlparser/vtgate propagation to populate it from SQL
directives is deferred to a follow-up PR.

AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Brett Wines <bwines@slack-corp.com>

* Gate DML with Snake load shedder (#871)

Acquire a Snake slot in TxPool.Begin (for fresh connections only,
not reserved conns) and release it in txComplete. This gates both
autocommit and explicit transaction DML through the same CoDel-based
load-shedding mechanism already used for the OLTP read pool.

Requests without a unique_id bypass the gate entirely, matching the
OLTP-read behavior.

AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :)

Signed-off-by: Brett Wines <bwines@slack-corp.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

---------

Signed-off-by: Brett Wines <bwines@slack-corp.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: bgwines

go/vt/proto/query/query.pb.go

@@ -46,6 +46,7 @@ import (
 	tacl "vitess.io/vitess/go/vt/tableacl/acl"
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/connpool"
+	"vitess.io/vitess/go/vt/vttablet/tabletserver/loadshed"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/planbuilder"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/rules"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/schema"
@@ -197,6 +198,9 @@ type QueryEngine struct {
 	accessCheckerLogger *logutil.ThrottledLogger
 
 	redactUIQuery bool
+
+	// snake is the CoDel-based load-shedding gate for the OLTP read pool.
+	snake *loadshed.Snake
 }
 
 // NewQueryEngine creates a new QueryEngine.
@@ -241,6 +245,20 @@ func NewQueryEngine(env tabletenv.Env, se *schema.Engine) *QueryEngine {
 	}
 	qe.txSerializer = txserializer.New(env)
 
+	if config.SnakeEnabled {
+		qe.snake = loadshed.NewSnake(loadshed.SnakeConfig{
+			Name: "oltp-read",
+			CoDel: loadshed.CoDelConfig{
+				TargetNs:       func() int64 { return config.SnakeTarget.Nanoseconds() },
+				IntervalNs:     func() int64 { return config.SnakeInterval.Nanoseconds() },
+				Exponent:       func() float64 { return 0.5 },
+				MinDropDelayNs: func() int64 { return int64(time.Millisecond) },
+			},
+			Capacity:            func() int { return config.OltpReadPool.Size },
+			LoadsheddingAllowed: func() bool { return true },
+		})
+	}
+
 	qe.strictTableACL = config.StrictTableACL
 	qe.enableTableACLDryRun = config.EnableTableACLDryRun
 

go/vt/proto/query/query_vtproto.pb.go

@@ -714,12 +714,13 @@ func (qre *QueryExecutor) execSelect() (*sqltypes.Result, error) {
 
 		if original {
 			defer q.Broadcast()
-			conn, err := qre.getConn()
+			conn, release, err := qre.getConn()
 
 			if err != nil {
 				q.SetErr(err)
 			} else {
 				defer conn.Recycle()
+				defer release()
 				res, err := qre.execDBConn(conn.Conn, sql, true)
 				if qre.tsv.config.ConsolidatorCacheProto3Rows && q.HasWaiters() {
 					res.CacheProto3Rows()
@@ -755,11 +756,12 @@ func (qre *QueryExecutor) execSelect() (*sqltypes.Result, error) {
 		}
 		// If waiter cap exceeded, fall through to independent execution
 	}
-	conn, err := qre.getConn()
+	conn, release, err := qre.getConn()
 	if err != nil {
 		return nil, err
 	}
 	defer conn.Recycle()
+	defer release()
 	res, err := qre.execDBConn(conn.Conn, sql, true)
 	if err != nil {
 		return nil, err
@@ -797,22 +799,45 @@ func (qre *QueryExecutor) verifyRowCount(count, maxrows int64) error {
 }
 
 func (qre *QueryExecutor) execOther() (*sqltypes.Result, error) {
-	conn, err := qre.getConn()
+	conn, release, err := qre.getConn()
 	if err != nil {
 		return nil, err
 	}
 	defer conn.Recycle()
+	defer release()
 	return qre.execDBConn(conn.Conn, qre.query, true)
 }
 
-func (qre *QueryExecutor) getConn() (*connpool.PooledConn, error) {
+func (qre *QueryExecutor) getConn() (*connpool.PooledConn, func(), error) {
 	span, ctx := trace.NewSpan(qre.ctx, "QueryExecutor.getConn")
 	defer span.Finish()
 
 	defer func(start time.Time) {
 		qre.logStats.WaitingForConnection += time.Since(start)
 	}(time.Now())
-	return qre.tsv.qe.conns.Get(ctx, qre.setting)
+
+	snake := qre.tsv.qe.snake
+	if snake != nil {
+		contentionID := qre.options.GetUniqueId()
+		if contentionID != "" {
+			unlock, err := snake.Acquire(ctx, contentionID)
+			if err != nil {
+				return nil, nil, vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED, "load shed: %v", err)
+			}
+			conn, err := qre.tsv.qe.conns.Get(ctx, qre.setting)
+			if err != nil {
+				unlock.Release()
+				return nil, nil, err
+			}
+			return conn, func() { unlock.Release() }, nil
+		}
+	}
+
+	conn, err := qre.tsv.qe.conns.Get(ctx, qre.setting)
+	if err != nil {
+		return nil, nil, err
+	}
+	return conn, func() {}, nil
 }
 
 func (qre *QueryExecutor) getStreamConn() (*connpool.PooledConn, error) {
@@ -927,11 +952,12 @@ func rewriteOUTParamError(err error) error {
 }
 
 func (qre *QueryExecutor) execCallProc() (*sqltypes.Result, error) {
-	conn, err := qre.getConn()
+	conn, release, err := qre.getConn()
 	if err != nil {
 		return nil, err
 	}
 	defer conn.Recycle()
+	defer release()
 	sql, _, err := qre.generateFinalSQL(qre.plan.FullQuery, qre.bindVars)
 	if err != nil {
 		return nil, err

go/vt/vttablet/tabletserver/query_engine.go

@@ -1621,10 +1621,11 @@ func TestGetConnectionLogStats(t *testing.T) {
 
 	// getConn() happy path
 	qre := newTestQueryExecutor(ctx, tsv, input, 0)
-	conn, err := qre.getConn()
+	conn, release, err := qre.getConn()
 	assert.NoError(t, err)
 	assert.NotNil(t, conn)
 	assert.True(t, qre.logStats.WaitingForConnection > 0)
+	release()
 
 	// getStreamConn() happy path
 	qre = newTestQueryExecutor(ctx, tsv, input, 0)
@@ -1638,7 +1639,7 @@ func TestGetConnectionLogStats(t *testing.T) {
 
 	// getConn() error path
 	qre = newTestQueryExecutor(ctx, tsv, input, 0)
-	_, err = qre.getConn()
+	_, _, err = qre.getConn()
 	assert.Error(t, err)
 	assert.True(t, qre.logStats.WaitingForConnection > 0)
 
@@ -1649,6 +1650,92 @@ func TestGetConnectionLogStats(t *testing.T) {
 	assert.True(t, qre.logStats.WaitingForConnection > 0)
 }
 
+func TestGetConnSnakeBypassed(t *testing.T) {
+	db := setUpQueryExecutorTest(t)
+	defer db.Close()
+
+	ctx := context.Background()
+	cfg := tabletenv.NewDefaultConfig()
+	cfg.OltpReadPool.Size = 2
+	cfg.TxPool.Size = 100
+	cfg.SnakeEnabled = true
+	cfg.SnakeTarget = 5 * time.Millisecond
+	cfg.SnakeInterval = 100 * time.Millisecond
+	cfg.DB = newDBConfigs(db)
+
+	srvTopoCounts := stats.NewCountersWithSingleLabel("", "Resilient srvtopo server operations", "type")
+	tsv := NewTabletServer(ctx, vtenv.NewTestEnv(), "TabletServerTest", cfg, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}, srvTopoCounts)
+	target := &querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
+	err := tsv.StartService(target, cfg.DB, nil)
+	require.NoError(t, err)
+	defer tsv.StopService()
+
+	require.NotNil(t, tsv.qe.snake)
+
+	input := "select * from test_table limit 1"
+
+	// Without UniqueId set, Snake gate is bypassed entirely
+	qre := newTestQueryExecutor(ctx, tsv, input, 0)
+	conn, release, err := qre.getConn()
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+	conn.Recycle()
+	release()
+}
+
+func TestGetConnWithSnake(t *testing.T) {
+	db := setUpQueryExecutorTest(t)
+	defer db.Close()
+
+	ctx := context.Background()
+	cfg := tabletenv.NewDefaultConfig()
+	cfg.OltpReadPool.Size = 2
+	cfg.TxPool.Size = 100
+	cfg.SnakeEnabled = true
+	cfg.SnakeTarget = 5 * time.Millisecond
+	cfg.SnakeInterval = 100 * time.Millisecond
+	cfg.DB = newDBConfigs(db)
+
+	srvTopoCounts := stats.NewCountersWithSingleLabel("", "Resilient srvtopo server operations", "type")
+	tsv := NewTabletServer(ctx, vtenv.NewTestEnv(), "TabletServerTest", cfg, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}, srvTopoCounts)
+	target := &querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
+	err := tsv.StartService(target, cfg.DB, nil)
+	require.NoError(t, err)
+	defer tsv.StopService()
+
+	require.NotNil(t, tsv.qe.snake, "snake should be initialized when SnakeEnabled=true")
+
+	input := "select * from test_table limit 1"
+
+	// With UniqueId set, Snake gate admits the request
+	qre := newTestQueryExecutor(ctx, tsv, input, 0)
+	qre.options = &querypb.ExecuteOptions{UniqueId: "test-request-123"}
+	conn, release, err := qre.getConn()
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+	conn.Recycle()
+	release()
+}
+
+func TestGetConnSnakeDisabled(t *testing.T) {
+	db := setUpQueryExecutorTest(t)
+	defer db.Close()
+
+	ctx := context.Background()
+	tsv := newTestTabletServer(ctx, noFlags, db)
+	defer tsv.StopService()
+
+	assert.Nil(t, tsv.qe.snake, "snake should be nil when SnakeEnabled=false")
+
+	input := "select * from test_table limit 1"
+	qre := newTestQueryExecutor(ctx, tsv, input, 0)
+	conn, release, err := qre.getConn()
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+	conn.Recycle()
+	release()
+}
+
 type executorFlags int64
 
 const (

go/vt/vttablet/tabletserver/query_executor.go

@@ -80,6 +80,7 @@ var (
 	unhealthyThreshold           time.Duration
 	transitionGracePeriod        time.Duration
 	enableReplicationReporter    bool
+	enableSnake                  bool
 )
 
 func init() {
@@ -223,6 +224,10 @@ func registerTabletEnvFlags(fs *pflag.FlagSet) {
 	fs.BoolVar(&currentConfig.EnablePerWorkloadTableMetrics, "enable-per-workload-table-metrics", defaultConfig.EnablePerWorkloadTableMetrics, "If true, query counts and query error metrics include a label that identifies the workload")
 
 	fs.BoolVar(&currentConfig.Unmanaged, "unmanaged", false, "Indicates an unmanaged tablet, i.e. using an external mysql-compatible database")
+
+	fs.BoolVar(&enableSnake, "snake-enabled", false, "If true, enables CoDel-based load shedding (Snake) on the OLTP read pool.")
+	fs.DurationVar(&currentConfig.SnakeTarget, "snake-target", 5*time.Millisecond, "CoDel target delay for the Snake load shedder.")
+	fs.DurationVar(&currentConfig.SnakeInterval, "snake-interval", 100*time.Millisecond, "CoDel interval for the Snake load shedder.")
 }
 
 var (
@@ -295,6 +300,7 @@ func Init() {
 	currentConfig.GracePeriods.Transition = transitionGracePeriod
 	currentConfig.SemiSyncMonitor.Interval = semiSyncMonitorInterval
 
+	currentConfig.SnakeEnabled = enableSnake
 	logFormat := streamlog.GetQueryLogConfig().Format
 	switch logFormat {
 	case streamlog.QueryLogFormatText:
@@ -387,6 +393,10 @@ type TabletConfig struct {
 	EnableViews bool `json:"-"`
 
 	EnablePerWorkloadTableMetrics bool `json:"-"`
+
+	SnakeEnabled  bool          `json:"-"`
+	SnakeTarget   time.Duration `json:"-"`
+	SnakeInterval time.Duration `json:"-"`
 }
 
 func (cfg *TabletConfig) MarshalJSON() ([]byte, error) {

go/vt/vttablet/tabletserver/query_executor_test.go

@@ -57,6 +57,8 @@ type (
 		LogToFile       bool
 
 		Stats *servenv.TimingsWrapper
+
+		SnakeRelease func()
 	}
 
 	// Query contains the query and involved tables executed inside transaction.

go/vt/vttablet/tabletserver/tabletenv/config.go

@@ -34,6 +34,7 @@ import (
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/connpool"
+	"vitess.io/vitess/go/vt/vttablet/tabletserver/loadshed"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/tx"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/txlimiter"
@@ -114,6 +115,19 @@ func NewTxEngine(env tabletenv.Env, dxNotifier func()) *TxEngine {
 	}
 	limiter := txlimiter.New(env)
 	te.txPool = NewTxPool(env, limiter)
+	if config.SnakeEnabled {
+		te.txPool.snake = loadshed.NewSnake(loadshed.SnakeConfig{
+			Name: "dml",
+			CoDel: loadshed.CoDelConfig{
+				TargetNs:       func() int64 { return config.SnakeTarget.Nanoseconds() },
+				IntervalNs:     func() int64 { return config.SnakeInterval.Nanoseconds() },
+				Exponent:       func() float64 { return 0.5 },
+				MinDropDelayNs: func() int64 { return int64(time.Millisecond) },
+			},
+			Capacity:            func() int { return config.TxPool.Size },
+			LoadsheddingAllowed: func() bool { return true },
+		})
+	}
 	// We initially allow twoPC (handles vttablet restarts).
 	// We will disallow them for a few reasons -
 	//	1. When a new tablet is promoted if semi-sync is turned off.

go/vt/vttablet/tabletserver/tx/api.go

@@ -31,6 +31,7 @@ import (
 	"vitess.io/vitess/go/vt/servenv"
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
+	"vitess.io/vitess/go/vt/vttablet/tabletserver/loadshed"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/tx"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/txlimiter"
@@ -68,6 +69,7 @@ type (
 		scp     *StatefulConnectionPool
 		ticks   *timer.Timer
 		limiter txlimiter.TxLimiter
+		snake   *loadshed.Snake
 
 		logMu   sync.Mutex
 		lastLog time.Time
@@ -235,6 +237,7 @@ func (tp *TxPool) Begin(ctx context.Context, options *querypb.ExecuteOptions, re
 
 	var conn *StatefulConnection
 	var err error
+	var snakeRelease func()
 	if reservedID != 0 {
 		conn, err = tp.scp.GetAndLock(reservedID, "start transaction on reserve conn")
 		if err != nil {
@@ -249,12 +252,27 @@ func (tp *TxPool) Begin(ctx context.Context, options *querypb.ExecuteOptions, re
 		if !tp.limiter.Get(immediateCaller, effectiveCaller) {
 			return nil, "", "", vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED, "per-user transaction pool connection limit exceeded")
 		}
+
+		if tp.snake != nil {
+			if contentionID := options.GetUniqueId(); contentionID != "" {
+				unlock, snakeErr := tp.snake.Acquire(ctx, contentionID)
+				if snakeErr != nil {
+					tp.limiter.Release(immediateCaller, effectiveCaller)
+					return nil, "", "", vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED, "dml load shed: %v", snakeErr)
+				}
+				snakeRelease = func() { unlock.Release() }
+			}
+		}
+
 		conn, err = tp.createConn(ctx, options, setting)
 		defer func() {
 			if err != nil {
 				// The transaction limiter frees transactions on rollback or commit. If we fail to create the transaction,
 				// release immediately since there will be no rollback or commit.
 				tp.limiter.Release(immediateCaller, effectiveCaller)
+				if snakeRelease != nil {
+					snakeRelease()
+				}
 			}
 		}()
 	}
@@ -272,6 +290,7 @@ func (tp *TxPool) Begin(ctx context.Context, options *querypb.ExecuteOptions, re
 	if setting != nil {
 		conn.TxProperties().RecordQueryDetail(setting.ApplyQuery(), nil)
 	}
+	conn.TxProperties().SnakeRelease = snakeRelease
 	return conn, sql, sessionStateChanges, nil
 }
 
@@ -444,6 +463,9 @@ func (tp *TxPool) LogActive() {
 
 func (tp *TxPool) txComplete(conn *StatefulConnection, reason tx.ReleaseReason) {
 	conn.LogTransaction(reason)
+	if release := conn.TxProperties().SnakeRelease; release != nil {
+		release()
+	}
 	tp.limiter.Release(conn.TxProperties().ImmediateCaller, conn.TxProperties().EffectiveCaller)
 	conn.CleanTxState()
 }