stream/protocol/fetch.go

package protocol

import (
	"time"

	log "github.com/hanzoai/kafka/logging"
	"github.com/hanzoai/kafka/serde"
	"github.com/hanzoai/kafka/types"
)

// FetchRequest represents the details of a FetchRequest (Version: 12).
type FetchRequest struct {
	ReplicaID           uint32
	MaxWaitMs           uint32
	MinBytes            uint32
	MaxBytes            uint32
	IsolationLevel      uint8
	SessionID           uint32
	SessionEpoch        uint32
	Topics              []FetchRequestTopic
	ForgottenTopicsData []FetchRequestForgottenTopic
	RackID              string `kafka:"CompactString"`
}

// FetchRequestTopic represents the topic-level data in a FetchRequest.
type FetchRequestTopic struct {
	Name       string `kafka:"CompactString"`
	Partitions []FetchRequestPartitionData
}

// FetchRequestPartitionData represents the partition-level data in a FetchRequest.
type FetchRequestPartitionData struct {
	PartitionIndex     uint32
	CurrentLeaderEpoch uint32
	FetchOffset        uint64
	LastFetchedEpoch   uint32
	LogStartOffset     uint64
	PartitionMaxBytes  uint32
}

// FetchRequestForgottenTopic represents the forgotten topic data in a FetchRequest.
type FetchRequestForgottenTopic struct {
	Topic      string `kafka:"CompactString"`
	Partitions []uint32
}

// FetchResponse represents the response to a fetch request.
type FetchResponse struct {
	ThrottleTimeMs uint32
	ErrorCode      uint16
	SessionID      uint32
	Responses      []FetchTopicResponse
}

// FetchTopicResponse represents the response for a topic in a fetch request.
type FetchTopicResponse struct {
	TopicName  string `kafka:"CompactString"`
	Partitions []FetchPartitionResponse
}

// FetchPartitionResponse represents the response for a partition in a fetch request.
type FetchPartitionResponse struct {
	PartitionIndex       uint32
	ErrorCode            uint16
	HighWatermark        uint64
	LastStableOffset     uint64
	LogStartOffset       uint64
	AbortedTransactions  []AbortedTransaction
	PreferredReadReplica uint32
	Records              []byte
}

// AbortedTransaction represents an aborted transaction in the fetch response.
type AbortedTransaction struct {
	ProducerID  uint64
	FirstOffset uint64
}

func decodeFetchRequest(d serde.Decoder, fetchRequest *FetchRequest, apiVersion uint16) {
	fetchRequest.ReplicaID = d.UInt32()
	fetchRequest.MaxWaitMs = d.UInt32()
	fetchRequest.MinBytes = d.UInt32()
	if apiVersion >= 3 {
		fetchRequest.MaxBytes = d.UInt32()
	}
	if apiVersion >= 4 {
		fetchRequest.IsolationLevel = d.UInt8()
	}
	if apiVersion >= 7 {
		fetchRequest.SessionID = d.UInt32()
		fetchRequest.SessionEpoch = d.UInt32()
	}

	if apiVersion >= 12 {
		// Flexible version: compact arrays/strings + tagged fields
		lenTopic := int(d.CompactArrayLen())
		for range lenTopic {
			topic := FetchRequestTopic{Name: d.CompactString()}
			lenPartitions := int(d.CompactArrayLen())
			for range lenPartitions {
				p := FetchRequestPartitionData{
					PartitionIndex:     d.UInt32(),
					CurrentLeaderEpoch: d.UInt32(),
					FetchOffset:        d.UInt64(),
					LastFetchedEpoch:   d.UInt32(),
					LogStartOffset:     d.UInt64(),
					PartitionMaxBytes:  d.UInt32(),
				}
				topic.Partitions = append(topic.Partitions, p)
				d.EndStruct()
			}
			fetchRequest.Topics = append(fetchRequest.Topics, topic)
			d.EndStruct()
		}
		// forgotten topics (compact array)
		lenForgotten := int(d.CompactArrayLen())
		for range lenForgotten {
			ft := FetchRequestForgottenTopic{Topic: d.CompactString()}
			lenParts := int(d.CompactArrayLen())
			for range lenParts {
				ft.Partitions = append(ft.Partitions, d.UInt32())
			}
			fetchRequest.ForgottenTopicsData = append(fetchRequest.ForgottenTopicsData, ft)
			d.EndStruct()
		}
		fetchRequest.RackID = d.CompactString()
	} else {
		// Non-flexible: int32 array counts, int16 strings
		lenTopic := int(int32(d.UInt32()))
		for range lenTopic {
			topic := FetchRequestTopic{Name: d.NullableString()}
			lenPartitions := int(int32(d.UInt32()))
			for range lenPartitions {
				p := FetchRequestPartitionData{
					PartitionIndex: d.UInt32(),
				}
				if apiVersion >= 9 {
					p.CurrentLeaderEpoch = d.UInt32()
				}
				p.FetchOffset = d.UInt64()
				if apiVersion >= 12 {
					p.LastFetchedEpoch = d.UInt32()
					p.LogStartOffset = d.UInt64()
				}
				p.PartitionMaxBytes = d.UInt32()
				topic.Partitions = append(topic.Partitions, p)
			}
			fetchRequest.Topics = append(fetchRequest.Topics, topic)
		}
		if apiVersion >= 7 {
			// forgotten topics (int32 array)
			lenForgotten := int(int32(d.UInt32()))
			for range lenForgotten {
				ft := FetchRequestForgottenTopic{Topic: d.NullableString()}
				lenParts := int(int32(d.UInt32()))
				for range lenParts {
					ft.Partitions = append(ft.Partitions, d.UInt32())
				}
				fetchRequest.ForgottenTopicsData = append(fetchRequest.ForgottenTopicsData, ft)
			}
		}
		if apiVersion >= 11 {
			fetchRequest.RackID = d.NullableString()
		}
	}
}

func (b *Broker) getFetchResponse(req types.Request) []byte {
	decoder := serde.NewDecoder(req.Body)
	fetchRequest := &FetchRequest{}
	decodeFetchRequest(decoder, fetchRequest, req.RequestAPIVersion)
	log.Debug("fetchRequest %+v", fetchRequest)

	numTotalRecordBytes := 0
	response := FetchResponse{}
	for _, tp := range fetchRequest.Topics {
		fetchTopicResponse := FetchTopicResponse{TopicName: tp.Name}
		for _, p := range tp.Partitions {
			pr := FetchPartitionResponse{
				PartitionIndex:       p.PartitionIndex,
				PreferredReadReplica: uint32(MinusOne), // -1 = no preferred replica
			}
			logStart, hw, err := b.partitionBounds(tp.Name, p.PartitionIndex)
			if err != nil {
				pr.ErrorCode = uint16(ErrUnknownTopicOrPartition.Code)
			} else {
				pr.HighWatermark = uint64(hw)
				pr.LastStableOffset = uint64(hw)
				pr.LogStartOffset = uint64(logStart)
				offset := int64(p.FetchOffset)
				switch {
				case offset < logStart || offset > hw:
					// Out of range is an answer, not a stall: the client resets
					// per its auto.offset.reset policy. Serving empty here
					// instead (the old behavior) left consumers polling a dead
					// position forever whenever their committed offset stopped
					// existing (purge, retention, an operator reset).
					pr.ErrorCode = uint16(ErrOffsetOutOfRange.Code)
				case offset == hw:
					// caught up: empty records, no error
				default:
					if msg := b.findRecordSet(tp.Name, p.PartitionIndex, offset); msg != nil {
						pr.Records = msg.Data
						numTotalRecordBytes += len(msg.Data)
					}
				}
			}
			fetchTopicResponse.Partitions = append(fetchTopicResponse.Partitions, pr)
		}
		response.Responses = append(response.Responses, fetchTopicResponse)
	}
	if numTotalRecordBytes == 0 {
		log.Debug("No data available for this fetch, waiting briefly")
		time.Sleep(300 * time.Millisecond)
	}
	return encodeFetchResponse(req, response)
}

// encodeFetchResponse manually encodes the Fetch response based on API version.
// v0-v11: non-flexible format (int32 arrays, int16 strings, no tagged fields).
// v12+: flexible format (compact arrays/strings, tagged field terminators).
func encodeFetchResponse(req types.Request, response FetchResponse) []byte {
	e := serde.NewEncoder()
	e.PutInt32(req.CorrelationID)

	if req.RequestAPIVersion >= 12 {
		// Flexible response header
		e.EndStruct()
		e.Encode(response)
	} else {
		// Non-flexible response
		if req.RequestAPIVersion >= 1 {
			e.PutInt32(response.ThrottleTimeMs)
		}
		if req.RequestAPIVersion >= 7 {
			e.PutInt16(response.ErrorCode)
			e.PutInt32(response.SessionID)
		}
		// responses array
		e.PutArrayLen(len(response.Responses))
		for _, topic := range response.Responses {
			e.PutString(topic.TopicName)
			e.PutArrayLen(len(topic.Partitions))
			for _, p := range topic.Partitions {
				e.PutInt32(p.PartitionIndex)
				e.PutInt16(p.ErrorCode)
				e.PutInt64(p.HighWatermark)
				if req.RequestAPIVersion >= 4 {
					e.PutInt64(p.LastStableOffset)
				}
				if req.RequestAPIVersion >= 5 {
					e.PutInt64(p.LogStartOffset)
				}
				if req.RequestAPIVersion >= 4 {
					// aborted transactions (int32 array, -1 = null)
					if len(p.AbortedTransactions) == 0 {
						e.PutInt32(uint32(MinusOne)) // null array = -1
					} else {
						e.PutArrayLen(len(p.AbortedTransactions))
						for _, at := range p.AbortedTransactions {
							e.PutInt64(at.ProducerID)
							e.PutInt64(at.FirstOffset)
						}
					}
				}
				if req.RequestAPIVersion >= 11 {
					e.PutInt32(p.PreferredReadReplica)
				}
				// records: int32 length prefix + bytes (-1 = null)
				if len(p.Records) == 0 {
					e.PutInt32(uint32(MinusOne)) // null
				} else {
					e.PutInt32(uint32(len(p.Records)))
					e.PutBytes(p.Records)
				}
			}
		}
	}

	e.PutLen()
	return e.Bytes()
}