fix kafka input partition ownership and autocommit on rebalance

e2e/kafka_auth/kafka_auth.go

@@ -181,4 +181,6 @@ type Consumer struct{}
 
 func (c Consumer) Assigned(_ context.Context, _ *kgo.Client, assigned map[string][]int32) {}
 
+func (c Consumer) Revoked(_ context.Context, _ *kgo.Client, revoked map[string][]int32) {}
+
 func (c Consumer) Lost(_ context.Context, _ *kgo.Client, lost map[string][]int32) {}

e2e/split_join/split_join.go

@@ -139,4 +139,6 @@ type Consumer struct{}
 
 func (c Consumer) Assigned(_ context.Context, _ *kgo.Client, assigned map[string][]int32) {}
 
+func (c Consumer) Revoked(_ context.Context, _ *kgo.Client, revoked map[string][]int32) {}
+
 func (c Consumer) Lost(_ context.Context, _ *kgo.Client, lost map[string][]int32) {}

plugin/input/kafka/client.go

@@ -24,7 +24,7 @@ func NewClient(ctx context.Context, c *Config, l *zap.Logger, s Consumer, tokenS
 		kgo.SessionTimeout(c.SessionTimeout_),
 		kgo.HeartbeatInterval(c.HeartbeatInterval_),
 		kgo.OnPartitionsAssigned(s.Assigned),
-		kgo.OnPartitionsRevoked(s.Lost),
+		kgo.OnPartitionsRevoked(s.Revoked),
 		kgo.OnPartitionsLost(s.Lost),
 		kgo.BlockRebalanceOnPoll(),
 	}...)

plugin/input/kafka/consumer.go

@@ -3,6 +3,7 @@ package kafka
 import (
 	"context"
 	"sync"
+	"time"
 
 	"github.com/ozontech/file.d/metric"
 	"github.com/ozontech/file.d/pipeline"
@@ -11,30 +12,40 @@ import (
 	"go.uber.org/zap"
 )
 
-type tp struct {
-	t string
-	p int32
+type partitionTopic struct {
+	topic     string
+	partition int32
 }
 
 type Consumer interface {
 	Assigned(_ context.Context, _ *kgo.Client, assigned map[string][]int32)
+	Revoked(_ context.Context, _ *kgo.Client, revoked map[string][]int32)
 	Lost(_ context.Context, _ *kgo.Client, lost map[string][]int32)
 }
 
 type splitConsume struct {
-	consumers              map[tp]*pconsumer
+	consumers              map[partitionTopic]*pconsumer
 	bufferSize             int
 	maxConcurrentConsumers int
 	idByTopic              map[string]int
 	controller             pipeline.InputPluginController
 	logger                 *zap.Logger
 	metaTemplater          *metadata.MetaTemplater
+	commitErrorsMetric     *metric.Counter
 	consumeErrorsMetric    *metric.Counter
+
+	mu    sync.RWMutex
+	owned map[partitionTopic]struct{}
 }
 
 func (s *splitConsume) Assigned(_ context.Context, _ *kgo.Client, assigned map[string][]int32) {
+	s.mu.Lock()
+	if s.owned == nil {
+		s.owned = make(map[partitionTopic]struct{})
+	}
 	for topic, partitions := range assigned {
 		for _, partition := range partitions {
+			s.owned[partitionTopic{topic, partition}] = struct{}{}
 			pc := &pconsumer{
 				topic:     topic,
 				partition: partition,
@@ -48,30 +59,58 @@ func (s *splitConsume) Assigned(_ context.Context, _ *kgo.Client, assigned map[s
 				logger:        s.logger,
 				metaTemplater: s.metaTemplater,
 			}
-			s.consumers[tp{topic, partition}] = pc
+			s.consumers[partitionTopic{topic, partition}] = pc
 			go pc.consume()
 		}
 	}
+	s.mu.Unlock()
+}
+
+func (s *splitConsume) Revoked(ctx context.Context, cl *kgo.Client, revoked map[string][]int32) {
+	commitCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	if err := cl.CommitMarkedOffsets(commitCtx); err != nil {
+		s.commitErrorsMetric.Inc()
+		s.logger.Error("can't commit marked offsets on revoke", zap.Error(err))
+	}
+	cancel()
+
+	s.stopConsumers(revoked)
 }
 
 func (s *splitConsume) Lost(_ context.Context, _ *kgo.Client, lost map[string][]int32) {
-	var wg sync.WaitGroup
-	defer wg.Wait()
+	s.stopConsumers(lost)
+}
 
+func (s *splitConsume) stopConsumers(lost map[string][]int32) {
+	s.mu.Lock()
+	var wg sync.WaitGroup
 	for topic, partitions := range lost {
 		for _, partition := range partitions {
-			tp := tp{topic, partition}
+			tp := partitionTopic{topic, partition}
 			pc := s.consumers[tp]
 			delete(s.consumers, tp)
-			close(pc.quit)
+			delete(s.owned, tp)
+			if pc == nil {
+				continue
+			}
 			pc.logger.Info("waiting for finish of consume", zap.String("topic", pc.topic), zap.Int32("partiton", pc.partition))
+			close(pc.quit)
 			wg.Add(1)
 			go func(pc *pconsumer) {
 				defer wg.Done()
 				<-pc.done
 			}(pc)
 		}
 	}
+	s.mu.Unlock()
+	wg.Wait()
+}
+
+func (s *splitConsume) Owns(topic string, partition int32) bool {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	_, ok := s.owned[partitionTopic{topic, partition}]
+	return ok
 }
 
 func (s *splitConsume) consume(ctx context.Context, cl *kgo.Client) {
@@ -94,13 +133,20 @@ func (s *splitConsume) consume(ctx context.Context, cl *kgo.Client) {
 		}
 
 		fetches.EachPartition(func(p kgo.FetchTopicPartition) {
-			tp := tp{p.Topic, p.Partition}
-			if consumer, ok := s.consumers[tp]; ok {
-				consumer.fetches <- p
-			} else {
+			tp := partitionTopic{p.Topic, p.Partition}
+			s.mu.RLock()
+			consumer, ok := s.consumers[tp]
+			s.mu.RUnlock()
+			if !ok {
 				s.logger.Error("consumer not ready yet", zap.String("topic", p.Topic), zap.Int32("partiton", p.Partition))
+				return
+			}
+			select {
+			case consumer.fetches <- p:
+			case <-ctx.Done():
 			}
 		})
+
 		cl.AllowRebalance()
 	}
 }
@@ -124,31 +170,51 @@ func (pc *pconsumer) consume() {
 	pc.logger.Info("starting consume", zap.String("topic", pc.topic), zap.Int32("partiton", pc.partition))
 	defer pc.logger.Info("closing consume", zap.String("topic", pc.topic), zap.Int32("partiton", pc.partition))
 	for {
+		var fetches kgo.FetchTopicPartition
+		var ok bool
 		select {
 		case <-pc.quit:
 			return
-		case fetches, ok := <-pc.fetches:
+		case fetches, ok = <-pc.fetches:
 			if !ok {
 				return // Channel closed
 			}
-			for i := range fetches.Records {
-				message := fetches.Records[i]
-				sourceID := assembleSourceID(
-					pc.topicID,
-					message.Partition,
-				)
-
-				offset := assembleOffset(message)
-				var metadataInfo metadata.MetaData
-				var err error
-				if pc.metaTemplater != nil {
-					metadataInfo, err = pc.metaTemplater.Render(newMetaInformation(message))
-					if err != nil {
-						pc.logger.Error("can't render meta data", zap.Error(err))
-					}
+		}
+
+		for i := range fetches.Records {
+			if !pc.active() {
+				return
+			}
+			message := fetches.Records[i]
+			sourceID := assembleSourceID(
+				pc.topicID,
+				message.Partition,
+			)
+
+			offset := assembleOffset(message)
+			var metadataInfo metadata.MetaData
+			var err error
+			if pc.metaTemplater != nil {
+				metadataInfo, err = pc.metaTemplater.Render(newMetaInformation(message))
+				if err != nil {
+					pc.logger.Error("can't render meta data", zap.Error(err))
 				}
-				_ = pc.controller.In(sourceID, "kafka", pipeline.NewOffsets(offset, nil), message.Value, true, metadataInfo)
 			}
+			select {
+			case <-pc.quit:
+				return
+			default:
+			}
+			pc.controller.In(sourceID, "kafka", pipeline.NewOffsets(offset, nil), message.Value, true, metadataInfo)
 		}
 	}
 }
+
+func (pc *pconsumer) active() bool {
+	select {
+	case <-pc.quit:
+		return false
+	default:
+		return true
+	}
+}

plugin/input/kafka/kafka.go

@@ -318,13 +318,14 @@ func (p *Plugin) Start(config pipeline.AnyConfig, params *pipeline.InputPluginPa
 	}
 
 	p.s = &splitConsume{
-		consumers:              make(map[tp]*pconsumer),
+		consumers:              make(map[partitionTopic]*pconsumer),
 		bufferSize:             p.config.ChannelBufferSize,
 		maxConcurrentConsumers: p.config.MaxConcurrentConsumers,
 		idByTopic:              idByTopic,
 		controller:             p.controller,
 		logger:                 p.logger,
 		metaTemplater:          p.metaTemplater,
+		commitErrorsMetric:     p.commitErrorsMetric,
 		consumeErrorsMetric:    p.consumeErrorsMetric,
 	}
 
@@ -371,9 +372,14 @@ func (p *Plugin) Stop() {
 func (p *Plugin) Commit(event *pipeline.Event) {
 	index, partition := disassembleSourceID(event.SourceID)
 
+	topic := p.config.Topics[index]
+	if !p.s.Owns(topic, partition) {
+		return
+	}
+
 	offset := disassembleOffset(event.Offset)
 	offsets := map[string]map[int32]kgo.EpochOffset{
-		p.config.Topics[index]: {partition: offset},
+		topic: {partition: offset},
 	}
 	p.client.MarkCommitOffsets(offsets)
 }
@@ -389,13 +395,18 @@ func disassembleSourceID(sourceID pipeline.SourceID) (index int, partition int32
 	return
 }
 
+// assembleOffset packs the record offset and leader epoch into a single int64
+// by shifting the offset left by 16 bits and placing the epoch in the low 16
+// bits. The epoch is stored offset by +1 so that an absent epoch (LeaderEpoch
+// == -1) is representable: it decodes back to -1 via disassembleOffset.
+// This encoding is valid for -1 <= LeaderEpoch <= 65534.
 func assembleOffset(message *kgo.Record) int64 {
-	return message.Offset<<16 + int64(message.LeaderEpoch)
+	return message.Offset<<16 + int64(message.LeaderEpoch+1)
 }
 
 func disassembleOffset(assembledOffset int64) kgo.EpochOffset {
 	offset := assembledOffset >> 16
-	epoch := int32(assembledOffset & 0xFFFF)
+	epoch := int32(assembledOffset&0xFFFF) - 1
 
 	return kgo.EpochOffset{
 		Offset: offset + 1,

plugin/input/kafka/kafka_test.go

@@ -21,16 +21,31 @@ func TestAssembleSourceID(t *testing.T) {
 }
 
 func TestAssembleOffset(t *testing.T) {
-	message := &kgo.Record{
-		LeaderEpoch: 27,
-		Offset:      237582035700,
+	tests := []struct {
+		name  string
+		epoch int32
+		off   int64
+	}{
+		{name: "regular epoch", epoch: 27, off: 237582035700},
+		{name: "no epoch", epoch: -1, off: 237582035700},
+		{name: "zero epoch", epoch: 0, off: 0},
+		{name: "max epoch", epoch: 65534, off: 12345},
+		{name: "high offset", epoch: 5, off: 1 << 40},
 	}
-	x := assembleOffset(message)
 
-	epochOffset := disassembleOffset(x)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			message := &kgo.Record{
+				LeaderEpoch: tt.epoch,
+				Offset:      tt.off,
+			}
 
-	assert.Equal(t, message.LeaderEpoch, epochOffset.Epoch, "values aren't equal")
-	assert.Equal(t, message.Offset+1, epochOffset.Offset, "values aren't equal")
+			epochOffset := disassembleOffset(assembleOffset(message))
+
+			assert.Equal(t, tt.epoch, epochOffset.Epoch, "epoch isn't equal")
+			assert.Equal(t, tt.off+1, epochOffset.Offset, "offset isn't equal")
+		})
+	}
 }
 
 func TestMetaInformationGetData(t *testing.T) {