[SCH-175] Update to buffer and pipeline commands (#16)

Author: danielstjules

broadcast/broadcast.go

@@ -2,6 +2,7 @@ package broadcast
 
 import (
 	"encoding/json"
+	"sync"
 	"time"
 
 	"github.com/bitly/go-nsq"
@@ -30,28 +31,47 @@ type Message struct {
 
 // Options for broadcast.
 type Options struct {
-	Redis        RedisPool
-	Metrics      *statsd.Client
-	Ratelimiter  *ratelimit.Ratelimiter
-	RatelimitKey string
-	Log          *log.Logger
+	Redis         RedisPool
+	Metrics       *statsd.Client
+	Ratelimiter   *ratelimit.Ratelimiter
+	RatelimitKey  string
+	Log           *log.Logger
+	FlushInterval time.Duration
 }
 
 // Broadcast consumer distributes messages to N handlers.
 type Broadcast struct {
+	Done chan struct{}
+	*Options
+
 	handlers []Handler
 	stats    *stats.Stats
-	*Options
+
+	// Single connection, channel and mutex are used when applying a flush interval
+	conn  Conn
+	mutex sync.Mutex
 }
 
 // New broadcast consumer.
 func New(o *Options) *Broadcast {
 	stats := stats.New()
 	go stats.TickEvery(10 * time.Second)
-	return &Broadcast{
+
+	broadcast := Broadcast{
 		stats:   stats,
 		Options: o,
+		Done:    make(chan struct{}, 1),
+	}
+
+	if o.FlushInterval < 0 {
+		panic("FlushInterval must not be a negative duration")
+	} else if o.FlushInterval > 0 {
+		conn := NewConn(o.Redis.Get())
+		broadcast.conn = conn
+		broadcast.flushOnInterval(conn)
 	}
+
+	return &broadcast
 }
 
 // Add handler.
@@ -80,9 +100,8 @@ func (b *Broadcast) HandleMessage(msg *nsq.Message) error {
 		return nil
 	}
 
-	db := b.Redis.Get()
-	defer db.Close()
-	conn := NewConn(db)
+	conn, done := b.getConn()
+	defer done()
 
 	for _, h := range b.handlers {
 		err := h.Handle(conn, m)
@@ -91,14 +110,63 @@ func (b *Broadcast) HandleMessage(msg *nsq.Message) error {
 		}
 	}
 
-	err = conn.Flush()
+	if b.FlushInterval == 0 {
+		if err := b.flush(conn); err != nil {
+			return err
+		}
+	}
+
+	b.Metrics.Duration("timers.broadcast", time.Since(start))
+	return nil
+}
+
+// Flushes all messages, then sends on the Done channel.
+func (b *Broadcast) Stop() {
+	if b.FlushInterval > 0 {
+		conn, done := b.getConn()
+		defer done()
+		b.flush(conn)
+	}
+
+	b.Done <- struct{}{}
+}
+
+func (b *Broadcast) getConn() (conn Conn, done func()) {
+	if b.FlushInterval == 0 {
+		db := b.Redis.Get()
+		conn = NewConn(db)
+		done = func() {
+			db.Close()
+		}
+	} else {
+		b.mutex.Lock()
+		conn = b.conn
+		done = func() {
+			b.mutex.Unlock()
+		}
+	}
+
+	return conn, done
+}
+
+func (b *Broadcast) flushOnInterval(conn Conn) {
+	go func() {
+		for range time.Tick(b.FlushInterval) {
+			b.mutex.Lock()
+			b.flush(conn)
+			b.mutex.Unlock()
+		}
+	}()
+}
+
+func (b *Broadcast) flush(conn Conn) error {
+	err := conn.Flush()
 	if err != nil {
 		b.Metrics.Incr("errors.flush")
 		b.Log.Error("flush: %s", err)
 		return err
 	}
 
-	b.Metrics.Duration("timers.broadcast", time.Since(start))
 	return nil
 }
 

broadcast/broadcast_test.go

@@ -2,6 +2,7 @@ package broadcast
 
 import (
 	"io/ioutil"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -84,6 +85,73 @@ func TestBroadcast(t *testing.T) {
 	h2.AssertExpectations(t)
 }
 
+func TestBroadcastInvalidFlushInterval(t *testing.T) {
+	assert.Panic(t, "FlushInterval must not be a negative duration", func() {
+		New(&Options{FlushInterval: -1 * time.Hour})
+	})
+}
+
+func TestBroadcastWithoutFlushInterval(t *testing.T) {
+	pool := getMockPool()
+	broadcast := New(&Options{
+		Redis:   pool,
+		Metrics: statsd.NewClient(ioutil.Discard),
+		Log:     log.Log,
+	})
+
+	sendMessages(broadcast)
+
+	mockConn := pool.Get().(*mocks.NoOpRedisConn)
+	assert.Equal(t, 2, int(atomic.LoadUint64(&mockConn.Flushes)))
+}
+
+func TestBroadcastWithFlushInterval(t *testing.T) {
+	pool := getMockPool()
+	broadcast := New(&Options{
+		Redis:         pool,
+		Metrics:       statsd.NewClient(ioutil.Discard),
+		Log:           log.Log,
+		FlushInterval: 2 * time.Millisecond,
+	})
+
+	sendMessages(broadcast)
+
+	mockConn := pool.Get().(*mocks.NoOpRedisConn)
+	assert.Equal(t, 0, int(atomic.LoadUint64(&mockConn.Flushes)))
+	<-time.After(3 * time.Millisecond)
+	assert.Equal(t, 1, int(atomic.LoadUint64(&mockConn.Flushes)))
+}
+
+func TestBroadcastWithFlushIntervalStop(t *testing.T) {
+	pool := getMockPool()
+	broadcast := New(&Options{
+		Redis:         pool,
+		Metrics:       statsd.NewClient(ioutil.Discard),
+		Log:           log.Log,
+		FlushInterval: 10 * time.Second,
+	})
+
+	sendMessages(broadcast)
+
+	mockConn := pool.Get().(*mocks.NoOpRedisConn)
+	assert.Equal(t, 0, int(atomic.LoadUint64(&mockConn.Flushes)))
+	broadcast.Stop()
+	<-broadcast.Done
+	assert.Equal(t, 1, int(atomic.LoadUint64(&mockConn.Flushes)))
+}
+
+func getMockPool() RedisPool {
+	pool := &mockRedisPool{}
+	pool.On("Get").Return(mocks.NewNoOpRedisConn())
+	return pool
+}
+
+func sendMessages(broadcast *Broadcast) {
+	nsqMsg := nsq.NewMessage(newNSQMessageId("nsq__message__id"), []byte(`{"projectId":"gy2d"}`))
+	broadcast.HandleMessage(nsqMsg)
+	broadcast.HandleMessage(nsqMsg)
+}
+
 func newNSQMessageId(id string) nsq.MessageID {
 	nsqId := [nsq.MsgIDLength]byte{}
 	copy(nsqId[:], id[:nsq.MsgIDLength])

broadcast/conn.go

@@ -41,5 +41,7 @@ func (c *conn) Flush() error {
 		}
 	}
 
+	c.pending = 0
+
 	return nil
 }

broadcast/mocks/RedisConn.go

@@ -1,34 +1,40 @@
 package mocks
 
-import "github.com/garyburd/redigo/redis"
+import (
+	"sync/atomic"
+
+	"github.com/garyburd/redigo/redis"
+)
 
 func NewNoOpRedisConn() redis.Conn {
-	return &noOpRedisConn{}
+	return &NoOpRedisConn{}
 }
 
-type noOpRedisConn struct{}
+type NoOpRedisConn struct {
+	Flushes uint64
+}
 
-func (c *noOpRedisConn) Close() error {
+func (c *NoOpRedisConn) Close() error {
 	return nil
 }
 
-func (c *noOpRedisConn) Err() error {
+func (c *NoOpRedisConn) Err() error {
 	return nil
 }
 
-func (c *noOpRedisConn) Do(commandName string, args ...interface{}) (reply interface{}, err error) {
+func (c *NoOpRedisConn) Do(commandName string, args ...interface{}) (reply interface{}, err error) {
 	return nil, nil
 }
 
-func (c *noOpRedisConn) Send(commandName string, args ...interface{}) error {
+func (c *NoOpRedisConn) Send(commandName string, args ...interface{}) error {
 	return nil
 }
 
-func (c *noOpRedisConn) Flush() error {
+func (c *NoOpRedisConn) Flush() error {
+	atomic.AddUint64(&c.Flushes, 1)
 	return nil
-
 }
 
-func (c *noOpRedisConn) Receive() (reply interface{}, err error) {
+func (c *NoOpRedisConn) Receive() (reply interface{}, err error) {
 	return nil, nil
 }

main.go

@@ -29,6 +29,7 @@ const usage = `
       [--lookupd-http-address addr...]
       [--nsqd-tcp-address addr...]
       [--redis-address addr]
+      [--flush-interval t]
       [--max-idle n]
       [--idle-timeout t]
       [--list name] [--list-size n]
@@ -47,6 +48,7 @@ const usage = `
     --redis-address addr         redis address [default: :6379]
     --max-attempts n             nsq max message attempts [default: 5]
     --max-in-flight n            nsq messages in-flight [default: 250]
+    --flush-interval t           time to buffer redis commands before flushing [default: 0s]
     --max-idle n                 redis max idle connections [default: 15]
     --idle-timeout t             idle connection timeout [default: 1m]
     --list-size n                redis list size [default: 100]
@@ -88,6 +90,14 @@ func main() {
 		log.Fatalf("error parsing idle timeout: %s", err)
 	}
 
+	flushInterval, err := time.ParseDuration(args["--flush-interval"].(string))
+	if err != nil {
+		log.Fatalf("error parsing flush-interval: %s", err)
+	}
+	if flushInterval < 0 {
+		log.Fatalf("flush-interval must not be a negative value")
+	}
+
 	maxIdle, err := strconv.Atoi(args["--max-idle"].(string))
 	if err != nil {
 		log.Fatalf("error parsing max-idle: %s", err)
@@ -105,11 +115,12 @@ func main() {
 	}
 
 	broadcast := broadcast.New(&broadcast.Options{
-		Redis:        pool,
-		Metrics:      metrics,
-		Log:          log.Log,
-		Ratelimiter:  ratelimiter(args),
-		RatelimitKey: args["--ratelimit-key"].(string),
+		Redis:         pool,
+		Metrics:       metrics,
+		Log:           log.Log,
+		Ratelimiter:   ratelimiter(args),
+		RatelimitKey:  args["--ratelimit-key"].(string),
+		FlushInterval: flushInterval,
 	})
 	config := config(args)
 
@@ -175,6 +186,8 @@ func main() {
 	log.Info("stopping")
 	consumer.Stop()
 	<-consumer.StopChan
+	broadcast.Stop()
+	<-broadcast.Done
 	log.Info("bye :)")
 }