diff --git a/README.md b/README.md index 6eeb64d..f32508a 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ installed. $ pip install hiredis [6] $ go get github.com/garyburd/redigo/redis [7] $ go get github.com/alecthomas/gozmq [8] + $ go get github.com/pebbe/zmq3 [9] 1. 2. @@ -35,6 +36,7 @@ installed. 6. 7. 8. +9. Contributions ============= diff --git a/pubsub3/broker.go b/pubsub3/broker.go new file mode 100644 index 0000000..b76fe9c --- /dev/null +++ b/pubsub3/broker.go @@ -0,0 +1,36 @@ +package pubsub + +import ( + "fmt" + zmq "github.com/pebbe/zmq3" + "time" +) + +func Serve(quiet bool) { + + context, _ := zmq.NewContext() + receiver, _ := context.NewSocket(zmq.PULL) + receiver.Bind("tcp://*:5562") + sender, _ := context.NewSocket(zmq.PUB) + sender.Bind("tcp://*:5561") + + last := time.Now() + messages := 0 + for { + message, err := receiver.Recv(0) + if err != nil { + fmt.Println(err) + } + sender.Send(message, 0) + if !quiet { + messages += 1 + now := time.Now() + if now.Sub(last).Seconds() > 1 { + println(messages, "msg/sec") + last = now + messages = 0 + } + } + } + +} diff --git a/pubsub3/clients.go b/pubsub3/clients.go new file mode 100644 index 0000000..3eaa812 --- /dev/null +++ b/pubsub3/clients.go @@ -0,0 +1,134 @@ +package pubsub + +import ( + "fmt" + "github.com/garyburd/redigo/redis" + zmq "github.com/pebbe/zmq3" + "strings" + "sync" + "time" +) + +// A pub-sub message - defined to support Redis receiving different +// message types, such as subscribe/unsubscribe info. +type Message struct { + Type string + Channel string + Data string +} + +// Client interface for both Redis and ZMQ pubsub clients. +type Client interface { + Subscribe(channels ...interface{}) (err error) + Unsubscribe(channels ...interface{}) (err error) + Publish(channel string, message string) error + Receive() (Message, error) +} + +// Redis client - defines the underlying connection and pub-sub +// connections, as well as a mutex for locking write access, +// since this occurs from multiple goroutines. +type RedisClient struct { + conn redis.Conn + redis.PubSubConn + sync.Mutex +} + +// ZMQ client - just defines the pub and sub ZMQ sockets. +type ZMQClient struct { + ctx *zmq.Context + pub *zmq.Socket + sub *zmq.Socket +} + +// Returns a new Redis client. The underlying redigo package uses +// Go's bufio package which will flush the connection when it contains +// enough data to send, but we still need to set up some kind of timed +// flusher, so it's done here with a goroutine. +func NewRedisClient(host string) *RedisClient { + host = fmt.Sprintf("%s:6379", host) + conn, _ := redis.Dial("tcp", host) + pubsub, _ := redis.Dial("tcp", host) + client := RedisClient{conn, redis.PubSubConn{pubsub}, sync.Mutex{}} + go func() { + for { + time.Sleep(200 * time.Millisecond) + client.Lock() + client.conn.Flush() + client.Unlock() + } + }() + return &client +} + +func (client *RedisClient) Publish(channel, message string) error { + client.Lock() + client.conn.Send("PUBLISH", channel, message) + client.Unlock() + + return nil +} + +func (client *RedisClient) Receive() (Message, error) { + switch message := client.PubSubConn.Receive().(type) { + case redis.Message: + return Message{"message", message.Channel, string(message.Data)}, nil + case redis.Subscription: + return Message{message.Kind, message.Channel, string(message.Count)}, nil + } + return Message{}, nil +} + +func NewZMQClient(host string) (*ZMQClient, error) { + var err error + var context *zmq.Context + context, err = zmq.NewContext() + if err != nil { + return nil, err + } + var pub *zmq.Socket + pub, err = context.NewSocket(zmq.PUSH) + if err != nil { + return nil, err + } + pub.Connect(fmt.Sprintf("tcp://%s:%d", host, 5562)) + var sub *zmq.Socket + sub, err = context.NewSocket(zmq.SUB) + if err != nil { + return nil, err + } + sub.Connect(fmt.Sprintf("tcp://%s:%d", host, 5561)) + return &ZMQClient{context, pub, sub}, nil +} + +func (client *ZMQClient) Subscribe(channels ...interface{}) error { + for _, channel := range channels { + if err := client.sub.SetSubscribe(channel.(string)); err != nil { + return err + } + } + return nil +} + +func (client *ZMQClient) Unsubscribe(channels ...interface{}) error { + for _, channel := range channels { + if err := client.sub.SetUnsubscribe(channel.(string)); err != nil { + return err + } + } + return nil +} + +func (client *ZMQClient) Publish(channel, message string) error { + _, err := client.pub.Send(channel+" "+message, 0) + return err +} + +func (client *ZMQClient) Receive() (Message, error) { + message, err := client.sub.Recv(0) + if err != nil { + return Message{}, err + } + parts := strings.SplitN(string(message), " ", 2) + return Message{Type: "message", Channel: parts[0], Data: parts[1]}, nil +} diff --git a/run_broker3.go b/run_broker3.go new file mode 100644 index 0000000..1a6ee8b --- /dev/null +++ b/run_broker3.go @@ -0,0 +1,13 @@ +package main + +import ( + "./pubsub3" + "flag" +) + +func main() { + var quiet bool + flag.BoolVar(&quiet, "quiet", false, "") + flag.Parse() + pubsub.Serve(quiet) +} diff --git a/test_client3.go b/test_client3.go new file mode 100644 index 0000000..ce00dd7 --- /dev/null +++ b/test_client3.go @@ -0,0 +1,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") + +}