upstream-centrifugo/internal/consuming/kafka.go

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
}