github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/clients/pkg/promtail/targets/kafka/consumer.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/Shopify/sarama"
    10  	"github.com/go-kit/log"
    11  	"github.com/go-kit/log/level"
    12  	"github.com/grafana/dskit/backoff"
    13  
    14  	"github.com/grafana/loki/clients/pkg/promtail/targets/target"
    15  )
    16  
    17  var defaultBackOff = backoff.Config{
    18  	MinBackoff: 1 * time.Second,
    19  	MaxBackoff: 60 * time.Second,
    20  	MaxRetries: 20,
    21  }
    22  
    23  type RunnableTarget interface {
    24  	target.Target
    25  	run()
    26  }
    27  
    28  type TargetDiscoverer interface {
    29  	NewTarget(sarama.ConsumerGroupSession, sarama.ConsumerGroupClaim) (RunnableTarget, error)
    30  }
    31  
    32  // consumer handle a group consumer instance.
    33  // It will create a new target for every consumer claim using the `TargetDiscoverer`.
    34  type consumer struct {
    35  	sarama.ConsumerGroup
    36  	discoverer TargetDiscoverer
    37  	logger     log.Logger
    38  
    39  	ctx    context.Context
    40  	cancel context.CancelFunc
    41  	wg     sync.WaitGroup
    42  
    43  	mutex          sync.Mutex // used during rebalancing setup and tear down
    44  	activeTargets  []target.Target
    45  	droppedTargets []target.Target
    46  }
    47  
    48  // start starts the consumer for a given list of topics.
    49  func (c *consumer) start(ctx context.Context, topics []string) {
    50  	c.wg.Wait()
    51  	c.wg.Add(1)
    52  
    53  	c.ctx, c.cancel = context.WithCancel(ctx)
    54  	level.Info(c.logger).Log("msg", "starting consumer", "topics", fmt.Sprintf("%+v", topics))
    55  
    56  	go func() {
    57  		defer c.wg.Done()
    58  		backoff := backoff.New(c.ctx, defaultBackOff)
    59  		for {
    60  			// Calling Consume in an infinite loop in case rebalancing is kicking in.
    61  			// In which case all claims will be renewed.
    62  			err := c.ConsumerGroup.Consume(c.ctx, topics, c)
    63  			if err != nil && err != context.Canceled {
    64  				level.Error(c.logger).Log("msg", "error from the consumer, retrying...", "err", err)
    65  				// backoff before re-trying.
    66  				backoff.Wait()
    67  				if backoff.Ongoing() {
    68  					continue
    69  				}
    70  				level.Error(c.logger).Log("msg", "maximun error from the consumer reached", "last_err", err)
    71  				return
    72  			}
    73  			if c.ctx.Err() != nil || err == context.Canceled {
    74  				level.Info(c.logger).Log("msg", "stopping consumer", "topics", fmt.Sprintf("%+v", topics))
    75  				return
    76  			}
    77  			backoff.Reset()
    78  		}
    79  	}()
    80  }
    81  
    82  // ConsumeClaim creates a target for the given received claim and start reading message from it.
    83  func (c *consumer) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
    84  	c.wg.Add(1)
    85  	defer c.wg.Done()
    86  
    87  	t, err := c.discoverer.NewTarget(session, claim)
    88  	if err != nil {
    89  		return err
    90  	}
    91  	if len(t.Labels()) == 0 {
    92  		c.addDroppedTarget(t)
    93  		t.run()
    94  		return nil
    95  	}
    96  	c.addTarget(t)
    97  	level.Info(c.logger).Log("msg", "consuming topic", "details", t.Details())
    98  	t.run()
    99  
   100  	return nil
   101  }
   102  
   103  // Setup is run at the beginning of a new session, before ConsumeClaim
   104  func (c *consumer) Setup(session sarama.ConsumerGroupSession) error {
   105  	c.resetTargets()
   106  	return nil
   107  }
   108  
   109  // Cleanup is run at the end of a session, once all ConsumeClaim goroutines have exited
   110  func (c *consumer) Cleanup(sarama.ConsumerGroupSession) error {
   111  	c.resetTargets()
   112  	return nil
   113  }
   114  
   115  // stop stops the consumer.
   116  func (c *consumer) stop() {
   117  	c.cancel()
   118  	c.wg.Wait()
   119  	c.resetTargets()
   120  }
   121  
   122  func (c *consumer) resetTargets() {
   123  	c.mutex.Lock()
   124  	defer c.mutex.Unlock()
   125  	c.activeTargets = nil
   126  	c.droppedTargets = nil
   127  }
   128  
   129  func (c *consumer) getActiveTargets() []target.Target {
   130  	c.mutex.Lock()
   131  	defer c.mutex.Unlock()
   132  	return c.activeTargets
   133  }
   134  
   135  func (c *consumer) getDroppedTargets() []target.Target {
   136  	c.mutex.Lock()
   137  	defer c.mutex.Unlock()
   138  	return c.droppedTargets
   139  }
   140  
   141  func (c *consumer) addTarget(t target.Target) {
   142  	c.mutex.Lock()
   143  	defer c.mutex.Unlock()
   144  	c.activeTargets = append(c.activeTargets, t)
   145  }
   146  
   147  func (c *consumer) addDroppedTarget(t target.Target) {
   148  	c.mutex.Lock()
   149  	defer c.mutex.Unlock()
   150  	c.droppedTargets = append(c.droppedTargets, t)
   151  }