Add method to fetch messages in batch

Since FetchMessage is already reading messages from a fetched
batch, this new method just hold the messages util the batchSize
number of messages are read.

Fixes #123

Author: krsoninikhil

go.mod

@@ -1,6 +1,6 @@
 module github.com/segmentio/kafka-go
 
-go 1.23
+go 1.23.0
 
 require (
 	github.com/klauspost/compress v1.15.9

reader.go

@@ -912,6 +912,68 @@ func (r *Reader) CommitMessages(ctx context.Context, msgs ...Message) error {
 	}
 }
 
+// FetchMessageBatch fetches a batch of messages from the reader. It is similar to
+// FetchMessage, except it blocks until no. of messages read reaches batchSize.
+func (r *Reader) FetchMessageBatch(ctx context.Context, batchSize int) ([]Message, error) {
+	r.activateReadLag()
+	msgBatch := make([]Message, 0, batchSize)
+
+	var i int
+	for i <= batchSize {
+		r.mutex.Lock()
+
+		if !r.closed && r.version == 0 {
+			r.start(r.getTopicPartitionOffset())
+		}
+
+		version := r.version
+		r.mutex.Unlock()
+
+		select {
+		case <-ctx.Done():
+			return []Message{}, ctx.Err()
+
+		case err := <-r.runError:
+			return []Message{}, err
+
+		case m, ok := <-r.msgs:
+			if !ok {
+				return []Message{}, io.EOF
+			}
+
+			if m.version < version {
+				continue
+			}
+
+			r.mutex.Lock()
+
+			switch {
+			case m.error != nil:
+			case version == r.version:
+				r.offset = m.message.Offset + 1
+				r.lag = m.watermark - r.offset
+			}
+
+			r.mutex.Unlock()
+
+			if errors.Is(m.error, io.EOF) {
+				// io.EOF is used as a marker to indicate that the stream
+				// has been closed, in case it was received from the inner
+				// reader we don't want to confuse the program and replace
+				// the error with io.ErrUnexpectedEOF.
+				m.error = io.ErrUnexpectedEOF
+			}
+			if m.error != nil {
+				return nil, m.error
+			}
+
+			msgBatch = append(msgBatch, m.message)
+		}
+		i++
+	}
+	return msgBatch, nil
+}
+
 // ReadLag returns the current lag of the reader by fetching the last offset of
 // the topic and partition and computing the difference between that value and
 // the offset of the last message returned by ReadMessage.
@@ -1487,7 +1549,7 @@ func (r *reader) initialize(ctx context.Context, offset int64) (conn *Conn, star
 	return
 }
 
-func (r *reader) read(ctx context.Context, offset int64, conn *Conn) (int64, error) {
+func (r *reader) read(ctx context.Context, offset int64, conn *Conn, batchSize int) (int64, error) {
 	r.stats.fetches.observe(1)
 	r.stats.offset.observe(offset)