github.com/QuangHoangHao/kafka-go@v0.4.36/protocol/encode.go (about)

     1  package protocol
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  	"fmt"
     7  	"hash/crc32"
     8  	"io"
     9  	"reflect"
    10  	"sync"
    11  	"sync/atomic"
    12  )
    13  
    14  type encoder struct {
    15  	writer io.Writer
    16  	err    error
    17  	table  *crc32.Table
    18  	crc32  uint32
    19  	buffer [32]byte
    20  }
    21  
    22  type encoderChecksum struct {
    23  	reader  io.Reader
    24  	encoder *encoder
    25  }
    26  
    27  func (e *encoderChecksum) Read(b []byte) (int, error) {
    28  	n, err := e.reader.Read(b)
    29  	if n > 0 {
    30  		e.encoder.update(b[:n])
    31  	}
    32  	return n, err
    33  }
    34  
    35  func (e *encoder) Reset(w io.Writer) {
    36  	e.writer = w
    37  	e.err = nil
    38  	e.table = nil
    39  	e.crc32 = 0
    40  	e.buffer = [32]byte{}
    41  }
    42  
    43  func (e *encoder) ReadFrom(r io.Reader) (int64, error) {
    44  	if e.table != nil {
    45  		r = &encoderChecksum{
    46  			reader:  r,
    47  			encoder: e,
    48  		}
    49  	}
    50  	return io.Copy(e.writer, r)
    51  }
    52  
    53  func (e *encoder) Write(b []byte) (int, error) {
    54  	if e.err != nil {
    55  		return 0, e.err
    56  	}
    57  	n, err := e.writer.Write(b)
    58  	if n > 0 {
    59  		e.update(b[:n])
    60  	}
    61  	if err != nil {
    62  		e.err = err
    63  	}
    64  	return n, err
    65  }
    66  
    67  func (e *encoder) WriteByte(b byte) error {
    68  	e.buffer[0] = b
    69  	_, err := e.Write(e.buffer[:1])
    70  	return err
    71  }
    72  
    73  func (e *encoder) WriteString(s string) (int, error) {
    74  	// This implementation is an optimization to avoid the heap allocation that
    75  	// would occur when converting the string to a []byte to call crc32.Update.
    76  	//
    77  	// Strings are rarely long in the kafka protocol, so the use of a 32 byte
    78  	// buffer is a good comprise between keeping the encoder value small and
    79  	// limiting the number of calls to Write.
    80  	//
    81  	// We introduced this optimization because memory profiles on the benchmarks
    82  	// showed that most heap allocations were caused by this code path.
    83  	n := 0
    84  
    85  	for len(s) != 0 {
    86  		c := copy(e.buffer[:], s)
    87  		w, err := e.Write(e.buffer[:c])
    88  		n += w
    89  		if err != nil {
    90  			return n, err
    91  		}
    92  		s = s[c:]
    93  	}
    94  
    95  	return n, nil
    96  }
    97  
    98  func (e *encoder) setCRC(table *crc32.Table) {
    99  	e.table, e.crc32 = table, 0
   100  }
   101  
   102  func (e *encoder) update(b []byte) {
   103  	if e.table != nil {
   104  		e.crc32 = crc32.Update(e.crc32, e.table, b)
   105  	}
   106  }
   107  
   108  func (e *encoder) encodeBool(v value) {
   109  	b := int8(0)
   110  	if v.bool() {
   111  		b = 1
   112  	}
   113  	e.writeInt8(b)
   114  }
   115  
   116  func (e *encoder) encodeInt8(v value) {
   117  	e.writeInt8(v.int8())
   118  }
   119  
   120  func (e *encoder) encodeInt16(v value) {
   121  	e.writeInt16(v.int16())
   122  }
   123  
   124  func (e *encoder) encodeInt32(v value) {
   125  	e.writeInt32(v.int32())
   126  }
   127  
   128  func (e *encoder) encodeInt64(v value) {
   129  	e.writeInt64(v.int64())
   130  }
   131  
   132  func (e *encoder) encodeString(v value) {
   133  	e.writeString(v.string())
   134  }
   135  
   136  func (e *encoder) encodeCompactString(v value) {
   137  	e.writeCompactString(v.string())
   138  }
   139  
   140  func (e *encoder) encodeNullString(v value) {
   141  	e.writeNullString(v.string())
   142  }
   143  
   144  func (e *encoder) encodeCompactNullString(v value) {
   145  	e.writeCompactNullString(v.string())
   146  }
   147  
   148  func (e *encoder) encodeBytes(v value) {
   149  	e.writeBytes(v.bytes())
   150  }
   151  
   152  func (e *encoder) encodeCompactBytes(v value) {
   153  	e.writeCompactBytes(v.bytes())
   154  }
   155  
   156  func (e *encoder) encodeNullBytes(v value) {
   157  	e.writeNullBytes(v.bytes())
   158  }
   159  
   160  func (e *encoder) encodeCompactNullBytes(v value) {
   161  	e.writeCompactNullBytes(v.bytes())
   162  }
   163  
   164  func (e *encoder) encodeArray(v value, elemType reflect.Type, encodeElem encodeFunc) {
   165  	a := v.array(elemType)
   166  	n := a.length()
   167  	e.writeInt32(int32(n))
   168  
   169  	for i := 0; i < n; i++ {
   170  		encodeElem(e, a.index(i))
   171  	}
   172  }
   173  
   174  func (e *encoder) encodeCompactArray(v value, elemType reflect.Type, encodeElem encodeFunc) {
   175  	a := v.array(elemType)
   176  	n := a.length()
   177  	e.writeUnsignedVarInt(uint64(n + 1))
   178  
   179  	for i := 0; i < n; i++ {
   180  		encodeElem(e, a.index(i))
   181  	}
   182  }
   183  
   184  func (e *encoder) encodeNullArray(v value, elemType reflect.Type, encodeElem encodeFunc) {
   185  	a := v.array(elemType)
   186  	if a.isNil() {
   187  		e.writeInt32(-1)
   188  		return
   189  	}
   190  
   191  	n := a.length()
   192  	e.writeInt32(int32(n))
   193  
   194  	for i := 0; i < n; i++ {
   195  		encodeElem(e, a.index(i))
   196  	}
   197  }
   198  
   199  func (e *encoder) encodeCompactNullArray(v value, elemType reflect.Type, encodeElem encodeFunc) {
   200  	a := v.array(elemType)
   201  	if a.isNil() {
   202  		e.writeUnsignedVarInt(0)
   203  		return
   204  	}
   205  
   206  	n := a.length()
   207  	e.writeUnsignedVarInt(uint64(n + 1))
   208  	for i := 0; i < n; i++ {
   209  		encodeElem(e, a.index(i))
   210  	}
   211  }
   212  
   213  func (e *encoder) writeInt8(i int8) {
   214  	writeInt8(e.buffer[:1], i)
   215  	e.Write(e.buffer[:1])
   216  }
   217  
   218  func (e *encoder) writeInt16(i int16) {
   219  	writeInt16(e.buffer[:2], i)
   220  	e.Write(e.buffer[:2])
   221  }
   222  
   223  func (e *encoder) writeInt32(i int32) {
   224  	writeInt32(e.buffer[:4], i)
   225  	e.Write(e.buffer[:4])
   226  }
   227  
   228  func (e *encoder) writeInt64(i int64) {
   229  	writeInt64(e.buffer[:8], i)
   230  	e.Write(e.buffer[:8])
   231  }
   232  
   233  func (e *encoder) writeString(s string) {
   234  	e.writeInt16(int16(len(s)))
   235  	e.WriteString(s)
   236  }
   237  
   238  func (e *encoder) writeVarString(s string) {
   239  	e.writeVarInt(int64(len(s)))
   240  	e.WriteString(s)
   241  }
   242  
   243  func (e *encoder) writeCompactString(s string) {
   244  	e.writeUnsignedVarInt(uint64(len(s)) + 1)
   245  	e.WriteString(s)
   246  }
   247  
   248  func (e *encoder) writeNullString(s string) {
   249  	if s == "" {
   250  		e.writeInt16(-1)
   251  	} else {
   252  		e.writeInt16(int16(len(s)))
   253  		e.WriteString(s)
   254  	}
   255  }
   256  
   257  func (e *encoder) writeCompactNullString(s string) {
   258  	if s == "" {
   259  		e.writeUnsignedVarInt(0)
   260  	} else {
   261  		e.writeUnsignedVarInt(uint64(len(s)) + 1)
   262  		e.WriteString(s)
   263  	}
   264  }
   265  
   266  func (e *encoder) writeBytes(b []byte) {
   267  	e.writeInt32(int32(len(b)))
   268  	e.Write(b)
   269  }
   270  
   271  func (e *encoder) writeCompactBytes(b []byte) {
   272  	e.writeUnsignedVarInt(uint64(len(b)) + 1)
   273  	e.Write(b)
   274  }
   275  
   276  func (e *encoder) writeNullBytes(b []byte) {
   277  	if b == nil {
   278  		e.writeInt32(-1)
   279  	} else {
   280  		e.writeInt32(int32(len(b)))
   281  		e.Write(b)
   282  	}
   283  }
   284  
   285  func (e *encoder) writeVarNullBytes(b []byte) {
   286  	if b == nil {
   287  		e.writeVarInt(-1)
   288  	} else {
   289  		e.writeVarInt(int64(len(b)))
   290  		e.Write(b)
   291  	}
   292  }
   293  
   294  func (e *encoder) writeCompactNullBytes(b []byte) {
   295  	if b == nil {
   296  		e.writeUnsignedVarInt(0)
   297  	} else {
   298  		e.writeUnsignedVarInt(uint64(len(b)) + 1)
   299  		e.Write(b)
   300  	}
   301  }
   302  
   303  func (e *encoder) writeNullBytesFrom(b Bytes) error {
   304  	if b == nil {
   305  		e.writeInt32(-1)
   306  		return nil
   307  	} else {
   308  		size := int64(b.Len())
   309  		e.writeInt32(int32(size))
   310  		n, err := io.Copy(e, b)
   311  		if err == nil && n != size {
   312  			err = fmt.Errorf("size of nullable bytes does not match the number of bytes that were written (size=%d, written=%d): %w", size, n, io.ErrUnexpectedEOF)
   313  		}
   314  		return err
   315  	}
   316  }
   317  
   318  func (e *encoder) writeVarNullBytesFrom(b Bytes) error {
   319  	if b == nil {
   320  		e.writeVarInt(-1)
   321  		return nil
   322  	} else {
   323  		size := int64(b.Len())
   324  		e.writeVarInt(size)
   325  		n, err := io.Copy(e, b)
   326  		if err == nil && n != size {
   327  			err = fmt.Errorf("size of nullable bytes does not match the number of bytes that were written (size=%d, written=%d): %w", size, n, io.ErrUnexpectedEOF)
   328  		}
   329  		return err
   330  	}
   331  }
   332  
   333  func (e *encoder) writeVarInt(i int64) {
   334  	e.writeUnsignedVarInt(uint64((i << 1) ^ (i >> 63)))
   335  }
   336  
   337  func (e *encoder) writeUnsignedVarInt(i uint64) {
   338  	b := e.buffer[:]
   339  	n := 0
   340  
   341  	for i >= 0x80 && n < len(b) {
   342  		b[n] = byte(i) | 0x80
   343  		i >>= 7
   344  		n++
   345  	}
   346  
   347  	if n < len(b) {
   348  		b[n] = byte(i)
   349  		n++
   350  	}
   351  
   352  	e.Write(b[:n])
   353  }
   354  
   355  type encodeFunc func(*encoder, value)
   356  
   357  var (
   358  	_ io.ReaderFrom   = (*encoder)(nil)
   359  	_ io.Writer       = (*encoder)(nil)
   360  	_ io.ByteWriter   = (*encoder)(nil)
   361  	_ io.StringWriter = (*encoder)(nil)
   362  
   363  	writerTo = reflect.TypeOf((*io.WriterTo)(nil)).Elem()
   364  )
   365  
   366  func encodeFuncOf(typ reflect.Type, version int16, flexible bool, tag structTag) encodeFunc {
   367  	if reflect.PtrTo(typ).Implements(writerTo) {
   368  		return writerEncodeFuncOf(typ)
   369  	}
   370  	switch typ.Kind() {
   371  	case reflect.Bool:
   372  		return (*encoder).encodeBool
   373  	case reflect.Int8:
   374  		return (*encoder).encodeInt8
   375  	case reflect.Int16:
   376  		return (*encoder).encodeInt16
   377  	case reflect.Int32:
   378  		return (*encoder).encodeInt32
   379  	case reflect.Int64:
   380  		return (*encoder).encodeInt64
   381  	case reflect.String:
   382  		return stringEncodeFuncOf(flexible, tag)
   383  	case reflect.Struct:
   384  		return structEncodeFuncOf(typ, version, flexible)
   385  	case reflect.Slice:
   386  		if typ.Elem().Kind() == reflect.Uint8 { // []byte
   387  			return bytesEncodeFuncOf(flexible, tag)
   388  		}
   389  		return arrayEncodeFuncOf(typ, version, flexible, tag)
   390  	default:
   391  		panic("unsupported type: " + typ.String())
   392  	}
   393  }
   394  
   395  func stringEncodeFuncOf(flexible bool, tag structTag) encodeFunc {
   396  	switch {
   397  	case flexible && tag.Nullable:
   398  		// In flexible messages, all strings are compact
   399  		return (*encoder).encodeCompactNullString
   400  	case flexible:
   401  		// In flexible messages, all strings are compact
   402  		return (*encoder).encodeCompactString
   403  	case tag.Nullable:
   404  		return (*encoder).encodeNullString
   405  	default:
   406  		return (*encoder).encodeString
   407  	}
   408  }
   409  
   410  func bytesEncodeFuncOf(flexible bool, tag structTag) encodeFunc {
   411  	switch {
   412  	case flexible && tag.Nullable:
   413  		// In flexible messages, all arrays are compact
   414  		return (*encoder).encodeCompactNullBytes
   415  	case flexible:
   416  		// In flexible messages, all arrays are compact
   417  		return (*encoder).encodeCompactBytes
   418  	case tag.Nullable:
   419  		return (*encoder).encodeNullBytes
   420  	default:
   421  		return (*encoder).encodeBytes
   422  	}
   423  }
   424  
   425  func structEncodeFuncOf(typ reflect.Type, version int16, flexible bool) encodeFunc {
   426  	type field struct {
   427  		encode encodeFunc
   428  		index  index
   429  		tagID  int
   430  	}
   431  
   432  	var fields []field
   433  	var taggedFields []field
   434  
   435  	forEachStructField(typ, func(typ reflect.Type, index index, tag string) {
   436  		if typ.Size() != 0 { // skip struct{}
   437  			forEachStructTag(tag, func(tag structTag) bool {
   438  				if tag.MinVersion <= version && version <= tag.MaxVersion {
   439  					f := field{
   440  						encode: encodeFuncOf(typ, version, flexible, tag),
   441  						index:  index,
   442  						tagID:  tag.TagID,
   443  					}
   444  
   445  					if tag.TagID < -1 {
   446  						// Normal required field
   447  						fields = append(fields, f)
   448  					} else {
   449  						// Optional tagged field (flexible messages only)
   450  						taggedFields = append(taggedFields, f)
   451  					}
   452  					return false
   453  				}
   454  				return true
   455  			})
   456  		}
   457  	})
   458  
   459  	return func(e *encoder, v value) {
   460  		for i := range fields {
   461  			f := &fields[i]
   462  			f.encode(e, v.fieldByIndex(f.index))
   463  		}
   464  
   465  		if flexible {
   466  			// See https://cwiki.apache.org/confluence/display/KAFKA/KIP-482%3A+The+Kafka+Protocol+should+Support+Optional+Tagged+Fields
   467  			// for details of tag buffers in "flexible" messages.
   468  			e.writeUnsignedVarInt(uint64(len(taggedFields)))
   469  
   470  			for i := range taggedFields {
   471  				f := &taggedFields[i]
   472  				e.writeUnsignedVarInt(uint64(f.tagID))
   473  
   474  				buf := &bytes.Buffer{}
   475  				se := &encoder{writer: buf}
   476  				f.encode(se, v.fieldByIndex(f.index))
   477  				e.writeUnsignedVarInt(uint64(buf.Len()))
   478  				e.Write(buf.Bytes())
   479  			}
   480  		}
   481  	}
   482  }
   483  
   484  func arrayEncodeFuncOf(typ reflect.Type, version int16, flexible bool, tag structTag) encodeFunc {
   485  	elemType := typ.Elem()
   486  	elemFunc := encodeFuncOf(elemType, version, flexible, tag)
   487  	switch {
   488  	case flexible && tag.Nullable:
   489  		// In flexible messages, all arrays are compact
   490  		return func(e *encoder, v value) { e.encodeCompactNullArray(v, elemType, elemFunc) }
   491  	case flexible:
   492  		// In flexible messages, all arrays are compact
   493  		return func(e *encoder, v value) { e.encodeCompactArray(v, elemType, elemFunc) }
   494  	case tag.Nullable:
   495  		return func(e *encoder, v value) { e.encodeNullArray(v, elemType, elemFunc) }
   496  	default:
   497  		return func(e *encoder, v value) { e.encodeArray(v, elemType, elemFunc) }
   498  	}
   499  }
   500  
   501  func writerEncodeFuncOf(typ reflect.Type) encodeFunc {
   502  	typ = reflect.PtrTo(typ)
   503  	return func(e *encoder, v value) {
   504  		// Optimization to write directly into the buffer when the encoder
   505  		// does no need to compute a crc32 checksum.
   506  		w := io.Writer(e)
   507  		if e.table == nil {
   508  			w = e.writer
   509  		}
   510  		_, err := v.iface(typ).(io.WriterTo).WriteTo(w)
   511  		if err != nil {
   512  			e.err = err
   513  		}
   514  	}
   515  }
   516  
   517  func writeInt8(b []byte, i int8) {
   518  	b[0] = byte(i)
   519  }
   520  
   521  func writeInt16(b []byte, i int16) {
   522  	binary.BigEndian.PutUint16(b, uint16(i))
   523  }
   524  
   525  func writeInt32(b []byte, i int32) {
   526  	binary.BigEndian.PutUint32(b, uint32(i))
   527  }
   528  
   529  func writeInt64(b []byte, i int64) {
   530  	binary.BigEndian.PutUint64(b, uint64(i))
   531  }
   532  
   533  func Marshal(version int16, value interface{}) ([]byte, error) {
   534  	typ := typeOf(value)
   535  	cache, _ := marshalers.Load().(map[versionedType]encodeFunc)
   536  	key := versionedType{typ: typ, version: version}
   537  	encode := cache[key]
   538  
   539  	if encode == nil {
   540  		encode = encodeFuncOf(reflect.TypeOf(value), version, false, structTag{
   541  			MinVersion: -1,
   542  			MaxVersion: -1,
   543  			TagID:      -2,
   544  			Compact:    true,
   545  			Nullable:   true,
   546  		})
   547  
   548  		newCache := make(map[versionedType]encodeFunc, len(cache)+1)
   549  		newCache[key] = encode
   550  
   551  		for typ, fun := range cache {
   552  			newCache[typ] = fun
   553  		}
   554  
   555  		marshalers.Store(newCache)
   556  	}
   557  
   558  	e, _ := encoders.Get().(*encoder)
   559  	if e == nil {
   560  		e = &encoder{writer: new(bytes.Buffer)}
   561  	}
   562  
   563  	b, _ := e.writer.(*bytes.Buffer)
   564  	defer func() {
   565  		b.Reset()
   566  		e.Reset(b)
   567  		encoders.Put(e)
   568  	}()
   569  
   570  	encode(e, nonAddressableValueOf(value))
   571  
   572  	if e.err != nil {
   573  		return nil, e.err
   574  	}
   575  
   576  	buf := b.Bytes()
   577  	out := make([]byte, len(buf))
   578  	copy(out, buf)
   579  	return out, nil
   580  }
   581  
   582  type versionedType struct {
   583  	typ     _type
   584  	version int16
   585  }
   586  
   587  var (
   588  	encoders   sync.Pool    // *encoder
   589  	marshalers atomic.Value // map[versionedType]encodeFunc
   590  )