zmq3 code

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. <http://www.zeromq.org>
 2. <http://redis.io>
@@ -35,6 +36,7 @@ installed.
 6. <https://github.com/pietern/hiredis-py>
 7. <https://github.com/garyburd/redigo>
 8. <https://github.com/alecthomas/gozmq>
+9. <https://github.com/pebbe/zmq3>
 
 Contributions
 =============

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
+			}
+		}
+	}
+
+}

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
+}

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)
+}

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")
+
+}