[1.13] Fix Kafka consumer cancellation (#3347)

Signed-off-by: yaron2 <schneider.yaron@live.com>

common/component/kafka/consumer.go

@@ -324,7 +324,7 @@ func (k *Kafka) Subscribe(ctx context.Context) error {
 			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 {
+				if k.consumer.consumeCtx.Err() != nil {
 					k.logger.Info("Consume context cancelled")
 					break
 				}
@@ -377,12 +377,14 @@ func (k *Kafka) Subscribe(ctx context.Context) error {
 }
 
 // Close down consumer group resources, refresh once.
-func (k *Kafka) closeSubscriptionResources() {
+func (k *Kafka) CloseSubscriptionResources() {
 	if k.cg != nil {
+		k.consumer.consumeCancel()
 		err := k.cg.Close()
 		if err != nil {
 			k.logger.Errorf("Error closing consumer group: %v", err)
 		}
+		k.cg = nil
 
 		k.consumer.once.Do(func() {
 			// Wait for shutdown to be complete

common/component/kafka/kafka.go

@@ -211,7 +211,7 @@ func (k *Kafka) Init(ctx context.Context, metadata map[string]string) error {
 }
 
 func (k *Kafka) Close() (err error) {
-	k.closeSubscriptionResources()
+	k.CloseSubscriptionResources()
 
 	if k.producer != nil {
 		err = k.producer.Close()

pubsub/kafka/kafka.go

@@ -90,6 +90,7 @@ func (p *PubSub) subscribeUtil(ctx context.Context, req pubsub.SubscribeRequest,
 		// Wait for context cancelation
 		select {
 		case <-ctx.Done():
+			p.kafka.CloseSubscriptionResources()
 		case <-p.closeCh:
 		}