[IncomingQueryThrottler] Implement Incoming Query Throttler with Strategy Pattern

go.mod

@@ -103,6 +103,7 @@ require (
 	github.com/kr/text v0.2.0
 	github.com/mitchellh/mapstructure v1.5.1-0.20231216201459-8508981c8b6c
 	github.com/nsf/jsondiff v0.0.0-20210926074059-1e845ec5d249
+	github.com/prashantv/gostub v1.1.0
 	github.com/shirou/gopsutil/v4 v4.25.4
 	github.com/spf13/afero v1.14.0
 	github.com/spf13/jwalterweatherman v1.1.0

go.sum

@@ -474,6 +474,8 @@ github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRI
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw=
 github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE=
+github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
+github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P1IAHhP5U=
 github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
 github.com/prometheus/client_golang v1.0.0/go.mod h1:db9x61etRT2tGnBNRi70OPL5FsnadC4Ky3P0J6CfImo=
 github.com/prometheus/client_golang v1.4.0/go.mod h1:e9GMxYsXl05ICDXkRhurwBS4Q3OK1iX/F2sw+iXX5zU=

go/vt/vttablet/tabletserver/incomingquerythrottler/cinnamon_strategy.go

@@ -0,0 +1,30 @@
+package incomingquerythrottler
+
+import (
+	"context"
+
+	querypb "vitess.io/vitess/go/vt/proto/query"
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+)
+
+var _ ThrottlingStrategyHandler = (*CinnamonStrategy)(nil)
+
+// CinnamonStrategy is a placeholder for the Cinnamon throttling strategy.
+type CinnamonStrategy struct{}
+
+func (s *CinnamonStrategy) ThrottleIfNeeded(ctx context.Context, targetTabletType topodatapb.TabletType, sql string, transactionID int64, options *querypb.ExecuteOptions) error {
+	// No-op for now
+	return nil
+}
+
+// Start is a placeholder for initializing the Cinnamon throttling strategy.
+// TODO: Implement actual Cinnamon strategy initialization when the strategy is fully developed.
+func (s *CinnamonStrategy) Start() {
+	// TODO: Initialize Cinnamon throttling strategy resources
+}
+
+// Stop is a placeholder for cleaning up the Cinnamon throttling strategy.
+// TODO: Implement actual Cinnamon strategy cleanup when the strategy is fully developed.
+func (s *CinnamonStrategy) Stop() {
+	// TODO: Clean up Cinnamon throttling strategy resources
+}

go/vt/vttablet/tabletserver/incomingquerythrottler/config.go

@@ -0,0 +1,31 @@
+package incomingquerythrottler
+
+// ThrottlingStrategy represents the strategy used to apply throttling
+// to incoming queries based on system load or external signals.
+type ThrottlingStrategy string
+
+// Predefined throttling strategies for the IncomingQueryThrottler.
+const (
+	// ThrottlingStrategyTabletThrottler uses Vitess Tablet Throttler to shed load
+	// from incoming queries when the tablet is under pressure.
+	// Reference: https://vitess.io/docs/21.0/reference/features/tablet-throttler/
+	ThrottlingStrategyTabletThrottler ThrottlingStrategy = "TabletThrottler"
+
+	// ThrottlingStrategyCinnamon uses Uber's Cinnamon load-shedding system
+	// to regulate incoming queries under high load conditions.
+	// Reference: https://www.uber.com/en-IN/blog/cinnamon-using-century-old-tech-to-build-a-mean-load-shedder/
+	ThrottlingStrategyCinnamon ThrottlingStrategy = "Cinnamon"
+
+	// ThrottlingStrategyUnknown is used when the strategy is not known.
+	ThrottlingStrategyUnknown ThrottlingStrategy = "Unknown"
+)
+
+// Config defines the runtime configuration for the IncomingQueryThrottler.
+// It specifies whether throttling is enabled and which strategy to use.
+type Config struct {
+	// Enabled indicates whether the throttler should actively apply throttling logic.
+	Enabled bool `json:"enabled"`
+
+	// Strategy selects which throttling strategy should be used.
+	Strategy ThrottlingStrategy `json:"strategy"`
+}

go/vt/vttablet/tabletserver/incomingquerythrottler/config_loader_interface.go

@@ -0,0 +1,8 @@
+package incomingquerythrottler
+
+import "context"
+
+type ConfigLoader interface {
+	// Load returns the latest throttler config (may come from file, topo, etc.)
+	Load(ctx context.Context) (Config, error)
+}

go/vt/vttablet/tabletserver/incomingquerythrottler/file_based_config_loader.go

@@ -0,0 +1,36 @@
+package incomingquerythrottler
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+)
+
+var (
+	_              ConfigLoader = (*FileBasedConfigLoader)(nil)
+	_osReadFile                 = os.ReadFile
+	_jsonUnmarshal              = json.Unmarshal
+	_configPath                 = "/config/throttler-config.json"
+)
+
+type FileBasedConfigLoader struct{}
+
+// NewFileBasedConfigLoader creates a new instance of FileBasedConfigLoader with the given file path.
+func NewFileBasedConfigLoader() *FileBasedConfigLoader {
+	return &FileBasedConfigLoader{}
+}
+
+// Load reads the configuration from a file at the specific config path.
+func (f *FileBasedConfigLoader) Load(ctx context.Context) (Config, error) {
+	data, err := _osReadFile(_configPath)
+	if err != nil {
+		return Config{}, err
+	}
+
+	var cfg Config
+	if unMarshalErr := _jsonUnmarshal(data, &cfg); unMarshalErr != nil {
+		return Config{}, unMarshalErr
+	}
+
+	return cfg, nil
+}

go/vt/vttablet/tabletserver/incomingquerythrottler/file_based_config_loader_test.go

@@ -0,0 +1,193 @@
+package incomingquerythrottler
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"testing"
+
+	"github.com/prashantv/gostub"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewFileBasedConfigLoader(t *testing.T) {
+	loader := NewFileBasedConfigLoader()
+	require.NotNil(t, loader)
+	require.IsType(t, &FileBasedConfigLoader{}, loader)
+}
+
+func TestFileBasedConfigLoader_Load(t *testing.T) {
+	tests := []struct {
+		name                string
+		stubOsReadFile      func(filename string) ([]byte, error)
+		stubJsonUnmarshal   func(data []byte, v interface{}) error
+		expectedConfig      Config
+		expectedError       string
+		expectedErrorNotNil bool
+	}{
+		{
+			name: "successful config load with minimal config",
+			stubOsReadFile: func(filename string) ([]byte, error) {
+				require.Equal(t, "/config/throttler-config.json", filename)
+				return []byte(`{"enabled": true, "strategy": "TabletThrottler"}`), nil
+			},
+			stubJsonUnmarshal: func(data []byte, v interface{}) error {
+				return json.Unmarshal(data, v)
+			},
+			expectedConfig: Config{
+				Enabled:  true,
+				Strategy: ThrottlingStrategyTabletThrottler,
+			},
+			expectedErrorNotNil: false,
+		},
+		{
+			name: "successful config load with disabled throttler",
+			stubOsReadFile: func(filename string) ([]byte, error) {
+				require.Equal(t, "/config/throttler-config.json", filename)
+				return []byte(`{"enabled": false, "strategy": "Unknown"}`), nil
+			},
+			stubJsonUnmarshal: func(data []byte, v interface{}) error {
+				return json.Unmarshal(data, v)
+			},
+			expectedConfig: Config{
+				Enabled:  false,
+				Strategy: ThrottlingStrategyUnknown,
+			},
+			expectedErrorNotNil: false,
+		},
+		{
+			name: "file read error - file not found",
+			stubOsReadFile: func(filename string) ([]byte, error) {
+				require.Equal(t, "/config/throttler-config.json", filename)
+				return nil, errors.New("no such file or directory")
+			},
+			stubJsonUnmarshal: func(data []byte, v interface{}) error {
+				// Should not be called
+				t.Fatal("jsonUnmarshal should not be called when file read fails")
+				return nil
+			},
+			expectedConfig:      Config{},
+			expectedError:       "no such file or directory",
+			expectedErrorNotNil: true,
+		},
+		{
+			name: "file read error - permission denied",
+			stubOsReadFile: func(filename string) ([]byte, error) {
+				require.Equal(t, "/config/throttler-config.json", filename)
+				return nil, errors.New("permission denied")
+			},
+			stubJsonUnmarshal: func(data []byte, v interface{}) error {
+				// Should not be called
+				t.Fatal("jsonUnmarshal should not be called when file read fails")
+				return nil
+			},
+			expectedConfig:      Config{},
+			expectedError:       "permission denied",
+			expectedErrorNotNil: true,
+		},
+		{
+			name: "json unmarshal error - invalid json",
+			stubOsReadFile: func(filename string) ([]byte, error) {
+				require.Equal(t, "/config/throttler-config.json", filename)
+				return []byte(`{"enabled": true, "strategy": `), nil // incomplete JSON
+			},
+			stubJsonUnmarshal: func(data []byte, v interface{}) error {
+				return errors.New("unexpected end of JSON input")
+			},
+			expectedConfig:      Config{},
+			expectedError:       "unexpected end of JSON input",
+			expectedErrorNotNil: true,
+		},
+		{
+			name: "json unmarshal error - invalid field type",
+			stubOsReadFile: func(filename string) ([]byte, error) {
+				require.Equal(t, "/config/throttler-config.json", filename)
+				return []byte(`{"enabled": "invalid_boolean", "strategy": "TabletThrottler"}`), nil
+			},
+			stubJsonUnmarshal: func(data []byte, v interface{}) error {
+				return errors.New("json: cannot unmarshal string into Go struct field Config.enabled of type bool")
+			},
+			expectedConfig:      Config{},
+			expectedError:       "json: cannot unmarshal string into Go struct field Config.enabled of type bool",
+			expectedErrorNotNil: true,
+		},
+		{
+			name: "empty file - should unmarshal to zero value config",
+			stubOsReadFile: func(filename string) ([]byte, error) {
+				require.Equal(t, "/config/throttler-config.json", filename)
+				return []byte(`{}`), nil
+			},
+			stubJsonUnmarshal: func(data []byte, v interface{}) error {
+				return json.Unmarshal(data, v)
+			},
+			expectedConfig: Config{
+				Enabled:  false, // zero value for bool
+				Strategy: "",    // zero value for ThrottlingStrategy
+			},
+			expectedErrorNotNil: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create stubs for the global variables
+			osReadFileStub := gostub.Stub(&_osReadFile, tt.stubOsReadFile)
+			jsonUnmarshalStub := gostub.Stub(&_jsonUnmarshal, tt.stubJsonUnmarshal)
+
+			// Ensure stubs are reset after the test
+			defer osReadFileStub.Reset()
+			defer jsonUnmarshalStub.Reset()
+
+			// Create loader and test Load method
+			loader := NewFileBasedConfigLoader()
+			config, err := loader.Load(context.Background())
+
+			// Verify error expectations
+			if tt.expectedErrorNotNil {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tt.expectedError)
+				require.Equal(t, tt.expectedConfig, config)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tt.expectedConfig, config)
+			}
+		})
+	}
+}
+
+func TestFileBasedConfigLoader_Load_ConfigPath(t *testing.T) {
+	// Test that the correct config path is being used
+	var capturedPath string
+
+	stubOsReadFile := func(filename string) ([]byte, error) {
+		capturedPath = filename
+		return []byte(`{"enabled": true, "strategy": "TabletThrottler"}`), nil
+	}
+
+	stubJsonUnmarshal := func(data []byte, v interface{}) error {
+		return json.Unmarshal(data, v)
+	}
+
+	// Create stubs
+	osReadFileStub := gostub.Stub(&_osReadFile, stubOsReadFile)
+	jsonUnmarshalStub := gostub.Stub(&_jsonUnmarshal, stubJsonUnmarshal)
+
+	defer osReadFileStub.Reset()
+	defer jsonUnmarshalStub.Reset()
+
+	// Test
+	loader := NewFileBasedConfigLoader()
+	_, err := loader.Load(context.Background())
+
+	require.NoError(t, err)
+	require.Equal(t, "/config/throttler-config.json", capturedPath)
+}
+
+func TestFileBasedConfigLoader_ImplementsConfigLoader(t *testing.T) {
+	// Verify that FileBasedConfigLoader implements ConfigLoader interface
+	var _ ConfigLoader = (*FileBasedConfigLoader)(nil)
+
+	// This should compile without issues, proving interface compliance
+	loader := NewFileBasedConfigLoader()
+	require.NotNil(t, loader)
+}