90 lines
1.9 KiB
Go
90 lines
1.9 KiB
Go
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)
|
|
}
|