Merge pull request nsqio#93 from zulily/nsqlookupd_error_on_connect

consumer: expose a way to retrieve runtime stats

consumer.go

@@ -55,6 +55,15 @@ type FailedMessageLogger interface {
 	LogFailedMessage(message *Message)
 }
 
+// ConsumerStats represents a snapshot of the state of a Consumer's connections and the messages
+// it has seen
+type ConsumerStats struct {
+	MessagesReceived uint64
+	MessagesFinished uint64
+	MessagesRequeued uint64
+	Connections      int
+}
+
 var instCount int64
 
 // Consumer is a high-level type to consume from NSQ.
@@ -168,6 +177,16 @@ func NewConsumer(topic string, channel string, config *Config) (*Consumer, error
 	return r, nil
 }
 
+// Stats retrieves the current connection and message statistics for a Consumer
+func (r *Consumer) Stats() *ConsumerStats {
+	return &ConsumerStats{
+		MessagesReceived: atomic.LoadUint64(&r.messagesReceived),
+		MessagesFinished: atomic.LoadUint64(&r.messagesFinished),
+		MessagesRequeued: atomic.LoadUint64(&r.messagesRequeued),
+		Connections:      len(r.conns()),
+	}
+}
+
 func (r *Consumer) conns() []*Conn {
 	r.mtx.RLock()
 	conns := make([]*Conn, 0, len(r.connections))

consumer_test.go

@@ -177,6 +177,11 @@ func consumerTest(t *testing.T, cb func(c *Config)) {
 		t.Fatal(err)
 	}
 
+	stats := q.Stats()
+	if stats.Connections == 0 {
+		t.Fatal("stats report 0 connections (should be > 0)")
+	}
+
 	err = q.ConnectToNSQD(addr)
 	if err == nil {
 		t.Fatal("should not be able to connect to the same NSQ twice")
@@ -199,6 +204,18 @@ func consumerTest(t *testing.T, cb func(c *Config)) {
 
 	<-q.StopChan
 
+	stats = q.Stats()
+	if stats.Connections != 0 {
+		t.Fatalf("stats report %d active connections (should be 0)", stats.Connections)
+	}
+
+	stats = q.Stats()
+	if stats.MessagesReceived != uint64(h.messagesReceived+h.messagesFailed) {
+		t.Fatalf("stats report %d messages received (should be %d)",
+			stats.MessagesReceived,
+			h.messagesReceived+h.messagesFailed)
+	}
+
 	if h.messagesReceived != 8 || h.messagesSent != 4 {
 		t.Fatalf("end of test. should have handled a diff number of messages (got %d, sent %d)", h.messagesReceived, h.messagesSent)
 	}