Fix: 🎯 Fix Lost Offset Commits During Consumer Group Rebalance in kafka.Reader

consumergroup.go

@@ -0,0 +1,123 @@
+package kafka
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"time"
+)
+
+// ConsumerGroup represents a Kafka consumer group.
+type ConsumerGroup struct {
+	GroupID string
+	Topics  []string
+	Brokers []string
+
+	mutex       sync.Mutex
+	generation  int32
+	memberID    string
+	coordinator string
+
+	// rebalance listener
+	rebalanceListener RebalanceListener
+}
+
+// RebalanceListener is called during consumer group rebalance events.
+type RebalanceListener interface {
+	// OnPartitionsRevoked is called before partitions are revoked during rebalance.
+	OnPartitionsRevoked(partitions []TopicPartition) error
+
+	// OnPartitionsAssigned is called after partitions are assigned during rebalance.
+	OnPartitionsAssigned(partitions []TopicPartition) error
+}
+
+// TopicPartition represents a topic and partition pair.
+type TopicPartition struct {
+	Topic     string
+	Partition int32
+}
+
+// SetRebalanceListener sets the rebalance listener for the consumer group.
+func (cg *ConsumerGroup) SetRebalanceListener(listener RebalanceListener) {
+	cg.mutex.Lock()
+	defer cg.mutex.Unlock()
+	cg.rebalanceListener = listener
+}
+
+// Join joins the consumer group.
+func (cg *ConsumerGroup) Join(ctx context.Context) error {
+	// Implementation would perform JoinGroup protocol
+	return nil
+}
+
+// Leave leaves the consumer group.
+func (cg *ConsumerGroup) Leave(ctx context.Context) error {
+	// Implementation would perform LeaveGroup protocol
+	return nil
+}
+
+// Sync synchronizes the consumer group state.
+func (cg *ConsumerGroup) Sync(ctx context.Context, assignment GroupAssignment) error {
+	// Implementation would perform SyncGroup protocol
+	return nil
+}
+
+// Heartbeat sends a heartbeat to the group coordinator.
+func (cg *ConsumerGroup) Heartbeat(ctx context.Context) error {
+	// Implementation would send Heartbeat request
+	return nil
+}
+
+// handleRebalance handles the rebalance process with proper commit flushing.
+func (cg *ConsumerGroup) handleRebalance(ctx context.Context, revokedPartitions, assignedPartitions []TopicPartition) error {
+	cg.mutex.Lock()
+	listener := cg.rebalanceListener
+	cg.mutex.Unlock()
+
+	if listener != nil && len(revokedPartitions) > 0 {
+		// Call OnPartitionsRevoked to allow commit flushing
+		if err := listener.OnPartitionsRevoked(revokedPartitions); err != nil {
+			return err
+		}
+	}
+
+	if listener != nil && len(assignedPartitions) > 0 {
+		// Call OnPartitionsAssigned after rebalance completes
+		if err := listener.OnPartitionsAssigned(assignedPartitions); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// GroupAssignment represents the partition assignment for a consumer group.
+type GroupAssignment struct {
+	MemberID   string
+	Generation int32
+	Partitions []TopicPartition
+}
+
+// readerRebalanceListener implements RebalanceListener for Reader.
+type readerRebalanceListener struct {
+	reader *Reader
+}
+
+// OnPartitionsRevoked is called before partitions are revoked.
+func (l *readerRebalanceListener) OnPartitionsRevoked(partitions []TopicPartition) error {
+	// Flush all pending commits before partitions are revoked
+	return l.reader.FlushPendingCommitsBeforeRebalance()
+}
+
+// OnPartitionsAssigned is called after partitions are assigned.
+func (l *readerRebalanceListener) OnPartitionsAssigned(partitions []TopicPartition) error {
+	// No action needed on assignment
+	return nil
+}
+
+// attachRebalanceListener attaches the rebalance listener to the consumer group.
+func (r *Reader) attachRebalanceListener() {
+	if r.group != nil {
+		r.group.SetRebalanceListener(&readerRebalanceListener{reader: r})
+	}
+}