/
robotApp.go
71 lines (55 loc) · 1.18 KB
/
robotApp.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
package app
import (
"fmt"
"sync"
"time"
"github.com/TheCoy/toolkit/routinepool"
"go.uber.org/ratelimit"
)
type RobotApp struct {
Times int64
QPS int64
Worker int
limiter ratelimit.Limiter
wg sync.WaitGroup
}
func (app *RobotApp) initApp() error {
app.limiter = ratelimit.New(int(app.QPS))
return nil
}
func (app *RobotApp) run() error {
workPool := routinepool.NewPool(app.Worker)
go func() {
defer close(workPool.EntryChannel)
for i := 0; i < int(app.Times); i++ {
app.limiter.Take()
app.wg.Add(1)
workPool.EntryChannel <- app.buildTask(int64(i))
}
}()
workPool.Run()
app.wg.Wait()
return nil
}
func (app *RobotApp) buildTask(i int64) *routinepool.Task {
task := routinepool.NewTask(func() error {
defer app.wg.Done()
maxTimes := 1000000
qps := int(app.QPS)
if qps > 1000 {
qps = 1000
}
sleep := maxTimes / qps / app.Worker
time.Sleep(time.Duration(sleep) * time.Microsecond)
fmt.Println("=======REQUESTED!=======", i)
return nil
})
return task
}
//SetNewLimiter changes built-in rateLimiter
func (app *RobotApp) SetNewLimiter(qps int) {
if qps != int(app.QPS) {
app.QPS = int64(qps)
}
app.limiter = ratelimit.New(qps)
}