*: add tidb_slow_log_max_per_sec variable to control the number of slow logs written per second (#63996)

close #63995

Author: zimulala

pkg/executor/adapter.go

@@ -1713,6 +1713,11 @@ func (a *ExecStmt) LogSlowQuery(txnTS uint64, succ bool, hasMoreResults bool) {
 		}
 	}
 
+	if !vardef.GlobalSlowLogRateLimiter.Allow() {
+		sampleLoggerFactory().Info("slow log skipped due to rate limiting", zap.Int64("tidb_slow_log_max_per_sec", int64(vardef.GlobalSlowLogRateLimiter.Limit())))
+		return
+	}
+
 	if slowItems == nil {
 		slowItems = &variable.SlowQueryLogItems{}
 	}

pkg/executor/adapter_slow_log.go

@@ -25,7 +25,9 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/sessionctx/slowlogrule"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
+	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/tikv/client-go/v2/util"
+	"go.uber.org/zap"
 )
 
 func init() {
@@ -41,6 +43,8 @@ func init() {
 	}
 }
 
+var sampleLoggerFactory = logutil.SampleLoggerFactory(time.Minute, 1, zap.String(logutil.LogFieldCategory, "slow log"))
+
 func mergeConditionFields(dst, src map[string]struct{}) {
 	for k := range src {
 		dst[k] = struct{}{}

pkg/executor/adapter_test.go

@@ -384,6 +384,61 @@ func TestWriteSlowLog(t *testing.T) {
 	checkWriteSlowLog(true)
 }
 
+func TestSlowLogMaxPerSec(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+
+	// default value
+	tk.MustQuery(`show variables like "tidb_slow_log_max_per_sec"`).Check(
+		testkit.Rows("tidb_slow_log_max_per_sec 0"),
+	)
+	_, err := tk.Exec(`select @@SESSION.tidb_slow_log_max_per_sec`)
+	require.Equal(t, "[variable:1238]Variable 'tidb_slow_log_max_per_sec' is a GLOBAL variable", err.Error())
+	tk.MustQuery(`select @@Global.tidb_slow_log_max_per_sec`).Check(
+		testkit.Rows("0"),
+	)
+
+	// test errors
+	_, err = tk.Exec(`set session tidb_slow_log_max_per_sec="0"`)
+	require.Equal(t, "[variable:1229]Variable 'tidb_slow_log_max_per_sec' is a GLOBAL variable and should be set with SET GLOBAL", err.Error())
+	_, err = tk.Exec(`set global tidb_slow_log_max_per_sec=""`)
+	require.Equal(t, "[variable:1232]Incorrect argument type to variable 'tidb_slow_log_max_per_sec'", err.Error())
+	_, err = tk.Exec(`set global tidb_slow_log_max_per_sec="1.23"`)
+	require.Equal(t, "[variable:1232]Incorrect argument type to variable 'tidb_slow_log_max_per_sec'", err.Error())
+
+	// test warnings
+	_, err = tk.Exec(`set global tidb_slow_log_max_per_sec="-1"`)
+	tk.MustQuery("SHOW WARNINGS").Check(testkit.Rows("Warning 1292 Truncated incorrect tidb_slow_log_max_per_sec value: '-1'"))
+	tk.MustQuery(`select @@Global.tidb_slow_log_max_per_sec`).Check(
+		testkit.Rows("0"),
+	)
+	tk.MustExec(`set global tidb_slow_log_max_per_sec="1234567"`)
+	tk.MustQuery("SHOW WARNINGS").Check(testkit.Rows("Warning 1292 Truncated incorrect tidb_slow_log_max_per_sec value: '1234567'"))
+	tk.MustQuery(`show variables like "tidb_slow_log_max_per_sec"`).Check(
+		testkit.Rows("tidb_slow_log_max_per_sec 1000000"),
+	)
+
+	// normal
+	tk.MustExec(`set global tidb_slow_log_max_per_sec="2"`)
+	require.True(t, vardef.GlobalSlowLogRateLimiter.Allow())
+	require.True(t, vardef.GlobalSlowLogRateLimiter.Allow())
+	require.False(t, vardef.GlobalSlowLogRateLimiter.Allow())
+	tk.MustQuery(`show variables like "tidb_slow_log_max_per_sec"`).Check(
+		testkit.Rows("tidb_slow_log_max_per_sec 2"),
+	)
+	tk.MustQuery(`select @@Global.tidb_slow_log_max_per_sec`).Check(
+		testkit.Rows("2"),
+	)
+	// no limit
+	tk.MustExec(`set global tidb_slow_log_max_per_sec="0"`)
+	require.True(t, vardef.GlobalSlowLogRateLimiter.Allow())
+	require.True(t, vardef.GlobalSlowLogRateLimiter.Allow())
+	require.True(t, vardef.GlobalSlowLogRateLimiter.Allow())
+	tk.MustQuery(`show variables like "tidb_slow_log_max_per_sec"`).Check(
+		testkit.Rows("tidb_slow_log_max_per_sec 0"),
+	)
+}
+
 func BenchmarkCheckSlowThreshold(b *testing.B) {
 	b.StopTimer()
 	b.ReportAllocs()

pkg/sessionctx/vardef/BUILD.bazel

@@ -19,6 +19,7 @@ go_library(
         "//pkg/util/paging",
         "//pkg/util/size",
         "@com_github_pingcap_tipb//go-tipb",
+        "@org_golang_x_time//rate",
         "@org_uber_go_atomic//:atomic",
     ],
 )
@@ -36,5 +37,6 @@ go_test(
     deps = [
         "//pkg/config/kerneltype",
         "@com_github_stretchr_testify//require",
+        "@org_golang_x_time//rate",
     ],
 )

pkg/sessionctx/vardef/tidb_vars.go

@@ -32,6 +32,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/size"
 	"github.com/pingcap/tipb/go-tipb"
 	"go.uber.org/atomic"
+	"golang.org/x/time/rate"
 )
 
 /*
@@ -234,6 +235,10 @@ const (
 	// TiDBSlowLogRules defines multi-dimensional trigger rules for flexible slow log control.
 	TiDBSlowLogRules = "tidb_slow_log_rules"
 
+	// TiDBSlowLogMaxPerSec is the maximum number of slow logs that can be recorded per second in the server.
+	// The default value is 0, which means no rate limiting is applied.
+	TiDBSlowLogMaxPerSec = "tidb_slow_log_max_per_sec"
+
 	// TiDBSlowTxnLogThreshold is used to set the slow transaction log threshold in the server.
 	TiDBSlowTxnLogThreshold = "tidb_slow_txn_log_threshold"
 
@@ -1759,19 +1764,20 @@ var (
 	//    the value of `tidb_analyze_column_options` determines the behavior of the analyze operation.
 	// 2. If `tidb_persist_analyze_options` is disabled, `tidb_analyze_column_options` is used directly to decide
 	//    whether to analyze all columns or just the predicate columns.
-	AnalyzeColumnOptions          = atomic.NewString(DefTiDBAnalyzeColumnOptions)
-	GlobalLogMaxDays              = atomic.NewInt32(int32(config.GetGlobalConfig().Log.File.MaxDays))
-	QueryLogMaxLen                = atomic.NewInt32(DefTiDBQueryLogMaxLen)
-	EnablePProfSQLCPU             = atomic.NewBool(false)
-	EnableBatchDML                = atomic.NewBool(false)
-	EnableTmpStorageOnOOM         = atomic.NewBool(DefTiDBEnableTmpStorageOnOOM)
-	DDLReorgWorkerCounter   int32 = DefTiDBDDLReorgWorkerCount
-	DDLReorgBatchSize       int32 = DefTiDBDDLReorgBatchSize
-	DDLFlashbackConcurrency int32 = DefTiDBDDLFlashbackConcurrency
-	DDLErrorCountLimit      int64 = DefTiDBDDLErrorCountLimit
-	DDLReorgRowFormat       int64 = DefTiDBRowFormatV2
-	DDLReorgMaxWriteSpeed         = atomic.NewInt64(DefTiDBDDLReorgMaxWriteSpeed)
-	MaxDeltaSchemaCount     int64 = DefTiDBMaxDeltaSchemaCount
+	AnalyzeColumnOptions           = atomic.NewString(DefTiDBAnalyzeColumnOptions)
+	GlobalLogMaxDays               = atomic.NewInt32(int32(config.GetGlobalConfig().Log.File.MaxDays))
+	QueryLogMaxLen                 = atomic.NewInt32(DefTiDBQueryLogMaxLen)
+	EnablePProfSQLCPU              = atomic.NewBool(false)
+	EnableBatchDML                 = atomic.NewBool(false)
+	EnableTmpStorageOnOOM          = atomic.NewBool(DefTiDBEnableTmpStorageOnOOM)
+	DDLReorgWorkerCounter    int32 = DefTiDBDDLReorgWorkerCount
+	DDLReorgBatchSize        int32 = DefTiDBDDLReorgBatchSize
+	DDLFlashbackConcurrency  int32 = DefTiDBDDLFlashbackConcurrency
+	DDLErrorCountLimit       int64 = DefTiDBDDLErrorCountLimit
+	DDLReorgRowFormat        int64 = DefTiDBRowFormatV2
+	DDLReorgMaxWriteSpeed          = atomic.NewInt64(DefTiDBDDLReorgMaxWriteSpeed)
+	MaxDeltaSchemaCount      int64 = DefTiDBMaxDeltaSchemaCount
+	GlobalSlowLogRateLimiter       = rate.NewLimiter(rate.Inf, 1)
 	// DDLSlowOprThreshold is the threshold for ddl slow operations, uint is millisecond.
 	DDLSlowOprThreshold = config.GetGlobalConfig().Instance.DDLSlowOprThreshold
 	GlobalSlowLogRules  = atomic.NewPointer[slowlogrule.GlobalSlowLogRules](

pkg/sessionctx/vardef/tidb_vars_test.go

@@ -15,10 +15,12 @@
 package vardef
 
 import (
+	"sync"
 	"testing"
 
 	"github.com/pingcap/tidb/pkg/config/kerneltype"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/time/rate"
 )
 
 func TestIsMDLEnabledInNextGen(t *testing.T) {
@@ -36,3 +38,48 @@ func TestIsMDLEnabledInNextGen(t *testing.T) {
 	SetEnableMDL(true)
 	require.True(t, IsMDLEnabled())
 }
+
+func runConcurrentTest(b *testing.B, limiter interface {
+	Allow() bool
+}, goroutines int) {
+	var wg sync.WaitGroup
+	startCh := make(chan struct{})
+
+	cnt := b.N / goroutines
+	for g := 0; g < goroutines; g++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			<-startCh
+			for i := 0; i < cnt; i++ {
+				limiter.Allow()
+			}
+		}()
+	}
+
+	b.ResetTimer()
+	close(startCh)
+	wg.Wait()
+	b.StopTimer()
+}
+
+const limit = 10000
+
+func BenchmarkRateLimiterSimple(b *testing.B) {
+	b.ReportAllocs()
+	rl := rate.NewLimiter(rate.Limit(limit), limit)
+	runConcurrentTest(b, rl, 1)
+}
+
+func BenchmarkRateLimiterCurrency100(b *testing.B) {
+	b.ReportAllocs()
+	rl := rate.NewLimiter(rate.Limit(limit), limit)
+	runConcurrentTest(b, rl, 100)
+}
+
+func BenchmarkRateLimiterCurrency1000(b *testing.B) {
+	b.ReportAllocs()
+	rl := rate.NewLimiter(rate.Limit(limit), limit)
+	runConcurrentTest(b, rl, 1000)
+}

pkg/sessionctx/variable/BUILD.bazel

@@ -90,6 +90,7 @@ go_library(
         "@com_github_tikv_client_go_v2//util",
         "@com_github_twmb_murmur3//:murmur3",
         "@org_golang_x_exp//maps",
+        "@org_golang_x_time//rate",
         "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
     ],

pkg/sessionctx/variable/sysvar.go

@@ -66,6 +66,7 @@ import (
 	"github.com/tikv/client-go/v2/oracle/oracles"
 	tikvcliutil "github.com/tikv/client-go/v2/util"
 	"go.uber.org/zap"
+	"golang.org/x/time/rate"
 )
 
 type concurrencySetter func(s *SessionVars, v int)
@@ -3804,6 +3805,30 @@ var defaultSysVars = []*SysVar{
 			return vardef.GlobalSlowLogRules.Load().RawRules, nil
 		},
 	},
+	{
+		Scope:    vardef.ScopeGlobal,
+		Name:     vardef.TiDBSlowLogMaxPerSec,
+		Value:    "0",
+		Type:     vardef.TypeInt,
+		MinValue: 0, MaxValue: 1000000,
+		SetGlobal: func(_ context.Context, sv *SessionVars, s string) error {
+			d := TidbOptInt(s, 0)
+			if d == int(vardef.GlobalSlowLogRateLimiter.Limit()) {
+				return nil
+			}
+
+			if d == 0 {
+				vardef.GlobalSlowLogRateLimiter.SetLimit(rate.Inf)
+				return nil
+			}
+			vardef.GlobalSlowLogRateLimiter.SetLimit(rate.Limit(d))
+			vardef.GlobalSlowLogRateLimiter.SetBurst(d)
+			return nil
+		},
+		GetGlobal: func(ctx context.Context, sv *SessionVars) (string, error) {
+			return strconv.Itoa(int(vardef.GlobalSlowLogRateLimiter.Limit())), nil
+		},
+	},
 }
 
 // GlobalSystemVariableInitialValue gets the default value for a system variable including ones that are dynamically set (e.g. based on the store)