PubSub Kafka: Respect Subscribe context (#3363)

Signed-off-by: joshvanl <[email protected]>

Author: JoshVanL

Diff for: bindings/kafka/kafka.go

@@ -100,29 +100,26 @@ func (b *Binding) Read(ctx context.Context, handler bindings.Handler) error {
 		return nil
 	}
 
-	handlerConfig := kafka.SubscriptionHandlerConfig{
-		IsBulkSubscribe: false,
-		Handler:         adaptHandler(handler),
-	}
-	for _, t := range b.topics {
-		b.kafka.AddTopicHandler(t, handlerConfig)
-	}
+	ctx, cancel := context.WithCancel(ctx)
+
 	b.wg.Add(1)
 	go func() {
-		defer b.wg.Done()
-		// Wait for context cancelation or closure.
 		select {
 		case <-ctx.Done():
 		case <-b.closeCh:
 		}
-
-		// Remove the topic handlers.
-		for _, t := range b.topics {
-			b.kafka.RemoveTopicHandler(t)
-		}
+		cancel()
+		b.wg.Done()
 	}()
 
-	return b.kafka.Subscribe(ctx)
+	handlerConfig := kafka.SubscriptionHandlerConfig{
+		IsBulkSubscribe: false,
+		Handler:         adaptHandler(handler),
+	}
+
+	b.kafka.Subscribe(ctx, handlerConfig, b.topics...)
+
+	return nil
 }
 
 func adaptHandler(handler bindings.Handler) kafka.EventHandler {

Diff for: common/component/kafka/consumer.go

@@ -14,12 +14,10 @@ limitations under the License.
 package kafka
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"strconv"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/IBM/sarama"
@@ -29,15 +27,8 @@ import (
 )
 
 type consumer struct {
-	k             *Kafka
-	ready         chan bool
-	running       chan struct{}
-	stopped       atomic.Bool
-	once          sync.Once
-	mutex         sync.Mutex
-	skipConsume   bool
-	consumeCtx    context.Context
-	consumeCancel context.CancelFunc
+	k     *Kafka
+	mutex sync.Mutex
 }
 
 func (consumer *consumer) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
@@ -233,27 +224,9 @@ func (consumer *consumer) Cleanup(sarama.ConsumerGroupSession) error {
 }
 
 func (consumer *consumer) Setup(sarama.ConsumerGroupSession) error {
-	consumer.once.Do(func() {
-		close(consumer.ready)
-	})
-
 	return nil
 }
 
-// AddTopicHandler adds a handler and configuration for a topic
-func (k *Kafka) AddTopicHandler(topic string, handlerConfig SubscriptionHandlerConfig) {
-	k.subscribeLock.Lock()
-	k.subscribeTopics[topic] = handlerConfig
-	k.subscribeLock.Unlock()
-}
-
-// RemoveTopicHandler removes a topic handler
-func (k *Kafka) RemoveTopicHandler(topic string) {
-	k.subscribeLock.Lock()
-	delete(k.subscribeTopics, topic)
-	k.subscribeLock.Unlock()
-}
-
 // checkBulkSubscribe checks if a bulk handler and config are correctly registered for provided topic
 func (k *Kafka) checkBulkSubscribe(topic string) bool {
 	if bulkHandlerConfig, ok := k.subscribeTopics[topic]; ok &&
@@ -275,120 +248,3 @@ func (k *Kafka) GetTopicHandlerConfig(topic string) (SubscriptionHandlerConfig,
 	return SubscriptionHandlerConfig{},
 		fmt.Errorf("any handler for messages of topic %s not found", topic)
 }
-
-// Subscribe to topic in the Kafka cluster, in a background goroutine
-func (k *Kafka) Subscribe(ctx context.Context) error {
-	if k.consumerGroup == "" {
-		return errors.New("kafka: consumerGroup must be set to subscribe")
-	}
-
-	k.subscribeLock.Lock()
-	defer k.subscribeLock.Unlock()
-
-	topics := k.subscribeTopics.TopicList()
-	if len(topics) == 0 {
-		// Nothing to subscribe to
-		return nil
-	}
-	k.consumer.skipConsume = true
-
-	ctxCreateFn := func() {
-		consumeCtx, cancel := context.WithCancel(context.Background())
-
-		k.consumer.consumeCtx = consumeCtx
-		k.consumer.consumeCancel = cancel
-
-		k.consumer.skipConsume = false
-	}
-
-	if k.cg == nil {
-		cg, err := sarama.NewConsumerGroup(k.brokers, k.consumerGroup, k.config)
-		if err != nil {
-			return err
-		}
-
-		k.cg = cg
-
-		ready := make(chan bool)
-		k.consumer = consumer{
-			k:       k,
-			ready:   ready,
-			running: make(chan struct{}),
-		}
-
-		ctxCreateFn()
-
-		go func() {
-			k.logger.Debugf("Subscribed and listening to topics: %s", topics)
-
-			for {
-				// If the context was cancelled, as is the case when handling SIGINT and SIGTERM below, then this pops
-				// us out of the consume loop
-				if ctx.Err() != nil {
-					k.logger.Info("Consume context cancelled")
-					break
-				}
-
-				k.logger.Debugf("Starting loop to consume.")
-
-				if k.consumer.skipConsume {
-					continue
-				}
-
-				topics = k.subscribeTopics.TopicList()
-
-				// Consume the requested topics
-				bo := backoff.WithContext(backoff.NewConstantBackOff(k.consumeRetryInterval), ctx)
-				innerErr := retry.NotifyRecover(func() error {
-					if ctxErr := ctx.Err(); ctxErr != nil {
-						return backoff.Permanent(ctxErr)
-					}
-					return k.cg.Consume(k.consumer.consumeCtx, topics, &(k.consumer))
-				}, bo, func(err error, t time.Duration) {
-					k.logger.Errorf("Error consuming %v. Retrying...: %v", topics, err)
-				}, func() {
-					k.logger.Infof("Recovered consuming %v", topics)
-				})
-				if innerErr != nil && !errors.Is(innerErr, context.Canceled) {
-					k.logger.Errorf("Permanent error consuming %v: %v", topics, innerErr)
-				}
-			}
-
-			k.logger.Debugf("Closing ConsumerGroup for topics: %v", topics)
-			err := k.cg.Close()
-			if err != nil {
-				k.logger.Errorf("Error closing consumer group: %v", err)
-			}
-
-			// Ensure running channel is only closed once.
-			if k.consumer.stopped.CompareAndSwap(false, true) {
-				close(k.consumer.running)
-			}
-		}()
-
-		<-ready
-	} else {
-		// The consumer group is already created and consuming topics. This means a new subscription is being added
-		k.consumer.consumeCancel()
-		ctxCreateFn()
-	}
-
-	return nil
-}
-
-// Close down consumer group resources, refresh once.
-func (k *Kafka) closeSubscriptionResources() {
-	if k.cg != nil {
-		err := k.cg.Close()
-		if err != nil {
-			k.logger.Errorf("Error closing consumer group: %v", err)
-		}
-
-		k.consumer.once.Do(func() {
-			// Wait for shutdown to be complete
-			<-k.consumer.running
-			close(k.consumer.ready)
-			k.consumer.once = sync.Once{}
-		})
-	}
-}

Diff for: common/component/kafka/kafka.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/IBM/sarama"
@@ -34,19 +35,24 @@ import (
 
 // Kafka allows reading/writing to a Kafka consumer group.
 type Kafka struct {
-	producer        sarama.SyncProducer
-	consumerGroup   string
-	brokers         []string
-	logger          logger.Logger
-	authType        string
-	saslUsername    string
-	saslPassword    string
-	initialOffset   int64
+	producer      sarama.SyncProducer
+	consumerGroup string
+	brokers       []string
+	logger        logger.Logger
+	authType      string
+	saslUsername  string
+	saslPassword  string
+	initialOffset int64
+	config        *sarama.Config
+
 	cg              sarama.ConsumerGroup
-	consumer        consumer
-	config          *sarama.Config
 	subscribeTopics TopicHandlerConfig
 	subscribeLock   sync.Mutex
+	consumerCancel  context.CancelFunc
+	consumerWG      sync.WaitGroup
+	closeCh         chan struct{}
+	closed          atomic.Bool
+	wg              sync.WaitGroup
 
 	// schema registry settings
 	srClient                   srclient.ISchemaRegistryClient
@@ -106,7 +112,7 @@ func NewKafka(logger logger.Logger) *Kafka {
 	return &Kafka{
 		logger:          logger,
 		subscribeTopics: make(TopicHandlerConfig),
-		subscribeLock:   sync.Mutex{},
+		closeCh:         make(chan struct{}),
 	}
 }
 
@@ -184,11 +190,11 @@ func (k *Kafka) Init(ctx context.Context, metadata map[string]string) error {
 
 	// Default retry configuration is used if no
 	// backOff properties are set.
-	if err := retry.DecodeConfigWithPrefix(
+	if rerr := retry.DecodeConfigWithPrefix(
 		&k.backOffConfig,
 		metadata,
-		"backOff"); err != nil {
-		return err
+		"backOff"); rerr != nil {
+		return rerr
 	}
 	k.consumeRetryEnabled = meta.ConsumeRetryEnabled
 	k.consumeRetryInterval = meta.ConsumeRetryInterval
@@ -207,22 +213,41 @@ func (k *Kafka) Init(ctx context.Context, metadata map[string]string) error {
 	}
 	k.logger.Debug("Kafka message bus initialization complete")
 
+	k.cg, err = sarama.NewConsumerGroup(k.brokers, k.consumerGroup, k.config)
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 
-func (k *Kafka) Close() (err error) {
-	k.closeSubscriptionResources()
+func (k *Kafka) Close() error {
+	defer k.wg.Wait()
+	defer k.consumerWG.Wait()
 
-	if k.producer != nil {
-		err = k.producer.Close()
-		k.producer = nil
-	}
+	errs := make([]error, 2)
+	if k.closed.CompareAndSwap(false, true) {
+		close(k.closeCh)
+
+		if k.producer != nil {
+			errs[0] = k.producer.Close()
+			k.producer = nil
+		}
+
+		if k.internalContext != nil {
+			k.internalContextCancel()
+		}
+
+		k.subscribeLock.Lock()
+		if k.consumerCancel != nil {
+			k.consumerCancel()
+		}
+		k.subscribeLock.Unlock()
 
-	if k.internalContext != nil {
-		k.internalContextCancel()
+		errs[1] = k.cg.Close()
 	}
 
-	return err
+	return errors.Join(errs...)
 }
 
 func getSchemaSubject(topic string) string {