253 lines
8.5 KiB
Go
253 lines
8.5 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
_ "github.com/go-sql-driver/mysql"
|
|
)
|
|
|
|
const (
|
|
temporaryMode = "temporary"
|
|
durableMode = "durable"
|
|
temporaryTable = "vehicle_alert_candidate_benchmark"
|
|
durableTable = "vehicle_alert_candidate_benchmark_durable"
|
|
durableLock = "lingniu.vehicle-alert-candidate-benchmark"
|
|
)
|
|
|
|
type benchmarkConfig struct {
|
|
Rows int
|
|
BatchSize int
|
|
Mode string
|
|
ConfirmDurable bool
|
|
}
|
|
|
|
type benchmarkResult struct {
|
|
Mode string `json:"mode"`
|
|
Rows int `json:"rows"`
|
|
BatchSize int `json:"batchSize"`
|
|
Batches int `json:"batches"`
|
|
Transactions int `json:"transactions"`
|
|
DurationMS int64 `json:"durationMs"`
|
|
RowsPerSec float64 `json:"rowsPerSec"`
|
|
VerifiedRows int `json:"verifiedRows"`
|
|
Temporary bool `json:"temporaryTable"`
|
|
DurableWrite bool `json:"durableWrite"`
|
|
CleanupMS int64 `json:"cleanupMs"`
|
|
CleanupVerified bool `json:"cleanupVerified"`
|
|
GlobalStatusDelta map[string]uint64 `json:"globalStatusDelta,omitempty"`
|
|
GlobalStatusAvailable bool `json:"globalStatusAvailable"`
|
|
}
|
|
|
|
func main() {
|
|
if err := run(); err != nil {
|
|
fmt.Fprintln(os.Stderr, "alert benchmark failed:", err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
func run() error {
|
|
config := benchmarkConfig{}
|
|
flag.IntVar(&config.Rows, "rows", 10_000, "candidate rows to write")
|
|
flag.IntVar(&config.BatchSize, "batch-size", 500, "rows per insert")
|
|
flag.StringVar(&config.Mode, "mode", temporaryMode, "storage mode: temporary or durable")
|
|
flag.BoolVar(&config.ConfirmDurable, "confirm-durable-write", false, "required acknowledgement for durable writes and DDL")
|
|
flag.Parse()
|
|
if err := config.validate(); err != nil {
|
|
return err
|
|
}
|
|
dsn := strings.TrimSpace(os.Getenv("MYSQL_DSN"))
|
|
if dsn == "" {
|
|
return fmt.Errorf("MYSQL_DSN is required")
|
|
}
|
|
db, err := sql.Open("mysql", dsn)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer db.Close()
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
|
defer cancel()
|
|
conn, err := db.Conn(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer conn.Close()
|
|
result, err := executeBenchmark(ctx, conn, config)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return json.NewEncoder(os.Stdout).Encode(result)
|
|
}
|
|
|
|
func (config benchmarkConfig) validate() error {
|
|
config.Mode = strings.ToLower(strings.TrimSpace(config.Mode))
|
|
if config.Rows < 1 || config.Rows > 100_000 || config.BatchSize < 1 || config.BatchSize > 1_000 {
|
|
return fmt.Errorf("rows must be 1..100000 and batch-size 1..1000")
|
|
}
|
|
if config.Mode != temporaryMode && config.Mode != durableMode {
|
|
return fmt.Errorf("mode must be temporary or durable")
|
|
}
|
|
if config.Mode == durableMode && !config.ConfirmDurable {
|
|
return fmt.Errorf("durable mode creates, commits to, and drops a dedicated physical table; pass --confirm-durable-write")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func executeBenchmark(ctx context.Context, conn *sql.Conn, config benchmarkConfig) (result benchmarkResult, err error) {
|
|
mode := strings.ToLower(strings.TrimSpace(config.Mode))
|
|
table := temporaryTable
|
|
createPrefix := "CREATE TEMPORARY TABLE"
|
|
dropPrefix := "DROP TEMPORARY TABLE IF EXISTS"
|
|
if mode == durableMode {
|
|
table = durableTable
|
|
createPrefix = "CREATE TABLE"
|
|
dropPrefix = "DROP TABLE IF EXISTS"
|
|
locked, lockErr := acquireBenchmarkLock(ctx, conn)
|
|
if lockErr != nil {
|
|
return result, fmt.Errorf("acquire durable benchmark lock: %w", lockErr)
|
|
}
|
|
if !locked {
|
|
return result, fmt.Errorf("another durable benchmark is already running")
|
|
}
|
|
defer releaseBenchmarkLock(conn)
|
|
if _, err = conn.ExecContext(ctx, dropPrefix+" "+table); err != nil {
|
|
return result, fmt.Errorf("remove stale durable benchmark table: %w", err)
|
|
}
|
|
}
|
|
if _, err = conn.ExecContext(ctx, createPrefix+" "+table+" LIKE vehicle_alert_candidate"); err != nil {
|
|
return result, fmt.Errorf("create %s benchmark table: %w", mode, err)
|
|
}
|
|
cleanupStarted := time.Time{}
|
|
cleanupComplete := false
|
|
defer func() {
|
|
cleanupStarted = time.Now()
|
|
cleanupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
_, cleanupErr := conn.ExecContext(cleanupCtx, dropPrefix+" "+table)
|
|
if cleanupErr == nil {
|
|
cleanupComplete = true
|
|
if mode == durableMode {
|
|
var remaining int
|
|
cleanupErr = conn.QueryRowContext(cleanupCtx, `SELECT COUNT(*) FROM information_schema.tables WHERE table_schema=DATABASE() AND table_name=?`, table).Scan(&remaining)
|
|
cleanupComplete = cleanupErr == nil && remaining == 0
|
|
}
|
|
}
|
|
result.CleanupMS = time.Since(cleanupStarted).Milliseconds()
|
|
result.CleanupVerified = cleanupComplete
|
|
if err == nil && cleanupErr != nil {
|
|
err = fmt.Errorf("cleanup benchmark table: %w", cleanupErr)
|
|
} else if err == nil && !cleanupComplete {
|
|
err = fmt.Errorf("benchmark table cleanup could not be verified")
|
|
}
|
|
}()
|
|
|
|
statusBefore, statusBeforeErr := readGlobalStatus(ctx, conn)
|
|
started := time.Now()
|
|
tx, err := conn.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return result, err
|
|
}
|
|
now := time.Now().UTC()
|
|
batches := 0
|
|
for start := 0; start < config.Rows; start += config.BatchSize {
|
|
end := min(start+config.BatchSize, config.Rows)
|
|
query, args := buildCandidateInsert(table, start, end, now)
|
|
if _, err = tx.ExecContext(ctx, query, args...); err != nil {
|
|
tx.Rollback()
|
|
return result, fmt.Errorf("insert batch %d: %w", batches+1, err)
|
|
}
|
|
batches++
|
|
}
|
|
if err = tx.Commit(); err != nil {
|
|
return result, err
|
|
}
|
|
duration := time.Since(started)
|
|
var verified int
|
|
if err = conn.QueryRowContext(ctx, `SELECT COUNT(*) FROM `+table).Scan(&verified); err != nil {
|
|
return result, err
|
|
}
|
|
if verified != config.Rows {
|
|
return result, fmt.Errorf("row verification failed: wrote %d, found %d", config.Rows, verified)
|
|
}
|
|
statusAfter, statusAfterErr := readGlobalStatus(ctx, conn)
|
|
statusDelta := map[string]uint64{}
|
|
statusAvailable := statusBeforeErr == nil && statusAfterErr == nil
|
|
if statusAvailable {
|
|
for key, after := range statusAfter {
|
|
if before, ok := statusBefore[key]; ok && after >= before {
|
|
statusDelta[key] = after - before
|
|
}
|
|
}
|
|
}
|
|
result = benchmarkResult{
|
|
Mode: mode,
|
|
Rows: config.Rows,
|
|
BatchSize: config.BatchSize,
|
|
Batches: batches,
|
|
Transactions: 1,
|
|
DurationMS: duration.Milliseconds(),
|
|
RowsPerSec: float64(config.Rows) / duration.Seconds(),
|
|
VerifiedRows: verified,
|
|
Temporary: mode == temporaryMode,
|
|
DurableWrite: mode == durableMode,
|
|
GlobalStatusDelta: statusDelta,
|
|
GlobalStatusAvailable: statusAvailable,
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func acquireBenchmarkLock(ctx context.Context, conn *sql.Conn) (bool, error) {
|
|
var acquired sql.NullInt64
|
|
if err := conn.QueryRowContext(ctx, `SELECT GET_LOCK(?,0)`, durableLock).Scan(&acquired); err != nil {
|
|
return false, err
|
|
}
|
|
return acquired.Valid && acquired.Int64 == 1, nil
|
|
}
|
|
|
|
func releaseBenchmarkLock(conn *sql.Conn) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
_, _ = conn.ExecContext(ctx, `SELECT RELEASE_LOCK(?)`, durableLock)
|
|
}
|
|
|
|
func readGlobalStatus(ctx context.Context, conn *sql.Conn) (map[string]uint64, error) {
|
|
rows, err := conn.QueryContext(ctx, `SHOW GLOBAL STATUS WHERE Variable_name IN ('Innodb_os_log_written','Binlog_cache_use','Binlog_cache_disk_use','Com_commit')`)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
values := map[string]uint64{}
|
|
for rows.Next() {
|
|
var key, raw string
|
|
if err := rows.Scan(&key, &raw); err != nil {
|
|
return nil, err
|
|
}
|
|
value, parseErr := strconv.ParseUint(raw, 10, 64)
|
|
if parseErr != nil {
|
|
return nil, parseErr
|
|
}
|
|
values[key] = value
|
|
}
|
|
return values, rows.Err()
|
|
}
|
|
|
|
func buildCandidateInsert(table string, start, end int, now time.Time) (string, []any) {
|
|
if table != temporaryTable && table != durableTable {
|
|
panic("unsupported benchmark table")
|
|
}
|
|
values := make([]string, 0, end-start)
|
|
args := make([]any, 0, (end-start)*7)
|
|
for index := start; index < end; index++ {
|
|
values = append(values, "(?,?,?,?,?,?,?)")
|
|
args = append(args, "benchmark-rule", fmt.Sprintf("BENCH%017d", index), "JT808", now, now, float64(index%120), fmt.Sprintf("benchmark-event-%d", index))
|
|
}
|
|
return `INSERT INTO ` + table + `(rule_id,vin,protocol,first_matched_at,last_matched_at,latest_value,source_event_id) VALUES ` + strings.Join(values, ","), args
|
|
}
|