forked from stephenmcd/two-queues
/
test_client3.go
143 lines (131 loc) · 3.15 KB
/
test_client3.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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
package main
import (
"./pubsub3"
"flag"
"fmt"
"log"
"math/rand"
"runtime"
"sort"
"strconv"
"strings"
"time"
)
var (
host string
numSeconds float64
numClients int
numChannels int
messageSize int
useRedis bool
quiet bool
channels []string
)
// Returns a new pubsub client instance - either the Redis or ZeroMQ
// client, based on command-line arg.
func NewClient() pubsub.Client {
var client pubsub.Client
if useRedis {
client = pubsub.NewRedisClient(host)
} else {
var err error
client, err = pubsub.NewZMQClient(host)
if err != nil {
log.Panicln(err)
}
}
return client
}
// Loops forever, publishing messages to random channels.
func Publisher() {
client := NewClient()
message := strings.Repeat("x", messageSize)
for {
channel := channels[rand.Intn(len(channels))]
if err := client.Publish(channel, message); err != nil {
log.Panicln(err)
}
}
}
// Subscribes to all channels, keeping a count of the number of
// messages received. Publishes and resets the total every second.
func Subscriber() {
client := NewClient()
for _, channel := range channels {
if err := client.Subscribe(channel); err != nil {
log.Panicln(err)
}
}
last := time.Now()
messages := 0
for {
if _, err := client.Receive(); err != nil {
log.Panicln(err)
}
messages += 1
now := time.Now()
if now.Sub(last).Seconds() > 1 {
if !quiet {
println(messages, "msg/sec")
}
if err := client.Publish("metrics", strconv.Itoa(messages)); err != nil {
log.Panicln(err)
}
last = now
messages = 0
}
}
}
// Creates goroutines * --num-clients, running the given target
// function for each.
func RunWorkers(target func()) {
for i := 0; i < numClients; i++ {
go target()
}
}
// Subscribes to the metrics channel and returns messages from
// it until --num-seconds has passed.
func GetMetrics() []int {
client := NewClient()
if err := client.Subscribe("metrics"); err != nil {
log.Panicln(err)
}
metrics := []int{}
start := time.Now()
for time.Now().Sub(start).Seconds() <= numSeconds {
message, err := client.Receive()
if err != nil {
log.Panicln(err)
}
if message.Type == "message" {
messages, _ := strconv.Atoi(message.Data)
metrics = append(metrics, messages)
}
}
return metrics
}
func main() {
// Set up and parse command-line args.
runtime.GOMAXPROCS(runtime.NumCPU())
flag.StringVar(&host, "host", "127.0.0.1", "")
flag.Float64Var(&numSeconds, "num-seconds", 10, "")
flag.IntVar(&numClients, "num-clients", 1, "")
flag.IntVar(&numChannels, "num-channels", 50, "")
flag.IntVar(&messageSize, "message-size", 20, "")
flag.BoolVar(&useRedis, "redis", false, "")
flag.BoolVar(&quiet, "quiet", false, "")
flag.Parse()
for i := 0; i < numChannels; i++ {
channels = append(channels, strconv.Itoa(i))
}
// Create publisher/subscriber goroutines, pausing to allow
// publishers to hit full throttle
RunWorkers(Publisher)
time.Sleep(1 * time.Second)
RunWorkers(Subscriber)
// Consume metrics until --num-seconds has passed, and display
// the median value.
metrics := GetMetrics()
sort.Ints(metrics)
fmt.Println("Num clients", numClients, "median", metrics[len(metrics)/2], "msg/sec")
}