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"
+	"fmt"
 	"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"
@@ -44,133 +46,246 @@ type Consumer struct {
 	internalConsumer consumer.Consumer
 	processorFactory ProcessorFactory
 
-	deadlockDetector deadlockDetector
+	deadlockDetector   deadlockDetector
+	deadlockNotifyCh   chan struct{}
+	deadlockNotifyOnce sync.Once
 
-	partitionIDToState  map[int32]*consumerState
-	partitionMapLock    sync.Mutex
-	partitionsHeld      int64
+	partitionsHeld      atomic.Int64
 	partitionsHeldGauge metrics.Gauge
 
+	consumerReady *consumerReady
+
 	doneWg sync.WaitGroup
+
+	topic  string
+	cancel context.CancelFunc
+}
+
+type consumerReady struct {
+	readyCh   chan struct{}
+	closeOnce sync.Once
+}
+
+func (c *consumerReady) waitReady() {
+	<-c.readyCh
+}
+
+func (c *consumerReady) markReady() {
+	c.close()
 }
 
-type consumerState struct {
-	partitionConsumer sc.PartitionConsumer
+func (c *consumerReady) close() {
+	c.closeOnce.Do(func() {
+		close(c.readyCh)
+	})
 }
 
 // New is a constructor for a Consumer
 func New(params Params) (*Consumer, error) {
 	deadlockDetector := newDeadlockDetector(params.MetricsFactory, params.Logger, params.DeadlockCheckInterval)
+	// new Consumer
 	return &Consumer{
 		metricsFactory:      params.MetricsFactory,
 		logger:              params.Logger,
 		internalConsumer:    params.InternalConsumer,
 		processorFactory:    params.ProcessorFactory,
 		deadlockDetector:    deadlockDetector,
-		partitionIDToState:  make(map[int32]*consumerState),
+		deadlockNotifyCh:    make(chan struct{}, 1),
 		partitionsHeldGauge: partitionsHeldGauge(params.MetricsFactory),
+		consumerReady: &consumerReady{
+			readyCh: make(chan struct{}, 1),
+		},
+		topic: params.ProcessorFactory.topic,
 	}, nil
 }
 
 // Start begins consuming messages in a go routine
 func (c *Consumer) Start() {
-	c.deadlockDetector.start()
+	c.doStart(true)
+}
+
+// Start begins consuming messages in a go routine and consumer is running
+func (c *Consumer) StartWithReady() {
+	c.Start()
+	c.consumerReady.waitReady()
+}
+
+// Start begins consuming and wait is running and disable global deadlock detector,
+// There are two deadlock detectors, one global and one specially partition
+// If the consumer is running, the specially partition deadlock detector takes effect
+func (c *Consumer) startWithReadyAndDisableGlobalDeadlockDetector() {
+	c.doStart(false)
+	c.consumerReady.waitReady()
+}
+
+// doStart begins consuming messages in a go routine
+func (c *Consumer) doStart(startAllPartitionDeadlockDetector bool) {
+	if startAllPartitionDeadlockDetector {
+		// all partition deadlock detector
+		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()
+		c.startConsume(ctx)
 	}()
+
+	go c.handleErrors(ctx)
+}
+
+func (c *Consumer) startConsume(ctx context.Context) {
+	topic := []string{c.topic}
+	// reference consumerReady and make sure to close it, because later consumerReady may reference a new object
+	ready := c.consumerReady
+	defer ready.close()
+
+	for {
+		select {
+		case <-ctx.Done():
+			c.logger.Error("ctx canceled")
+			return
+		default:
+			c.logger.Info("Topic", zap.Strings("topic", topic))
+			if err := c.internalConsumer.Consume(ctx, topic, c); err != nil {
+				c.logger.Error("Error from consumer", zap.Error(err))
+			}
+
+			c.consumerReady.close()
+
+			c.consumerReady = &consumerReady{
+				readyCh: make(chan struct{}, 1),
+			}
+
+			// check if context was cancelled, signaling that the consumer should stop
+			if ctx.Err() != nil {
+				if c.cancel != nil {
+					c.cancel()
+				}
+			}
+		}
+	}
+}
+
+// markDeadlock mark consume is deadlocked
+func (c *Consumer) markDeadlock() {
+	c.deadlockNotifyOnce.Do(func() {
+		close(c.deadlockNotifyCh)
+	})
+}
+
+// Deadlock return a consume deadlock notify chan
+func (c *Consumer) Deadlock() <-chan struct{} {
+	return c.deadlockNotifyCh
 }
 
 // Close closes the Consumer and underlying sarama consumer
 func (c *Consumer) Close() error {
 	// Close the internal consumer, which will close each partition consumers' message and error channels.
 	c.logger.Info("Closing parent consumer")
 	err := c.internalConsumer.Close()
-
-	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()
-	defer func() {
-		c.closePartition(pc)
-		c.partitionMapLock.Lock()
-		c.partitionsHeld--
-		c.partitionsHeldGauge.Update(c.partitionsHeld)
-		c.partitionMapLock.Unlock()
-		c.doneWg.Done()
-	}()
+// 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.markReady()
+	return nil
+}
 
-	msgMetrics := c.newMsgMetrics(pc.Partition())
+// 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)
+
+	return c.startMessagesLoop(session, claim)
+}
+
+func (c *Consumer) startMessagesLoop(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
+	return c.handleMessages(session, claim)
+}
+
+// handleMessages starting message handler
+func (c *Consumer) handleMessages(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
+	c.logger.Info("Starting message handler", zap.Int32("partition", claim.Partition()))
 
-	var msgProcessor processor.SpanProcessor
+	c.partitionsHeld.Add(1)
+	c.partitionsHeldGauge.Update(c.partitionsHeld.Load())
 
-	deadlockDetector := c.deadlockDetector.startMonitoringForPartition(pc.Partition())
-	defer deadlockDetector.close()
+	msgMetrics := c.newMsgMetrics(claim.Partition())
+	var msgProcessor processor.SpanProcessor = nil
+
+	partitionDeadlockDetector := c.deadlockDetector.startMonitoringForPartition(claim.Partition())
+
+	defer func() {
+		partitionDeadlockDetector.close()
+		c.closePartition(claim)
+		c.partitionsHeld.Add(-1)
+		c.partitionsHeldGauge.Update(c.partitionsHeld.Load())
+	}()
 
 	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)
-			deadlockDetector.incrementMsgCount()
+			msgMetrics.lagGauge.Update(claim.HighWaterMarkOffset() - msg.Offset - 1)
+			partitionDeadlockDetector.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})
 
-		case <-deadlockDetector.closePartitionChannel():
-			c.logger.Info("Closing partition due to inactivity", zap.Int32("partition", pc.Partition()))
-			return
+		// 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 session.Context().Err()
+		case <-partitionDeadlockDetector.closePartitionChannel():
+			c.logger.Info("Closing partition due to inactivity", zap.Int32("partition", claim.Partition()))
+			c.markDeadlock()
+			return fmt.Errorf("closing partition[%d] due to inactivity", claim.Partition())
 		}
 	}
 }
 
-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) {
-	c.logger.Info("Starting error handler", zap.Int32("partition", partition))
-	defer c.doneWg.Done()
+func (c *Consumer) handleErrors(ctx context.Context) {
+	c.logger.Info("Starting error handler")
 
-	errMetrics := c.newErrMetrics(partition)
+	errChan := c.internalConsumer.Errors()
+	errMetrics := c.newErrMetrics()
 	for err := range errChan {
 		errMetrics.errCounter.Inc(1)
 		c.logger.Error("Error consuming from Kafka", zap.Error(err))
 	}
-	c.logger.Info("Finished handling errors", zap.Int32("partition", partition))
 }

cmd/ingester/app/consumer/consumer_metrics.go

@@ -37,25 +37,29 @@ type partitionMetrics struct {
 	closeCounter metrics.Counter
 }
 
-func (c *Consumer) namespace(partition int32) metrics.Factory {
+func (c *Consumer) namespace() metrics.Factory {
+	return c.metricsFactory.Namespace(metrics.NSOptions{Name: consumerNamespace, Tags: map[string]string{}})
+}
+
+func (c *Consumer) namespaceWithPartition(partition int32) metrics.Factory {
 	return c.metricsFactory.Namespace(metrics.NSOptions{Name: consumerNamespace, Tags: map[string]string{"partition": strconv.Itoa(int(partition))}})
 }
 
 func (c *Consumer) newMsgMetrics(partition int32) msgMetrics {
-	f := c.namespace(partition)
+	f := c.namespaceWithPartition(partition)
 	return msgMetrics{
 		counter:     f.Counter(metrics.Options{Name: "messages", Tags: nil}),
 		offsetGauge: f.Gauge(metrics.Options{Name: "current-offset", Tags: nil}),
 		lagGauge:    f.Gauge(metrics.Options{Name: "offset-lag", Tags: nil}),
 	}
 }
 
-func (c *Consumer) newErrMetrics(partition int32) errMetrics {
-	return errMetrics{errCounter: c.namespace(partition).Counter(metrics.Options{Name: "errors", Tags: nil})}
+func (c *Consumer) newErrMetrics() errMetrics {
+	return errMetrics{errCounter: c.namespace().Counter(metrics.Options{Name: "errors", Tags: nil})}
 }
 
 func (c *Consumer) partitionMetrics(partition int32) partitionMetrics {
-	f := c.namespace(partition)
+	f := c.namespaceWithPartition(partition)
 	return partitionMetrics{
 		closeCounter: f.Counter(metrics.Options{Name: "partition-close", Tags: nil}),
 		startCounter: f.Counter(metrics.Options{Name: "partition-start", Tags: nil}),