Merge pull request #160 from dmarkham/empty_channel_testing

nsqd: Empty Queue should reset the InFlight count of consumers

nsqd/channel.go

@@ -24,6 +24,7 @@ type Consumer interface {
 	Close() error
 	TimedOutMessage()
 	Stats() ClientStats
+	Empty()
 }
 
 // Channel represents the concrete type for a NSQ channel (and also
@@ -202,6 +203,9 @@ func (c *Channel) Empty() error {
 	defer c.Unlock()
 
 	c.initPQ()
+	for _, client := range c.clients {
+		client.Empty()
+	}
 
 	for {
 		select {

nsqd/channel_test.go

@@ -118,3 +118,38 @@ func TestChannelEmpty(t *testing.T) {
 	assert.Equal(t, len(channel.deferredPQ), 0)
 	assert.Equal(t, channel.Depth(), int64(0))
 }
+
+func TestChannelEmptyConsumer(t *testing.T) {
+	log.SetOutput(ioutil.Discard)
+	defer log.SetOutput(os.Stdout)
+
+	tcpAddr, _ := mustStartNSQd(NewNsqdOptions())
+	defer nsqd.Exit()
+	conn, _ := mustConnectNSQd(tcpAddr)
+
+	topicName := "test_channel_empty" + strconv.Itoa(int(time.Now().Unix()))
+	topic := nsqd.GetTopic(topicName)
+	channel := topic.GetChannel("channel")
+	client := NewClientV2(conn)
+	client.SetReadyCount(25)
+	channel.AddClient(client)
+
+	for i := 0; i < 25; i++ {
+		msg := nsq.NewMessage(<-nsqd.idChan, []byte("test"))
+		channel.StartInFlightTimeout(msg, client)
+		client.SendingMessage()
+	}
+
+	for _, cl := range channel.clients {
+		stats := cl.Stats()
+		assert.Equal(t, stats.InFlightCount, int64(25))
+	}
+
+	channel.Empty()
+
+	for _, cl := range channel.clients {
+		stats := cl.Stats()
+		assert.Equal(t, stats.InFlightCount, int64(0))
+	}
+
+}

nsqd/client_v2.go

@@ -123,6 +123,11 @@ func (c *ClientV2) FinishedMessage() {
 	c.tryUpdateReadyState()
 }
 
+func (c *ClientV2) Empty() {
+	atomic.StoreInt64(&c.InFlightCount, 0)
+	c.tryUpdateReadyState()
+}
+
 func (c *ClientV2) SendingMessage() {
 	atomic.AddInt64(&c.ReadyCount, -1)
 	atomic.AddInt64(&c.InFlightCount, 1)