Wait for Kafka partition to shutdown

example/daemons/number_printer.go

@@ -19,7 +19,7 @@ func (n *NumberPrinter) Startup() {
 		log.Printf("Oh, crap! There was a panic, take a look: %v", err)
 	})
 
-	n.LimitRate(1, 2*time.Second)
+	n.LimitRate(3, 1*time.Second)
 	n.SystemProcess("Random Number Generator", n.generateNumbers)
 }
 

example/daemons/price_consumer.go

@@ -19,10 +19,10 @@ type PriceUpdate struct {
 
 // Startup creates a new subscription for ProductPriceUpdates topic.
 func (p *PriceConsumer) Startup() {
-	p.LimitRate(1, 500*time.Millisecond)
 	b.Subscribe("ProductPriceUpdates", func(u PriceUpdate) {
 		log.Printf("Price for %q is now $%.2f", u.Product, u.Amount)
 	})
+	p.LimitRate(5, 1*time.Second)
 }
 
 // Shutdown is empty because PriceConsumer requires no cleanup upon exiting.

example/kafka/kafka.go

@@ -6,6 +6,7 @@ import (
 	"io/ioutil"
 	"log"
 	"os"
+	"sync"
 
 	"github.com/Shopify/sarama"
 	"github.com/localhots/satan"
@@ -21,6 +22,7 @@ type ConsumerState struct {
 type Stream struct {
 	messages chan []byte
 	shutdown chan struct{}
+	wg       sync.WaitGroup
 }
 
 const (
@@ -80,15 +82,21 @@ func Subscribe(consumer, topic string) satan.Streamer {
 		shutdown: make(chan struct{}),
 	}
 	go func() {
+		stream.wg.Add(1)
+		defer stream.wg.Done()
+		defer pc.Close()
 		for {
 			select {
 			case msg := <-pc.Messages():
-				stream.messages <- msg.Value
-				t.Offset = msg.Offset
+				select {
+				case stream.messages <- msg.Value:
+					t.Offset = msg.Offset
+				case <-stream.shutdown:
+					return
+				}
 			case err := <-pc.Errors():
 				log.Println("Kafka error:", err.Error())
 			case <-stream.shutdown:
-				pc.Close()
 				return
 			}
 		}
@@ -105,6 +113,7 @@ func (s *Stream) Messages() <-chan []byte {
 // Close stops Kafka partition consumer.
 func (s *Stream) Close() {
 	close(s.shutdown)
+	s.wg.Wait()
 }
 
 func loadConsumerConfig() {