nsqd: drop per-channel queue workers

This moves to a single goroutine to process in-flight and deferred
priority queues.  It manages a pool of workers (configurable max) that
process channels concurrently.

It is a copy of redis's probabilistic algorithm that wakes up every 100ms
to select a random 20 channels from a locally cached list (refreshed
every 5s).

If either of the queues had work to do the channel is considered "dirty".

If 25% of the selected channels were dirty, the loop continues without
sleep.

For 1000 topics and channels with no clients connected - idle cpu usage
dropped to ~8%.

Author: mreiferson

internal/util/rand.go

@@ -0,0 +1,22 @@
+package util
+
+import (
+	"math/rand"
+)
+
+func UniqRands(l int, n int) []int {
+	set := make(map[int]struct{})
+	nums := make([]int, 0, l)
+	for {
+		num := rand.Intn(n)
+		if _, ok := set[num]; !ok {
+			set[num] = struct{}{}
+			nums = append(nums, num)
+		}
+		if len(nums) == l {
+			goto exit
+		}
+	}
+exit:
+	return nums
+}

nsqd/channel.go

@@ -12,12 +12,8 @@ import (
 
 	"github.com/bitly/nsq/internal/pqueue"
 	"github.com/bitly/nsq/internal/quantile"
-	"github.com/bitly/nsq/internal/util"
 )
 
-// the amount of time a worker will wait when idle
-const defaultWorkerWait = 100 * time.Millisecond
-
 type Consumer interface {
 	UnPause()
 	Pause()
@@ -52,8 +48,8 @@ type Channel struct {
 	memoryMsgChan chan *Message
 	clientMsgChan chan *Message
 	exitChan      chan int
-	waitGroup     util.WaitGroupWrapper
 	exitFlag      int32
+	exitMutex     sync.RWMutex
 
 	// state tracking
 	clients        map[int64]Consumer
@@ -116,9 +112,6 @@ func NewChannel(topicName string, channelName string, ctx *context,
 
 	go c.messagePump()
 
-	c.waitGroup.Wrap(func() { c.deferredWorker() })
-	c.waitGroup.Wrap(func() { c.inFlightWorker() })
-
 	c.ctx.nsqd.Notify(c)
 
 	return c
@@ -155,6 +148,9 @@ func (c *Channel) Close() error {
 }
 
 func (c *Channel) exit(deleted bool) error {
+	c.exitMutex.Lock()
+	defer c.exitMutex.Unlock()
+
 	if !atomic.CompareAndSwapInt32(&c.exitFlag, 0, 1) {
 		return errors.New("exiting")
 	}
@@ -178,9 +174,6 @@ func (c *Channel) exit(deleted bool) error {
 
 	close(c.exitChan)
 
-	// synchronize the close of router() and pqWorkers (2)
-	c.waitGroup.Wait()
-
 	if deleted {
 		// empty the queue (deletes the backend files, too)
 		c.Empty()
@@ -587,80 +580,70 @@ exit:
 	close(c.clientMsgChan)
 }
 
-func (c *Channel) deferredWorker() {
-	c.pqWorker(&c.deferredPQ, &c.deferredMutex, func(item *pqueue.Item) {
-		msg := item.Value.(*Message)
-		_, err := c.popDeferredMessage(msg.ID)
-		if err != nil {
-			return
-		}
-		c.doRequeue(msg)
-	})
-}
+func (c *Channel) processDeferredQueue(t int64) bool {
+	c.exitMutex.RLock()
+	defer c.exitMutex.RUnlock()
+
+	if c.Exiting() {
+		return false
+	}
 
-func (c *Channel) inFlightWorker() {
-	ticker := time.NewTicker(defaultWorkerWait)
+	dirty := false
 	for {
-		select {
-		case <-ticker.C:
-		case <-c.exitChan:
+		c.deferredMutex.Lock()
+		item, _ := c.deferredPQ.PeekAndShift(t)
+		c.deferredMutex.Unlock()
+
+		if item == nil {
 			goto exit
 		}
-		now := time.Now().UnixNano()
-		for {
-			c.inFlightMutex.Lock()
-			msg, _ := c.inFlightPQ.PeekAndShift(now)
-			c.inFlightMutex.Unlock()
-
-			if msg == nil {
-				break
-			}
+		dirty = true
 
-			_, err := c.popInFlightMessage(msg.clientID, msg.ID)
-			if err != nil {
-				break
-			}
-			atomic.AddUint64(&c.timeoutCount, 1)
-			c.RLock()
-			client, ok := c.clients[msg.clientID]
-			c.RUnlock()
-			if ok {
-				client.TimedOutMessage()
-			}
-			c.doRequeue(msg)
+		msg := item.Value.(*Message)
+		_, err := c.popDeferredMessage(msg.ID)
+		if err != nil {
+			goto exit
 		}
+		c.doRequeue(msg)
 	}
 
 exit:
-	c.ctx.nsqd.logf("CHANNEL(%s): closing ... inFlightWorker", c.name)
-	ticker.Stop()
+	return dirty
 }
 
-// generic loop (executed in a goroutine) that periodically wakes up to walk
-// the priority queue and call the callback
-func (c *Channel) pqWorker(pq *pqueue.PriorityQueue, mutex *sync.Mutex, callback func(item *pqueue.Item)) {
-	ticker := time.NewTicker(defaultWorkerWait)
+func (c *Channel) processInFlightQueue(t int64) bool {
+	c.exitMutex.RLock()
+	defer c.exitMutex.RUnlock()
+
+	if c.Exiting() {
+		return false
+	}
+
+	dirty := false
 	for {
-		select {
-		case <-ticker.C:
-		case <-c.exitChan:
+		c.inFlightMutex.Lock()
+		msg, _ := c.inFlightPQ.PeekAndShift(t)
+		c.inFlightMutex.Unlock()
+
+		if msg == nil {
 			goto exit
 		}
-		now := time.Now().UnixNano()
-		for {
-			mutex.Lock()
-			item, _ := pq.PeekAndShift(now)
-			mutex.Unlock()
-
-			if item == nil {
-				break
-			}
+		dirty = true
 
-			callback(item)
+		_, err := c.popInFlightMessage(msg.clientID, msg.ID)
+		if err != nil {
+			goto exit
 		}
+		atomic.AddUint64(&c.timeoutCount, 1)
+		c.RLock()
+		client, ok := c.clients[msg.clientID]
+		c.RUnlock()
+		if ok {
+			client.TimedOutMessage()
+		}
+		c.doRequeue(msg)
 	}
 
 exit:
-	c.ctx.nsqd.logf("CHANNEL(%s): closing ... pqueue worker", c.name)
-	ticker.Stop()
+	return dirty
 }

nsqd/channel_test.go

@@ -60,6 +60,7 @@ func TestInFlightWorker(t *testing.T) {
 	opts := NewNSQDOptions()
 	opts.Logger = newTestLogger(t)
 	opts.MsgTimeout = 100 * time.Millisecond
+	opts.QueueScanRefreshInterval = 100 * time.Millisecond
 	_, _, nsqd := mustStartNSQD(opts)
 	defer os.RemoveAll(opts.DataPath)
 	defer nsqd.Exit()

nsqd/nsqd.go

@@ -59,6 +59,8 @@ type NSQD struct {
 	httpsListener net.Listener
 	tlsConfig     *tls.Config
 
+	poolSize int
+
 	idChan     chan MessageID
 	notifyChan chan interface{}
 	exitChan   chan int
@@ -245,6 +247,7 @@ func (n *NSQD) Main() {
 		http_api.Serve(n.httpListener, httpServer, n.opts.Logger, "HTTP")
 	})
 
+	n.waitGroup.Wrap(func() { n.queueScanLoop() })
 	n.waitGroup.Wrap(func() { n.idPump() })
 	n.waitGroup.Wrap(func() { n.lookupLoop() })
 	if n.opts.StatsdAddress != "" {
@@ -551,6 +554,137 @@ func (n *NSQD) Notify(v interface{}) {
 	})
 }
 
+// channels returns a flat slice of all channels in all topics
+func (n *NSQD) channels() []*Channel {
+	var channels []*Channel
+	n.RLock()
+	for _, t := range n.topicMap {
+		t.RLock()
+		for _, c := range t.channelMap {
+			channels = append(channels, c)
+		}
+		t.RUnlock()
+	}
+	n.RUnlock()
+	return channels
+}
+
+// resizePool adjusts the size of the pool of queueScanWorker goroutines
+//
+// 	1 <= pool <= min(num * 0.25, QueueScanWorkerPoolMax)
+//
+func (n *NSQD) resizePool(num int, workCh chan *Channel, responseCh chan bool, closeCh chan int) {
+	idealPoolSize := int(float64(num) * 0.25)
+	if idealPoolSize < 1 {
+		idealPoolSize = 1
+	} else if idealPoolSize > n.opts.QueueScanWorkerPoolMax {
+		idealPoolSize = n.opts.QueueScanWorkerPoolMax
+	}
+	for {
+		if idealPoolSize == n.poolSize {
+			break
+		} else if idealPoolSize < n.poolSize {
+			// contract
+			closeCh <- 1
+			n.poolSize--
+		} else {
+			// expand
+			n.waitGroup.Wrap(func() {
+				n.queueScanWorker(workCh, responseCh, closeCh)
+			})
+			n.poolSize++
+		}
+	}
+}
+
+// queueScanWorker receives work (in the form of a channel) from queueScanLoop
+// and processes the deferred and in-flight queues
+func (n *NSQD) queueScanWorker(workCh chan *Channel, responseCh chan bool, closeCh chan int) {
+	for {
+		select {
+		case c := <-workCh:
+			now := time.Now().UnixNano()
+			dirty := false
+			if c.processInFlightQueue(now) {
+				dirty = true
+			}
+			if c.processDeferredQueue(now) {
+				dirty = true
+			}
+			responseCh <- dirty
+		case <-closeCh:
+			return
+		}
+	}
+}
+
+// queueScanLoop runs in a single goroutine to process in-flight and deferred
+// priority queues. It manages a pool of queueScanWorker (configurable max of
+// QueueScanWorkerPoolMax (default: 4)) that process channels concurrently.
+//
+// It copies Redis's probabilistic expiration algorithm: it wakes up every
+// QueueScanInterval (default: 100ms) to select a random QueueScanSelectionCount
+// (default: 20) channels from a locally cached list (refreshed every
+// QueueScanRefreshInterval (default: 5s)).
+//
+// If either of the queues had work to do the channel is considered "dirty".
+//
+// If QueueScanDirtyPercent (default: 25%) of the selected channels were dirty,
+// the loop continues without sleep.
+func (n *NSQD) queueScanLoop() {
+	workCh := make(chan *Channel, n.opts.QueueScanSelectionCount)
+	responseCh := make(chan bool, n.opts.QueueScanSelectionCount)
+	closeCh := make(chan int)
+
+	workTicker := time.NewTicker(n.opts.QueueScanInterval)
+	refreshTicker := time.NewTicker(n.opts.QueueScanRefreshInterval)
+
+	channels := n.channels()
+	n.resizePool(len(channels), workCh, responseCh, closeCh)
+
+	for {
+		select {
+		case <-workTicker.C:
+			if len(channels) == 0 {
+				continue
+			}
+		case <-refreshTicker.C:
+			channels = n.channels()
+			n.resizePool(len(channels), workCh, responseCh, closeCh)
+			continue
+		case <-n.exitChan:
+			goto exit
+		}
+
+		num := n.opts.QueueScanSelectionCount
+		if num > len(channels) {
+			num = len(channels)
+		}
+
+	loop:
+		for _, i := range util.UniqRands(num, len(channels)) {
+			workCh <- channels[i]
+		}
+
+		numDirty := 0
+		for i := 0; i < num; i++ {
+			if <-responseCh {
+				numDirty++
+			}
+		}
+
+		if float64(numDirty)/float64(num) > n.opts.QueueScanDirtyPercent {
+			goto loop
+		}
+	}
+
+exit:
+	n.logf("QUEUESCAN: closing")
+	close(closeCh)
+	workTicker.Stop()
+	refreshTicker.Stop()
+}
+
 func buildTLSConfig(opts *nsqdOptions) (*tls.Config, error) {
 	var tlsConfig *tls.Config
 

nsqd/options.go

@@ -28,6 +28,12 @@ type nsqdOptions struct {
 	SyncEvery       int64         `flag:"sync-every"`
 	SyncTimeout     time.Duration `flag:"sync-timeout"`
 
+	QueueScanInterval        time.Duration
+	QueueScanRefreshInterval time.Duration
+	QueueScanSelectionCount  int
+	QueueScanWorkerPoolMax   int
+	QueueScanDirtyPercent    float64
+
 	// msg and command options
 	MsgTimeout    time.Duration `flag:"msg-timeout" arg:"1ms"`
 	MaxMsgTimeout time.Duration `flag:"max-msg-timeout"`
@@ -85,6 +91,12 @@ func NewNSQDOptions() *nsqdOptions {
 		SyncEvery:       2500,
 		SyncTimeout:     2 * time.Second,
 
+		QueueScanInterval:        100 * time.Millisecond,
+		QueueScanRefreshInterval: 5 * time.Second,
+		QueueScanSelectionCount:  20,
+		QueueScanWorkerPoolMax:   4,
+		QueueScanDirtyPercent:    0.25,
+
 		MsgTimeout:    60 * time.Second,
 		MaxMsgTimeout: 15 * time.Minute,
 		MaxMsgSize:    1024768,