diff --git a/e2e/kafka_auth/kafka_auth.go b/e2e/kafka_auth/kafka_auth.go index 6d0e87050..6a5f72e51 100644 --- a/e2e/kafka_auth/kafka_auth.go +++ b/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) {} diff --git a/e2e/split_join/split_join.go b/e2e/split_join/split_join.go index 577784b93..92306abae 100644 --- a/e2e/split_join/split_join.go +++ b/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) {} diff --git a/plugin/input/kafka/client.go b/plugin/input/kafka/client.go index a570aeca8..0ed9f35f0 100644 --- a/plugin/input/kafka/client.go +++ b/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(), }...) diff --git a/plugin/input/kafka/consumer.go b/plugin/input/kafka/consumer.go index 8337bc4e3..73f909780 100644 --- a/plugin/input/kafka/consumer.go +++ b/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,23 +59,42 @@ 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() @@ -72,6 +102,15 @@ func (s *splitConsume) Lost(_ context.Context, _ *kgo.Client, lost map[string][] }(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 + } +} diff --git a/plugin/input/kafka/kafka.go b/plugin/input/kafka/kafka.go index 0d57ee54f..74e72220d 100644 --- a/plugin/input/kafka/kafka.go +++ b/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, diff --git a/plugin/input/kafka/kafka_test.go b/plugin/input/kafka/kafka_test.go index 5fb1b7edc..216a686c2 100644 --- a/plugin/input/kafka/kafka_test.go +++ b/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) {