github.com/hack0072008/kafka-go@v1.0.1/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) encodeVarString(v value) { 137 e.writeVarString(v.string()) 138 } 139 140 func (e *encoder) encodeCompactString(v value) { 141 e.writeCompactString(v.string()) 142 } 143 144 func (e *encoder) encodeNullString(v value) { 145 e.writeNullString(v.string()) 146 } 147 148 func (e *encoder) encodeVarNullString(v value) { 149 e.writeVarNullString(v.string()) 150 } 151 152 func (e *encoder) encodeCompactNullString(v value) { 153 e.writeCompactNullString(v.string()) 154 } 155 156 func (e *encoder) encodeBytes(v value) { 157 e.writeBytes(v.bytes()) 158 } 159 160 func (e *encoder) encodeVarBytes(v value) { 161 e.writeVarBytes(v.bytes()) 162 } 163 164 func (e *encoder) encodeCompactBytes(v value) { 165 e.writeCompactBytes(v.bytes()) 166 } 167 168 func (e *encoder) encodeNullBytes(v value) { 169 e.writeNullBytes(v.bytes()) 170 } 171 172 func (e *encoder) encodeVarNullBytes(v value) { 173 e.writeVarNullBytes(v.bytes()) 174 } 175 176 func (e *encoder) encodeCompactNullBytes(v value) { 177 e.writeCompactNullBytes(v.bytes()) 178 } 179 180 func (e *encoder) encodeArray(v value, elemType reflect.Type, encodeElem encodeFunc) { 181 a := v.array(elemType) 182 n := a.length() 183 e.writeInt32(int32(n)) 184 185 for i := 0; i < n; i++ { 186 encodeElem(e, a.index(i)) 187 } 188 } 189 190 func (e *encoder) encodeCompactArray(v value, elemType reflect.Type, encodeElem encodeFunc) { 191 a := v.array(elemType) 192 n := a.length() 193 e.writeUnsignedVarInt(uint64(n + 1)) 194 195 for i := 0; i < n; i++ { 196 encodeElem(e, a.index(i)) 197 } 198 } 199 200 func (e *encoder) encodeNullArray(v value, elemType reflect.Type, encodeElem encodeFunc) { 201 a := v.array(elemType) 202 if a.isNil() { 203 e.writeInt32(-1) 204 return 205 } 206 207 n := a.length() 208 e.writeInt32(int32(n)) 209 210 for i := 0; i < n; i++ { 211 encodeElem(e, a.index(i)) 212 } 213 } 214 215 func (e *encoder) encodeCompactNullArray(v value, elemType reflect.Type, encodeElem encodeFunc) { 216 a := v.array(elemType) 217 if a.isNil() { 218 e.writeUnsignedVarInt(0) 219 return 220 } 221 222 n := a.length() 223 e.writeUnsignedVarInt(uint64(n + 1)) 224 for i := 0; i < n; i++ { 225 encodeElem(e, a.index(i)) 226 } 227 } 228 229 func (e *encoder) writeInt8(i int8) { 230 writeInt8(e.buffer[:1], i) 231 e.Write(e.buffer[:1]) 232 } 233 234 func (e *encoder) writeInt16(i int16) { 235 writeInt16(e.buffer[:2], i) 236 e.Write(e.buffer[:2]) 237 } 238 239 func (e *encoder) writeInt32(i int32) { 240 writeInt32(e.buffer[:4], i) 241 e.Write(e.buffer[:4]) 242 } 243 244 func (e *encoder) writeInt64(i int64) { 245 writeInt64(e.buffer[:8], i) 246 e.Write(e.buffer[:8]) 247 } 248 249 func (e *encoder) writeString(s string) { 250 e.writeInt16(int16(len(s))) 251 e.WriteString(s) 252 } 253 254 func (e *encoder) writeVarString(s string) { 255 e.writeVarInt(int64(len(s))) 256 e.WriteString(s) 257 } 258 259 func (e *encoder) writeCompactString(s string) { 260 e.writeUnsignedVarInt(uint64(len(s)) + 1) 261 e.WriteString(s) 262 } 263 264 func (e *encoder) writeNullString(s string) { 265 if s == "" { 266 e.writeInt16(-1) 267 } else { 268 e.writeInt16(int16(len(s))) 269 e.WriteString(s) 270 } 271 } 272 273 func (e *encoder) writeVarNullString(s string) { 274 if s == "" { 275 e.writeVarInt(-1) 276 } else { 277 e.writeVarInt(int64(len(s))) 278 e.WriteString(s) 279 } 280 } 281 282 func (e *encoder) writeCompactNullString(s string) { 283 if s == "" { 284 e.writeUnsignedVarInt(0) 285 } else { 286 e.writeUnsignedVarInt(uint64(len(s)) + 1) 287 e.WriteString(s) 288 } 289 } 290 291 func (e *encoder) writeBytes(b []byte) { 292 e.writeInt32(int32(len(b))) 293 e.Write(b) 294 } 295 296 func (e *encoder) writeVarBytes(b []byte) { 297 e.writeVarInt(int64(len(b))) 298 e.Write(b) 299 } 300 301 func (e *encoder) writeCompactBytes(b []byte) { 302 e.writeUnsignedVarInt(uint64(len(b)) + 1) 303 e.Write(b) 304 } 305 306 func (e *encoder) writeNullBytes(b []byte) { 307 if b == nil { 308 e.writeInt32(-1) 309 } else { 310 e.writeInt32(int32(len(b))) 311 e.Write(b) 312 } 313 } 314 315 func (e *encoder) writeVarNullBytes(b []byte) { 316 if b == nil { 317 e.writeVarInt(-1) 318 } else { 319 e.writeVarInt(int64(len(b))) 320 e.Write(b) 321 } 322 } 323 324 func (e *encoder) writeCompactNullBytes(b []byte) { 325 if b == nil { 326 e.writeUnsignedVarInt(0) 327 } else { 328 e.writeUnsignedVarInt(uint64(len(b)) + 1) 329 e.Write(b) 330 } 331 } 332 333 func (e *encoder) writeBytesFrom(b Bytes) error { 334 size := int64(b.Len()) 335 e.writeInt32(int32(size)) 336 n, err := io.Copy(e, b) 337 if err == nil && n != size { 338 err = fmt.Errorf("size of bytes does not match the number of bytes that were written (size=%d, written=%d): %w", size, n, io.ErrUnexpectedEOF) 339 } 340 return err 341 } 342 343 func (e *encoder) writeNullBytesFrom(b Bytes) error { 344 if b == nil { 345 e.writeInt32(-1) 346 return nil 347 } else { 348 size := int64(b.Len()) 349 e.writeInt32(int32(size)) 350 n, err := io.Copy(e, b) 351 if err == nil && n != size { 352 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) 353 } 354 return err 355 } 356 } 357 358 func (e *encoder) writeVarNullBytesFrom(b Bytes) error { 359 if b == nil { 360 e.writeVarInt(-1) 361 return nil 362 } else { 363 size := int64(b.Len()) 364 e.writeVarInt(size) 365 n, err := io.Copy(e, b) 366 if err == nil && n != size { 367 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) 368 } 369 return err 370 } 371 } 372 373 func (e *encoder) writeCompactNullBytesFrom(b Bytes) error { 374 if b == nil { 375 e.writeUnsignedVarInt(0) 376 return nil 377 } else { 378 size := int64(b.Len()) 379 e.writeUnsignedVarInt(uint64(size + 1)) 380 n, err := io.Copy(e, b) 381 if err == nil && n != size { 382 err = fmt.Errorf("size of compact nullable bytes does not match the number of bytes that were written (size=%d, written=%d): %w", size, n, io.ErrUnexpectedEOF) 383 } 384 return err 385 } 386 } 387 388 func (e *encoder) writeVarInt(i int64) { 389 e.writeUnsignedVarInt(uint64((i << 1) ^ (i >> 63))) 390 } 391 392 func (e *encoder) writeUnsignedVarInt(i uint64) { 393 b := e.buffer[:] 394 n := 0 395 396 for i >= 0x80 && n < len(b) { 397 b[n] = byte(i) | 0x80 398 i >>= 7 399 n++ 400 } 401 402 if n < len(b) { 403 b[n] = byte(i) 404 n++ 405 } 406 407 e.Write(b[:n]) 408 } 409 410 type encodeFunc func(*encoder, value) 411 412 var ( 413 _ io.ReaderFrom = (*encoder)(nil) 414 _ io.Writer = (*encoder)(nil) 415 _ io.ByteWriter = (*encoder)(nil) 416 _ io.StringWriter = (*encoder)(nil) 417 418 writerTo = reflect.TypeOf((*io.WriterTo)(nil)).Elem() 419 ) 420 421 func encodeFuncOf(typ reflect.Type, version int16, flexible bool, tag structTag) encodeFunc { 422 if reflect.PtrTo(typ).Implements(writerTo) { 423 return writerEncodeFuncOf(typ) 424 } 425 switch typ.Kind() { 426 case reflect.Bool: 427 return (*encoder).encodeBool 428 case reflect.Int8: 429 return (*encoder).encodeInt8 430 case reflect.Int16: 431 return (*encoder).encodeInt16 432 case reflect.Int32: 433 return (*encoder).encodeInt32 434 case reflect.Int64: 435 return (*encoder).encodeInt64 436 case reflect.String: 437 return stringEncodeFuncOf(flexible, tag) 438 case reflect.Struct: 439 return structEncodeFuncOf(typ, version, flexible) 440 case reflect.Slice: 441 if typ.Elem().Kind() == reflect.Uint8 { // []byte 442 return bytesEncodeFuncOf(flexible, tag) 443 } 444 return arrayEncodeFuncOf(typ, version, flexible, tag) 445 default: 446 panic("unsupported type: " + typ.String()) 447 } 448 } 449 450 func stringEncodeFuncOf(flexible bool, tag structTag) encodeFunc { 451 switch { 452 case flexible && tag.Nullable: 453 // In flexible messages, all strings are compact 454 return (*encoder).encodeCompactNullString 455 case flexible: 456 // In flexible messages, all strings are compact 457 return (*encoder).encodeCompactString 458 case tag.Nullable: 459 return (*encoder).encodeNullString 460 default: 461 return (*encoder).encodeString 462 } 463 } 464 465 func bytesEncodeFuncOf(flexible bool, tag structTag) encodeFunc { 466 switch { 467 case flexible && tag.Nullable: 468 // In flexible messages, all arrays are compact 469 return (*encoder).encodeCompactNullBytes 470 case flexible: 471 // In flexible messages, all arrays are compact 472 return (*encoder).encodeCompactBytes 473 case tag.Nullable: 474 return (*encoder).encodeNullBytes 475 default: 476 return (*encoder).encodeBytes 477 } 478 } 479 480 func structEncodeFuncOf(typ reflect.Type, version int16, flexible bool) encodeFunc { 481 type field struct { 482 encode encodeFunc 483 index index 484 tagID int 485 } 486 487 var fields []field 488 var taggedFields []field 489 490 forEachStructField(typ, func(typ reflect.Type, index index, tag string) { 491 if typ.Size() != 0 { // skip struct{} 492 forEachStructTag(tag, func(tag structTag) bool { 493 if tag.MinVersion <= version && version <= tag.MaxVersion { 494 f := field{ 495 encode: encodeFuncOf(typ, version, flexible, tag), 496 index: index, 497 tagID: tag.TagID, 498 } 499 500 if tag.TagID < -1 { 501 // Normal required field 502 fields = append(fields, f) 503 } else { 504 // Optional tagged field (flexible messages only) 505 taggedFields = append(taggedFields, f) 506 } 507 return false 508 } 509 return true 510 }) 511 } 512 }) 513 514 return func(e *encoder, v value) { 515 for i := range fields { 516 f := &fields[i] 517 f.encode(e, v.fieldByIndex(f.index)) 518 } 519 520 if flexible { 521 // See https://cwiki.apache.org/confluence/display/KAFKA/KIP-482%3A+The+Kafka+Protocol+should+Support+Optional+Tagged+Fields 522 // for details of tag buffers in "flexible" messages. 523 e.writeUnsignedVarInt(uint64(len(taggedFields))) 524 525 for i := range taggedFields { 526 f := &taggedFields[i] 527 e.writeUnsignedVarInt(uint64(f.tagID)) 528 529 buf := &bytes.Buffer{} 530 se := &encoder{writer: buf} 531 f.encode(se, v.fieldByIndex(f.index)) 532 e.writeUnsignedVarInt(uint64(buf.Len())) 533 e.Write(buf.Bytes()) 534 } 535 } 536 } 537 } 538 539 func arrayEncodeFuncOf(typ reflect.Type, version int16, flexible bool, tag structTag) encodeFunc { 540 elemType := typ.Elem() 541 elemFunc := encodeFuncOf(elemType, version, flexible, tag) 542 switch { 543 case flexible && tag.Nullable: 544 // In flexible messages, all arrays are compact 545 return func(e *encoder, v value) { e.encodeCompactNullArray(v, elemType, elemFunc) } 546 case flexible: 547 // In flexible messages, all arrays are compact 548 return func(e *encoder, v value) { e.encodeCompactArray(v, elemType, elemFunc) } 549 case tag.Nullable: 550 return func(e *encoder, v value) { e.encodeNullArray(v, elemType, elemFunc) } 551 default: 552 return func(e *encoder, v value) { e.encodeArray(v, elemType, elemFunc) } 553 } 554 } 555 556 func writerEncodeFuncOf(typ reflect.Type) encodeFunc { 557 typ = reflect.PtrTo(typ) 558 return func(e *encoder, v value) { 559 // Optimization to write directly into the buffer when the encoder 560 // does no need to compute a crc32 checksum. 561 w := io.Writer(e) 562 if e.table == nil { 563 w = e.writer 564 } 565 _, err := v.iface(typ).(io.WriterTo).WriteTo(w) 566 if err != nil { 567 e.err = err 568 } 569 } 570 } 571 572 func writeInt8(b []byte, i int8) { 573 b[0] = byte(i) 574 } 575 576 func writeInt16(b []byte, i int16) { 577 binary.BigEndian.PutUint16(b, uint16(i)) 578 } 579 580 func writeInt32(b []byte, i int32) { 581 binary.BigEndian.PutUint32(b, uint32(i)) 582 } 583 584 func writeInt64(b []byte, i int64) { 585 binary.BigEndian.PutUint64(b, uint64(i)) 586 } 587 588 func Marshal(version int16, value interface{}) ([]byte, error) { 589 typ := typeOf(value) 590 cache, _ := marshalers.Load().(map[versionedType]encodeFunc) 591 key := versionedType{typ: typ, version: version} 592 encode := cache[key] 593 594 if encode == nil { 595 encode = encodeFuncOf(reflect.TypeOf(value), version, false, structTag{ 596 MinVersion: -1, 597 MaxVersion: -1, 598 TagID: -2, 599 Compact: true, 600 Nullable: true, 601 }) 602 603 newCache := make(map[versionedType]encodeFunc, len(cache)+1) 604 newCache[key] = encode 605 606 for typ, fun := range cache { 607 newCache[typ] = fun 608 } 609 610 marshalers.Store(newCache) 611 } 612 613 e, _ := encoders.Get().(*encoder) 614 if e == nil { 615 e = &encoder{writer: new(bytes.Buffer)} 616 } 617 618 b, _ := e.writer.(*bytes.Buffer) 619 defer func() { 620 b.Reset() 621 e.Reset(b) 622 encoders.Put(e) 623 }() 624 625 encode(e, nonAddressableValueOf(value)) 626 627 if e.err != nil { 628 return nil, e.err 629 } 630 631 buf := b.Bytes() 632 out := make([]byte, len(buf)) 633 copy(out, buf) 634 return out, nil 635 } 636 637 type versionedType struct { 638 typ _type 639 version int16 640 } 641 642 var ( 643 encoders sync.Pool // *encoder 644 marshalers atomic.Value // map[versionedType]encodeFunc 645 )