rework of the event consumer

clients/highwater/highwaterMock_test.go

@@ -27,8 +27,8 @@ func TestMock(t *testing.T) {
 	client.PostWithUser(USERID, EVENT_NAME, TOKEN, p)
 }
 
-//log.Panic is called when not all required args are passed.
-//This test fails if the panic and subseqent recover are not called
+// log.Panic is called when not all required args are passed.
+// This test fails if the panic and subseqent recover are not called
 func TestMock_Fails(t *testing.T) {
 
 	defer func() {

clients/mocks.go

@@ -12,7 +12,7 @@ type (
 	SeagullMock struct{}
 )
 
-//A mock of the Gatekeeper interface
+// A mock of the Gatekeeper interface
 func NewGatekeeperMock(expectedPermissions Permissions, expectedError error) *GatekeeperMock {
 	return &GatekeeperMock{expectedPermissions, expectedError}
 }
@@ -45,7 +45,7 @@ func (mock *GatekeeperMock) SetPermissions(userID, groupID string, permissions P
 	return Permissions{userID: Allowed}, nil
 }
 
-//A mock of the Seagull interface
+// A mock of the Seagull interface
 func NewSeagullMock() *SeagullMock {
 	return &SeagullMock{}
 }

events/config.go

@@ -1,8 +1,6 @@
 package events
 
 import (
-	"errors"
-
 	"github.com/IBM/sarama"
 	"github.com/kelseyhightower/envconfig"
 )
@@ -71,10 +69,3 @@ func (k *CloudEventsConfig) GetDeadLettersTopic() string {
 func (k *CloudEventsConfig) IsDeadLettersEnabled() bool {
 	return k.GetDeadLettersTopic() != ""
 }
-
-func validateConsumerConfig(config *CloudEventsConfig) error {
-	if config.KafkaConsumerGroup == "" {
-		return errors.New("consumer group cannot be empty")
-	}
-	return nil
-}

events/consumer_group.go

@@ -2,146 +2,134 @@ package events
 
 import (
 	"context"
-	"fmt"
-	"github.com/IBM/sarama"
-	"github.com/tidepool-org/go-common/errors"
+	stderrors "errors"
 	"log"
 	"sync"
+
+	"github.com/IBM/sarama"
+
+	"github.com/tidepool-org/go-common/errors"
 )
 
 var ErrConsumerStopped = errors.New("consumer has been stopped")
 
-type SaramaConsumerGroup struct {
-	config        *CloudEventsConfig
-	consumerGroup sarama.ConsumerGroup
-	consumer      MessageConsumer
-	ready         chan bool
-	stop          chan struct{}
-	stopOnce      *sync.Once
-	topic         string
-	wg            *sync.WaitGroup
+// SaramaEventConsumer implements EventConsumer to consume messages received
+// via Sarama consumer groups.
+type SaramaEventConsumer struct {
+	config   *CloudEventsConfig
+	consumer MessageConsumer
+
+	cancelFuncMu sync.Mutex
+	cancelFunc   context.CancelFunc
 }
 
 func NewSaramaConsumerGroup(config *CloudEventsConfig, consumer MessageConsumer) (EventConsumer, error) {
 	if err := validateConsumerConfig(config); err != nil {
 		return nil, err
 	}
+	if err := consumer.Initialize(config); err != nil {
+		return nil, err
+	}
 
-	return &SaramaConsumerGroup{
+	return &SaramaEventConsumer{
 		config:   config,
 		consumer: consumer,
-		topic:    config.GetPrefixedTopic(),
-		wg:       &sync.WaitGroup{},
-		ready:    make(chan bool, 1),
-		stop:     make(chan struct{}),
-		stopOnce: &sync.Once{},
 	}, nil
 }
 
-func (s *SaramaConsumerGroup) Setup(session sarama.ConsumerGroupSession) error {
-	// Mark the consumer as ready
-	close(s.ready)
+func validateConsumerConfig(config *CloudEventsConfig) error {
+	if config.KafkaConsumerGroup == "" {
+		return errors.New("consumer group cannot be empty")
+	}
 	return nil
 }
 
-func (s *SaramaConsumerGroup) Cleanup(session sarama.ConsumerGroupSession) error {
-	return nil
-}
+func (s *SaramaEventConsumer) Start() error {
 
-func (s *SaramaConsumerGroup) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
-	for message := range claim.Messages() {
-		if err := s.consumer.HandleKafkaMessage(message); err != nil {
-			log.Printf("failed to process kafka message: %v", err)
-			return err
-		}
-		session.MarkMessage(message, "")
-	}
+	ctx, cancel := s.newContext()
+	defer cancel()
 
-	return nil
-}
-
-func (s *SaramaConsumerGroup) Start() error {
-	if err := s.initialize(); err != nil {
+	cg, err := sarama.NewConsumerGroup(
+		s.config.KafkaBrokers,
+		s.config.KafkaConsumerGroup,
+		s.config.SaramaConfig,
+	)
+	if err != nil {
 		return err
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-	go func() {
-		<-s.stop
-		cancel()
-	}()
-
-	errChan := make(chan error)
-	s.wg.Add(1)
-	go func() {
-		defer s.wg.Done()
-		for {
-			// `Consume` should be called inside an infinite loop, when a
-			// server-side rebalance happens, the consumer session will need to be
-			// recreated to get the new claims
-			if err := s.consumerGroup.Consume(ctx, []string{s.topic}, s); err != nil {
-				log.Printf("Error from consumer: %v", err)
-				// It's not clear whether this condition can be true
-				if err == context.Canceled {
-					err = ErrConsumerStopped
-				}
-				errChan <- err
-				return
+	handler := &SaramaMessageConsumer{s.consumer}
+	topics := []string{s.config.GetPrefixedTopic()}
+	for {
+		// `Consume` should be called inside an infinite loop, when a
+		// server-side rebalance happens, the consumer session will need to be
+		// recreated to get the new claims
+		if err := cg.Consume(ctx, topics, handler); err != nil {
+			log.Printf("Error from consumer: %v", err)
+			if stderrors.Is(err, context.Canceled) {
+				return ErrConsumerStopped
 			}
-			// check if context was cancelled, signaling that the consumer should stop
-			if ctx.Err() != nil {
-				errChan <- ErrConsumerStopped
-				return
+			return err
+		}
+		// Double check the context isn't canceled before looping. This is necessary as
+		// Consume() sometimes returns nil when the context is canceled.
+		if err := ctx.Err(); err != nil {
+			if stderrors.Is(err, context.Canceled) {
+				return ErrConsumerStopped
 			}
-			s.ready = make(chan bool)
+			return err
 		}
-	}()
-
-	err := <-errChan
-	if err == ErrConsumerStopped {
-		return err
 	}
+}
 
-	// The consumer group was terminated with an unexpected error.
-	// We need to call stop, so we cancel the context and stop the
-	// go routine so it doesn't leak on restart.
-	if e := s.Stop(); e != nil {
-		err = fmt.Errorf("%w: %s", err, e.Error())
-	}
+func (s *SaramaEventConsumer) Stop() error {
+	s.cancelFuncMu.Lock()
+	defer s.cancelFuncMu.Unlock()
 
-	return err
+	if s.cancelFunc != nil {
+		s.cancelFunc()
+		s.cancelFunc = nil
+	}
+	return nil
 }
 
-func (s *SaramaConsumerGroup) Stop() error {
-	// Initialization failed
-	if s.consumerGroup == nil {
-		return nil
+func (s *SaramaEventConsumer) newContext() (context.Context, context.CancelFunc) {
+	s.cancelFuncMu.Lock()
+	defer s.cancelFuncMu.Unlock()
+	if s.cancelFunc != nil {
+		s.cancelFunc()
 	}
+	ctx, cancel := context.WithCancel(context.Background())
+	s.cancelFunc = cancel
+	return ctx, cancel
+}
 
-	// Signal that the consumer group should be terminated
-	s.stopOnce.Do(func() {
-		s.stop <- struct{}{}
-	})
+// SaramaMessageConsumer implements sarama.ConsumerGroupHandler.
+//
+// It adapts a MessageConsumer for this purpose.
+type SaramaMessageConsumer struct {
+	MessageConsumer
+}
 
-	// Wait for the consumer group to exit
-	s.wg.Wait()
-	return s.consumerGroup.Close()
+// Cleanup implements sarama.ConsumerGroupHandler.
+func (c *SaramaMessageConsumer) Cleanup(_ sarama.ConsumerGroupSession) error {
+	return nil
 }
 
-func (s *SaramaConsumerGroup) initialize() error {
-	cg, err := sarama.NewConsumerGroup(
-		s.config.KafkaBrokers,
-		s.config.KafkaConsumerGroup,
-		s.config.SaramaConfig,
-	)
-	if err != nil {
-		return err
+// ConsumeClaim implements sarama.ConsumerGroupHandler.
+func (c *SaramaMessageConsumer) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
+	for message := range claim.Messages() {
+		if err := c.HandleKafkaMessage(message); err != nil {
+			log.Printf("failed to process kafka message: %v", err)
+			return err
+		}
+		session.MarkMessage(message, "")
 	}
 
-	if err := s.consumer.Initialize(s.config); err != nil {
-		return err
-	}
+	return nil
+}
 
-	s.consumerGroup = cg
+// Setup implements sarama.ConsumerGroupHandler.
+func (c *SaramaMessageConsumer) Setup(_ sarama.ConsumerGroupSession) error {
 	return nil
 }