Throttler code: some refactoring and cleanup (vitessio#16368)

Signed-off-by: Shlomi Noach <[email protected]>
Signed-off-by: Florent Poinsard <[email protected]>

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)
+	})
+}