diff --git a/go/vt/vttablet/tabletserver/querythrottler/config.go b/go/vt/vttablet/tabletserver/querythrottler/config.go index f81f7b7fc89..a424184abe5 100644 --- a/go/vt/vttablet/tabletserver/querythrottler/config.go +++ b/go/vt/vttablet/tabletserver/querythrottler/config.go @@ -27,6 +27,11 @@ type Config struct { // Enabled indicates whether the throttler should actively apply throttling logic. Enabled bool `json:"enabled"` + // DryRun indicates whether throttling decisions should be logged but not enforced. + // When true, queries that would be throttled are allowed to proceed, but the + // throttling decision is logged for observability. + DryRun bool `json:"dry_run"` + // Strategy selects which throttling strategy should be used. Strategy registry.ThrottlingStrategy `json:"strategy"` } diff --git a/go/vt/vttablet/tabletserver/querythrottler/file_based_config_loader_test.go b/go/vt/vttablet/tabletserver/querythrottler/file_based_config_loader_test.go index 19ea02c3e2b..67904e32da6 100644 --- a/go/vt/vttablet/tabletserver/querythrottler/file_based_config_loader_test.go +++ b/go/vt/vttablet/tabletserver/querythrottler/file_based_config_loader_test.go @@ -90,6 +90,23 @@ func TestFileBasedConfigLoader_Load(t *testing.T) { expectedError: "no such file or directory", expectedErrorNotNil: true, }, + { + name: "successful config load with dry run as enabled", + configPath: "/config/throttler-config.json", + mockReadFile: func(filename string) ([]byte, error) { + require.Equal(t, "/config/throttler-config.json", filename) + return []byte(`{"enabled": true, "strategy": "TabletThrottler", "dry_run": true}`), nil + }, + mockJsonUnmarshal: func(data []byte, v interface{}) error { + return json.Unmarshal(data, v) + }, + expectedConfig: Config{ + Enabled: true, + Strategy: registry.ThrottlingStrategyTabletThrottler, + DryRun: true, + }, + expectedErrorNotNil: false, + }, { name: "file read error - permission denied", configPath: "/config/throttler-config.json", diff --git a/go/vt/vttablet/tabletserver/querythrottler/query_throttler.go b/go/vt/vttablet/tabletserver/querythrottler/query_throttler.go index 3d7920f48f1..08e6b7bdba1 100644 --- a/go/vt/vttablet/tabletserver/querythrottler/query_throttler.go +++ b/go/vt/vttablet/tabletserver/querythrottler/query_throttler.go @@ -103,6 +103,12 @@ func (qt *QueryThrottler) Throttle(ctx context.Context, tabletType topodatapb.Ta return nil } + // If dry-run mode is enabled, log the decision but don't throttle + if qt.cfg.DryRun { + log.Warningf("[DRY-RUN] %s", decision.Message) + return nil + } + // Normal throttling: return an error to reject the query return vterrors.New(vtrpcpb.Code_RESOURCE_EXHAUSTED, decision.Message) } diff --git a/go/vt/vttablet/tabletserver/querythrottler/query_throttler_test.go b/go/vt/vttablet/tabletserver/querythrottler/query_throttler_test.go index 5824edaf6cc..6f1caeaddef 100644 --- a/go/vt/vttablet/tabletserver/querythrottler/query_throttler_test.go +++ b/go/vt/vttablet/tabletserver/querythrottler/query_throttler_test.go @@ -18,9 +18,15 @@ package querythrottler import ( "context" + "fmt" "testing" "time" + "vitess.io/vitess/go/vt/log" + querypb "vitess.io/vitess/go/vt/proto/query" + topodatapb "vitess.io/vitess/go/vt/proto/topodata" + "vitess.io/vitess/go/vt/sqlparser" + "vitess.io/vitess/go/vt/vttablet/tabletserver/querythrottler/registry" "vitess.io/vitess/go/vt/vtenv" @@ -150,3 +156,181 @@ func TestQueryThrottler_Shutdown(t *testing.T) { iqt.mu.RUnlock() require.NotNil(t, strategy) } + +// TestIncomingQueryThrottler_DryRunMode tests that dry-run mode logs decisions but doesn't throttle queries. +func TestIncomingQueryThrottler_DryRunMode(t *testing.T) { + tests := []struct { + name string + enabled bool + dryRun bool + throttleDecision registry.ThrottleDecision + expectError bool + expectDryRunLog bool + expectedLogMsg string + }{ + { + name: "Disabled throttler - no checks performed", + enabled: false, + dryRun: false, + throttleDecision: registry.ThrottleDecision{ + Throttle: true, + Message: "Should not be evaluated", + }, + expectError: false, + expectDryRunLog: false, + }, + { + name: "Disabled throttler with dry-run - no checks performed", + enabled: false, + dryRun: true, + throttleDecision: registry.ThrottleDecision{ + Throttle: true, + Message: "Should not be evaluated", + }, + expectError: false, + expectDryRunLog: false, + }, + { + name: "Normal mode - query allowed", + enabled: true, + dryRun: false, + throttleDecision: registry.ThrottleDecision{ + Throttle: false, + Message: "Query allowed", + }, + expectError: false, + expectDryRunLog: false, + }, + { + name: "Normal mode - query throttled", + enabled: true, + dryRun: false, + throttleDecision: registry.ThrottleDecision{ + Throttle: true, + Message: "Query throttled: metric=cpu value=90.0 threshold=80.0", + MetricName: "cpu", + MetricValue: 90.0, + Threshold: 80.0, + ThrottlePercentage: 1.0, + }, + expectError: true, + expectDryRunLog: false, + }, + { + name: "Dry-run mode - query would be throttled but allowed", + enabled: true, + dryRun: true, + throttleDecision: registry.ThrottleDecision{ + Throttle: true, + Message: "Query throttled: metric=cpu value=95.0 threshold=80.0", + MetricName: "cpu", + MetricValue: 95.0, + Threshold: 80.0, + ThrottlePercentage: 1.0, + }, + expectError: false, + expectDryRunLog: true, + expectedLogMsg: "[DRY-RUN] Query throttled: metric=cpu value=95.0 threshold=80.0", + }, + { + name: "Dry-run mode - query allowed normally", + enabled: true, + dryRun: true, + throttleDecision: registry.ThrottleDecision{ + Throttle: false, + Message: "Query allowed", + }, + expectError: false, + expectDryRunLog: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create a mock strategy with controlled decision + mockStrategy := &mockThrottlingStrategy{ + decision: tt.throttleDecision, + } + + // Create throttler with controlled config + iqt := &QueryThrottler{ + ctx: context.Background(), + cfg: Config{ + Enabled: tt.enabled, + DryRun: tt.dryRun, + }, + strategy: mockStrategy, + } + + // Capture log output + logCapture := &testLogCapture{} + originalLogWarningf := log.Warningf + defer func() { + // Restore original logging function + log.Warningf = originalLogWarningf + }() + + // Mock log.Warningf to capture output + log.Warningf = logCapture.captureLog + + // Test the enforcement + err := iqt.Throttle( + context.Background(), + topodatapb.TabletType_REPLICA, + &sqlparser.ParsedQuery{Query: "SELECT * FROM test_table WHERE id = 1"}, + 12345, + &querypb.ExecuteOptions{ + WorkloadName: "test-workload", + Priority: "50", + }, + ) + + // Verify error expectation + if tt.expectError { + require.EqualError(t, err, tt.throttleDecision.Message, "Error should match the throttle message exactly") + } else { + require.NoError(t, err, "Expected no throttling error") + } + + // Verify log expectation + if tt.expectDryRunLog { + require.Len(t, logCapture.logs, 1, "Expected exactly one log message") + require.Equal(t, tt.expectedLogMsg, logCapture.logs[0], "Log message should match expected") + } else { + require.Empty(t, logCapture.logs, "Expected no log messages") + } + }) + } +} + +// mockThrottlingStrategy is a test strategy that allows us to control throttling decisions +type mockThrottlingStrategy struct { + decision registry.ThrottleDecision + started bool + stopped bool +} + +func (m *mockThrottlingStrategy) Evaluate(ctx context.Context, targetTabletType topodatapb.TabletType, parsedQuery *sqlparser.ParsedQuery, transactionID int64, options *querypb.ExecuteOptions) registry.ThrottleDecision { + return m.decision +} + +func (m *mockThrottlingStrategy) Start() { + m.started = true +} + +func (m *mockThrottlingStrategy) Stop() { + m.stopped = true +} + +func (m *mockThrottlingStrategy) GetStrategyName() string { + return "MockStrategy" +} + +// testLogCapture captures log output for testing +type testLogCapture struct { + logs []string +} + +func (lc *testLogCapture) captureLog(msg string, args ...interface{}) { + lc.logs = append(lc.logs, fmt.Sprintf(msg, args...)) +}