etcd-io
/
dbtester
mirror of https://github.com/etcd-io/dbtester.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

control: implement rate-limiting

This commit is contained in:
Gyu-Ho Lee 2016-12-20 12:53:28 -08:00
parent 1cd6cf534a
commit 3ccd4b7a98
No known key found for this signature in database
GPG Key ID: 1DDD39C7EB70C24C
1 changed files with 20 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
control/control.go
View File

@ -27,6 +27,7 @@ import (
"github.com/coreos/dbtester/remotestorage"
"github.com/spf13/cobra"
"golang.org/x/net/context"
"golang.org/x/time/rate"
"google.golang.org/grpc"
"github.com/cheggaaa/pb"
@ -586,7 +587,16 @@ func newReadOneshotHandlers(cfg Config) []ReqHandler {
func generateReads(cfg Config, key string, requests chan<- request) {
defer close(requests)
var rateLimiter *rate.Limiter
if cfg.Step2.RequestsPerSecond > 0 {
rateLimiter = rate.NewLimiter(rate.Limit(cfg.Step2.RequestsPerSecond), cfg.Step2.RequestsPerSecond)
}
for i := 0; i < cfg.Step2.TotalRequests; i++ {
if rateLimiter != nil {
rateLimiter.Wait(context.TODO())
}
switch cfg.Database {
case "etcdv2":
// serializable read by default
@ -613,18 +623,21 @@ func generateReads(cfg Config, key string, requests chan<- request) {
}
requests <- request{consulOp: op}
}
if cfg.Step2.RequestIntervalMs > 0 {
time.Sleep(time.Duration(cfg.Step2.RequestIntervalMs) * time.Millisecond)
}
}
}
func generateWrites(cfg Config, vals values, requests chan<- request) {
var rateLimiter *rate.Limiter
if cfg.Step2.RequestsPerSecond > 0 {
rateLimiter = rate.NewLimiter(rate.Limit(cfg.Step2.RequestsPerSecond), cfg.Step2.RequestsPerSecond)
}
var wg sync.WaitGroup
defer func() {
close(requests)
wg.Wait()
}()
for i := 0; i < cfg.Step2.TotalRequests; i++ {
k := sequentialKey(cfg.Step2.KeySize, i)
if cfg.Step2.SameKey {
@ -634,6 +647,10 @@ func generateWrites(cfg Config, vals values, requests chan<- request) {
v := vals.bytes[i%vals.sampleSize]
vs := vals.strings[i%vals.sampleSize]
if rateLimiter != nil {
rateLimiter.Wait(context.TODO())
}
switch cfg.Database {
case "etcdv2":
requests <- request{etcdv2Op: etcdv2Op{key: k, value: vs}}
@ -644,8 +661,5 @@ func generateWrites(cfg Config, vals values, requests chan<- request) {
case "consul", "cetcd":
requests <- request{consulOp: consulOp{key: k, value: v}}
}
if cfg.Step2.RequestIntervalMs > 0 {
time.Sleep(time.Duration(cfg.Step2.RequestIntervalMs) * time.Millisecond)
}
}
}