Skip to content

Instantly share code, notes, and snippets.

@siddontang
Created March 20, 2018 14:06
Show Gist options
  • Save siddontang/f0ba1d29c6b6989e680ed224a18df350 to your computer and use it in GitHub Desktop.
Save siddontang/f0ba1d29c6b6989e680ed224a18df350 to your computer and use it in GitHub Desktop.
package main
import (
"database/sql"
"flag"
"fmt"
"math/rand"
"os"
"sync"
"time"
_ "github.com/go-sql-driver/mysql"
)
var (
concurrent = flag.Int("C", 100, "concurrent")
number = flag.Int("N", 12000000, "number")
)
func perror(err error) {
if err != nil {
fmt.Printf("%s\n", err)
os.Exit(1)
}
}
var letters = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
func randString(r *rand.Rand, n int) string {
b := make([]byte, n)
for i := range b {
b[i] = letters[r.Intn(len(letters))]
}
return string(b)
}
func main() {
db, err := sql.Open("mysql", "root@tcp(172.16.10.1:4000)/test")
perror(err)
db.SetMaxIdleConns(*concurrent + 1)
defer db.Close()
_, err = db.Exec("drop table if exists account")
perror(err)
_, err = db.Exec(
`CREATE TABLE account (
accountId int(20) NOT NULL AUTO_INCREMENT,
name varchar(255) DEFAULT NULL,
balance int(20) NOT NULL,
PRIMARY KEY (accountId),
KEY INDEX_balance (balance)
) ENGINE=InnoDB DEFAULT CHARSET=utf8 COLLATE=utf8_bin;`)
perror(err)
var wg sync.WaitGroup
wg.Add(*concurrent)
for i := 0; i < *concurrent; i++ {
r := rand.New(rand.NewSource(time.Now().UnixNano()))
step := *number / *concurrent
go func(i int) {
defer wg.Done()
for j := 0; j < step; j++ {
n := r.Intn(4) + 6
id := i*step + j + 1
_, err1 := db.Exec("insert ignore into account (accountId, name, balance) values (?, ?, ?)", id, randString(r, n), 10000)
perror(err1)
}
}(i)
}
wg.Wait()
total, err := sumBalance(db)
perror(err)
for i := 0; i < *concurrent; i++ {
go func() {
r := rand.New(rand.NewSource(time.Now().UnixNano()))
for {
if err := doTransaction(db, r); err != nil {
fmt.Printf("do transaction failed %v\n", err)
}
}
}()
}
for {
checkSum(db, total)
}
}
func sumBalance(db *sql.DB) (int64, error) {
var total int64
err := db.QueryRow("select sum(balance) from account where id < 10000000").Scan(&total)
return total, err
}
func checkSum(db *sql.DB, total int64) {
n, err := sumBalance(db)
if err != nil {
fmt.Printf("get total err %v\n", err)
return
}
if n != total {
fmt.Printf("check sum failed, want %d, but got %d", total, n)
os.Exit(1)
}
}
func doTransaction(db *sql.DB, r *rand.Rand) error {
txn, err := db.Begin()
if err != nil {
return err
}
defer txn.Rollback()
_, err = txn.Exec("update account set balance = balance + 100 where accountId = ?", r.Intn(10000000-1)+1)
if err != nil {
return err
}
_, err = txn.Exec("update account set balance = balance - 100 where accountId = ?", r.Intn(10000000-1)+1)
if err != nil {
return err
}
return txn.Commit()
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment