github.com/hack0072008/kafka-go@v1.0.1/protocol/record_batch.go (about)

     1  package protocol
     2  
     3  import (
     4  	"errors"
     5  	"io"
     6  	"time"
     7  )
     8  
     9  // RecordReader is an interface representing a sequence of records. Record sets
    10  // are used in both produce and fetch requests to represent the sequence of
    11  // records that are sent to or receive from kafka brokers.
    12  //
    13  // RecordSet values are not safe to use concurrently from multiple goroutines.
    14  type RecordReader interface {
    15  	// Returns the next record in the set, or io.EOF if the end of the sequence
    16  	// has been reached.
    17  	//
    18  	// The returned Record is guaranteed to be valid until the next call to
    19  	// ReadRecord. If the program needs to retain the Record value it must make
    20  	// a copy.
    21  	ReadRecord() (*Record, error)
    22  }
    23  
    24  // NewRecordReader constructs a reader exposing the records passed as arguments.
    25  func NewRecordReader(records ...Record) RecordReader {
    26  	switch len(records) {
    27  	case 0:
    28  		return emptyRecordReader{}
    29  	default:
    30  		r := &recordReader{records: make([]Record, len(records))}
    31  		copy(r.records, records)
    32  		return r
    33  	}
    34  }
    35  
    36  // MultiRecordReader merges multiple record batches into one.
    37  func MultiRecordReader(batches ...RecordReader) RecordReader {
    38  	switch len(batches) {
    39  	case 0:
    40  		return emptyRecordReader{}
    41  	case 1:
    42  		return batches[0]
    43  	default:
    44  		m := &multiRecordReader{batches: make([]RecordReader, len(batches))}
    45  		copy(m.batches, batches)
    46  		return m
    47  	}
    48  }
    49  
    50  func forEachRecord(r RecordReader, f func(int, *Record) error) error {
    51  	for i := 0; ; i++ {
    52  		rec, err := r.ReadRecord()
    53  
    54  		if err != nil {
    55  			if errors.Is(err, io.EOF) {
    56  				err = nil
    57  			}
    58  			return err
    59  		}
    60  
    61  		if err := handleRecord(i, rec, f); err != nil {
    62  			return err
    63  		}
    64  	}
    65  }
    66  
    67  func handleRecord(i int, r *Record, f func(int, *Record) error) error {
    68  	if r.Key != nil {
    69  		defer r.Key.Close()
    70  	}
    71  	if r.Value != nil {
    72  		defer r.Value.Close()
    73  	}
    74  	return f(i, r)
    75  }
    76  
    77  type recordReader struct {
    78  	records []Record
    79  	index   int
    80  }
    81  
    82  func (r *recordReader) ReadRecord() (*Record, error) {
    83  	if i := r.index; i >= 0 && i < len(r.records) {
    84  		r.index++
    85  		return &r.records[i], nil
    86  	}
    87  	return nil, io.EOF
    88  }
    89  
    90  type multiRecordReader struct {
    91  	batches []RecordReader
    92  	index   int
    93  }
    94  
    95  func (m *multiRecordReader) ReadRecord() (*Record, error) {
    96  	for {
    97  		if m.index == len(m.batches) {
    98  			return nil, io.EOF
    99  		}
   100  		r, err := m.batches[m.index].ReadRecord()
   101  		if err == nil {
   102  			return r, nil
   103  		}
   104  		if !errors.Is(err, io.EOF) {
   105  			return nil, err
   106  		}
   107  		m.index++
   108  	}
   109  }
   110  
   111  func concatRecordReader(head RecordReader, tail RecordReader) RecordReader {
   112  	if head == nil {
   113  		return tail
   114  	}
   115  	if m, _ := head.(*multiRecordReader); m != nil {
   116  		m.batches = append(m.batches, tail)
   117  		return m
   118  	}
   119  	return MultiRecordReader(head, tail)
   120  }
   121  
   122  // optimizedRecordReader is an implementation of a RecordReader which exposes a
   123  // sequence
   124  type optimizedRecordReader struct {
   125  	records []optimizedRecord
   126  	index   int
   127  	buffer  Record
   128  	headers [][]Header
   129  }
   130  
   131  func (r *optimizedRecordReader) ReadRecord() (*Record, error) {
   132  	if i := r.index; i >= 0 && i < len(r.records) {
   133  		rec := &r.records[i]
   134  		r.index++
   135  		r.buffer = Record{
   136  			Offset: rec.offset,
   137  			Time:   rec.time(),
   138  			Key:    rec.key(),
   139  			Value:  rec.value(),
   140  		}
   141  		if i < len(r.headers) {
   142  			r.buffer.Headers = r.headers[i]
   143  		}
   144  		return &r.buffer, nil
   145  	}
   146  	return nil, io.EOF
   147  }
   148  
   149  type optimizedRecord struct {
   150  	offset    int64
   151  	timestamp int64
   152  	keyRef    *pageRef
   153  	valueRef  *pageRef
   154  }
   155  
   156  func (r *optimizedRecord) time() time.Time {
   157  	return makeTime(r.timestamp)
   158  }
   159  
   160  func (r *optimizedRecord) key() Bytes {
   161  	return makeBytes(r.keyRef)
   162  }
   163  
   164  func (r *optimizedRecord) value() Bytes {
   165  	return makeBytes(r.valueRef)
   166  }
   167  
   168  func makeBytes(ref *pageRef) Bytes {
   169  	if ref == nil {
   170  		return nil
   171  	}
   172  	return ref
   173  }
   174  
   175  type emptyRecordReader struct{}
   176  
   177  func (emptyRecordReader) ReadRecord() (*Record, error) { return nil, io.EOF }
   178  
   179  // ControlRecord represents a record read from a control batch.
   180  type ControlRecord struct {
   181  	Offset  int64
   182  	Time    time.Time
   183  	Version int16
   184  	Type    int16
   185  	Data    []byte
   186  	Headers []Header
   187  }
   188  
   189  func ReadControlRecord(r *Record) (*ControlRecord, error) {
   190  	if r.Key != nil {
   191  		defer r.Key.Close()
   192  	}
   193  	if r.Value != nil {
   194  		defer r.Value.Close()
   195  	}
   196  
   197  	k, err := ReadAll(r.Key)
   198  	if err != nil {
   199  		return nil, err
   200  	}
   201  	if k == nil {
   202  		return nil, Error("invalid control record with nil key")
   203  	}
   204  	if len(k) != 4 {
   205  		return nil, Errorf("invalid control record with key of size %d", len(k))
   206  	}
   207  
   208  	v, err := ReadAll(r.Value)
   209  	if err != nil {
   210  		return nil, err
   211  	}
   212  
   213  	c := &ControlRecord{
   214  		Offset:  r.Offset,
   215  		Time:    r.Time,
   216  		Version: readInt16(k[:2]),
   217  		Type:    readInt16(k[2:]),
   218  		Data:    v,
   219  		Headers: r.Headers,
   220  	}
   221  
   222  	return c, nil
   223  }
   224  
   225  func (cr *ControlRecord) Key() Bytes {
   226  	k := make([]byte, 4)
   227  	writeInt16(k[:2], cr.Version)
   228  	writeInt16(k[2:], cr.Type)
   229  	return NewBytes(k)
   230  }
   231  
   232  func (cr *ControlRecord) Value() Bytes {
   233  	return NewBytes(cr.Data)
   234  }
   235  
   236  func (cr *ControlRecord) Record() Record {
   237  	return Record{
   238  		Offset:  cr.Offset,
   239  		Time:    cr.Time,
   240  		Key:     cr.Key(),
   241  		Value:   cr.Value(),
   242  		Headers: cr.Headers,
   243  	}
   244  }
   245  
   246  // ControlBatch is an implementation of the RecordReader interface representing
   247  // control batches returned by kafka brokers.
   248  type ControlBatch struct {
   249  	Attributes           Attributes
   250  	PartitionLeaderEpoch int32
   251  	BaseOffset           int64
   252  	ProducerID           int64
   253  	ProducerEpoch        int16
   254  	BaseSequence         int32
   255  	Records              RecordReader
   256  }
   257  
   258  // NewControlBatch constructs a control batch from the list of records passed as
   259  // arguments.
   260  func NewControlBatch(records ...ControlRecord) *ControlBatch {
   261  	rawRecords := make([]Record, len(records))
   262  	for i, cr := range records {
   263  		rawRecords[i] = cr.Record()
   264  	}
   265  	return &ControlBatch{
   266  		Records: NewRecordReader(rawRecords...),
   267  	}
   268  }
   269  
   270  func (c *ControlBatch) ReadRecord() (*Record, error) {
   271  	return c.Records.ReadRecord()
   272  }
   273  
   274  func (c *ControlBatch) ReadControlRecord() (*ControlRecord, error) {
   275  	r, err := c.ReadRecord()
   276  	if err != nil {
   277  		return nil, err
   278  	}
   279  	if r.Key != nil {
   280  		defer r.Key.Close()
   281  	}
   282  	if r.Value != nil {
   283  		defer r.Value.Close()
   284  	}
   285  	return ReadControlRecord(r)
   286  }
   287  
   288  func (c *ControlBatch) Offset() int64 {
   289  	return c.BaseOffset
   290  }
   291  
   292  func (c *ControlBatch) Version() int {
   293  	return 2
   294  }
   295  
   296  // RecordBatch is an implementation of the RecordReader interface representing
   297  // regular record batches (v2).
   298  type RecordBatch struct {
   299  	Attributes           Attributes
   300  	PartitionLeaderEpoch int32
   301  	BaseOffset           int64
   302  	ProducerID           int64
   303  	ProducerEpoch        int16
   304  	BaseSequence         int32
   305  	Records              RecordReader
   306  }
   307  
   308  func (r *RecordBatch) ReadRecord() (*Record, error) {
   309  	return r.Records.ReadRecord()
   310  }
   311  
   312  func (r *RecordBatch) Offset() int64 {
   313  	return r.BaseOffset
   314  }
   315  
   316  func (r *RecordBatch) Version() int {
   317  	return 2
   318  }
   319  
   320  // MessageSet is an implementation of the RecordReader interface representing
   321  // regular message sets (v1).
   322  type MessageSet struct {
   323  	Attributes Attributes
   324  	BaseOffset int64
   325  	Records    RecordReader
   326  }
   327  
   328  func (m *MessageSet) ReadRecord() (*Record, error) {
   329  	return m.Records.ReadRecord()
   330  }
   331  
   332  func (m *MessageSet) Offset() int64 {
   333  	return m.BaseOffset
   334  }
   335  
   336  func (m *MessageSet) Version() int {
   337  	return 1
   338  }
   339  
   340  // RecordStream is an implementation of the RecordReader interface which
   341  // combines multiple underlying RecordReader and only expose records that
   342  // are not from control batches.
   343  type RecordStream struct {
   344  	Records []RecordReader
   345  	index   int
   346  }
   347  
   348  func (s *RecordStream) ReadRecord() (*Record, error) {
   349  	for {
   350  		if s.index < 0 || s.index >= len(s.Records) {
   351  			return nil, io.EOF
   352  		}
   353  
   354  		if _, isControl := s.Records[s.index].(*ControlBatch); isControl {
   355  			s.index++
   356  			continue
   357  		}
   358  
   359  		r, err := s.Records[s.index].ReadRecord()
   360  		if err != nil {
   361  			if errors.Is(err, io.EOF) {
   362  				s.index++
   363  				continue
   364  			}
   365  		}
   366  
   367  		return r, err
   368  	}
   369  }