bank1.go

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()
}