kafka/protocol/fetch.go

package protocol

import (
	"time"

	log "github.com/hanzoai/kafka/logging"
	"github.com/hanzoai/kafka/pubsub"
	"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 i := 0; i < lenTopic; i++ {
			topic := FetchRequestTopic{Name: d.CompactString()}
			lenPartitions := int(d.CompactArrayLen())
			for j := 0; j < lenPartitions; j++ {
				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 i := 0; i < lenForgotten; i++ {
			ft := FetchRequestForgottenTopic{Topic: d.CompactString()}
			lenParts := int(d.CompactArrayLen())
			for j := 0; j < lenParts; j++ {
				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 i := 0; i < lenTopic; i++ {
			topic := FetchRequestTopic{Name: d.NullableString()}
			lenPartitions := int(int32(d.UInt32()))
			for j := 0; j < lenPartitions; j++ {
				p := FetchRequestPartitionData{
					PartitionIndex: d.UInt32(),
				}
				if apiVersion >= 9 {
					p.CurrentLeaderEpoch = d.UInt32()
				}
				p.FetchOffset = d.UInt64()
				if apiVersion >= 5 {
					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 i := 0; i < lenForgotten; i++ {
				ft := FetchRequestForgottenTopic{Topic: d.NullableString()}
				lenParts := int(int32(d.UInt32()))
				for j := 0; j < lenParts; j++ {
					ft.Partitions = append(ft.Partitions, d.UInt32())
				}
				fetchRequest.ForgottenTopicsData = append(fetchRequest.ForgottenTopicsData, ft)
			}
		}
		if apiVersion >= 11 {
			fetchRequest.RackID = d.NullableString()
		}
	}
}

// maxFetchWait caps how long a fetch is held for data. A client waits for the
// response for at least the MaxWaitMs it sent, so answering earlier is always
// valid; the cap bounds what a misconfigured or hostile client can pin.
const maxFetchWait = 30 * time.Second

// fetchWait is how long a fetch with nothing to return is held. MaxWaitMs is
// an INT32 on the wire: zero or negative means answer at once.
func fetchWait(maxWaitMs uint32) time.Duration {
	ms := int32(maxWaitMs)
	if ms <= 0 {
		return 0
	}
	return min(time.Duration(ms)*time.Millisecond, maxFetchWait)
}

// getFetchResponse answers a Fetch the way Kafka's delayed fetch does: at once
// when the partitions hold MinBytes of records, a partition holds more than
// this response can carry, or any partition answers with an error; otherwise
// once a requested partition gets a new tail that satisfies it or MaxWaitMs
// runs out, whichever comes first. Reading an
// unchanged partition costs one StreamInfo, on arrival and on each wake; a
// parked fetch costs nothing while it waits.
func (b *Broker) getFetchResponse(req types.Request) []byte {
	start := time.Now()
	decoder := serde.NewDecoder(req.Body)
	fetchRequest := &FetchRequest{}
	decodeFetchRequest(decoder, fetchRequest, req.RequestAPIVersion)
	log.Debug("fetchRequest %+v", fetchRequest)

	minBytes := int(int32(fetchRequest.MinBytes))
	deadline := start.Add(fetchWait(fetchRequest.MaxWaitMs))
	var streams []string
	for _, tp := range fetchRequest.Topics {
		for _, p := range tp.Partitions {
			streams = append(streams, pubsub.StreamName(tp.Name, p.PartitionIndex))
		}
	}
	// Park before the first read: a tail announced while the partitions are
	// being read is then held in wake rather than lost.
	wake := make(chan struct{}, 1)
	if time.Until(deadline) > 0 && len(streams) > 0 {
		defer b.park(streams, wake)()
	}

	for {
		response, size, now := b.fetchOnce(fetchRequest)
		wait := time.Until(deadline)
		if now || size >= minBytes || wait <= 0 || len(streams) == 0 {
			return encodeFetchResponse(req, response)
		}
		timer := time.NewTimer(wait)
		select {
		case <-wake:
			timer.Stop()
			continue
		case <-timer.C:
		case <-b.ShutDownSignal:
			timer.Stop()
		}
		return encodeFetchResponse(req, response)
	}
}

// fetchOnce reads every requested partition as the store holds it now. It
// returns the response, the record bytes in it, and whether to answer
// without waiting for MinBytes: a partition answered with an error, or a
// partition holds records past the one record set a fetch serves it, so
// waiting cannot make this response any fuller.
func (b *Broker) fetchOnce(fetchRequest *FetchRequest) (response FetchResponse, size int, now bool) {
	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
			}
			bd, err := b.partitionBounds(tp.Name, p.PartitionIndex)
			if err != nil {
				pr.ErrorCode = uint16(ErrUnknownTopicOrPartition.Code)
			} else {
				pr.HighWatermark = uint64(bd.next)
				pr.LastStableOffset = uint64(bd.next)
				pr.LogStartOffset = uint64(bd.logStart)
				offset := int64(p.FetchOffset)
				switch {
				case offset < bd.logStart || offset > bd.next:
					// 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 == bd.next:
					// caught up: empty records, no error
				default:
					if msg := b.findRecordSet(tp.Name, p.PartitionIndex, offset, bd); msg != nil {
						pr.Records = msg.Data
						size += len(msg.Data)
						if _, last, ok := batchSpan(msg.Data); ok && last+1 < bd.next {
							now = true
						}
					}
				}
			}
			now = now || pr.ErrorCode != 0
			fetchTopicResponse.Partitions = append(fetchTopicResponse.Partitions, pr)
		}
		response.Responses = append(response.Responses, fetchTopicResponse)
	}
	return response, size, now
}

// 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()
}