bench.go

package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"sync"
)

var settings string

func GenInsert(start, end int) string {
		insert := "insert into t_async_insert_dedup values"
		for v := start; v < end; v++ {
			insert += fmt.Sprintf(" (%d)", v)
			if v+1 < end {
				insert += ","
			} else {
				insert += ";"
			}
		}
	return insert
}

func GenData(queue chan string, gStart, size, dupRate int) {
	counter := 0
	chunkSize := 200
	gEnd := gStart + size
	for i := gStart; i < gEnd; i += chunkSize {
		insert := GenInsert(i, i + chunkSize)
		queue <- insert
		if (counter % 100) < dupRate {
			pos := gStart + (counter / 2) * chunkSize
			insert := GenInsert(pos, pos + chunkSize)
			queue <- insert
		}
		counter ++
	}
	close(queue)
}

func ConsumeData(queue chan string, wg *sync.WaitGroup) {
	defer wg.Done()
	for insert := range queue {
		cmd := exec.Command("clickhouse-client", "-nq", settings+insert)
		err := cmd.Run()
		if err != nil {
			log.Fatal(err)
		}
	}
}

func main() {
	queue := make(chan string, 1000000)
	if len(os.Args) < 5 {
		log.Fatal("command should be: bench start end concurrency dupRate [settings...]")
	}
	start, err := strconv.Atoi(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	size, err := strconv.Atoi(os.Args[2])
	if err != nil {
		log.Fatal(err)
	}
	con, err := strconv.Atoi(os.Args[3])
	if err != nil {
		log.Fatal(err)
	}
	dupRate, err := strconv.Atoi(os.Args[4])
	if err != nil {
		log.Fatal(err)
	}
	for i, setting := range os.Args {
		if i > 4 {
			settings += "set " + setting + ";"
		}
	}
	go GenData(queue, start, size, dupRate)
	var wg sync.WaitGroup
	wg.Add(con)
	for i := 0; i < con; i++ {
		go ConsumeData(queue, &wg)
	}
	wg.Wait()
}