insert_mysql.go

package main

import (
    "database/sql"
    "flag"
    "fmt"
    "log"
    "math/rand"
    "os"
    "os/signal"
    "sync"
    "sync/atomic"
    "syscall"
    "time"

    _ "github.com/go-sql-driver/mysql"
    "github.com/twinj/uuid"
)

func main() {
    port := flag.String("p", "4000", "port")
    count := flag.Int("c", 1, "number goroutine")
    flag.Parse()
    log.Println("start")
    var n int64
    lastTime := time.Now()
    var done int32
    wg := sync.WaitGroup{}
    wg.Add(*count)
    db, err := sql.Open("mysql", fmt.Sprintf("root@tcp(localhost:%s)/test", *port))
    if err != nil {
        log.Fatal(err)
    }
    defer db.Close()

    _, err = db.Exec(`create table if not exists tbl (
    id VARCHAR(1000) NOT NULL,
    filed1 BIGINT(20),
    filed2 VARCHAR(1000),
    filed INT(10) DEFAULT NULL,
    filed3 VARCHAR(1000),
    filed4 BIGINT(20),
    filed5 VARCHAR(100),
    PRIMARY KEY (id))`)

    if err != nil {
        log.Fatal(err)
    }

    for i := 0; i < *count; i++ {
        go func() {
            for atomic.LoadInt32(&done) == 0 {
                u := uuid.NewV4().Bytes()
                _, err = db.Exec("insert into tbl (id, filed1, filed2, filed3, filed4, filed5) values (?, ?, ?, ?, ?, ?)",
                    u, rand.Int(), u, u, rand.Int63(), u)
                if err != nil {
                    log.Println(err)
                    continue
                }
                current := atomic.AddInt64(&n, 1)
                if current%10000 == 0 {
                    now := time.Now()
                    useTime := now.Sub(lastTime)
                    lastTime = now
                    log.Printf("total: %v speed: %.2f", current, 10000.0/useTime.Seconds())
                }
            }
            wg.Done()
        }()
    }
    sigchan := make(chan os.Signal, 1)
    signal.Notify(sigchan, os.Interrupt)
    signal.Notify(sigchan, syscall.SIGTERM)
    sig := <-sigchan
    atomic.AddInt32(&done, 1)
    log.Printf("receive signal %v and close\n", sig)
    wg.Wait()
    log.Println("total", n)
}