Merge pull request #118 from codecrafters-io/refactor-record-batches

Refactor record batch and children to use kafka values

Author: UdeshyaDhungana

protocol/kafka_files_generator/cluster_metadata_generator.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/codecrafters-io/kafka-tester/protocol/encoder"
 	"github.com/codecrafters-io/kafka-tester/protocol/kafkaapi"
+	"github.com/codecrafters-io/kafka-tester/protocol/value"
 )
 
 type ClusterMetadataGenerator struct {
@@ -50,21 +51,21 @@ func (g *ClusterMetadataGenerator) writeLogFile() error {
 	}
 
 	recordBatch1 := kafkaapi.RecordBatch{
-		BaseOffset:           baseOffset,
-		PartitionLeaderEpoch: 1,
-		Attributes:           0,
-		LastOffsetDelta:      0, // len(records) - 1
-		FirstTimestamp:       1726045943832,
-		MaxTimestamp:         1726045943832,
-		ProducerId:           -1,
-		ProducerEpoch:        -1,
-		BaseSequence:         -1,
+		BaseOffset:           value.Int64{Value: baseOffset},
+		PartitionLeaderEpoch: value.Int32{Value: 1},
+		Attributes:           value.Int16{Value: 0},
+		LastOffsetDelta:      value.Int32{Value: 0}, // len(records) - 1
+		FirstTimestamp:       value.Int64{Value: 1726045943832},
+		MaxTimestamp:         value.Int64{Value: 1726045943832},
+		ProducerId:           value.Int64{Value: -1},
+		ProducerEpoch:        value.Int16{Value: -1},
+		BaseSequence:         value.Int32{Value: -1},
 		Records: []kafkaapi.Record{
 			{
-				Attributes:     0,
-				TimestampDelta: 0,
-				Key:            nil,
-				Value:          GetEncodedBytes(featureLevelRecord),
+				Attributes:     value.Int8{Value: 0},
+				TimestampDelta: value.Int64{Value: 0},
+				Key:            value.RawBytes{Value: nil},
+				Value:          value.RawBytes{Value: GetEncodedBytes(featureLevelRecord)},
 				Headers:        []kafkaapi.RecordHeader{},
 			},
 		},
@@ -113,34 +114,34 @@ func (g *ClusterMetadataGenerator) writeLogFile() error {
 
 		// Add topic record
 		records = append(records, kafkaapi.Record{
-			Attributes:     0,
-			TimestampDelta: 0,
-			Key:            nil,
-			Value:          GetEncodedBytes(topicRecord),
+			Attributes:     value.Int8{Value: 0},
+			TimestampDelta: value.Int64{Value: 0},
+			Key:            value.RawBytes{Value: nil},
+			Value:          value.RawBytes{Value: GetEncodedBytes(topicRecord)},
 			Headers:        []kafkaapi.RecordHeader{},
 		})
 
 		// Add all partition records
 		for _, partitionRecord := range partitionRecords {
 			records = append(records, kafkaapi.Record{
-				Attributes:     0,
-				TimestampDelta: 0,
-				Key:            nil,
-				Value:          GetEncodedBytes(partitionRecord),
+				Attributes:     value.Int8{Value: 0},
+				TimestampDelta: value.Int64{Value: 0},
+				Key:            value.RawBytes{Value: nil},
+				Value:          value.RawBytes{Value: GetEncodedBytes(partitionRecord)},
 				Headers:        []kafkaapi.RecordHeader{},
 			})
 		}
 
 		recordBatch := kafkaapi.RecordBatch{
-			BaseOffset:           baseOffset,
-			PartitionLeaderEpoch: 1,
-			Attributes:           0,
-			LastOffsetDelta:      int32(len(records) - 1),
-			FirstTimestamp:       1726045957397,
-			MaxTimestamp:         1726045957397,
-			ProducerId:           -1,
-			ProducerEpoch:        -1,
-			BaseSequence:         -1,
+			BaseOffset:           value.Int64{Value: baseOffset},
+			PartitionLeaderEpoch: value.Int32{Value: 1},
+			Attributes:           value.Int16{Value: 0},
+			LastOffsetDelta:      value.Int32{Value: int32(len(records) - 1)},
+			FirstTimestamp:       value.Int64{Value: 1726045957397},
+			MaxTimestamp:         value.Int64{Value: 1726045957397},
+			ProducerId:           value.Int64{Value: -1},
+			ProducerEpoch:        value.Int16{Value: -1},
+			BaseSequence:         value.Int32{Value: -1},
 			Records:              records,
 		}
 		recordBatches = append(recordBatches, recordBatch)

protocol/kafka_files_generator/partition_generator.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/codecrafters-io/kafka-tester/protocol/encoder"
 	"github.com/codecrafters-io/kafka-tester/protocol/kafkaapi"
+	"github.com/codecrafters-io/kafka-tester/protocol/value"
 	"github.com/codecrafters-io/tester-utils/logger"
 )
 
@@ -95,21 +96,21 @@ func (c *PartitionGenerationConfig) generateRecordBatchesFromLogs(logs []string)
 
 	for i, message := range logs {
 		recordBatches = append(recordBatches, kafkaapi.RecordBatch{
-			BaseOffset:           int64(i),
-			PartitionLeaderEpoch: 0,
-			Attributes:           0,
-			LastOffsetDelta:      0,
-			FirstTimestamp:       1726045973899,
-			MaxTimestamp:         1726045973899,
-			ProducerId:           0,
-			ProducerEpoch:        0,
-			BaseSequence:         0,
+			BaseOffset:           value.Int64{Value: int64(i)},
+			PartitionLeaderEpoch: value.Int32{Value: 0},
+			Attributes:           value.Int16{Value: 0},
+			LastOffsetDelta:      value.Int32{Value: 0},
+			FirstTimestamp:       value.Int64{Value: 1726045973899},
+			MaxTimestamp:         value.Int64{Value: 1726045973899},
+			ProducerId:           value.Int64{Value: 0},
+			ProducerEpoch:        value.Int16{Value: 0},
+			BaseSequence:         value.Int32{Value: 0},
 			Records: []kafkaapi.Record{
 				{
-					Attributes:     0,
-					TimestampDelta: 0,
-					Key:            nil,
-					Value:          []byte(message),
+					Attributes:     value.Int8{Value: 0},
+					TimestampDelta: value.Int64{Value: 0},
+					Key:            value.RawBytes{},
+					Value:          value.RawBytes{Value: []byte(message)},
 					Headers:        []kafkaapi.RecordHeader{},
 				},
 			},

protocol/kafkaapi/record.go

@@ -1,35 +1,38 @@
 package kafkaapi
 
-import "github.com/codecrafters-io/kafka-tester/protocol/encoder"
+import (
+	"github.com/codecrafters-io/kafka-tester/protocol/encoder"
+	"github.com/codecrafters-io/kafka-tester/protocol/value"
+)
 
 type Record struct {
-	Length         int32
-	Attributes     int8
-	TimestampDelta int64
-	OffsetDelta    int32
-	Key            []byte
-	Value          []byte
+	Length         value.Int32
+	Attributes     value.Int8
+	TimestampDelta value.Int64
+	OffsetDelta    value.Int32
+	Key            value.RawBytes
+	Value          value.RawBytes
 	Headers        []RecordHeader
 }
 
 func (r Record) Encode(pe *encoder.Encoder) {
 	propertiesEncoder := encoder.NewEncoder()
 
-	propertiesEncoder.WriteInt8(r.Attributes)
-	propertiesEncoder.WriteVarint(r.TimestampDelta)
-	propertiesEncoder.WriteVarint(int64(r.OffsetDelta))
+	propertiesEncoder.WriteInt8(r.Attributes.Value)
+	propertiesEncoder.WriteVarint(r.TimestampDelta.Value)
+	propertiesEncoder.WriteVarint(int64(r.OffsetDelta.Value))
 
 	// Special encoding that does not belong to any data type and is only present inside Records
 	// similar to protobuf encoding. It is mentioned in the Kafka docs here:  https://kafka.apache.org/documentation/#recordheader
-	if r.Key == nil {
+	if r.Key.Value == nil {
 		propertiesEncoder.WriteVarint(-1)
 	} else {
-		propertiesEncoder.WriteVarint(int64(len(r.Key)))
-		propertiesEncoder.WriteRawBytes(r.Key)
+		propertiesEncoder.WriteVarint(int64(len(r.Key.Value)))
+		propertiesEncoder.WriteRawBytes(r.Key.Value)
 	}
 
-	propertiesEncoder.WriteVarint(int64(len(r.Value)))
-	propertiesEncoder.WriteRawBytes(r.Value)
+	propertiesEncoder.WriteVarint(int64(len(r.Value.Value)))
+	propertiesEncoder.WriteRawBytes(r.Value.Value)
 	propertiesEncoder.WriteVarint(int64(len(r.Headers)))
 	for _, header := range r.Headers {
 		header.Encode(propertiesEncoder)
@@ -42,13 +45,13 @@ func (r Record) Encode(pe *encoder.Encoder) {
 }
 
 type RecordHeader struct {
-	Key   string
-	Value []byte
+	Key   value.RawBytes
+	Value value.RawBytes
 }
 
 func (rh RecordHeader) Encode(pe *encoder.Encoder) {
-	pe.WriteVarint(int64(len(rh.Key)))
-	pe.WriteRawBytes([]byte(rh.Key))
-	pe.WriteVarint(int64(len(rh.Value)))
-	pe.WriteRawBytes(rh.Value)
+	pe.WriteVarint(int64(len(rh.Key.Value)))
+	pe.WriteRawBytes(rh.Key.Value)
+	pe.WriteVarint(int64(len(rh.Value.Value)))
+	pe.WriteRawBytes(rh.Value.Value)
 }

protocol/kafkaapi/record_batch.go

@@ -4,59 +4,64 @@ import (
 	"hash/crc32"
 
 	"github.com/codecrafters-io/kafka-tester/protocol/encoder"
+	"github.com/codecrafters-io/kafka-tester/protocol/value"
 )
 
 type RecordBatches []RecordBatch
 
 func (rbs RecordBatches) Encode(pe *encoder.Encoder) {
-	for _, rb := range rbs {
-		rb.Encode(pe)
+	for i := range rbs {
+		rbs[i].Encode(pe)
 	}
 }
 
 type RecordBatch struct {
-	BaseOffset           int64
-	BatchLength          int32
-	PartitionLeaderEpoch int32
-	Magic                int8
-	CRC                  int32
-	Attributes           int16
-	LastOffsetDelta      int32
-	FirstTimestamp       int64
-	MaxTimestamp         int64
-	ProducerId           int64
-	ProducerEpoch        int16
-	BaseSequence         int32
+	BaseOffset           value.Int64
+	BatchLength          value.Int32
+	PartitionLeaderEpoch value.Int32
+	Magic                value.Int8
+	CRC                  value.Int32
+	Attributes           value.Int16
+	LastOffsetDelta      value.Int32
+	FirstTimestamp       value.Int64
+	MaxTimestamp         value.Int64
+	ProducerId           value.Int64
+	ProducerEpoch        value.Int16
+	BaseSequence         value.Int32
 	Records              []Record
 }
 
-func (rb RecordBatch) Encode(pe *encoder.Encoder) {
+func (rb *RecordBatch) Encode(pe *encoder.Encoder) {
 	propertiesEncoder := encoder.NewEncoder()
-	propertiesEncoder.WriteInt16(rb.Attributes)
-	propertiesEncoder.WriteInt32(rb.LastOffsetDelta)
-	propertiesEncoder.WriteInt64(rb.FirstTimestamp)
-	propertiesEncoder.WriteInt64(rb.MaxTimestamp)
-	propertiesEncoder.WriteInt64(rb.ProducerId)
-	propertiesEncoder.WriteInt16(rb.ProducerEpoch)
-	propertiesEncoder.WriteInt32(rb.BaseSequence)
+	propertiesEncoder.WriteInt16(rb.Attributes.Value)
+	propertiesEncoder.WriteInt32(rb.LastOffsetDelta.Value)
+	propertiesEncoder.WriteInt64(rb.FirstTimestamp.Value)
+	propertiesEncoder.WriteInt64(rb.MaxTimestamp.Value)
+	propertiesEncoder.WriteInt64(rb.ProducerId.Value)
+	propertiesEncoder.WriteInt16(rb.ProducerEpoch.Value)
+	propertiesEncoder.WriteInt32(rb.BaseSequence.Value)
 	propertiesEncoder.WriteInt32(int32(len(rb.Records)))
 
 	for i, record := range rb.Records {
-		record.OffsetDelta = int32(i) // Offset Deltas are consecutive numerals from 0 to N-1
+		record.OffsetDelta = value.Int32{Value: int32(i)} // Offset Deltas are consecutive numerals from 0 to N-1
 		// We can set them programmatically as we know the order of the records
 		record.Encode(propertiesEncoder)
 	}
 
 	propertiesEncoderBytes := propertiesEncoder.Bytes()
 	computedChecksum := crc32.Checksum(propertiesEncoderBytes, crc32.MakeTable(crc32.Castagnoli))
-	rb.CRC = int32(computedChecksum)
-	rb.BatchLength = int32(len(propertiesEncoderBytes) + 4 + 1 + 4) // partitionLeaderEpoch + magic value + CRC
+	rb.CRC = value.Int32{
+		Value: int32(computedChecksum),
+	}
+	rb.BatchLength = value.Int32{
+		Value: int32(len(propertiesEncoderBytes) + 4 + 1 + 4), // partitionLeaderEpoch + magic value + CRC
+	}
 
 	// Encode everything now
-	pe.WriteInt64(rb.BaseOffset)
-	pe.WriteInt32(rb.BatchLength)
-	pe.WriteInt32(rb.PartitionLeaderEpoch)
-	pe.WriteInt8(2)       // Magic value is 2
-	pe.WriteInt32(rb.CRC) // CRC placeholder
+	pe.WriteInt64(rb.BaseOffset.Value)
+	pe.WriteInt32(rb.BatchLength.Value)
+	pe.WriteInt32(rb.PartitionLeaderEpoch.Value)
+	pe.WriteInt8(2)             // Magic value is 2
+	pe.WriteInt32(rb.CRC.Value) // CRC placeholder
 	pe.WriteRawBytes(propertiesEncoderBytes)
 }

protocol/value/int64.go

@@ -0,0 +1,15 @@
+package value
+
+import "fmt"
+
+type Int64 struct {
+	Value int64
+}
+
+func (v Int64) String() string {
+	return fmt.Sprintf("%d", v.Value)
+}
+
+func (v Int64) GetType() string {
+	return "INT64"
+}