snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)
This commit is contained in:
upstream-mirror
934
internal/consuming/kafka.go
Normal file
934
internal/consuming/kafka.go
Normal file
@@ -0,0 +1,934 @@
|
||||
package consuming
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/centrifugal/centrifugo/v6/internal/api"
|
||||
"github.com/centrifugal/centrifugo/v6/internal/configtypes"
|
||||
"github.com/centrifugal/centrifugo/v6/internal/logging"
|
||||
"github.com/centrifugal/centrifugo/v6/internal/metrics"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
awsv2cfg "github.com/aws/aws-sdk-go-v2/config"
|
||||
"github.com/aws/aws-sdk-go-v2/service/sts"
|
||||
"github.com/rs/zerolog"
|
||||
"github.com/twmb/franz-go/pkg/kerr"
|
||||
"github.com/twmb/franz-go/pkg/kgo"
|
||||
"github.com/twmb/franz-go/pkg/kmsg"
|
||||
mskaws "github.com/twmb/franz-go/pkg/sasl/aws"
|
||||
"github.com/twmb/franz-go/pkg/sasl/plain"
|
||||
"github.com/twmb/franz-go/pkg/sasl/scram"
|
||||
)
|
||||
|
||||
type KafkaConfig = configtypes.KafkaConsumerConfig
|
||||
|
||||
type testOnlyConfig struct {
|
||||
topicPartitionBeforePauseCh chan topicPartition
|
||||
topicPartitionPauseProceedCh chan struct{}
|
||||
fetchTopicPartitionSubmittedCh chan kgo.FetchTopicPartition
|
||||
// leaveGroupOnShutdown when true sends a manual LeaveGroup request on shutdown.
|
||||
// Used only in tests to simulate the old behavior with dynamic instance IDs.
|
||||
leaveGroupOnShutdown bool
|
||||
// injectFatalPollErrorCh, when a value is received from this channel (non-blocking
|
||||
// check on each pollUntilFatal entry), causes pollUntilFatal to return it as a fatal
|
||||
// error instead of actually polling — used to exercise the client re-init path in tests.
|
||||
injectFatalPollErrorCh chan error
|
||||
}
|
||||
|
||||
type topicPartition struct {
|
||||
topic string
|
||||
partition int32
|
||||
}
|
||||
|
||||
type KafkaConsumer struct {
|
||||
name string
|
||||
client *kgo.Client
|
||||
dispatcher Dispatcher
|
||||
config KafkaConfig
|
||||
// consumersMu protects the consumers map. During normal operation, concurrent access is already
|
||||
// prevented: franz-go serializes assigned/revoked/lost callbacks, and BlockRebalanceOnPoll
|
||||
// prevents callbacks from firing while pollUntilFatal reads the map. The mutex is needed
|
||||
// for shutdown, when the main goroutine accesses the map (via allTopicPartitions/killConsumers)
|
||||
// while the group management goroutine may still be running callbacks. Must not be held
|
||||
// during killConsumers' wg.Wait() to avoid blocking concurrent callbacks.
|
||||
consumersMu sync.Mutex
|
||||
consumers map[topicPartition]*partitionConsumer
|
||||
doneCh chan struct{}
|
||||
common *consumerCommon
|
||||
testOnlyConfig testOnlyConfig
|
||||
}
|
||||
|
||||
func NewKafkaConsumer(
|
||||
config KafkaConfig, dispatcher Dispatcher, common *consumerCommon,
|
||||
) (*KafkaConsumer, error) {
|
||||
if len(config.Brokers) == 0 {
|
||||
return nil, errors.New("brokers required")
|
||||
}
|
||||
if len(config.Topics) == 0 {
|
||||
return nil, errors.New("topics required")
|
||||
}
|
||||
if len(config.ConsumerGroup) == 0 {
|
||||
return nil, errors.New("consumer_group required")
|
||||
}
|
||||
if config.MaxPollRecords == 0 {
|
||||
config.MaxPollRecords = 100
|
||||
}
|
||||
if config.FetchMaxWait == 0 {
|
||||
config.FetchMaxWait = configtypes.Duration(500 * time.Millisecond)
|
||||
}
|
||||
if config.PartitionQueueMaxSize == -1 {
|
||||
config.PartitionQueueMaxSize = 0
|
||||
} else if config.PartitionQueueMaxSize == 0 {
|
||||
config.PartitionQueueMaxSize = 1000
|
||||
}
|
||||
consumer := &KafkaConsumer{
|
||||
name: common.name,
|
||||
dispatcher: dispatcher,
|
||||
config: config,
|
||||
consumers: make(map[topicPartition]*partitionConsumer),
|
||||
doneCh: make(chan struct{}),
|
||||
common: common,
|
||||
}
|
||||
cl, err := consumer.initClient()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error init Kafka client: %w", err)
|
||||
}
|
||||
consumer.client = cl
|
||||
return consumer, nil
|
||||
}
|
||||
|
||||
const kafkaClientID = "centrifugo"
|
||||
|
||||
// kgoZerologAdapter adapts zerolog.Logger to kgo.Logger interface.
|
||||
type kgoZerologAdapter struct {
|
||||
log zerolog.Logger
|
||||
}
|
||||
|
||||
func (l *kgoZerologAdapter) Level() kgo.LogLevel {
|
||||
switch {
|
||||
case logging.Enabled(logging.TraceLevel):
|
||||
return kgo.LogLevelDebug // Kafka client debug level is more like our trace.
|
||||
case logging.Enabled(logging.InfoLevel):
|
||||
return kgo.LogLevelInfo
|
||||
case logging.Enabled(logging.WarnLevel):
|
||||
return kgo.LogLevelWarn
|
||||
case logging.Enabled(logging.ErrorLevel):
|
||||
return kgo.LogLevelError
|
||||
default:
|
||||
return kgo.LogLevelNone
|
||||
}
|
||||
}
|
||||
|
||||
func (l *kgoZerologAdapter) Log(level kgo.LogLevel, msg string, keyvals ...any) {
|
||||
var e *zerolog.Event
|
||||
switch level {
|
||||
case kgo.LogLevelError:
|
||||
e = l.log.Error()
|
||||
case kgo.LogLevelWarn:
|
||||
e = l.log.Warn()
|
||||
case kgo.LogLevelInfo:
|
||||
e = l.log.Info()
|
||||
case kgo.LogLevelDebug:
|
||||
e = l.log.Trace() // Kafka client debug level is more like our trace.
|
||||
default:
|
||||
return
|
||||
}
|
||||
for i := 0; i+1 < len(keyvals); i += 2 {
|
||||
if key, ok := keyvals[i].(string); ok {
|
||||
e = e.Interface(key, keyvals[i+1])
|
||||
}
|
||||
}
|
||||
e.Msg(msg)
|
||||
}
|
||||
|
||||
// leaveGroup sends a manual LeaveGroup request for static members. franz-go's Close does not
|
||||
// send LeaveGroup for static members by design. This method is only used in tests to simulate
|
||||
// the old behavior where dynamic instance IDs were combined with a manual LeaveGroup.
|
||||
func (c *KafkaConsumer) leaveGroup(ctx context.Context, client *kgo.Client) error {
|
||||
req := kmsg.NewPtrLeaveGroupRequest()
|
||||
req.Group = c.config.ConsumerGroup
|
||||
instanceID := c.config.InstanceID
|
||||
reason := "shutdown"
|
||||
req.Members = []kmsg.LeaveGroupRequestMember{
|
||||
{
|
||||
InstanceID: &instanceID,
|
||||
Reason: &reason,
|
||||
},
|
||||
}
|
||||
_, err := req.RequestWith(ctx, client)
|
||||
return err
|
||||
}
|
||||
|
||||
// mskAssumeRoleSessionName is the STS RoleSessionName for AssumeRole calls.
|
||||
const mskAssumeRoleSessionName = "centrifugo-msk"
|
||||
|
||||
// mskAssumeRoleAuth obtains fresh STS credentials for each MSK IAM SASL handshake.
|
||||
// Call AssumeRole directly on every callback; do not use CredentialsCache, which can
|
||||
// return the same credentials MSK already considers expired (see franz-go #731).
|
||||
// RoleSessionName stays constant so the broker principal is stable across re-auth
|
||||
// (see aws-msk-iam-auth README).
|
||||
type mskAssumeRoleAuth struct {
|
||||
stsClient *sts.Client
|
||||
roleARN string
|
||||
}
|
||||
|
||||
func newMSKAssumeRoleAuth(ctx context.Context, roleARN string) (*mskAssumeRoleAuth, error) {
|
||||
awsCfg, err := awsv2cfg.LoadDefaultConfig(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &mskAssumeRoleAuth{
|
||||
stsClient: sts.NewFromConfig(awsCfg),
|
||||
roleARN: roleARN,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (a *mskAssumeRoleAuth) auth(ctx context.Context) (mskaws.Auth, error) {
|
||||
out, err := a.stsClient.AssumeRole(ctx, &sts.AssumeRoleInput{
|
||||
RoleArn: aws.String(a.roleARN),
|
||||
RoleSessionName: aws.String(mskAssumeRoleSessionName),
|
||||
})
|
||||
if err != nil {
|
||||
return mskaws.Auth{}, err
|
||||
}
|
||||
if out.Credentials == nil {
|
||||
return mskaws.Auth{}, errors.New("assume role returned no credentials")
|
||||
}
|
||||
creds := out.Credentials
|
||||
if creds.AccessKeyId == nil || creds.SecretAccessKey == nil {
|
||||
return mskaws.Auth{}, errors.New("assume role returned incomplete credentials")
|
||||
}
|
||||
sessionToken := ""
|
||||
if creds.SessionToken != nil {
|
||||
sessionToken = *creds.SessionToken
|
||||
}
|
||||
return mskaws.Auth{
|
||||
AccessKey: *creds.AccessKeyId,
|
||||
SecretKey: *creds.SecretAccessKey,
|
||||
SessionToken: sessionToken,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// defaultKafkaDialTimeout is the fallback TCP dial timeout to a single broker.
|
||||
const defaultKafkaDialTimeout = 3 * time.Second
|
||||
|
||||
func effectiveDialTimeout(d configtypes.Duration) time.Duration {
|
||||
if d <= 0 {
|
||||
return defaultKafkaDialTimeout
|
||||
}
|
||||
return time.Duration(d)
|
||||
}
|
||||
|
||||
// pingTimeout returns a context deadline long enough for Ping to try every
|
||||
// seed broker. franz-go retries each broker's connection once on retryable
|
||||
// errors (see broker.handleReq), so each broker may cost up to 2 × dialTimeout.
|
||||
func pingTimeout(dialTimeout time.Duration, numBrokers int) time.Duration {
|
||||
const maxAttemptsPerBroker = 2
|
||||
return dialTimeout * time.Duration(numBrokers) * maxAttemptsPerBroker
|
||||
}
|
||||
|
||||
func (c *KafkaConsumer) initClient() (*kgo.Client, error) {
|
||||
dialTimeout := effectiveDialTimeout(c.config.DialTimeout)
|
||||
opts := []kgo.Opt{
|
||||
kgo.SeedBrokers(c.config.Brokers...),
|
||||
kgo.DialTimeout(dialTimeout),
|
||||
kgo.ConsumeTopics(c.config.Topics...),
|
||||
kgo.ConsumerGroup(c.config.ConsumerGroup),
|
||||
kgo.OnPartitionsAssigned(c.assigned),
|
||||
kgo.OnPartitionsRevoked(c.revoked),
|
||||
kgo.OnPartitionsLost(c.lost),
|
||||
kgo.AutoCommitMarks(),
|
||||
kgo.BlockRebalanceOnPoll(),
|
||||
kgo.ClientID(kafkaClientID),
|
||||
kgo.FetchMaxWait(time.Duration(c.config.FetchMaxWait)),
|
||||
kgo.WithLogger(&kgoZerologAdapter{log: c.common.log}),
|
||||
}
|
||||
if c.config.InstanceID != "" { // Important to keep separate, otherwise pointer to empty string is used in kgo.InstanceID.
|
||||
opts = append(opts, kgo.InstanceID(c.config.InstanceID))
|
||||
}
|
||||
if c.config.FetchMaxBytes > 0 {
|
||||
opts = append(opts, kgo.FetchMaxBytes(c.config.FetchMaxBytes))
|
||||
}
|
||||
if !c.config.FetchReadUncommitted {
|
||||
opts = append(opts, kgo.FetchIsolationLevel(kgo.ReadCommitted()))
|
||||
} else {
|
||||
opts = append(opts, kgo.FetchIsolationLevel(kgo.ReadUncommitted()))
|
||||
}
|
||||
if c.config.TLS.Enabled {
|
||||
tlsConfig, err := c.config.TLS.ToGoTLSConfig("kafka:" + c.name)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error making TLS configuration: %w", err)
|
||||
}
|
||||
dialer := &tls.Dialer{
|
||||
NetDialer: &net.Dialer{Timeout: 10 * time.Second},
|
||||
Config: tlsConfig,
|
||||
}
|
||||
opts = append(opts, kgo.Dialer(dialer.DialContext))
|
||||
}
|
||||
|
||||
if c.config.SASLMechanism != "" {
|
||||
switch c.config.SASLMechanism {
|
||||
case "plain":
|
||||
opts = append(opts, kgo.SASL(plain.Auth{
|
||||
User: c.config.SASLUser,
|
||||
Pass: c.config.SASLPassword,
|
||||
}.AsMechanism()))
|
||||
case "scram-sha-256":
|
||||
opts = append(opts, kgo.SASL(scram.Auth{
|
||||
User: c.config.SASLUser,
|
||||
Pass: c.config.SASLPassword,
|
||||
}.AsSha256Mechanism()))
|
||||
case "scram-sha-512":
|
||||
opts = append(opts, kgo.SASL(scram.Auth{
|
||||
User: c.config.SASLUser,
|
||||
Pass: c.config.SASLPassword,
|
||||
}.AsSha512Mechanism()))
|
||||
case "aws-msk-iam":
|
||||
if c.config.AssumeRoleARN != "" {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
mskAuth, err := newMSKAssumeRoleAuth(ctx, c.config.AssumeRoleARN)
|
||||
cancel()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("load AWS config for MSK IAM assume role: %w", err)
|
||||
}
|
||||
opts = append(opts, kgo.SASL(mskaws.ManagedStreamingIAM(func(ctx context.Context) (mskaws.Auth, error) {
|
||||
return mskAuth.auth(ctx)
|
||||
})))
|
||||
} else {
|
||||
opts = append(opts, kgo.SASL(mskaws.Auth{
|
||||
AccessKey: c.config.SASLUser,
|
||||
SecretKey: c.config.SASLPassword,
|
||||
}.AsManagedStreamingIAMMechanism()))
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported SASL mechanism: %s", c.config.SASLMechanism)
|
||||
}
|
||||
}
|
||||
|
||||
client, err := kgo.NewClient(opts...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error initializing client: %w", err)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), pingTimeout(dialTimeout, len(c.config.Brokers)))
|
||||
defer cancel()
|
||||
if err := client.Ping(ctx); err != nil {
|
||||
client.Close()
|
||||
return nil, fmt.Errorf("error ping Kafka: %w", err)
|
||||
}
|
||||
return client, nil
|
||||
}
|
||||
|
||||
// fetchPartitionErrorRequiresClientReinit reports whether a PollRecords fetch error
|
||||
// should trigger a full consumer client re-initialization. franz-go classifies Kafka
|
||||
// protocol errors with kerr.IsRetriable; broker-level retryability uses
|
||||
// kgo.IsRetryableBrokerErr — these differ, so both are consulted. *kgo.ErrDataLoss
|
||||
// is handled by offset reset inside the client and does not require restarting it.
|
||||
func fetchPartitionErrorRequiresClientReinit(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var dataLoss *kgo.ErrDataLoss
|
||||
if errors.As(err, &dataLoss) {
|
||||
return false
|
||||
}
|
||||
if kerr.IsRetriable(err) || kgo.IsRetryableBrokerErr(err) {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (c *KafkaConsumer) Run(ctx context.Context) error {
|
||||
started := time.Now()
|
||||
defer func() {
|
||||
c.common.log.Info().Str("was_running_for", time.Since(started).String()).Msg("stopping Kafka consumer")
|
||||
stopStarted := time.Now()
|
||||
close(c.doneCh)
|
||||
if c.client != nil {
|
||||
// Stop all partition consumers first to ensure all offsets are marked.
|
||||
c.killConsumers(c.allTopicPartitions())
|
||||
if c.testOnlyConfig.leaveGroupOnShutdown && c.config.InstanceID != "" {
|
||||
// Test-only path simulating old behavior with dynamic instance IDs.
|
||||
// Commit before manual LeaveGroup — the coordinator rejects commits
|
||||
// from members that already left the group.
|
||||
closeCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := c.client.CommitMarkedOffsets(closeCtx); err != nil {
|
||||
c.common.log.Error().Err(err).Msg("error committing marked offsets on shutdown")
|
||||
}
|
||||
if err := c.leaveGroup(closeCtx, c.client); err != nil {
|
||||
c.common.log.Error().Err(err).Msg("error leaving consumer group")
|
||||
}
|
||||
}
|
||||
// CloseAllowingRebalance calls AllowRebalance (required since we use
|
||||
// BlockRebalanceOnPoll) and then Close. Close triggers the manage loop
|
||||
// exit which calls onRevoked with a valid client context and no
|
||||
// joinAndSync WLock contention — so CommitMarkedOffsets in onRevoked
|
||||
// succeeds without racing. For non-static members (no InstanceID),
|
||||
// Close sends LeaveGroup automatically. For static members (InstanceID
|
||||
// is set), Close does not send LeaveGroup – so the replacement consumer
|
||||
// with the same instance ID can take over partitions seamlessly.
|
||||
c.client.CloseAllowingRebalance()
|
||||
}
|
||||
c.common.log.Info().Str("elapsed", time.Since(stopStarted).String()).Msg("stopped Kafka consumer")
|
||||
}()
|
||||
for {
|
||||
pollStarted := time.Now()
|
||||
err := c.pollUntilFatal(ctx)
|
||||
if err != nil {
|
||||
// franz-go may return context.DeadlineExceeded instead of context.Canceled on
|
||||
// shutdown, so check the external context to detect shutdown — avoids logging
|
||||
// the unwrapped error as if it were a real polling failure.
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
c.common.log.Error().Err(err).Msg("fatal error polling Kafka, need client re-init")
|
||||
}
|
||||
// Upon returning from polling loop with fatal error we re-initialize consumer client.
|
||||
c.client.CloseAllowingRebalance()
|
||||
c.client = nil
|
||||
c.common.log.Info().Str("was_polling_for", time.Since(pollStarted).String()).Msg("start re-initializing Kafka consumer client")
|
||||
reInitStarted := time.Now()
|
||||
err = c.reInitClient(ctx)
|
||||
if err != nil {
|
||||
// Only context.Canceled may be returned.
|
||||
return err
|
||||
}
|
||||
c.common.log.Info().Str("elapsed", time.Since(reInitStarted).String()).Msg("Kafka consumer client re-initialized")
|
||||
}
|
||||
}
|
||||
|
||||
func (c *KafkaConsumer) pollUntilFatal(ctx context.Context) error {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
if ch := c.testOnlyConfig.injectFatalPollErrorCh; ch != nil {
|
||||
select {
|
||||
case err := <-ch:
|
||||
return err
|
||||
default:
|
||||
}
|
||||
}
|
||||
fetches := c.client.PollRecords(ctx, c.config.MaxPollRecords)
|
||||
if fetches.IsClientClosed() {
|
||||
// Defensive check: in our code the client is only closed by Run after this loop
|
||||
// exits, so this should not normally happen. Trigger re-init via Run to recover.
|
||||
c.common.log.Warn().Msg("Kafka client unexpectedly closed during poll, will re-init")
|
||||
return errors.New("client closed during poll")
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
// Check fetch errors. Only return fatal error if at least one partition error needs a
|
||||
// full client re-init. Kafka-retriable protocol errors (kerr) and retryable broker
|
||||
// errors (kgo) are left to franz-go; ErrDataLoss is handled via offset reset inside the
|
||||
// client. Otherwise we log and continue so EachPartition still processes fetched
|
||||
// records for partitions without errors.
|
||||
fetchErrors := fetches.Errors()
|
||||
if len(fetchErrors) > 0 {
|
||||
var errs []error
|
||||
var fatal bool
|
||||
for _, fetchErr := range fetchErrors {
|
||||
errs = append(errs, fetchErr.Err)
|
||||
reinit := fetchPartitionErrorRequiresClientReinit(fetchErr.Err)
|
||||
c.common.log.Error().
|
||||
Err(fetchErr.Err).
|
||||
Str("topic", fetchErr.Topic).
|
||||
Int32("partition", fetchErr.Partition).
|
||||
Bool("fatal", reinit).
|
||||
Msg("error while polling topic partition")
|
||||
if reinit {
|
||||
fatal = true
|
||||
}
|
||||
}
|
||||
if fatal {
|
||||
// Client will be re-initialized in Run. This poll returns before EachPartition,
|
||||
// so records from this cycle are not enqueued yet; recovery is at-least-once
|
||||
// via refetch after re-init.
|
||||
return fmt.Errorf("poll error: %w", errors.Join(errs...))
|
||||
}
|
||||
}
|
||||
|
||||
// Track which partitions we've paused in this poll cycle.
|
||||
pausedTopicPartitions := map[topicPartition]struct{}{}
|
||||
|
||||
fetches.EachPartition(func(p kgo.FetchTopicPartition) {
|
||||
if len(p.Records) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
tp := topicPartition{p.Topic, p.Partition}
|
||||
|
||||
consumer := c.consumers[tp]
|
||||
if consumer == nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-consumer.partitionCtx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
partitionsToPause := map[string][]int32{p.Topic: {p.Partition}}
|
||||
if _, paused := pausedTopicPartitions[tp]; !paused {
|
||||
// Only pause if queue threshold would be exceeded after adding records, or if threshold is 0 (always pause).
|
||||
shouldPause := consumer.queue.NumRecords()+len(p.Records) >= c.config.PartitionQueueMaxSize
|
||||
if shouldPause {
|
||||
// Pause partition BEFORE submitting to queue to avoid race condition
|
||||
// between pause and resume operations.
|
||||
if c.testOnlyConfig.topicPartitionBeforePauseCh != nil {
|
||||
c.testOnlyConfig.topicPartitionBeforePauseCh <- tp
|
||||
}
|
||||
if c.testOnlyConfig.topicPartitionPauseProceedCh != nil {
|
||||
<-c.testOnlyConfig.topicPartitionPauseProceedCh
|
||||
}
|
||||
c.client.PauseFetchPartitions(partitionsToPause)
|
||||
pausedTopicPartitions[tp] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
// Clone p.Records to a fresh, tightly-sized backing array before enqueueing.
|
||||
// The []*kgo.Record slice returned by franz-go points into the per-partition
|
||||
// array franz-go grew via append while decoding the fetch response. That array
|
||||
// can hold many more *Record pointers than the window we take here (up to
|
||||
// MaxPollRecords); Go's GC tracks the whole backing allocation, so retaining
|
||||
// the sub-slice keeps every pointer in the original array reachable — including
|
||||
// records semantically stripped when the partition gets paused. That in turn
|
||||
// keeps their Record structs and the batch buffers their Values point into
|
||||
// alive, amplifying heap usage under lag. Cloning breaks the reference so the
|
||||
// queue only pins the records we actually intend to process.
|
||||
p.Records = slices.Clone(p.Records)
|
||||
|
||||
// Now submit to unbounded queue.
|
||||
if !consumer.queue.Push(p) {
|
||||
// Queue is closed, partition consumer is shutting down
|
||||
// Resume the partition if we paused it but won't process
|
||||
if _, wasPaused := pausedTopicPartitions[tp]; wasPaused {
|
||||
c.client.ResumeFetchPartitions(partitionsToPause)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if c.testOnlyConfig.fetchTopicPartitionSubmittedCh != nil {
|
||||
c.testOnlyConfig.fetchTopicPartitionSubmittedCh <- p
|
||||
}
|
||||
})
|
||||
|
||||
c.client.AllowRebalance()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *KafkaConsumer) reInitClient(ctx context.Context) error {
|
||||
var backoffDuration time.Duration = 0
|
||||
retries := 0
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
}
|
||||
client, err := c.initClient()
|
||||
if err != nil {
|
||||
retries++
|
||||
backoffDuration = getNextBackoffDuration(backoffDuration, retries)
|
||||
c.common.log.Error().Err(err).Msg("error initializing Kafka client")
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(backoffDuration):
|
||||
continue
|
||||
}
|
||||
}
|
||||
c.client = client
|
||||
break
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *KafkaConsumer) assigned(ctx context.Context, cl *kgo.Client, assigned map[string][]int32) {
|
||||
c.consumersMu.Lock()
|
||||
defer c.consumersMu.Unlock()
|
||||
select {
|
||||
case <-c.doneCh:
|
||||
return // Shutdown in progress, skip creating consumers that will be immediately killed.
|
||||
default:
|
||||
}
|
||||
for topic, partitions := range assigned {
|
||||
for _, partition := range partitions {
|
||||
partitionCtx, cancel := context.WithCancel(ctx)
|
||||
pc := &partitionConsumer{
|
||||
partitionCtx: partitionCtx,
|
||||
cancelFunc: cancel,
|
||||
dispatcher: c.dispatcher,
|
||||
cl: cl,
|
||||
topic: topic,
|
||||
partition: partition,
|
||||
config: c.config,
|
||||
name: c.name,
|
||||
common: c.common,
|
||||
|
||||
done: make(chan struct{}),
|
||||
queue: newUnboundedQueue(),
|
||||
}
|
||||
c.consumers[topicPartition{topic, partition}] = pc
|
||||
go pc.consume()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *KafkaConsumer) revoked(ctx context.Context, cl *kgo.Client, revoked map[string][]int32) {
|
||||
c.killConsumers(revoked)
|
||||
// Always commit marked offsets on revoke.
|
||||
if err := cl.CommitMarkedOffsets(ctx); err != nil {
|
||||
// On shutdown the app context is already canceled; commit will fail with
|
||||
// context.Canceled. That is expected, not an error — log at info level instead.
|
||||
if errors.Is(err, context.Canceled) {
|
||||
c.common.log.Info().Err(err).Msg("skipping commit on revoke due to context cancellation")
|
||||
} else {
|
||||
c.common.log.Error().Err(err).Msg("error committing marked offsets on revoke")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *KafkaConsumer) lost(_ context.Context, _ *kgo.Client, lost map[string][]int32) {
|
||||
c.killConsumers(lost)
|
||||
// Losing means we cannot commit: an error happened.
|
||||
}
|
||||
|
||||
func (c *KafkaConsumer) allTopicPartitions() map[string][]int32 {
|
||||
c.consumersMu.Lock()
|
||||
defer c.consumersMu.Unlock()
|
||||
result := make(map[string][]int32)
|
||||
for tp := range c.consumers {
|
||||
result[tp.topic] = append(result[tp.topic], tp.partition)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (c *KafkaConsumer) killConsumers(lost map[string][]int32) {
|
||||
c.consumersMu.Lock()
|
||||
var pcs []*partitionConsumer
|
||||
for topic, partitions := range lost {
|
||||
for _, partition := range partitions {
|
||||
tp := topicPartition{topic, partition}
|
||||
pc := c.consumers[tp]
|
||||
if pc == nil {
|
||||
continue
|
||||
}
|
||||
delete(c.consumers, tp)
|
||||
pcs = append(pcs, pc)
|
||||
}
|
||||
}
|
||||
c.consumersMu.Unlock()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for _, pc := range pcs {
|
||||
pc.queue.Close()
|
||||
pc.cancelFunc()
|
||||
wg.Add(1)
|
||||
go func() { <-pc.done; wg.Done() }()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
type partitionConsumer struct {
|
||||
partitionCtx context.Context
|
||||
cancelFunc context.CancelFunc
|
||||
dispatcher Dispatcher
|
||||
cl *kgo.Client
|
||||
topic string
|
||||
partition int32
|
||||
config KafkaConfig
|
||||
name string
|
||||
common *consumerCommon
|
||||
|
||||
done chan struct{}
|
||||
queue *unboundedQueue
|
||||
}
|
||||
|
||||
func getUint64HeaderValue(record *kgo.Record, headerKey string) (uint64, error) {
|
||||
if headerKey == "" {
|
||||
return 0, nil
|
||||
}
|
||||
headerValue := getHeaderValue(record, headerKey)
|
||||
if headerValue == "" {
|
||||
return 0, nil
|
||||
}
|
||||
value, err := strconv.ParseUint(headerValue, 10, 64)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("error parsing header %s value %s: %w", headerKey, headerValue, err)
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func getBoolHeaderValue(record *kgo.Record, headerKey string) (bool, error) {
|
||||
if headerKey == "" {
|
||||
return false, nil
|
||||
}
|
||||
headerValue := getHeaderValue(record, headerKey)
|
||||
if headerValue == "" {
|
||||
return false, nil
|
||||
}
|
||||
value, err := strconv.ParseBool(headerValue)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("error parsing header %s value %s: %w", headerKey, headerValue, err)
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func getHeaderValue(record *kgo.Record, headerKey string) string {
|
||||
if headerKey == "" {
|
||||
return ""
|
||||
}
|
||||
for _, header := range record.Headers {
|
||||
if header.Key == headerKey {
|
||||
return string(header.Value)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func publicationTagsFromKafkaRecord(record *kgo.Record, tagsHeaderPrefix string) map[string]string {
|
||||
if tagsHeaderPrefix == "" {
|
||||
return nil
|
||||
}
|
||||
var tags map[string]string
|
||||
for _, header := range record.Headers {
|
||||
if strings.HasPrefix(header.Key, tagsHeaderPrefix) {
|
||||
if tags == nil {
|
||||
tags = make(map[string]string)
|
||||
}
|
||||
tags[header.Key[len(tagsHeaderPrefix):]] = string(header.Value)
|
||||
}
|
||||
}
|
||||
return tags
|
||||
}
|
||||
|
||||
func (pc *partitionConsumer) processPublicationDataRecord(ctx context.Context, record *kgo.Record) error {
|
||||
data := record.Value
|
||||
idempotencyKey := getHeaderValue(record, pc.config.PublicationDataMode.IdempotencyKeyHeader)
|
||||
delta, err := getBoolHeaderValue(record, pc.config.PublicationDataMode.DeltaHeader)
|
||||
if err != nil {
|
||||
pc.common.log.Error().Err(err).Str("topic", record.Topic).Int32("partition", record.Partition).Msg("error parsing delta header value, skip message")
|
||||
return nil
|
||||
}
|
||||
channelsHeader := getHeaderValue(record, pc.config.PublicationDataMode.ChannelsHeader)
|
||||
if channelsHeader == "" {
|
||||
pc.common.log.Info().Str("topic", record.Topic).Int32("partition", record.Partition).Msg("no channels found, skip message")
|
||||
return nil
|
||||
}
|
||||
channels := strings.Split(channelsHeader, ",")
|
||||
if len(channels) == 0 || (len(channels) == 1 && channels[0] == "") {
|
||||
pc.common.log.Info().Str("topic", record.Topic).Int32("partition", record.Partition).Msg("no channels found, skip message")
|
||||
return nil
|
||||
}
|
||||
version, err := getUint64HeaderValue(record, pc.config.PublicationDataMode.VersionHeader)
|
||||
if err != nil {
|
||||
pc.common.log.Error().Err(err).Str("topic", record.Topic).Int32("partition", record.Partition).Msg("error parsing version header value, skip message")
|
||||
return nil
|
||||
}
|
||||
return pc.dispatcher.DispatchPublication(
|
||||
ctx,
|
||||
channels,
|
||||
api.ConsumedPublication{
|
||||
Data: data,
|
||||
IdempotencyKey: idempotencyKey,
|
||||
Delta: delta,
|
||||
Tags: publicationTagsFromKafkaRecord(record, pc.config.PublicationDataMode.TagsHeaderPrefix),
|
||||
Version: version,
|
||||
VersionEpoch: getHeaderValue(record, pc.config.PublicationDataMode.VersionEpochHeader),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func (pc *partitionConsumer) processRecord(ctx context.Context, record *kgo.Record) error {
|
||||
if pc.config.PublicationDataMode.Enabled {
|
||||
return pc.processPublicationDataRecord(ctx, record)
|
||||
}
|
||||
method := getHeaderValue(record, pc.config.MethodHeader)
|
||||
return pc.dispatcher.DispatchCommand(ctx, method, record.Value)
|
||||
}
|
||||
|
||||
func (pc *partitionConsumer) processRecords(records []*kgo.Record) {
|
||||
for _, record := range records {
|
||||
select {
|
||||
case <-pc.partitionCtx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
var backoffDuration time.Duration = 0
|
||||
retries := 0
|
||||
for {
|
||||
err := pc.processRecord(pc.partitionCtx, record)
|
||||
if err == nil {
|
||||
if retries > 0 {
|
||||
pc.common.log.Info().Str("topic", record.Topic).Int32("partition", record.Partition).Msg("OK processing message after errors")
|
||||
}
|
||||
if logging.Enabled(logging.TraceLevel) {
|
||||
pc.common.log.Trace().Str("topic", record.Topic).Int32("partition", record.Partition).Int64("offset", record.Offset).Msg("record processed")
|
||||
}
|
||||
metrics.ConsumerProcessedTotal.WithLabelValues(pc.name).Inc()
|
||||
pc.cl.MarkCommitRecords(record)
|
||||
break
|
||||
}
|
||||
// Only stop on our own cancellation: an unrelated context.Canceled bubbling up
|
||||
// from the dispatcher must be retried, not treated as shutdown. Bailing out here
|
||||
// would skip the rest of the batch and let a later batch commit an offset past
|
||||
// this record — losing it for good.
|
||||
if pc.partitionCtx.Err() != nil {
|
||||
return
|
||||
}
|
||||
retries++
|
||||
backoffDuration = getNextBackoffDuration(backoffDuration, retries)
|
||||
metrics.ConsumerErrorsTotal.WithLabelValues(pc.name).Inc()
|
||||
pc.common.log.Error().Err(err).Str("topic", record.Topic).Int32("partition", record.Partition).Str("next_attempt_in", backoffDuration.String()).Msg("error processing consumed record")
|
||||
select {
|
||||
case <-time.After(backoffDuration):
|
||||
case <-pc.partitionCtx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (pc *partitionConsumer) consume() {
|
||||
if logging.Enabled(logging.DebugLevel) {
|
||||
pc.common.log.Debug().Str("topic", pc.topic).Int32("partition", pc.partition).Msg("starting partition consumer")
|
||||
}
|
||||
defer close(pc.done)
|
||||
defer pc.queue.Close()
|
||||
defer func() {
|
||||
if logging.Enabled(logging.DebugLevel) {
|
||||
// Queue is closed, partition consumer is shutting down.
|
||||
pc.common.log.Debug().Str("topic", pc.topic).Int32("partition", pc.partition).Msg("partition consumer shutting down")
|
||||
}
|
||||
}()
|
||||
|
||||
partitionsToResume := map[string][]int32{pc.topic: {pc.partition}}
|
||||
resumeConsuming := func() {
|
||||
pc.cl.ResumeFetchPartitions(partitionsToResume)
|
||||
}
|
||||
defer resumeConsuming()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-pc.partitionCtx.Done():
|
||||
return
|
||||
case <-pc.queue.NotEmpty():
|
||||
if pc.queue.IsClosed() {
|
||||
return
|
||||
}
|
||||
// Process all available items in the queue.
|
||||
for {
|
||||
select {
|
||||
case <-pc.partitionCtx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
p, ok := pc.queue.Pop()
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
pc.processRecords(p.Records)
|
||||
// Resume partition after processing this batch. Only one additional batch will
|
||||
// be added to the queue before we pause again.
|
||||
resumeConsuming()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// unboundedQueue implements an unbounded queue for FetchTopicPartition
|
||||
type unboundedQueue struct {
|
||||
mu sync.RWMutex
|
||||
items []kgo.FetchTopicPartition
|
||||
numRecords int
|
||||
notEmpty chan struct{}
|
||||
closed bool
|
||||
}
|
||||
|
||||
func newUnboundedQueue() *unboundedQueue {
|
||||
return &unboundedQueue{
|
||||
items: make([]kgo.FetchTopicPartition, 0),
|
||||
notEmpty: make(chan struct{}, 1),
|
||||
}
|
||||
}
|
||||
|
||||
func (q *unboundedQueue) Push(item kgo.FetchTopicPartition) bool {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
|
||||
if q.closed {
|
||||
return false
|
||||
}
|
||||
|
||||
q.items = append(q.items, item)
|
||||
q.numRecords += len(item.Records)
|
||||
|
||||
// Signal that queue is not empty.
|
||||
select {
|
||||
case q.notEmpty <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
func (q *unboundedQueue) Pop() (kgo.FetchTopicPartition, bool) {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
|
||||
if q.closed || len(q.items) == 0 {
|
||||
return kgo.FetchTopicPartition{}, false
|
||||
}
|
||||
|
||||
// We do not expect many items in the queue, so we can use a simple slice pop operation and
|
||||
// not worry about memory being reserved for the underlying slice.
|
||||
item := q.items[0]
|
||||
q.items[0] = kgo.FetchTopicPartition{} // Clear the first item for faster GC.
|
||||
q.items = q.items[1:]
|
||||
q.numRecords -= len(item.Records)
|
||||
|
||||
return item, true
|
||||
}
|
||||
|
||||
func (q *unboundedQueue) NumRecords() int {
|
||||
q.mu.RLock()
|
||||
defer q.mu.RUnlock()
|
||||
return q.numRecords
|
||||
}
|
||||
|
||||
func (q *unboundedQueue) Close() {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
|
||||
if !q.closed {
|
||||
q.closed = true
|
||||
close(q.notEmpty)
|
||||
}
|
||||
}
|
||||
|
||||
func (q *unboundedQueue) IsClosed() bool {
|
||||
q.mu.RLock()
|
||||
defer q.mu.RUnlock()
|
||||
return q.closed
|
||||
}
|
||||
|
||||
func (q *unboundedQueue) NotEmpty() <-chan struct{} {
|
||||
return q.notEmpty
|
||||
}
|
||||
Reference in New Issue
Block a user