《后端架构设计·高性能与高可用》11.1 索引、查询与慢 SQL

// file: internal/perf/slowquery/slowquery.go
// tier: L1
// verify: go test ./internal/perf/slowquery/ -count=1 -v
package slowquery

import (
	"encoding/csv"
	"errors"
	"io"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

// Entry mirrors the five pg_stat_statements columns this heuristic needs.
// Rows is the total across calls, not rows per call.
type Entry struct {
	QueryID    int64
	Calls      int64
	MeanExecMs float64
	Rows       int64
	Query      string
}

func (e Entry) TotalMs() float64 { return float64(e.Calls) * e.MeanExecMs }

func (e Entry) RowsPerCall() float64 {
	if e.Calls == 0 {
		return 0
	}
	return float64(e.Rows) / float64(e.Calls)
}

// Parse reads tab-separated, pg_stat_statements-like rows. A leading # marks
// a fixture comment; quoted fields are accepted.
func Parse(r io.Reader) ([]Entry, error) {
	cr := csv.NewReader(r)
	cr.Comma, cr.LazyQuotes, cr.FieldsPerRecord = '\t', true, -1
	var out []Entry
	header := false
	for {
		rec, err := cr.Read()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return nil, err
		}
		if len(rec) == 0 || strings.HasPrefix(strings.TrimSpace(rec[0]), "#") {
			continue
		}
		if !header {
			header = len(rec) >= 5 && strings.TrimSpace(rec[0]) == "queryid"
			if !header {
				return nil, errors.New("slowquery: bad header")
			}
			continue
		}
		if len(rec) < 5 {
			return nil, errors.New("slowquery: short row")
		}
		id, e1 := strconv.ParseInt(strings.TrimSpace(rec[0]), 10, 64)
		calls, e2 := strconv.ParseInt(strings.TrimSpace(rec[1]), 10, 64)
		mean, e3 := strconv.ParseFloat(strings.TrimSpace(rec[2]), 64)
		rows, e4 := strconv.ParseInt(strings.TrimSpace(rec[3]), 10, 64)
		if err := errors.Join(e1, e2, e3, e4); err != nil {
			return nil, err
		}
		out = append(out, Entry{id, calls, mean, rows, strings.TrimSpace(rec[4])})
	}
	if !header {
		return nil, errors.New("slowquery: empty input")
	}
	return out, nil
}

// Rank orders by total execution time, not mean: 3ms x 40k calls costs more
// than 150ms x 800 calls.
func Rank(entries []Entry) []Entry {
	out := append([]Entry(nil), entries...)
	sort.SliceStable(out, func(i, j int) bool {
		ti, tj := out[i].TotalMs(), out[j].TotalMs()
		if ti == tj {
			return out[i].QueryID < out[j].QueryID
		}
		return ti > tj
	})
	return out
}

// ShapeOf is a heuristic over normalized SQL, not a parser: OR, BETWEEN and
// CTEs are out of scope and need human review.
var (
	reTable = regexp.MustCompile(`(?i)\bfrom\s+([a-z_][a-z0-9_]*)`)
	reWhere = regexp.MustCompile(`(?i)\bwhere\s+(.*?)(?:\bgroup\s+by\b|\border\s+by\b|\blimit\b|;|$)`)
	reOrder = regexp.MustCompile(`(?i)\border\s+by\s+([^;]*?)(?:\blimit\b|;|$)`)
	reLimit = regexp.MustCompile(`(?i)\blimit\s+(\d+)`)
	reAnd   = regexp.MustCompile(`(?i)\s+and\s+`)
	rePred  = regexp.MustCompile(`(?i)^\s*(?:[a-z_][a-z0-9_]*\.)?([a-z_][a-z0-9_]*)\s*(=|>=|<=|<>|!=|>|<|\bin\b)`)
)

type Shape struct {
	Table string
	Equal []string
	Range []string
	Sort  []string
	Limit int64
}

func ShapeOf(query string) (Shape, error) {
	var s Shape
	if m := reTable.FindStringSubmatch(query); m != nil {
		s.Table = strings.ToLower(m[1])
	}
	if s.Table == "" {
		return s, errors.New("slowquery: no FROM table")
	}
	if m := reLimit.FindStringSubmatch(query); m != nil {
		s.Limit, _ = strconv.ParseInt(m[1], 10, 64)
	}
	if m := reWhere.FindStringSubmatch(query); m != nil {
		for _, pred := range reAnd.Split(m[1], -1) {
			pm := rePred.FindStringSubmatch(pred)
			if pm == nil {
				continue
			}
			col, op := strings.ToLower(pm[1]), strings.ToLower(pm[2])
			if op == "=" || op == "in" {
				s.Equal = addCol(s.Equal, col)
			} else {
				s.Range = addCol(s.Range, col)
			}
		}
	}
	if m := reOrder.FindStringSubmatch(query); m != nil {
		for _, part := range strings.Split(m[1], ",") {
			if f := strings.Fields(part); len(f) > 0 {
				col := strings.ToLower(f[0])
				if i := strings.LastIndexByte(col, '.'); i >= 0 {
					col = col[i+1:]
				}
				s.Sort = addCol(s.Sort, col)
			}
		}
	}
	return s, nil
}

func addCol(dst []string, c string) []string {
	for _, x := range dst {
		if x == c {
			return dst
		}
	}
	return append(dst, c)
}

func uniq(cols []string) []string {
	seen := make(map[string]bool, len(cols))
	out := cols[:0]
	for _, c := range cols {
		if c != "" && !seen[c] {
			seen[c] = true
			out = append(out, c)
		}
	}
	return out
}

type Candidate struct {
	Columns []string
	Reason  string
}

// Candidates follows equality, then range, then sort. A range column and a
// different sort column cannot both be served by one index, so two
// alternatives are returned.
func Candidates(s Shape) []Candidate {
	eq := append([]string(nil), s.Equal...)
	var out []Candidate
	add := func(cols []string, reason string) {
		cols = uniq(cols)
		if len(cols) == 0 {
			return
		}
		for _, c := range out {
			if strings.Join(c.Columns, ",") == strings.Join(cols, ",") {
				return
			}
		}
		out = append(out, Candidate{cols, reason})
	}
	if len(s.Range) > 0 && len(s.Sort) > 0 && s.Range[0] != s.Sort[0] {
		add(append(append([]string(nil), eq...), s.Range[0]), "equality+range")
		add(append(append([]string(nil), eq...), s.Sort...), "equality+sort")
		return out
	}
	cols, reason := append([]string(nil), eq...), "equality"
	if len(s.Range) > 0 {
		cols = append(cols, s.Range[0])
		reason = "equality+range"
	}
	cols = append(cols, s.Sort...)
	if len(s.Sort) > 0 {
		reason += "+sort"
	}
	add(cols, reason)
	return out
}

// TableStats is fixture input. Rows is the table row count; write/read rates
// are optional and only drive the write-amplification guard.
type TableStats struct {
	Rows         int64
	WritesPerSec float64
	ReadsPerSec  float64
	Existing     [][]string
}

const (
	Recommend = "recommend"
	Covered   = "covered_by_existing"
	Defer     = "defer_write_cost"
	PreferSeq = "prefer_seq_scan"
)

type Assessment struct {
	Columns     []string
	Reason      string
	QueryCount  int
	TotalMs     float64
	Selectivity float64
	Amplify     float64
	Decision    string
}

type Advisor struct{ Stats TableStats }

// Recommend aggregates one column shape across queries, then applies the rules
// in a fixed order so the output is reproducible.
func (a Advisor) Recommend(entries []Entry) []Assessment {
	if a.Stats.Rows <= 0 {
		a.Stats.Rows = 1
	}
	byKey := make(map[string]*Assessment)
	var order []string
	for _, e := range entries {
		s, err := ShapeOf(e.Query)
		if err != nil {
			continue
		}
		sel := e.RowsPerCall() / float64(a.Stats.Rows)
		if s.Limit > 0 && e.RowsPerCall() <= float64(s.Limit) {
			// pg_stat_statements reports post-LIMIT rows: a floor, not an
			// estimate. Real selectivity needs pg_stats or EXPLAIN.
			sel = 0
		}
		for _, c := range Candidates(s) {
			k := strings.Join(c.Columns, ",")
			as := byKey[k]
			if as == nil {
				as = &Assessment{Columns: c.Columns, Reason: c.Reason}
				byKey[k] = as
				order = append(order, k)
			}
			as.QueryCount++
			as.TotalMs += e.TotalMs()
			if sel > as.Selectivity {
				as.Selectivity = sel
			}
		}
	}
	out := make([]Assessment, 0, len(order))
	for _, k := range order {
		as := *byKey[k]
		as.Amplify = float64(len(a.Stats.Existing)+2) / float64(len(a.Stats.Existing)+1)
		as.Decision = decide(as, a.Stats)
		out = append(out, as)
	}
	sort.SliceStable(out, func(i, j int) bool { return out[i].TotalMs > out[j].TotalMs })
	return out
}

func decide(a Assessment, st TableStats) string {
	if a.Selectivity >= 0.20 {
		return PreferSeq
	}
	prefix := strings.Join(a.Columns, ",") + ","
	for _, idx := range st.Existing {
		if strings.HasPrefix(strings.Join(idx, ",")+",", prefix) {
			return Covered
		}
	}
	if a.Amplify > 1.50 && a.TotalMs < 1000 {
		return Defer
	}
	if st.WritesPerSec > 0.25*st.ReadsPerSec && a.QueryCount < 2 {
		return Defer
	}
	return Recommend
}

top queryid=104 total_ms=120000
cols=status reason=equality queries=1 total_ms=120000 sel=0.250000 amp=1.333 decision=prefer_seq_scan
cols=status,created_at reason=equality+range+sort total_ms=102000 sel=0.000200 amp=1.333 decision=recommend
cols=customer_id,created_at reason=equality+sort total_ms=62000 sel=0.000000 amp=1.333 decision=recommend
cols=order_no reason=equality queries=1 total_ms=12000 sel=0.000001 amp=1.333 decision=covered_by_existing
cols=customer_id,status reason=equality queries=1 total_ms=3600 sel=0.000001 amp=1.333 decision=recommend