Throttler code: some refactoring and cleanup (#16368)

Signed-off-by: Shlomi Noach <2607934+shlomi-noach@users.noreply.github.com>

go/vt/vttablet/tabletserver/throttle/mysql/mysql_inventory.go → go/vt/vttablet/tabletserver/throttle/base/inventory.go

@@ -39,22 +39,7 @@ limitations under the License.
  SOFTWARE.
 */
 
-package mysql
-
-import (
- "vitess.io/vitess/go/vt/vttablet/tabletserver/throttle/base"
-)
-
-// TabletResultMap maps a tablet to a result
-type TabletResultMap map[string]base.MetricResultMap
-
-func (m TabletResultMap) Split(alias string) (withAlias TabletResultMap, all TabletResultMap) {
- withAlias = make(TabletResultMap)
- if val, ok := m[alias]; ok {
- withAlias[alias] = val
- }
- return withAlias, m
-}
+package base
 
 // Inventory has the operational data about probes, their metrics, and relevant configuration
 type Inventory struct {

go/vt/vttablet/tabletserver/throttle/mysql/mysql_inventory_test.go → go/vt/vttablet/tabletserver/throttle/base/inventory_test.go

@@ -14,22 +14,20 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package mysql
+package base
 
 import (
  "testing"
 
  "github.com/stretchr/testify/assert"
  "golang.org/x/exp/maps"
-
- "vitess.io/vitess/go/vt/vttablet/tabletserver/throttle/base"
 )
 
 func TestTabletResultMapSplit(t *testing.T) {
  tabletResultMap := TabletResultMap{
- "a": make(base.MetricResultMap),
- "b": make(base.MetricResultMap),
- "c": make(base.MetricResultMap),
+ "a": make(MetricResultMap),
+ "b": make(MetricResultMap),
+ "c": make(MetricResultMap),
  }
  {
  withAlias, all := tabletResultMap.Split("b")

go/vt/vttablet/tabletserver/throttle/mysql/probe.go → go/vt/vttablet/tabletserver/throttle/base/probe.go

@@ -39,7 +39,7 @@ limitations under the License.
  SOFTWARE.
 */
 
-package mysql
+package base
 
 import (
  "fmt"

go/vt/vttablet/tabletserver/throttle/mysql/probe_test.go → go/vt/vttablet/tabletserver/throttle/base/probe_test.go

@@ -39,7 +39,7 @@ limitations under the License.
  SOFTWARE.
 */
 
-package mysql
+package base
 
 import (
  "testing"

go/vt/vttablet/tabletserver/throttle/mysql.go → go/vt/vttablet/tabletserver/throttle/base/tablet_results.go

@@ -14,63 +14,42 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-// This codebase originates from https://github.com/github/freno, See https://github.com/github/freno/blob/master/LICENSE
-/*
- MIT License
-
- Copyright (c) 2017 GitHub
-
- Permission is hereby granted, free of charge, to any person obtaining a copy
- of this software and associated documentation files (the "Software"), to deal
- in the Software without restriction, including without limitation the rights
- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
- copies of the Software, and to permit persons to whom the Software is
- furnished to do so, subject to the following conditions:
-
- The above copyright notice and this permission notice shall be included in all
- copies or substantial portions of the Software.
+package base
 
- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
- SOFTWARE.
-*/
-
-package throttle
+import "sort"
 
-import (
- "context"
- "sort"
+// TabletResultMap maps a tablet to a result
+type TabletResultMap map[string]MetricResultMap
 
- "vitess.io/vitess/go/vt/vttablet/tabletserver/throttle/base"
- "vitess.io/vitess/go/vt/vttablet/tabletserver/throttle/mysql"
-)
+func (m TabletResultMap) Split(alias string) (withAlias TabletResultMap, all TabletResultMap) {
+ withAlias = make(TabletResultMap)
+ if val, ok := m[alias]; ok {
+ withAlias[alias] = val
+ }
+ return withAlias, m
+}
 
-func aggregateMySQLProbes(
- ctx context.Context,
- metricName base.MetricName,
- tabletResultsMap mysql.TabletResultMap,
+func AggregateTabletMetricResults(
+ metricName MetricName,
+ tabletResultsMap TabletResultMap,
  ignoreHostsCount int,
  IgnoreDialTCPErrors bool,
  ignoreHostsThreshold float64,
-) (worstMetric base.MetricResult) {
+) (worstMetric MetricResult) {
  // probes is known not to change. It can be *replaced*, but not changed.
  // so it's safe to iterate it
  probeValues := []float64{}
  for _, tabletMetricResults := range tabletResultsMap {
  tabletMetricResult, ok := tabletMetricResults[metricName]
  if !ok {
- return base.NoSuchMetric
+ return NoSuchMetric
  }
  if tabletMetricResult == nil {
- return base.NoMetricResultYet
+ return NoMetricResultYet
  }
  value, err := tabletMetricResult.Get()
  if err != nil {
- if IgnoreDialTCPErrors && base.IsDialTCPError(err) {
+ if IgnoreDialTCPErrors && IsDialTCPError(err) {
  continue
  }
  if ignoreHostsCount > 0 {
@@ -85,7 +64,7 @@ func aggregateMySQLProbes(
  probeValues = append(probeValues, value)
  }
  if len(probeValues) == 0 {
- return base.NoHostsMetricResult
+ return NoHostsMetricResult
  }
 
  // If we got here, that means no errors (or good-to-skip errors)
@@ -115,6 +94,6 @@ func aggregateMySQLProbes(
  ignoreHostsCount = ignoreHostsCount - 1
  }
  worstValue := probeValues[len(probeValues)-1]
- worstMetric = base.NewSimpleMetricResult(worstValue)
+ worstMetric = NewSimpleMetricResult(worstValue)
  return worstMetric
 }

go/vt/vttablet/tabletserver/throttle/base/tablet_results_test.go

@@ -0,0 +1,199 @@
+/*
+Copyright 2024 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 base
+
+import (
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+)
+
+var (
+ alias1 = "zone1-0001"
+ alias2 = "zone1-0002"
+ alias3 = "zone1-0003"
+ alias4 = "zone1-0004"
+ alias5 = "zone1-0005"
+)
+
+const (
+ nonexistentMetricName MetricName = "nonexistent"
+)
+
+func newMetricResultMap(val float64) MetricResultMap {
+ return MetricResultMap{
+ DefaultMetricName: NewSimpleMetricResult(val),
+ LagMetricName: NewSimpleMetricResult(val),
+ LoadAvgMetricName: NewSimpleMetricResult(3.14),
+ }
+}
+func noSuchMetricMap() MetricResultMap {
+ result := make(MetricResultMap)
+ for _, metricName := range KnownMetricNames {
+ result[metricName] = NoSuchMetric
+ }
+ return result
+}
+
+func TestAggregateTabletMetricResultsNoErrors(t *testing.T) {
+ tabletResultsMap := TabletResultMap{
+ alias1: newMetricResultMap(1.2),
+ alias2: newMetricResultMap(1.7),
+ alias3: newMetricResultMap(0.3),
+ alias4: newMetricResultMap(0.6),
+ alias5: newMetricResultMap(1.1),
+ }
+
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 0, false, 0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 1.7)
+ }
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 1, false, 0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 1.2)
+ }
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 2, false, 0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 1.1)
+ }
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 3, false, 0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 0.6)
+ }
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 4, false, 0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 0.3)
+ }
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 5, false, 0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 0.3)
+ }
+}
+
+func TestAggregateTabletMetricResultsNoErrorsIgnoreHostsThreshold(t *testing.T) {
+ tabletResultsMap := TabletResultMap{
+ alias1: newMetricResultMap(1.2),
+ alias2: newMetricResultMap(1.7),
+ alias3: newMetricResultMap(0.3),
+ alias4: newMetricResultMap(0.6),
+ alias5: newMetricResultMap(1.1),
+ }
+
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 0, false, 1.0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 1.7)
+ }
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 1, false, 1.0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 1.2)
+ }
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 2, false, 1.0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 1.1)
+ }
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 3, false, 1.0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 0.6)
+ }
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 4, false, 1.0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 0.6)
+ }
+ {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 5, false, 1.0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 0.6)
+ }
+}
+
+func TestAggregateTabletMetricResultsWithErrors(t *testing.T) {
+ tabletResultsMap := TabletResultMap{
+ alias1: newMetricResultMap(1.2),
+ alias2: newMetricResultMap(1.7),
+ alias3: newMetricResultMap(0.3),
+ alias4: noSuchMetricMap(),
+ alias5: newMetricResultMap(1.1),
+ }
+
+ t.Run("nonexistent", func(t *testing.T) {
+ worstMetric := AggregateTabletMetricResults(nonexistentMetricName, tabletResultsMap, 0, false, 0)
+ _, err := worstMetric.Get()
+ assert.Error(t, err)
+ assert.Equal(t, ErrNoSuchMetric, err)
+ })
+ t.Run("no ignore", func(t *testing.T) {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 0, false, 0)
+ _, err := worstMetric.Get()
+ assert.Error(t, err)
+ assert.Equal(t, ErrNoSuchMetric, err)
+ })
+ t.Run("ignore 1", func(t *testing.T) {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 1, false, 0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, 1.7, value)
+ })
+ t.Run("ignore 2", func(t *testing.T) {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 2, false, 0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, 1.2, value)
+ })
+
+ tabletResultsMap[alias1][DefaultMetricName] = NoSuchMetric
+ t.Run("no such metric", func(t *testing.T) {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 0, false, 0)
+ _, err := worstMetric.Get()
+ assert.Error(t, err)
+ assert.Equal(t, ErrNoSuchMetric, err)
+ })
+ t.Run("no such metric, ignore 1", func(t *testing.T) {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 1, false, 0)
+ _, err := worstMetric.Get()
+ assert.Error(t, err)
+ assert.Equal(t, ErrNoSuchMetric, err)
+ })
+ t.Run("metric found", func(t *testing.T) {
+ worstMetric := AggregateTabletMetricResults(DefaultMetricName, tabletResultsMap, 2, false, 0)
+ value, err := worstMetric.Get()
+ assert.NoError(t, err)
+ assert.Equal(t, value, 1.7)
+ })
+}