domain: record runaway and quarantine query

domain/domain.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"math"
 	"math/rand"
+	"net"
 	"sort"
 	"strconv"
 	"strings"
@@ -1219,6 +1220,7 @@ func (do *Domain) Init(
 	do.wg.Run(do.topNSlowQueryLoop, "topNSlowQueryLoop")
 	do.wg.Run(do.infoSyncerKeeper, "infoSyncerKeeper")
 	do.wg.Run(do.globalConfigSyncerKeeper, "globalConfigSyncerKeeper")
+	do.wg.Run(do.runawayRecordFlushLoop, "runawayRecordFlushLoop")
 	if !skipRegisterToDashboard {
 		do.wg.Run(do.topologySyncerKeeper, "topologySyncerKeeper")
 	}
@@ -1246,6 +1248,121 @@ func (do *Domain) SetOnClose(onClose func()) {
 	do.onClose = onClose
 }
 
+func (do *Domain) runawayRecordFlushLoop() {
+	defer util.Recover(metrics.LabelDomain, "runawayRecordFlushLoop", nil, false)
+	// this times is used to batch flushing rocords, with 1s duration,
+	// we can guarantee a watch record can be seen by the user within 1s.
+	timer := time.NewTimer(time.Second)
+	fired := false
+	recordCh := do.RunawayManager().RunawayRecordChan()
+	quarantineRecordCh := do.RunawayManager().QuarantineRecordChan()
+	flushThrehold := do.runawayManager.FlushThreshold()
+	records := make([]*resourcegroup.RunawayRecord, 0, flushThrehold)
+	quarantineRecords := make([]*resourcegroup.QuarantineRecord, 0)
+
+	flushRunawayRecords := func() {
+		if len(records) == 0 {
+			return
+		}
+		sql, params := genRunawayQueriesStmt(records)
+		if err := do.execFlushSQL(sql, params); err != nil {
+			logutil.BgLogger().Info("flush runaway records failed", zap.Error(err), zap.Int("count", len(records)))
+		}
+		records = records[:0]
+	}
+	flushQuarantineRecords := func() {
+		if len(quarantineRecords) == 0 {
+			return
+		}
+		sql, params := genQuarantineQueriesStmt(quarantineRecords)
+		if err := do.execFlushSQL(sql, params); err != nil {
+			logutil.BgLogger().Info("flush quarantine records failed", zap.Error(err), zap.Int("count", len(records)))
+		}
+		quarantineRecords = quarantineRecords[:0]
+	}
+
+	for {
+		select {
+		case <-do.exit:
+			return
+		case <-timer.C:
+			flushRunawayRecords()
+			fired = true
+		case r := <-quarantineRecordCh:
+			quarantineRecords = append(quarantineRecords, r)
+			// we expect quarantine record should not be triggered very often, so always
+			// flush as soon as possible.
+			if len(quarantineRecordCh) == 0 || len(quarantineRecords) >= flushThrehold {
+				flushQuarantineRecords()
+			}
+		case r := <-recordCh:
+			records = append(records, r)
+			if len(records) >= flushThrehold {
+				flushRunawayRecords()
+			} else if fired {
+				fired = false
+				// meet a new record, reset the timer.
+				timer.Reset(time.Second)
+			}
+		}
+	}
+}
+
+func (do *Domain) execFlushSQL(sql string, params []interface{}) error {
+	se, err := do.sysSessionPool.Get()
+	defer func() {
+		do.sysSessionPool.Put(se)
+	}()
+	if err != nil {
+		return errors.Annotate(err, "get session failed")
+	}
+	exec := se.(sqlexec.RestrictedSQLExecutor)
+	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)
+	_, _, err = exec.ExecRestrictedSQL(ctx, []sqlexec.OptionFuncAlias{sqlexec.ExecOptionUseCurSession},
+		sql, params...,
+	)
+	return err
+}
+
+func genRunawayQueriesStmt(records []*resourcegroup.RunawayRecord) (string, []interface{}) {
+	var builder strings.Builder
+	params := make([]interface{}, 0, len(records)*7)
+	builder.WriteString("insert into mysql.tidb_runaway_queries VALUES ")
+	for count, r := range records {
+		if count > 0 {
+			builder.WriteByte(',')
+		}
+		builder.WriteString("(%?, %?, %?, %?, %?, %?, %?)")
+		params = append(params, r.ResourceGroupName)
+		params = append(params, r.Time)
+		params = append(params, r.Match)
+		params = append(params, r.Action)
+		params = append(params, r.SQLText)
+		params = append(params, r.PlanDigest)
+		params = append(params, r.From)
+	}
+	return builder.String(), params
+}
+
+func genQuarantineQueriesStmt(records []*resourcegroup.QuarantineRecord) (string, []interface{}) {
+	var builder strings.Builder
+	params := make([]interface{}, 0, len(records)*7)
+	builder.WriteString("insert into mysql.tidb_runaway_quarantined_watch VALUES ")
+	for count, r := range records {
+		if count > 0 {
+			builder.WriteByte(',')
+		}
+		builder.WriteString("(%?, %?, %?, %?, %?, %?)")
+		params = append(params, r.ResourceGroupName)
+		params = append(params, r.StartTime)
+		params = append(params, r.EndTime)
+		params = append(params, r.Watch)
+		params = append(params, r.WatchText)
+		params = append(params, r.From)
+	}
+	return builder.String(), params
+}
+
 func (do *Domain) initResourceGroupsController(ctx context.Context, pdClient pd.Client) error {
 	if pdClient == nil {
 		logutil.BgLogger().Warn("cannot setup up resource controller, not using tikv storage")
@@ -1258,7 +1375,12 @@ func (do *Domain) initResourceGroupsController(ctx context.Context, pdClient pd.
 		return err
 	}
 	control.Start(ctx)
-	do.runawayManager = resourcegroup.NewRunawayManager(control)
+	serverInfo, err := infosync.GetServerInfo()
+	if err != nil {
+		return err
+	}
+	serverAddr := net.JoinHostPort(serverInfo.IP, strconv.Itoa(int(serverInfo.Port)))
+	do.runawayManager = resourcegroup.NewRunawayManager(control, serverAddr)
 	do.resourceGroupsController = control
 	return nil
 }

domain/resourcegroup/runaway.go

@@ -15,6 +15,8 @@
 package resourcegroup
 
 import (
+	"strings"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -32,22 +34,73 @@ const (
 	// DefaultResourceGroupName is the default resource group name.
 	DefaultResourceGroupName = "default"
 	// MaxWaitDuration is the max duration to wait for acquiring token buckets.
-	MaxWaitDuration = time.Second * 30
+	MaxWaitDuration           = time.Second * 30
+	maxWatchListCap           = 10000
+	maxWatchRecordChannelSize = 1024
 )
 
+// RunawayMatchType is used to indicates whether qurey was interrupted by runaway identification or quarantine watch.
+type RunawayMatchType uint
+
+const (
+	// RunawayMatchTypeWatch shows quarantine watch.
+	RunawayMatchTypeWatch RunawayMatchType = iota
+	// RunawayMatchTypeIdentify shows identification.
+	RunawayMatchTypeIdentify
+)
+
+func (t RunawayMatchType) String() string {
+	switch t {
+	case RunawayMatchTypeWatch:
+		return "watch"
+	case RunawayMatchTypeIdentify:
+		return "identify"
+	default:
+		panic("unknown type")
+	}
+}
+
+// RunawayRecord is used to save records which will be insert into mysql.tidb_runaway_queries.
+type RunawayRecord struct {
+	ResourceGroupName string
+	Time              time.Time
+	Match             string
+	Action            string
+	SQLText           string
+	PlanDigest        string
+	From              string
+}
+
+// QuarantineRecord is used to save records which will be insert into mysql.tidb_runaway_quarantined_watch.
+type QuarantineRecord struct {
+	ResourceGroupName string
+	StartTime         time.Time
+	EndTime           time.Time
+	Watch             string
+	WatchText         string
+	From              string
+}
+
 // RunawayManager is used to detect and record runaway queries.
 type RunawayManager struct {
-	resourceGroupCtl *rmclient.ResourceGroupsController
-	watchList        *ttlcache.Cache[string, struct{}]
+	queryLock          sync.Mutex
+	resourceGroupCtl   *rmclient.ResourceGroupsController
+	watchList          *ttlcache.Cache[string, struct{}]
+	serverID           string
+	runawayQueriesChan chan *RunawayRecord
+	quarantineChan     chan *QuarantineRecord
 }
 
 // NewRunawayManager creates a new RunawayManager.
-func NewRunawayManager(resourceGroupCtl *rmclient.ResourceGroupsController) *RunawayManager {
-	watchList := ttlcache.New[string, struct{}]()
+func NewRunawayManager(resourceGroupCtl *rmclient.ResourceGroupsController, serverAddr string) *RunawayManager {
+	watchList := ttlcache.New[string, struct{}](ttlcache.WithCapacity[string, struct{}](maxWatchListCap))
 	go watchList.Start()
 	return &RunawayManager{
-		resourceGroupCtl: resourceGroupCtl,
-		watchList:        watchList,
+		resourceGroupCtl:   resourceGroupCtl,
+		watchList:          watchList,
+		serverID:           serverAddr,
+		runawayQueriesChan: make(chan *RunawayRecord, maxWatchRecordChannelSize),
+		quarantineChan:     make(chan *QuarantineRecord, maxWatchRecordChannelSize),
 	}
 }
 
@@ -61,17 +114,64 @@ func (rm *RunawayManager) DeriveChecker(resourceGroupName string, originalSQL st
 	if group.RunawaySettings == nil {
 		return nil
 	}
-
 	return newRunawayChecker(rm, resourceGroupName, group.RunawaySettings, originalSQL, planDigest)
 }
 
-// MarkRunaway marks the query as runaway.
-func (rm *RunawayManager) MarkRunaway(resourceGroupName string, convict string, ttl time.Duration) {
-	rm.watchList.Set(resourceGroupName+"/"+convict, struct{}{}, ttl)
+func (rm *RunawayManager) markQuarantine(resourceGroupName, convict, watchType string, ttl time.Duration, action string, now *time.Time) {
+	key := resourceGroupName + "/" + convict
+	if rm.watchList.Get(key) == nil {
+		rm.queryLock.Lock()
+		if rm.watchList.Get(key) == nil {
+			rm.watchList.Set(key, struct{}{}, ttl)
+		}
+		rm.queryLock.Unlock()
+	}
+	select {
+	case rm.quarantineChan <- &QuarantineRecord{
+		ResourceGroupName: resourceGroupName,
+		StartTime:         *now,
+		EndTime:           now.Add(ttl),
+		Watch:             watchType,
+		WatchText:         convict,
+		From:              rm.serverID,
+	}:
+	default:
+		// TODO: add warning for discard flush records
+	}
+}
+func (rm *RunawayManager) markRunaway(resourceGroupName, originalSQL, planDigest string, action string, matchType RunawayMatchType, now *time.Time) {
+	select {
+	case rm.runawayQueriesChan <- &RunawayRecord{
+		ResourceGroupName: resourceGroupName,
+		Time:              *now,
+		Match:             matchType.String(),
+		Action:            action,
+		SQLText:           originalSQL,
+		PlanDigest:        planDigest,
+		From:              rm.serverID,
+	}:
+	default:
+		// TODO: add warning for discard flush records
+	}
+}
+
+// FlushThreshold specifies the threshold for the number of records in trigger flush
+func (rm *RunawayManager) FlushThreshold() int {
+	return maxWatchRecordChannelSize / 2
 }
 
-// ExamineWatchList check whether the query is in watch list.
-func (rm *RunawayManager) ExamineWatchList(resourceGroupName string, convict string) bool {
+// RunawayRecordChan returns the channel of RunawayRecord
+func (rm *RunawayManager) RunawayRecordChan() <-chan *RunawayRecord {
+	return rm.runawayQueriesChan
+}
+
+// QuarantineRecordChan returns the channel of QuarantineRecord
+func (rm *RunawayManager) QuarantineRecordChan() <-chan *QuarantineRecord {
+	return rm.quarantineChan
+}
+
+// examineWatchList check whether the query is in watch list.
+func (rm *RunawayManager) examineWatchList(resourceGroupName string, convict string) bool {
 	return rm.watchList.Get(resourceGroupName+"/"+convict) != nil
 }
 
@@ -91,6 +191,7 @@ type RunawayChecker struct {
 
 	deadline time.Time
 	setting  *rmpb.RunawaySettings
+	action   string
 
 	marked atomic.Bool
 }
@@ -104,6 +205,7 @@ func newRunawayChecker(manager *RunawayManager, resourceGroupName string, settin
 		deadline:          time.Now().Add(time.Duration(setting.Rule.ExecElapsedTimeMs) * time.Millisecond),
 		setting:           setting,
 		marked:            atomic.Bool{},
+		action:            strings.ToLower(setting.Action.String()),
 	}
 }
 
@@ -112,9 +214,13 @@ func (r *RunawayChecker) BeforeExecutor() error {
 	if r == nil {
 		return nil
 	}
-	result := r.manager.ExamineWatchList(r.resourceGroupName, r.getConvictIdentifier())
+	result := r.manager.examineWatchList(r.resourceGroupName, r.getConvictIdentifier())
 	if result {
 		r.marked.Store(result)
+		if result {
+			now := time.Now()
+			r.markRunaway(RunawayMatchTypeWatch, &now)
+		}
 		switch r.setting.Action {
 		case rmpb.RunawayAction_Kill:
 			return exeerrors.ErrResourceGroupQueryRunawayQuarantine
@@ -144,7 +250,9 @@ func (r *RunawayChecker) BeforeCopRequest(req *tikvrpc.Request) error {
 		}
 		// execution time exceeds the threshold, mark the query as runaway
 		if r.marked.CompareAndSwap(false, true) {
-			r.markRunaway()
+			now := time.Now()
+			r.markRunaway(RunawayMatchTypeIdentify, &now)
+			r.markQuarantine(&now)
 		}
 	}
 	switch r.setting.Action {
@@ -166,16 +274,25 @@ func (r *RunawayChecker) AfterCopRequest() {
 	// Here only marks the query as runaway
 	if !r.marked.Load() && r.deadline.Before(time.Now()) {
 		if r.marked.CompareAndSwap(false, true) {
-			r.markRunaway()
+			now := time.Now()
+			r.markRunaway(RunawayMatchTypeIdentify, &now)
+			r.markQuarantine(&now)
 		}
 	}
 }
 
-func (r *RunawayChecker) markRunaway() {
+func (r *RunawayChecker) markQuarantine(now *time.Time) {
 	if r.setting.Watch == nil {
 		return
 	}
-	r.manager.MarkRunaway(r.resourceGroupName, r.getConvictIdentifier(), time.Duration(r.setting.Watch.LastingDurationMs)*time.Millisecond)
+	watchType := strings.ToLower(r.setting.Watch.Type.String())
+	ttl := time.Duration(r.setting.Watch.LastingDurationMs) * time.Millisecond
+
+	r.manager.markQuarantine(r.resourceGroupName, r.getConvictIdentifier(), watchType, ttl, r.action, now)
+}
+
+func (r *RunawayChecker) markRunaway(matchType RunawayMatchType, now *time.Time) {
+	r.manager.markRunaway(r.resourceGroupName, r.originalSQL, r.planDigest, r.action, matchType, now)
 }
 
 func (r *RunawayChecker) getConvictIdentifier() string {

executor/infoschema_cluster_table_test.go

@@ -290,7 +290,7 @@ func TestTableStorageStats(t *testing.T) {
 		"test 2",
 	))
 	rows := tk.MustQuery("select TABLE_NAME from information_schema.TABLE_STORAGE_STATS where TABLE_SCHEMA = 'mysql';").Rows()
-	result := 49
+	result := 51
 	require.Len(t, rows, result)
 
 	// More tests about the privileges.