mock-data.go

package main

import (
	"database/sql"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/ngaut/log"
	"golang.org/x/net/context"
)

var (
	host        = flag.String("host", "10.3.3.48", "TiDB host")
	port        = flag.Int("port", 4000, "TiDB port")
	username    = flag.String("username", "root", "TiDB username")
	password    = flag.String("password", "", "TiDB password")
	concurrency = flag.Int("concurrent", 10, "Concurrency")
	numRows     = flag.Int("num-rows", 10000, "Number of Rows")
)

func openDB() (*sql.DB, error) {
	dbDSN := fmt.Sprintf("%s:%s@tcp(%s:%d)/ns_test", *username, *password, *host, *port)
	db, err := sql.Open("mysql", dbDSN)
	if err != nil {
		return nil, err
	}
	db.SetMaxIdleConns(*concurrency)
	return db, nil
}

func mustExec(db *sql.DB, query string, args ...interface{}) sql.Result {
	r, err := db.Exec(query, args...)
	if err != nil {
		log.Fatalf("exec %s err %v", query, err)
	}
	return r
}

func main() {
	flag.Parse()

	db, err := openDB()
	if err != nil {
		log.Fatal(err)
	}

	ctx, cancel := context.WithCancel(context.Background())

	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)

	go func() {
		sig := <-sc
		log.Infof("Got signal [%d] to exit.", sig)

		db.Close()
		cancel()
	}()

	jobChan := make(chan int, *concurrency)

	batchSize := 100
	jobCount := *numRows / batchSize

	var wg sync.WaitGroup
	wg.Add(*concurrency)

	for i := 0; i < *concurrency; i++ {
		go func() {
			defer wg.Done()
			// do job
			for {
				select {
				case <-ctx.Done():
					return
				case startIndex, ok := <-jobChan:
					if !ok {
						return
					}
					args := make([]string, batchSize)

					start := time.Now()

					var i int
					for i = 0; i < batchSize; i++ {
						cur := startIndex + i
						if cur >= *numRows {
							break
						}
						args[i] = fmt.Sprintf("('num', '%d')", startIndex+i)
					}
					// shrink slice to avoid bad sql
					if i < batchSize {
						args = args[:i]
					}
					if len(args) > 0 {
						mustExec(db, fmt.Sprintf("INSERT INTO ns2(name, value) VALUES %s", strings.Join(args, ",")))
						log.Infof("insert %d rows, takes %s", i, time.Now().Sub(start))
					}
				}
			}
		}()
	}

	// add jobs
	go func() {
		// avoid empty job if the number of rows is smaller than the batch size
		for i := 0; i < jobCount+1; i++ {
			jobChan <- i * batchSize
		}
		close(jobChan)
	}()

	wg.Wait()
	log.Info("Done!")
}