-
Notifications
You must be signed in to change notification settings - Fork 309
/
tpcc.go
111 lines (94 loc) · 2.91 KB
/
tpcc.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
package main
import (
"context"
"fmt"
"net/http"
_ "net/http/pprof"
"os"
"runtime"
"github.com/pingcap/go-tpc/pkg/workload"
"github.com/pingcap/go-tpc/tpcc"
"github.com/spf13/cobra"
)
var tpccConfig tpcc.Config
func executeTpcc(action string) {
if pprofAddr != "" {
go func() {
err := http.ListenAndServe(pprofAddr, http.DefaultServeMux)
if err != nil {
fmt.Printf("failed to ListenAndServe: %s\n", err.Error())
}
}()
}
runtime.GOMAXPROCS(maxProcs)
openDB()
defer closeDB()
tpccConfig.DBName = dbName
tpccConfig.Threads = threads
tpccConfig.Isolation = isolationLevel
var (
w workload.Workloader
err error
)
switch tpccConfig.OutputType {
case "csv", "CSV":
if tpccConfig.OutputDir == "" {
fmt.Printf("Output Directory cannot be empty when generating files")
os.Exit(1)
}
w, err = tpcc.NewCSVWorkloader(globalDB, &tpccConfig)
default:
w, err = tpcc.NewWorkloader(globalDB, &tpccConfig)
}
if err != nil {
fmt.Printf("Failed to init work loader: %v\n", err)
os.Exit(1)
}
timeoutCtx, cancel := context.WithTimeout(globalCtx, totalTime)
defer cancel()
executeWorkload(timeoutCtx, w, action)
}
func registerTpcc(root *cobra.Command) {
cmd := &cobra.Command{
Use: "tpcc",
}
cmd.PersistentFlags().IntVar(&tpccConfig.Parts, "parts", 1, "Number to partition warehouses")
cmd.PersistentFlags().IntVar(&tpccConfig.Warehouses, "warehouses", 10, "Number of warehouses")
cmd.PersistentFlags().BoolVar(&tpccConfig.CheckAll, "check-all", false, "Run all consistency checks")
var cmdPrepare = &cobra.Command{
Use: "prepare",
Short: "Prepare data for TPCC",
Run: func(cmd *cobra.Command, _ []string) {
executeTpcc("prepare")
},
}
cmdPrepare.PersistentFlags().StringVar(&tpccConfig.OutputType, "output-type", "", "Output file type."+
" If empty, then load data to db. Current only support csv")
cmdPrepare.PersistentFlags().StringVar(&tpccConfig.OutputDir, "output-dir", "", "Output directory for generating file if specified")
cmdPrepare.PersistentFlags().StringVar(&tpccConfig.SpecifiedTables, "tables", "", "Specified tables for "+
"generating file, separated by ','. Valid only if output is set. If this flag is not set, generate all tables by default")
var cmdRun = &cobra.Command{
Use: "run",
Short: "Run workload",
Run: func(cmd *cobra.Command, _ []string) {
executeTpcc("run")
},
}
cmdRun.PersistentFlags().BoolVar(&tpccConfig.Wait, "wait", false, "including keying & thinking time described on TPC-C Standard Specification")
var cmdCleanup = &cobra.Command{
Use: "cleanup",
Short: "Cleanup data for the workload",
Run: func(cmd *cobra.Command, _ []string) {
executeTpcc("cleanup")
},
}
var cmdCheck = &cobra.Command{
Use: "check",
Short: "Check data consistency for the workload",
Run: func(cmd *cobra.Command, _ []string) {
executeTpcc("check")
},
}
cmd.AddCommand(cmdRun, cmdPrepare, cmdCleanup, cmdCheck)
root.AddCommand(cmd)
}