WIP fix.

Author: nachogiljaldo

batch.go

@@ -232,6 +232,7 @@ func (batch *Batch) ReadMessage() (Message, error) {
 	msg.HighWaterMark = batch.highWaterMark
 	msg.Time = makeTime(timestamp)
 	msg.Headers = headers
+	msg.GenerationId = batch.conn.generationId
 
 	return msg, err
 }

commit.go

@@ -3,18 +3,20 @@ package kafka
 // A commit represents the instruction of publishing an update of the last
 // offset read by a program for a topic and partition.
 type commit struct {
-	topic     string
-	partition int
-	offset    int64
+	topic        string
+	partition    int
+	offset       int64
+	generationId int32
 }
 
 // makeCommit builds a commit value from a message, the resulting commit takes
 // its topic, partition, and offset from the message.
 func makeCommit(msg Message) commit {
 	return commit{
-		topic:     msg.Topic,
-		partition: msg.Partition,
-		offset:    msg.Offset + 1,
+		topic:        msg.Topic,
+		partition:    msg.Partition,
+		offset:       msg.Offset + 1,
+		generationId: msg.GenerationId,
 	}
 }
 

conn.go

@@ -65,6 +65,8 @@ type Conn struct {
 	apiVersions atomic.Value // apiVersionMap
 
 	transactionalID *string
+
+	generationId int32
 }
 
 type apiVersionMap map[apiKey]ApiVersion
@@ -182,6 +184,7 @@ func NewConnWith(conn net.Conn, config ConnConfig) *Conn {
 		offset:          FirstOffset,
 		requiredAcks:    -1,
 		transactionalID: emptyToNullable(config.TransactionalID),
+		generationId:    -1,
 	}
 
 	c.wb.w = &c.wbuf
@@ -388,6 +391,7 @@ func (c *Conn) joinGroup(request joinGroupRequestV1) (joinGroupResponseV1, error
 		return joinGroupResponseV1{}, Error(response.ErrorCode)
 	}
 
+	c.generationId = response.GenerationID
 	return response, nil
 }
 

message.go

@@ -20,6 +20,10 @@ type Message struct {
 	Value         []byte
 	Headers       []Header
 
+	// If the message has been sent by a consumer group, it contains the
+	// generation's id. Value is -1 if not using consumer groups.
+	GenerationId int32
+
 	// This field is used to hold arbitrary data you wish to include, so it
 	// will be available when handle it on the Writer's `Completion` method,
 	// this support the application can do any post operation on each message.

reader.go

@@ -150,35 +150,49 @@ func (r *Reader) commitOffsetsWithRetry(gen *Generation, offsetStash offsetStash
 		backoffDelayMax = 5 * time.Second
 	)
 
+	messagesToSend := make(map[string]map[int]int64)
+	for topic, partitionsInfo := range offsetStash {
+		messagesToSend[topic] = make(map[int]int64)
+		for partition, commitInfo := range partitionsInfo {
+			messagesToSend[topic][partition] = commitInfo.offset
+		}
+	}
 	for attempt := 0; attempt < retries; attempt++ {
 		if attempt != 0 {
 			if !sleep(r.stctx, backoff(attempt, backoffDelayMin, backoffDelayMax)) {
 				return
 			}
 		}
 
-		if err = gen.CommitOffsets(offsetStash); err == nil {
+		if err = gen.CommitOffsets(messagesToSend); err == nil {
 			return
 		}
 	}
 
 	return // err will not be nil
 }
 
-// offsetStash holds offsets by topic => partition => offset.
-type offsetStash map[string]map[int]int64
+// offsetStash holds offsets by topic => partition => offsetEntry.
+type offsetEntry struct {
+	offset       int64
+	generationId int32
+}
+type offsetStash map[string]map[int]offsetEntry
 
 // merge updates the offsetStash with the offsets from the provided messages.
 func (o offsetStash) merge(commits []commit) {
 	for _, c := range commits {
 		offsetsByPartition, ok := o[c.topic]
 		if !ok {
-			offsetsByPartition = map[int]int64{}
+			offsetsByPartition = map[int]offsetEntry{}
 			o[c.topic] = offsetsByPartition
 		}
 
-		if offset, ok := offsetsByPartition[c.partition]; !ok || c.offset > offset {
-			offsetsByPartition[c.partition] = c.offset
+		if offset, ok := offsetsByPartition[c.partition]; !ok || c.offset > offset.offset {
+			offsetsByPartition[c.partition] = offsetEntry{
+				offset:       c.offset,
+				generationId: c.generationId,
+			}
 		}
 	}
 }

reader_test.go

@@ -897,6 +897,36 @@ func TestReaderConsumerGroup(t *testing.T) {
 	}
 }
 
+func TestReaderConsumerGroup2(t *testing.T) {
+	// It appears that some of the tests depend on all these tests being
+	// run concurrently to pass... this is brittle and should be fixed
+	// at some point.
+	t.Parallel()
+
+	topic := makeTopic()
+	createTopic(t, topic, 2)
+	defer deleteTopic(t, topic)
+
+	groupID := makeGroupID()
+	r := NewReader(ReaderConfig{
+		Brokers:           []string{"localhost:9092"},
+		Topic:             topic,
+		GroupID:           groupID,
+		HeartbeatInterval: 2 * time.Second,
+		CommitInterval:    0,
+		RebalanceTimeout:  2 * time.Second,
+		RetentionTime:     time.Hour,
+		MinBytes:          1,
+		MaxBytes:          1e6,
+	})
+	defer r.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	testReaderConsumerGroupRebalanceDoesNotCommitNotOwnedPartitions(t, ctx, r)
+}
+
 func testReaderConsumerGroupHandshake(t *testing.T, ctx context.Context, r *Reader) {
 	prepareReader(t, context.Background(), r, makeTestSequence(5)...)
 
@@ -1228,27 +1258,7 @@ func testReaderConsumerGroupRebalanceDoesNotCommitNotOwnedPartitions(t *testing.
 	}
 
 	require.NoError(t, secondReader.CommitMessages(ctx, msgsForSecondReader[len(msgsForSecondReader)-1]))
-	require.NoError(t, firstReader.CommitMessages(ctx, msgsForSecondReader[0]))
-	resp, err := client.OffsetFetch(
-		ctx,
-		&OffsetFetchRequest{
-			GroupID: firstReader.config.GroupID,
-			Topics:  map[string][]int{firstReader.config.Topic: {0, 1}},
-		},
-	)
-	require.NoError(t, err)
-	require.NotNil(t, resp)
-	if topics, ok := resp.Topics[firstReader.config.Topic]; !ok {
-		require.True(t, ok, "Topic not found")
-	} else {
-		for _, topic := range topics {
-			if offset, ok := topicsToCommit[topic.Partition]; ok {
-				assert.Equal(t, offset+1, topic.CommittedOffset, "committed partition %d had committed offset %d instead of %d", topic.Partition, topic.CommittedOffset, offset)
-			} else {
-				assert.Equal(t, int64(-1), topic.CommittedOffset, "not-committed partition %d had committed offset %d instead of 0", topic.Partition, topic.CommittedOffset)
-			}
-		}
-	}
+	require.ErrorIs(t, errInvalidWritePartition, firstReader.CommitMessages(ctx, msgsForSecondReader[0]))
 }
 
 func TestOffsetStash(t *testing.T) {
@@ -1272,16 +1282,16 @@ func TestOffsetStash(t *testing.T) {
 			Given:    offsetStash{},
 			Messages: []Message{newMessage(0, 0)},
 			Expected: offsetStash{
-				topic: {0: 1},
+				topic: {0: {1, 1}},
 			},
 		},
 		"ignores earlier offsets": {
 			Given: offsetStash{
-				topic: {0: 2},
+				topic: {0: {2, 1}},
 			},
 			Messages: []Message{newMessage(0, 0)},
 			Expected: offsetStash{
-				topic: {0: 2},
+				topic: {0: {2, 1}},
 			},
 		},
 		"uses latest offset": {
@@ -1292,7 +1302,7 @@ func TestOffsetStash(t *testing.T) {
 				newMessage(0, 1),
 			},
 			Expected: offsetStash{
-				topic: {0: 4},
+				topic: {0: {4, 1}},
 			},
 		},
 		"uses latest offset, across multiple topics": {
@@ -1306,8 +1316,8 @@ func TestOffsetStash(t *testing.T) {
 			},
 			Expected: offsetStash{
 				topic: {
-					0: 4,
-					1: 7,
+					0: {4, 1},
+					1: {7, 1},
 				},
 			},
 		},
@@ -1359,10 +1369,11 @@ func TestCommitLoopImmediateFlushOnGenerationEnd(t *testing.T) {
 				return offsetCommitResponseV2{}, nil
 			},
 		},
-		done:     make(chan struct{}),
-		log:      func(func(Logger)) {},
-		logError: func(func(Logger)) {},
-		joined:   make(chan struct{}),
+		done:        make(chan struct{}),
+		log:         func(func(Logger)) {},
+		logError:    func(func(Logger)) {},
+		joined:      make(chan struct{}),
+		Assignments: map[string][]PartitionAssignment{"topic": {{0, 1}}},
 	}
 
 	// initialize commits so that the commitLoopImmediate select statement blocks
@@ -1396,7 +1407,7 @@ func TestCommitLoopImmediateFlushOnGenerationEnd(t *testing.T) {
 }
 
 func TestCommitOffsetsWithRetry(t *testing.T) {
-	offsets := offsetStash{"topic": {0: 0}}
+	offsets := offsetStash{"topic": {0: {0, 1}}}
 
 	tests := map[string]struct {
 		Fails       int
@@ -1430,9 +1441,10 @@ func TestCommitOffsetsWithRetry(t *testing.T) {
 						return offsetCommitResponseV2{}, nil
 					},
 				},
-				done:     make(chan struct{}),
-				log:      func(func(Logger)) {},
-				logError: func(func(Logger)) {},
+				done:        make(chan struct{}),
+				log:         func(func(Logger)) {},
+				logError:    func(func(Logger)) {},
+				Assignments: map[string][]PartitionAssignment{"topic": {{0, 1}}},
 			}
 
 			r := &Reader{stctx: context.Background()}
@@ -1636,7 +1648,7 @@ func TestConsumerGroupWithGroupTopicsSingle(t *testing.T) {
 	}
 }
 
-func TestConsumerGroupWithGroupTopicsMultple(t *testing.T) {
+func TestConsumerGroupWithGroupTopicsMultiple(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()