File: dqlite-benchmark.go

package info (click to toggle)
golang-github-canonical-go-dqlite 2.0.0-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, trixie
  • size: 720 kB
  • sloc: sh: 380; makefile: 5
file content (153 lines) | stat: -rw-r--r-- 5,383 bytes parent folder | download | duplicates (2)
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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"time"

	"github.com/canonical/go-dqlite/v2/app"
	"github.com/canonical/go-dqlite/v2/benchmark"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"golang.org/x/sys/unix"
)

const (
	defaultClusterTimeout = 120
	defaultDir            = "/tmp/dqlite-benchmark"
	defaultDiskMode       = false
	defaultDriver         = false
	defaultDurationS      = 60
	defaultKvKeySize      = 32
	defaultKvValueSize    = 1024
	defaultWorkers        = 1
	defaultWorkload       = "kvwrite"
	docString             = "For benchmarking dqlite.\n\n" +
		"Run a 1 node benchmark:\n" +
		"dqlite-benchmark -d 127.0.0.1:9001 --driver --cluster 127.0.0.1:9001\n\n" +
		"Run a multi-node benchmark, the first node will self-elect and become leader,\n" +
		"the driver flag results in the workload being run from the first, leader node.\n" +
		"dqlite-benchmark --db 127.0.0.1:9001 --driver --cluster 127.0.0.1:9001,127.0.0.1:9002,127.0.0.1:9003 &\n" +
		"dqlite-benchmark --db 127.0.0.1:9002 --join 127.0.0.1:9001 &\n" +
		"dqlite-benchmark --db 127.0.0.1:9003 --join 127.0.0.1:9001 &\n\n" +
		"Run a multi-node benchmark, the first node will self-elect and become leader,\n" +
		"the driver flag results in the workload being run from the third, non-leader node.\n" +
		"dqlite-benchmark --db 127.0.0.1:9001 &\n" +
		"dqlite-benchmark --db 127.0.0.1:9002 --join 127.0.0.1:9001 &\n" +
		"dqlite-benchmark --db 127.0.0.1:9003 --join 127.0.0.1:9001 --driver --cluster 127.0.0.1:9001,127.0.0.1:9002,127.0.0.1:9003 &\n\n" +
		"The results can be found on the `driver` node in " + defaultDir + "/results or in the directory provided to the tool.\n" +
		"Benchmark results are files named `n-q-timestamp` where `n` is the number of the worker,\n" +
		"`q` is the type of query that was tracked. All results in the file are in milliseconds.\n"
)

func signalChannel() chan os.Signal {
	ch := make(chan os.Signal, 32)
	signal.Notify(ch, unix.SIGPWR)
	signal.Notify(ch, unix.SIGINT)
	signal.Notify(ch, unix.SIGQUIT)
	signal.Notify(ch, unix.SIGTERM)
	return ch
}

func main() {
	var cluster *[]string
	var clusterTimeout int
	var db string
	var dir string
	var driver bool
	var duration int
	var join *[]string
	var kvKeySize int
	var kvValueSize int
	var workers int
	var workload string
	var diskMode bool

	cmd := &cobra.Command{
		Use:   "dqlite-benchmark",
		Short: "For benchmarking dqlite",
		Long:  docString,
		RunE: func(cmd *cobra.Command, args []string) error {
			dir := filepath.Join(dir, db)
			if err := os.MkdirAll(dir, 0755); err != nil {
				return errors.Wrapf(err, "can't create %s", dir)
			}

			app, err := app.New(dir, app.WithDiskMode(diskMode), app.WithAddress(db), app.WithCluster(*join))
			if err != nil {
				return err
			}

			readyCtx, cancel := context.WithTimeout(context.Background(), time.Duration(clusterTimeout)*time.Second)
			defer cancel()
			if err := app.Ready(readyCtx); err != nil {
				return errors.Wrap(err, "App not ready in time")
			}

			ch := signalChannel()
			if !driver {
				fmt.Println("Benchmark client ready. Send signal to abort or when done.")
				<-ch
				return nil
			}

			if len(*cluster) == 0 {
				return fmt.Errorf("driver node, `--cluster` flag must be provided")
			}

			db, err := app.Open(context.Background(), "benchmark")
			if err != nil {
				return err
			}
			db.SetMaxOpenConns(500)
			db.SetMaxIdleConns(500)

			bm, err := benchmark.New(
				app,
				db,
				dir,
				benchmark.WithWorkload(workload),
				benchmark.WithDuration(duration),
				benchmark.WithWorkers(workers),
				benchmark.WithKvKeySize(kvKeySize),
				benchmark.WithKvValueSize(kvValueSize),
				benchmark.WithCluster(*cluster),
				benchmark.WithClusterTimeout(clusterTimeout),
			)
			if err != nil {
				return err
			}

			if err := bm.Run(ch); err != nil {
				return err
			}

			db.Close()
			app.Close()
			return nil
		},
	}

	flags := cmd.Flags()
	flags.StringVarP(&db, "db", "d", "", "Address used for internal database replication.")
	join = flags.StringSliceP("join", "j", nil, "Database addresses of existing nodes.")
	cluster = flags.StringSliceP("cluster", "c", nil, "Database addresses of all nodes taking part in the benchmark.\n"+
		"The driver will wait for all nodes to be online before running the benchmark.")
	flags.IntVar(&clusterTimeout, "cluster-timeout", defaultClusterTimeout, "How long the benchmark should wait in seconds for the whole cluster to be online.")
	flags.StringVarP(&dir, "dir", "D", defaultDir, "Data directory.")
	flags.StringVarP(&workload, "workload", "w", defaultWorkload, "The workload to run: \"kvwrite\" or \"kvreadwrite\".")
	flags.BoolVar(&driver, "driver", defaultDriver, "Set this flag to run the benchmark from this instance. Must be set on 1 node.")
	flags.IntVar(&duration, "duration", defaultDurationS, "Run duration in seconds.")
	flags.IntVar(&workers, "workers", defaultWorkers, "Number of workers executing the workload.")
	flags.IntVar(&kvKeySize, "key-size", defaultKvKeySize, "Size of the KV keys in bytes.")
	flags.IntVar(&kvValueSize, "value-size", defaultKvValueSize, "Size of the KV values in bytes.")
	flags.BoolVar(&diskMode, "disk", defaultDiskMode, "Warning: Unstable, Experimental. Set this flag to enable dqlite's disk-mode.")

	cmd.MarkFlagRequired("db")
	if err := cmd.Execute(); err != nil {
		os.Exit(1)
	}
}