[iman] graceful shutdown

BaritoLog · Jun 12, 2018 · 77ad1b5 · 77ad1b5
1 parent 0d7da92
commit 77ad1b5
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 0 deletions.
diff --git a/cmds/consumer.go b/cmds/consumer.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	"github.com/BaritoLog/barito-flow/flow"
+	"github.com/BaritoLog/go-boilerplate/srvkit"
 	"github.com/BaritoLog/go-boilerplate/timekit"
 	"github.com/BaritoLog/instru"
 	cluster "github.com/bsm/sarama-cluster"
@@ -55,6 +56,8 @@ func Consumer(c *cli.Context) (err error) {
 		},
 	}
 
+	srvkit.AsyncGracefulShutdown(agent.Close)
+
 	return agent.Start()
 
 }

diff --git a/cmds/producer.go b/cmds/producer.go
@@ -2,6 +2,7 @@ package cmds
 
 import (
 	"github.com/BaritoLog/barito-flow/flow"
+	"github.com/BaritoLog/go-boilerplate/srvkit"
 	"github.com/Shopify/sarama"
 	log "github.com/sirupsen/logrus"
 
@@ -41,5 +42,7 @@ func Producer(c *cli.Context) (err error) {
 		producerMaxTps,
 	)
 
+	srvkit.AsyncGracefulShutdown(agent.Close)
+
 	return agent.Start()
 }
diff --git a/flow/kafka_agent.go b/flow/kafka_agent.go
@@ -20,6 +20,12 @@ func (a *KafkaAgent) Start() (err error) {
 	return
 }
 
+func (a *KafkaAgent) Close() {
+	if a.Consumer != nil {
+		a.Consumer.Close()
+	}
+}
+
 func (a *KafkaAgent) loopMain() {
 	for {
 		select {

diff --git a/flow/kafka_agent_test.go b/flow/kafka_agent_test.go
@@ -43,6 +43,7 @@ func TestKafkaAgent(t *testing.T) {
 			gotNotification = notification
 		},
 	}
+	defer agent.Close()
 
 	go agent.Start()
 	timekit.Sleep("1ms")
@@ -74,6 +75,7 @@ func TestKafkaAgent_StoreError(t *testing.T) {
 			err0 = err
 		},
 	}
+	defer agent.Close()
 
 	go agent.Start()
 	timekit.Sleep("1ms")
@@ -103,6 +105,7 @@ func TestKafkaAgent_KafkaError(t *testing.T) {
 			err0 = err
 		},
 	}
+	defer agent.Close()
 
 	go agent.Start()
 	timekit.Sleep("1ms")

diff --git a/flow/kafka_consumer.go b/flow/kafka_consumer.go
@@ -11,6 +11,7 @@ type KafkaConsumer interface {
 	Notifications() <-chan *cluster.Notification
 	Errors() <-chan error
 	MarkOffset(msg *sarama.ConsumerMessage, metadata string)
+	Close() error
 }
 
 type dummyKafkaConsumer struct {
@@ -33,3 +34,7 @@ func (c *dummyKafkaConsumer) Notifications() <-chan *cluster.Notification {
 
 func (c *dummyKafkaConsumer) MarkOffset(msg *sarama.ConsumerMessage, metadata string) {
 }
+
+func (c *dummyKafkaConsumer) Close() error {
+	return nil
+}