forked from influxdata/influxdb
/
util.go
89 lines (78 loc) · 1.9 KB
/
util.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
package stressClient
import (
"crypto/rand"
"fmt"
"log"
"strconv"
"sync"
)
// ###########################################
// ConcurrencyLimiter and associated methods #
// ###########################################
// ConcurrencyLimiter ensures that no more than a specified
// max number of goroutines are running.
type ConcurrencyLimiter struct {
inc chan chan struct{}
dec chan struct{}
max int
count int
sync.Mutex
}
// NewConcurrencyLimiter returns a configured limiter that will
// ensure that calls to Increment will block if the max is hit.
func NewConcurrencyLimiter(max int) *ConcurrencyLimiter {
c := &ConcurrencyLimiter{
inc: make(chan chan struct{}),
dec: make(chan struct{}, max),
max: max,
}
go c.handleLimits()
return c
}
// Increment will increase the count of running goroutines by 1.
// if the number is currently at the max, the call to Increment
// will block until another goroutine decrements.
func (c *ConcurrencyLimiter) Increment() {
r := make(chan struct{})
c.inc <- r
<-r
}
// Decrement will reduce the count of running goroutines by 1
func (c *ConcurrencyLimiter) Decrement() {
c.dec <- struct{}{}
}
// NewMax resets the max of a ConcurrencyLimiter.
func (c *ConcurrencyLimiter) NewMax(i int) {
c.Lock()
defer c.Unlock()
c.max = i
}
// handleLimits runs in a goroutine to manage the count of
// running goroutines.
func (c *ConcurrencyLimiter) handleLimits() {
for {
r := <-c.inc
c.Lock()
if c.count >= c.max {
<-c.dec
c.count--
}
c.Unlock()
c.count++
r <- struct{}{}
}
}
// Utility interger parsing function
func parseInt(s string) int {
i, err := strconv.ParseInt(s, 10, 64)
if err != nil {
log.Fatalf("Error parsing integer:\n String: %v\n Error: %v\n", s, err)
}
return int(i)
}
// Utility for making random strings of length n
func randStr(n int) string {
b := make([]byte, n/2)
_, _ = rand.Read(b)
return fmt.Sprintf("%x", b)
}