-
Notifications
You must be signed in to change notification settings - Fork 0
/
bench_reader.go
122 lines (112 loc) · 2.9 KB
/
bench_reader.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
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
package main
import (
"bufio"
"flag"
"log"
"net"
"runtime"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/nsqio/go-nsq"
)
var (
runfor = flag.Duration("runfor", 10*time.Second, "duration of time to run")
tcpAddress = flag.String("nsqd-tcp-address", "127.0.0.1:4150", "<addr>:<port> to connect to nsqd")
size = flag.Int("size", 200, "size of messages")
topic = flag.String("topic", "sub_bench", "topic to receive messages on")
channel = flag.String("channel", "ch", "channel to receive messages on")
deadline = flag.String("deadline", "", "deadline to start the benchmark run")
rdy = flag.Int("rdy", 2500, "RDY count to use")
)
var totalMsgCount int64
func main() {
flag.Parse()
var wg sync.WaitGroup
log.SetPrefix("[bench_reader] ")
goChan := make(chan int)
rdyChan := make(chan int)
workers := runtime.GOMAXPROCS(0)
for j := 0; j < workers; j++ {
wg.Add(1)
go func(id int) {
subWorker(*runfor, workers, *tcpAddress, *topic, *channel, rdyChan, goChan, id)
wg.Done()
}(j)
<-rdyChan
}
if *deadline != "" {
t, err := time.Parse("2006-01-02 15:04:05", *deadline)
if err != nil {
log.Fatal(err)
}
d := t.Sub(time.Now())
log.Printf("sleeping until %s (%s)", t, d)
time.Sleep(d)
}
start := time.Now()
close(goChan)
wg.Wait()
end := time.Now()
duration := end.Sub(start)
tmc := atomic.LoadInt64(&totalMsgCount)
log.Printf("duration: %s - %.03fmb/s - %.03fops/s - %.03fus/op",
duration,
float64(tmc*int64(*size))/duration.Seconds()/1024/1024,
float64(tmc)/duration.Seconds(),
float64(duration/time.Microsecond)/float64(tmc))
}
func subWorker(td time.Duration, workers int, tcpAddr string, topic string, channel string, rdyChan chan int, goChan chan int, id int) {
conn, err := net.DialTimeout("tcp", tcpAddr, time.Second)
if err != nil {
panic(err.Error())
}
conn.Write(nsq.MagicV2)
rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
ci := make(map[string]interface{})
ci["short_id"] = "test"
ci["long_id"] = "test"
cmd, _ := nsq.Identify(ci)
cmd.WriteTo(rw)
nsq.Subscribe(topic, channel).WriteTo(rw)
rdyChan <- 1
<-goChan
nsq.Ready(*rdy).WriteTo(rw)
rw.Flush()
nsq.ReadResponse(rw)
nsq.ReadResponse(rw)
var msgCount int64
go func() {
time.Sleep(td)
conn.Close()
}()
for {
resp, err := nsq.ReadResponse(rw)
if err != nil {
if strings.Contains(err.Error(), "use of closed network connection") {
break
}
panic(err.Error())
}
frameType, data, err := nsq.UnpackResponse(resp)
if err != nil {
panic(err.Error())
}
if frameType == nsq.FrameTypeError {
panic(string(data))
} else if frameType == nsq.FrameTypeResponse {
continue
}
msg, err := nsq.DecodeMessage(data)
if err != nil {
panic(err.Error())
}
nsq.Finish(msg.ID).WriteTo(rw)
msgCount++
if float64(msgCount%int64(*rdy)) > float64(*rdy)*0.75 {
rw.Flush()
}
}
atomic.AddInt64(&totalMsgCount, msgCount)
}