|
- package protocol
-
- import (
- "bytes"
- "encoding/binary"
- "fmt"
- "hash/crc32"
- "io"
- "math"
- "reflect"
- "sync"
- "sync/atomic"
- )
-
- type encoder struct {
- writer io.Writer
- err error
- table *crc32.Table
- crc32 uint32
- buffer [32]byte
- }
-
- type encoderChecksum struct {
- reader io.Reader
- encoder *encoder
- }
-
- func (e *encoderChecksum) Read(b []byte) (int, error) {
- n, err := e.reader.Read(b)
- if n > 0 {
- e.encoder.update(b[:n])
- }
- return n, err
- }
-
- func (e *encoder) Reset(w io.Writer) {
- e.writer = w
- e.err = nil
- e.table = nil
- e.crc32 = 0
- e.buffer = [32]byte{}
- }
-
- func (e *encoder) ReadFrom(r io.Reader) (int64, error) {
- if e.table != nil {
- r = &encoderChecksum{
- reader: r,
- encoder: e,
- }
- }
- return io.Copy(e.writer, r)
- }
-
- func (e *encoder) Write(b []byte) (int, error) {
- if e.err != nil {
- return 0, e.err
- }
- n, err := e.writer.Write(b)
- if n > 0 {
- e.update(b[:n])
- }
- if err != nil {
- e.err = err
- }
- return n, err
- }
-
- func (e *encoder) WriteByte(b byte) error {
- e.buffer[0] = b
- _, err := e.Write(e.buffer[:1])
- return err
- }
-
- func (e *encoder) WriteString(s string) (int, error) {
- // This implementation is an optimization to avoid the heap allocation that
- // would occur when converting the string to a []byte to call crc32.Update.
- //
- // Strings are rarely long in the kafka protocol, so the use of a 32 byte
- // buffer is a good comprise between keeping the encoder value small and
- // limiting the number of calls to Write.
- //
- // We introduced this optimization because memory profiles on the benchmarks
- // showed that most heap allocations were caused by this code path.
- n := 0
-
- for len(s) != 0 {
- c := copy(e.buffer[:], s)
- w, err := e.Write(e.buffer[:c])
- n += w
- if err != nil {
- return n, err
- }
- s = s[c:]
- }
-
- return n, nil
- }
-
- func (e *encoder) setCRC(table *crc32.Table) {
- e.table, e.crc32 = table, 0
- }
-
- func (e *encoder) update(b []byte) {
- if e.table != nil {
- e.crc32 = crc32.Update(e.crc32, e.table, b)
- }
- }
-
- func (e *encoder) encodeBool(v value) {
- b := int8(0)
- if v.bool() {
- b = 1
- }
- e.writeInt8(b)
- }
-
- func (e *encoder) encodeInt8(v value) {
- e.writeInt8(v.int8())
- }
-
- func (e *encoder) encodeInt16(v value) {
- e.writeInt16(v.int16())
- }
-
- func (e *encoder) encodeInt32(v value) {
- e.writeInt32(v.int32())
- }
-
- func (e *encoder) encodeInt64(v value) {
- e.writeInt64(v.int64())
- }
-
- func (e *encoder) encodeFloat64(v value) {
- e.writeFloat64(v.float64())
- }
-
- func (e *encoder) encodeString(v value) {
- e.writeString(v.string())
- }
-
- func (e *encoder) encodeCompactString(v value) {
- e.writeCompactString(v.string())
- }
-
- func (e *encoder) encodeNullString(v value) {
- e.writeNullString(v.string())
- }
-
- func (e *encoder) encodeCompactNullString(v value) {
- e.writeCompactNullString(v.string())
- }
-
- func (e *encoder) encodeBytes(v value) {
- e.writeBytes(v.bytes())
- }
-
- func (e *encoder) encodeCompactBytes(v value) {
- e.writeCompactBytes(v.bytes())
- }
-
- func (e *encoder) encodeNullBytes(v value) {
- e.writeNullBytes(v.bytes())
- }
-
- func (e *encoder) encodeCompactNullBytes(v value) {
- e.writeCompactNullBytes(v.bytes())
- }
-
- func (e *encoder) encodeArray(v value, elemType reflect.Type, encodeElem encodeFunc) {
- a := v.array(elemType)
- n := a.length()
- e.writeInt32(int32(n))
-
- for i := 0; i < n; i++ {
- encodeElem(e, a.index(i))
- }
- }
-
- func (e *encoder) encodeCompactArray(v value, elemType reflect.Type, encodeElem encodeFunc) {
- a := v.array(elemType)
- n := a.length()
- e.writeUnsignedVarInt(uint64(n + 1))
-
- for i := 0; i < n; i++ {
- encodeElem(e, a.index(i))
- }
- }
-
- func (e *encoder) encodeNullArray(v value, elemType reflect.Type, encodeElem encodeFunc) {
- a := v.array(elemType)
- if a.isNil() {
- e.writeInt32(-1)
- return
- }
-
- n := a.length()
- e.writeInt32(int32(n))
-
- for i := 0; i < n; i++ {
- encodeElem(e, a.index(i))
- }
- }
-
- func (e *encoder) encodeCompactNullArray(v value, elemType reflect.Type, encodeElem encodeFunc) {
- a := v.array(elemType)
- if a.isNil() {
- e.writeUnsignedVarInt(0)
- return
- }
-
- n := a.length()
- e.writeUnsignedVarInt(uint64(n + 1))
- for i := 0; i < n; i++ {
- encodeElem(e, a.index(i))
- }
- }
-
- func (e *encoder) writeInt8(i int8) {
- writeInt8(e.buffer[:1], i)
- e.Write(e.buffer[:1])
- }
-
- func (e *encoder) writeInt16(i int16) {
- writeInt16(e.buffer[:2], i)
- e.Write(e.buffer[:2])
- }
-
- func (e *encoder) writeInt32(i int32) {
- writeInt32(e.buffer[:4], i)
- e.Write(e.buffer[:4])
- }
-
- func (e *encoder) writeInt64(i int64) {
- writeInt64(e.buffer[:8], i)
- e.Write(e.buffer[:8])
- }
-
- func (e *encoder) writeFloat64(f float64) {
- writeFloat64(e.buffer[:8], f)
- e.Write(e.buffer[:8])
- }
-
- func (e *encoder) writeString(s string) {
- e.writeInt16(int16(len(s)))
- e.WriteString(s)
- }
-
- func (e *encoder) writeVarString(s string) {
- e.writeVarInt(int64(len(s)))
- e.WriteString(s)
- }
-
- func (e *encoder) writeCompactString(s string) {
- e.writeUnsignedVarInt(uint64(len(s)) + 1)
- e.WriteString(s)
- }
-
- func (e *encoder) writeNullString(s string) {
- if s == "" {
- e.writeInt16(-1)
- } else {
- e.writeInt16(int16(len(s)))
- e.WriteString(s)
- }
- }
-
- func (e *encoder) writeCompactNullString(s string) {
- if s == "" {
- e.writeUnsignedVarInt(0)
- } else {
- e.writeUnsignedVarInt(uint64(len(s)) + 1)
- e.WriteString(s)
- }
- }
-
- func (e *encoder) writeBytes(b []byte) {
- e.writeInt32(int32(len(b)))
- e.Write(b)
- }
-
- func (e *encoder) writeCompactBytes(b []byte) {
- e.writeUnsignedVarInt(uint64(len(b)) + 1)
- e.Write(b)
- }
-
- func (e *encoder) writeNullBytes(b []byte) {
- if b == nil {
- e.writeInt32(-1)
- } else {
- e.writeInt32(int32(len(b)))
- e.Write(b)
- }
- }
-
- func (e *encoder) writeVarNullBytes(b []byte) {
- if b == nil {
- e.writeVarInt(-1)
- } else {
- e.writeVarInt(int64(len(b)))
- e.Write(b)
- }
- }
-
- func (e *encoder) writeCompactNullBytes(b []byte) {
- if b == nil {
- e.writeUnsignedVarInt(0)
- } else {
- e.writeUnsignedVarInt(uint64(len(b)) + 1)
- e.Write(b)
- }
- }
-
- func (e *encoder) writeNullBytesFrom(b Bytes) error {
- if b == nil {
- e.writeInt32(-1)
- return nil
- } else {
- size := int64(b.Len())
- e.writeInt32(int32(size))
- n, err := io.Copy(e, b)
- if err == nil && n != size {
- 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)
- }
- return err
- }
- }
-
- func (e *encoder) writeVarNullBytesFrom(b Bytes) error {
- if b == nil {
- e.writeVarInt(-1)
- return nil
- } else {
- size := int64(b.Len())
- e.writeVarInt(size)
- n, err := io.Copy(e, b)
- if err == nil && n != size {
- 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)
- }
- return err
- }
- }
-
- func (e *encoder) writeVarInt(i int64) {
- e.writeUnsignedVarInt(uint64((i << 1) ^ (i >> 63)))
- }
-
- func (e *encoder) writeUnsignedVarInt(i uint64) {
- b := e.buffer[:]
- n := 0
-
- for i >= 0x80 && n < len(b) {
- b[n] = byte(i) | 0x80
- i >>= 7
- n++
- }
-
- if n < len(b) {
- b[n] = byte(i)
- n++
- }
-
- e.Write(b[:n])
- }
-
- type encodeFunc func(*encoder, value)
-
- var (
- _ io.ReaderFrom = (*encoder)(nil)
- _ io.Writer = (*encoder)(nil)
- _ io.ByteWriter = (*encoder)(nil)
- _ io.StringWriter = (*encoder)(nil)
-
- writerTo = reflect.TypeOf((*io.WriterTo)(nil)).Elem()
- )
-
- func encodeFuncOf(typ reflect.Type, version int16, flexible bool, tag structTag) encodeFunc {
- if reflect.PtrTo(typ).Implements(writerTo) {
- return writerEncodeFuncOf(typ)
- }
- switch typ.Kind() {
- case reflect.Bool:
- return (*encoder).encodeBool
- case reflect.Int8:
- return (*encoder).encodeInt8
- case reflect.Int16:
- return (*encoder).encodeInt16
- case reflect.Int32:
- return (*encoder).encodeInt32
- case reflect.Int64:
- return (*encoder).encodeInt64
- case reflect.Float64:
- return (*encoder).encodeFloat64
- case reflect.String:
- return stringEncodeFuncOf(flexible, tag)
- case reflect.Struct:
- return structEncodeFuncOf(typ, version, flexible)
- case reflect.Slice:
- if typ.Elem().Kind() == reflect.Uint8 { // []byte
- return bytesEncodeFuncOf(flexible, tag)
- }
- return arrayEncodeFuncOf(typ, version, flexible, tag)
- default:
- panic("unsupported type: " + typ.String())
- }
- }
-
- func stringEncodeFuncOf(flexible bool, tag structTag) encodeFunc {
- switch {
- case flexible && tag.Nullable:
- // In flexible messages, all strings are compact
- return (*encoder).encodeCompactNullString
- case flexible:
- // In flexible messages, all strings are compact
- return (*encoder).encodeCompactString
- case tag.Nullable:
- return (*encoder).encodeNullString
- default:
- return (*encoder).encodeString
- }
- }
-
- func bytesEncodeFuncOf(flexible bool, tag structTag) encodeFunc {
- switch {
- case flexible && tag.Nullable:
- // In flexible messages, all arrays are compact
- return (*encoder).encodeCompactNullBytes
- case flexible:
- // In flexible messages, all arrays are compact
- return (*encoder).encodeCompactBytes
- case tag.Nullable:
- return (*encoder).encodeNullBytes
- default:
- return (*encoder).encodeBytes
- }
- }
-
- func structEncodeFuncOf(typ reflect.Type, version int16, flexible bool) encodeFunc {
- type field struct {
- encode encodeFunc
- index index
- tagID int
- }
-
- var fields []field
- var taggedFields []field
-
- forEachStructField(typ, func(typ reflect.Type, index index, tag string) {
- if typ.Size() != 0 { // skip struct{}
- forEachStructTag(tag, func(tag structTag) bool {
- if tag.MinVersion <= version && version <= tag.MaxVersion {
- f := field{
- encode: encodeFuncOf(typ, version, flexible, tag),
- index: index,
- tagID: tag.TagID,
- }
-
- if tag.TagID < -1 {
- // Normal required field
- fields = append(fields, f)
- } else {
- // Optional tagged field (flexible messages only)
- taggedFields = append(taggedFields, f)
- }
- return false
- }
- return true
- })
- }
- })
-
- return func(e *encoder, v value) {
- for i := range fields {
- f := &fields[i]
- f.encode(e, v.fieldByIndex(f.index))
- }
-
- if flexible {
- // See https://cwiki.apache.org/confluence/display/KAFKA/KIP-482%3A+The+Kafka+Protocol+should+Support+Optional+Tagged+Fields
- // for details of tag buffers in "flexible" messages.
- e.writeUnsignedVarInt(uint64(len(taggedFields)))
-
- for i := range taggedFields {
- f := &taggedFields[i]
- e.writeUnsignedVarInt(uint64(f.tagID))
-
- buf := &bytes.Buffer{}
- se := &encoder{writer: buf}
- f.encode(se, v.fieldByIndex(f.index))
- e.writeUnsignedVarInt(uint64(buf.Len()))
- e.Write(buf.Bytes())
- }
- }
- }
- }
-
- func arrayEncodeFuncOf(typ reflect.Type, version int16, flexible bool, tag structTag) encodeFunc {
- elemType := typ.Elem()
- elemFunc := encodeFuncOf(elemType, version, flexible, tag)
- switch {
- case flexible && tag.Nullable:
- // In flexible messages, all arrays are compact
- return func(e *encoder, v value) { e.encodeCompactNullArray(v, elemType, elemFunc) }
- case flexible:
- // In flexible messages, all arrays are compact
- return func(e *encoder, v value) { e.encodeCompactArray(v, elemType, elemFunc) }
- case tag.Nullable:
- return func(e *encoder, v value) { e.encodeNullArray(v, elemType, elemFunc) }
- default:
- return func(e *encoder, v value) { e.encodeArray(v, elemType, elemFunc) }
- }
- }
-
- func writerEncodeFuncOf(typ reflect.Type) encodeFunc {
- typ = reflect.PtrTo(typ)
- return func(e *encoder, v value) {
- // Optimization to write directly into the buffer when the encoder
- // does no need to compute a crc32 checksum.
- w := io.Writer(e)
- if e.table == nil {
- w = e.writer
- }
- _, err := v.iface(typ).(io.WriterTo).WriteTo(w)
- if err != nil {
- e.err = err
- }
- }
- }
-
- func writeInt8(b []byte, i int8) {
- b[0] = byte(i)
- }
-
- func writeInt16(b []byte, i int16) {
- binary.BigEndian.PutUint16(b, uint16(i))
- }
-
- func writeInt32(b []byte, i int32) {
- binary.BigEndian.PutUint32(b, uint32(i))
- }
-
- func writeInt64(b []byte, i int64) {
- binary.BigEndian.PutUint64(b, uint64(i))
- }
-
- func writeFloat64(b []byte, f float64) {
- binary.BigEndian.PutUint64(b, math.Float64bits(f))
- }
-
- func Marshal(version int16, value interface{}) ([]byte, error) {
- typ := typeOf(value)
- cache, _ := marshalers.Load().(map[versionedType]encodeFunc)
- key := versionedType{typ: typ, version: version}
- encode := cache[key]
-
- if encode == nil {
- encode = encodeFuncOf(reflect.TypeOf(value), version, false, structTag{
- MinVersion: -1,
- MaxVersion: -1,
- TagID: -2,
- Compact: true,
- Nullable: true,
- })
-
- newCache := make(map[versionedType]encodeFunc, len(cache)+1)
- newCache[key] = encode
-
- for typ, fun := range cache {
- newCache[typ] = fun
- }
-
- marshalers.Store(newCache)
- }
-
- e, _ := encoders.Get().(*encoder)
- if e == nil {
- e = &encoder{writer: new(bytes.Buffer)}
- }
-
- b, _ := e.writer.(*bytes.Buffer)
- defer func() {
- b.Reset()
- e.Reset(b)
- encoders.Put(e)
- }()
-
- encode(e, nonAddressableValueOf(value))
-
- if e.err != nil {
- return nil, e.err
- }
-
- buf := b.Bytes()
- out := make([]byte, len(buf))
- copy(out, buf)
- return out, nil
- }
-
- type versionedType struct {
- typ _type
- version int16
- }
-
- var (
- encoders sync.Pool // *encoder
- marshalers atomic.Value // map[versionedType]encodeFunc
- )
|