Bump github.com/Shopify/sarama from v1.33.0 to v1.38.1

Signed-off-by: axfor <aixiaoxiang2009@hotmail.com>

cmd/ingester/app/consumer/consumer.go

@@ -15,11 +15,13 @@
 package consumer
 
 import (
+	"context"
+	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/Shopify/sarama"
-	sc "github.com/bsm/sarama-cluster"
 	"go.uber.org/zap"
 
 	"github.com/jaegertracing/jaeger/cmd/ingester/app/processor"
@@ -46,16 +48,30 @@ type Consumer struct {
 
 	deadlockDetector deadlockDetector
 
-	partitionIDToState  map[int32]*consumerState
-	partitionMapLock    sync.Mutex
-	partitionsHeld      int64
+	partitionsHeld      atomic.Int64
 	partitionsHeldGauge metrics.Gauge
 
 	doneWg sync.WaitGroup
+
+	consumerReady *ConsumerReady
+
+	topic  string
+	cancel context.CancelFunc
+}
+
+type ConsumerReady struct {
+	ready     chan struct{}
+	closeOnce sync.Once
 }
 
-type consumerState struct {
-	partitionConsumer sc.PartitionConsumer
+func (c *ConsumerReady) Ready() {
+	<-c.ready
+}
+
+func (c *ConsumerReady) Close() {
+	c.closeOnce.Do(func() {
+		close(c.ready)
+	})
 }
 
 // New is a constructor for a Consumer
@@ -67,25 +83,41 @@ func New(params Params) (*Consumer, error) {
 		internalConsumer:    params.InternalConsumer,
 		processorFactory:    params.ProcessorFactory,
 		deadlockDetector:    deadlockDetector,
-		partitionIDToState:  make(map[int32]*consumerState),
 		partitionsHeldGauge: partitionsHeldGauge(params.MetricsFactory),
+		consumerReady: &ConsumerReady{
+			ready: make(chan struct{}),
+		},
+		topic: params.ProcessorFactory.topic,
 	}, nil
 }
 
 // Start begins consuming messages in a go routine
 func (c *Consumer) Start() {
 	c.deadlockDetector.start()
+	ctx, cancel := context.WithCancel(context.Background())
+	c.cancel = cancel
+	c.doneWg.Add(1)
+
 	go func() {
-		c.logger.Info("Starting main loop")
-		for pc := range c.internalConsumer.Partitions() {
-			c.partitionMapLock.Lock()
-			c.partitionIDToState[pc.Partition()] = &consumerState{partitionConsumer: pc}
-			c.partitionMapLock.Unlock()
-			c.partitionMetrics(pc.Partition()).startCounter.Inc(1)
-
-			c.doneWg.Add(2)
-			go c.handleMessages(pc)
-			go c.handleErrors(pc.Partition(), pc.Errors())
+		defer c.doneWg.Done()
+		for {
+			select {
+			case <-ctx.Done():
+				c.logger.Error("ctx canceld")
+				return
+			default:
+				c.logger.Info("Topic", zap.Strings("topic", strings.Split(c.topic, ",")))
+				if err := c.internalConsumer.Consume(ctx, strings.Split(c.topic, ","), c); err != nil {
+					c.logger.Error("Error from consumer", zap.Error(err))
+				}
+				// check if context was cancelled, signaling that the consumer should stop
+				if ctx.Err() != nil {
+					return
+				}
+				c.consumerReady = &ConsumerReady{
+					ready: make(chan struct{}),
+				}
+			}
 		}
 	}()
 }
@@ -98,79 +130,109 @@ func (c *Consumer) Close() error {
 
 	c.logger.Debug("Closing deadlock detector")
 	c.deadlockDetector.close()
+	if c.cancel != nil {
+		c.cancel()
+	}
 
 	c.logger.Debug("Waiting for messages and errors to be handled")
 	c.doneWg.Wait()
 
 	return err
 }
 
-// handleMessages handles incoming Kafka messages on a channel
-func (c *Consumer) handleMessages(pc sc.PartitionConsumer) {
-	c.logger.Info("Starting message handler", zap.Int32("partition", pc.Partition()))
-	c.partitionMapLock.Lock()
-	c.partitionsHeld++
-	c.partitionsHeldGauge.Update(c.partitionsHeld)
-	c.partitionMapLock.Unlock()
+// Ready is consumer running
+func (c *Consumer) Ready() {
+	c.consumerReady.Ready()
+}
+
+// Setup is run at the beginning of a new session, before ConsumeClaim
+func (c *Consumer) Setup(sarama.ConsumerGroupSession) error {
+	// Mark the consumer as ready
+	c.consumerReady.Close()
+	return nil
+}
+
+// Cleanup is run at the end of a session, once all ConsumeClaim goroutines have exited
+func (c *Consumer) Cleanup(sarama.ConsumerGroupSession) error {
+	return nil
+}
+
+// ConsumeClaim must start a consumer loop of ConsumerGroupClaim's Messages().
+func (c *Consumer) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
+	c.partitionMetrics(claim.Partition()).startCounter.Inc(1)
+
+	c.doneWg.Add(2)
+	go c.handleErrors(claim.Partition(), c.internalConsumer.Errors())
+	c.handleMessages(session, claim)
+	return nil
+}
+
+// handleMessages starting message handler
+func (c *Consumer) handleMessages(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) {
+	c.logger.Info("Starting message handler", zap.Int32("partition", claim.Partition()))
+
+	c.partitionsHeld.Add(1)
+	c.partitionsHeldGauge.Update(c.partitionsHeld.Load())
+
 	defer func() {
-		c.closePartition(pc)
-		c.partitionMapLock.Lock()
-		c.partitionsHeld--
-		c.partitionsHeldGauge.Update(c.partitionsHeld)
-		c.partitionMapLock.Unlock()
+		c.closePartition(claim)
+		c.partitionsHeld.Add(-1)
+		c.partitionsHeldGauge.Update(c.partitionsHeld.Load())
 		c.doneWg.Done()
 	}()
 
-	msgMetrics := c.newMsgMetrics(pc.Partition())
+	msgMetrics := c.newMsgMetrics(claim.Partition())
 
 	var msgProcessor processor.SpanProcessor
 
-	deadlockDetector := c.deadlockDetector.startMonitoringForPartition(pc.Partition())
+	deadlockDetector := c.deadlockDetector.startMonitoringForPartition(claim.Partition())
 	defer deadlockDetector.close()
 
 	for {
 		select {
-		case msg, ok := <-pc.Messages():
-			if !ok {
-				c.logger.Info("Message channel closed. ", zap.Int32("partition", pc.Partition()))
-				return
-			}
+		case msg := <-claim.Messages():
+
 			c.logger.Debug("Got msg", zap.Any("msg", msg))
 			msgMetrics.counter.Inc(1)
 			msgMetrics.offsetGauge.Update(msg.Offset)
-			msgMetrics.lagGauge.Update(pc.HighWaterMarkOffset() - msg.Offset - 1)
+			msgMetrics.lagGauge.Update(claim.HighWaterMarkOffset() - msg.Offset - 1)
 			deadlockDetector.incrementMsgCount()
 
 			if msgProcessor == nil {
-				msgProcessor = c.processorFactory.new(pc.Partition(), msg.Offset-1)
+				msgProcessor = c.processorFactory.new(session, claim, msg.Offset-1)
 				defer msgProcessor.Close()
 			}
 
 			msgProcessor.Process(saramaMessageWrapper{msg})
 
+		// Should return when `session.Context()` is done.
+		// If not, will raise `ErrRebalanceInProgress` or `read tcp <ip>:<port>: i/o timeout` when kafka rebalance. see:
+		// https://github.com/Shopify/sarama/issues/1192
+		case <-session.Context().Done():
+			c.logger.Info("Session done", zap.Int32("partition", claim.Partition()))
+			return
 		case <-deadlockDetector.closePartitionChannel():
-			c.logger.Info("Closing partition due to inactivity", zap.Int32("partition", pc.Partition()))
+			c.logger.Info("Closing partition due to inactivity", zap.Int32("partition", claim.Partition()))
 			return
 		}
 	}
 }
 
-func (c *Consumer) closePartition(partitionConsumer sc.PartitionConsumer) {
-	c.logger.Info("Closing partition consumer", zap.Int32("partition", partitionConsumer.Partition()))
-	partitionConsumer.Close() // blocks until messages channel is drained
-	c.partitionMetrics(partitionConsumer.Partition()).closeCounter.Inc(1)
-	c.logger.Info("Closed partition consumer", zap.Int32("partition", partitionConsumer.Partition()))
+// closePartition close partition of consumer
+func (c *Consumer) closePartition(claim sarama.ConsumerGroupClaim) {
+	c.logger.Info("Closing partition consumer", zap.Int32("partition", claim.Partition()))
+	c.partitionMetrics(claim.Partition()).closeCounter.Inc(1)
 }
 
 // handleErrors handles incoming Kafka consumer errors on a channel
-func (c *Consumer) handleErrors(partition int32, errChan <-chan *sarama.ConsumerError) {
+func (c *Consumer) handleErrors(partition int32, errChan <-chan error) {
 	c.logger.Info("Starting error handler", zap.Int32("partition", partition))
 	defer c.doneWg.Done()
 
 	errMetrics := c.newErrMetrics(partition)
 	for err := range errChan {
 		errMetrics.errCounter.Inc(1)
-		c.logger.Error("Error consuming from Kafka", zap.Error(err))
+		c.logger.Error("Error consuming from kafka", zap.Error(err))
 	}
 	c.logger.Info("Finished handling errors", zap.Int32("partition", partition))
 }