update pubsub consume with stop signal

Author: ochom

pubsub/consumer.go

@@ -1,6 +1,8 @@
 package pubsub
 
-import "fmt"
+import (
+	"fmt"
+)
 
 // Consumer ...
 type Consumer struct {
@@ -40,19 +42,18 @@ func NewConsumer(rabbitURL, exchange, queue string, config ...Config) *Consumer
 }
 
 // Consume consume messages from the channels
-func (c *Consumer) Consume(workerFunc func([]byte)) error {
-	conn, ch, err := initQ(c.url)
+func (c *Consumer) Consume(stop chan bool, workerFunc func([]byte)) error {
+	ps, err := initQ(c.url)
 	if err != nil {
 		return fmt.Errorf("failed to initialize a connection: %s", err.Error())
 	}
-	defer ch.Close()
-	defer conn.Close()
+	defer ps.Close()
 
-	if err := initPubSub(ch, c.exchange, c.queue, string(c.config.Type)); err != nil {
+	if err := initPubSub(ps.ch, c.exchange, c.queue, string(c.config.Type)); err != nil {
 		return fmt.Errorf("failed to initialize a pubsub: %s", err.Error())
 	}
 
-	deliveries, err := ch.Consume(
+	deliveries, err := ps.ch.Consume(
 		c.queue,            // queue
 		c.config.Tag,       // consumerTag
 		c.config.AutoAck,   // auto-ack
@@ -66,9 +67,12 @@ func (c *Consumer) Consume(workerFunc func([]byte)) error {
 		return fmt.Errorf("failed to consume messages: %s", err.Error())
 	}
 
-	for d := range deliveries {
-		workerFunc(d.Body)
+	for {
+		select {
+		case <-stop:
+			return fmt.Errorf("stop signal received")
+		case message := <-deliveries:
+			workerFunc(message.Body)
+		}
 	}
-
-	return nil
 }

pubsub/publisher.go

@@ -22,19 +22,18 @@ func NewPublisher(rabbitURL, exchange, queue string, exchangeType ExchangeType)
 
 // publish ...
 func (p *publisher) publish(body []byte, delay time.Duration) error {
-	conn, ch, err := initQ(p.url)
+	ps, err := initQ(p.url)
 	if err != nil {
 		return fmt.Errorf("failed to initialize a connection: %s", err.Error())
 	}
-	defer ch.Close()
-	defer conn.Close()
+	defer ps.Close()
 
-	if err := initPubSub(ch, p.exchange, p.queue, p.exchangeType); err != nil {
+	if err := initPubSub(ps.ch, p.exchange, p.queue, p.exchangeType); err != nil {
 		return fmt.Errorf("failed to initialize a pubsub: %s", err.Error())
 	}
 
 	// publish message to exchange
-	err = ch.Publish(
+	err = ps.ch.Publish(
 		p.exchange, // exchange
 		p.queue,    // routing key
 		true,       // mandatory

pubsub/pubsub.go

@@ -17,23 +17,40 @@ var (
 	Delayed ExchangeType = "x-delayed-message"
 )
 
-func initQ(url string) (*amqp.Connection, *amqp.Channel, error) {
+// Pubsub ...
+type Pubsub struct {
+	conn *amqp.Connection
+	ch   *amqp.Channel
+}
+
+// Close ...
+func (p *Pubsub) Close() {
+	if err := p.conn.Close(); err != nil {
+		fmt.Println("failed to close connection: ", err.Error())
+	}
+
+	if err := p.ch.Close(); err != nil {
+		fmt.Println("failed to close channel: ", err.Error())
+	}
+}
+
+func initQ(url string) (*Pubsub, error) {
 	conn, err := amqp.Dial(url)
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to connect to RabbitMQ: %s", err.Error())
+		return nil, fmt.Errorf("failed to connect to RabbitMQ: %s", err.Error())
 	}
 
 	ch, err := conn.Channel()
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to open a channel: %s", err.Error())
+		return nil, fmt.Errorf("failed to open a channel: %s", err.Error())
 	}
 
 	err = ch.Qos(10, 0, false) // fair dispatch
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to set QoS: %s", err.Error())
+		return nil, fmt.Errorf("failed to set QoS: %s", err.Error())
 	}
 
-	return conn, ch, nil
+	return &Pubsub{conn: conn, ch: ch}, nil
 }
 
 // initPubSub ...