-
Notifications
You must be signed in to change notification settings - Fork 6
/
main.go
107 lines (88 loc) · 3.5 KB
/
main.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
// load_timescaledb loads a TimescaleDB instance with data from stdin.
//
// If the database exists beforehand, it will be *DROPPED*.
package main
import (
"fmt"
"sync"
"github.com/blagojts/viper"
"github.com/spf13/pflag"
"github.com/cnosdb/tsdb-comparisons/internal/utils"
"github.com/cnosdb/tsdb-comparisons/load"
"github.com/cnosdb/tsdb-comparisons/pkg/data/source"
"github.com/cnosdb/tsdb-comparisons/pkg/targets/timescaledb"
)
// Parse args:
func initProgramOptions() (*timescaledb.LoadingOptions, load.BenchmarkRunner, *load.BenchmarkRunnerConfig) {
target := timescaledb.NewTarget()
loaderConf := load.BenchmarkRunnerConfig{}
loaderConf.AddToFlagSet(pflag.CommandLine)
target.TargetSpecificFlags("", pflag.CommandLine)
pflag.Parse()
err := utils.SetupConfigFile()
if err != nil {
panic(fmt.Errorf("fatal error config file: %s", err))
}
if err := viper.Unmarshal(&loaderConf); err != nil {
panic(fmt.Errorf("unable to decode config: %s", err))
}
opts := timescaledb.LoadingOptions{}
viper.SetTypeByDefaultValue(true)
opts.PostgresConnect = viper.GetString("postgres")
opts.Host = viper.GetString("host")
opts.Port = viper.GetString("port")
opts.User = viper.GetString("user")
opts.Pass = viper.GetString("pass")
opts.ConnDB = viper.GetString("admin-db-name")
opts.LogBatches = viper.GetBool("log-batches")
opts.UseHypertable = viper.GetBool("use-hypertable")
opts.ChunkTime = viper.GetDuration("chunk-time")
opts.UseJSON = viper.GetBool("use-jsonb-tags")
// This must be set to 'true' if you are going to test
// distributed hypertable queries and insert. Replication
// factor must also be set to true for distributed hypertables
opts.InTableTag = viper.GetBool("in-table-partition-tag")
// We currently use `create_hypertable` for all variations. When
// `replication-factor`>=1, we automatically create a distributed
// hypertable.
opts.ReplicationFactor = viper.GetInt("replication-factor")
// Currently ignored for distributed hypertables. We assume all
// data nodes will be used based on the partition-column above
opts.NumberPartitions = viper.GetInt("partitions")
opts.TimeIndex = viper.GetBool("time-index")
opts.TimePartitionIndex = viper.GetBool("time-partition-index")
opts.PartitionIndex = viper.GetBool("partition-index")
opts.FieldIndex = viper.GetString("field-index")
opts.FieldIndexCount = viper.GetInt("field-index-count")
opts.ProfileFile = viper.GetString("write-profile")
opts.ReplicationStatsFile = viper.GetString("write-replication-stats")
opts.CreateMetricsTable = viper.GetBool("create-metrics-table")
opts.ForceTextFormat = viper.GetBool("force-text-format")
opts.UseInsert = viper.GetBool("use-insert")
loader := load.GetBenchmarkRunner(loaderConf)
return &opts, loader, &loaderConf
}
func main() {
opts, loader, loaderConf := initProgramOptions()
// If specified, generate a performance profile
if len(opts.ProfileFile) > 0 {
go profileCPUAndMem(opts.ProfileFile)
}
var replicationStatsWaitGroup sync.WaitGroup
if len(opts.ReplicationStatsFile) > 0 {
go OutputReplicationStats(
opts.GetConnectString(loader.DatabaseName()), opts.ReplicationStatsFile, &replicationStatsWaitGroup,
)
}
benchmark, err := timescaledb.NewBenchmark(loaderConf.DBName, opts, &source.DataSourceConfig{
Type: source.FileDataSourceType,
File: &source.FileDataSourceConfig{Location: loaderConf.FileName},
})
if err != nil {
panic(err)
}
loader.RunBenchmark(benchmark)
if len(opts.ReplicationStatsFile) > 0 {
replicationStatsWaitGroup.Wait()
}
}