github.com/hack0072008/kafka-go@v1.0.1/protocol/buffer.go (about) 1 package protocol 2 3 import ( 4 "bytes" 5 "fmt" 6 "io" 7 "math" 8 "sync" 9 "sync/atomic" 10 ) 11 12 // Bytes is an interface implemented by types that represent immutable 13 // sequences of bytes. 14 // 15 // Bytes values are used to abstract the location where record keys and 16 // values are read from (e.g. in-memory buffers, network sockets, files). 17 // 18 // The Close method should be called to release resources held by the object 19 // when the program is done with it. 20 // 21 // Bytes values are generally not safe to use concurrently from multiple 22 // goroutines. 23 type Bytes interface { 24 io.ReadCloser 25 // Returns the number of bytes remaining to be read from the payload. 26 Len() int 27 } 28 29 // NewBytes constructs a Bytes value from b. 30 // 31 // The returned value references b, it does not make a copy of the backing 32 // array. 33 // 34 // If b is nil, nil is returned to represent a null BYTES value in the kafka 35 // protocol. 36 func NewBytes(b []byte) Bytes { 37 if b == nil { 38 return nil 39 } 40 r := new(bytesReader) 41 r.Reset(b) 42 return r 43 } 44 45 // ReadAll is similar to ioutil.ReadAll, but it takes advantage of knowing the 46 // length of b to minimize the memory footprint. 47 // 48 // The function returns a nil slice if b is nil. 49 func ReadAll(b Bytes) ([]byte, error) { 50 if b == nil { 51 return nil, nil 52 } 53 s := make([]byte, b.Len()) 54 _, err := io.ReadFull(b, s) 55 return s, err 56 } 57 58 type bytesReader struct{ bytes.Reader } 59 60 func (*bytesReader) Close() error { return nil } 61 62 type refCount uintptr 63 64 func (rc *refCount) ref() { atomic.AddUintptr((*uintptr)(rc), 1) } 65 66 func (rc *refCount) unref(onZero func()) { 67 if atomic.AddUintptr((*uintptr)(rc), ^uintptr(0)) == 0 { 68 onZero() 69 } 70 } 71 72 const ( 73 // Size of the memory buffer for a single page. We use a farily 74 // large size here (64 KiB) because batches exchanged with kafka 75 // tend to be multiple kilobytes in size, sometimes hundreds. 76 // Using large pages amortizes the overhead of the page metadata 77 // and algorithms to manage the pages. 78 pageSize = 65536 79 ) 80 81 type page struct { 82 refc refCount 83 offset int64 84 length int 85 buffer *[pageSize]byte 86 } 87 88 func newPage(offset int64) *page { 89 p, _ := pagePool.Get().(*page) 90 if p != nil { 91 p.offset = offset 92 p.length = 0 93 p.ref() 94 } else { 95 p = &page{ 96 refc: 1, 97 offset: offset, 98 buffer: &[pageSize]byte{}, 99 } 100 } 101 return p 102 } 103 104 func (p *page) ref() { p.refc.ref() } 105 106 func (p *page) unref() { p.refc.unref(func() { pagePool.Put(p) }) } 107 108 func (p *page) slice(begin, end int64) []byte { 109 i, j := begin-p.offset, end-p.offset 110 111 if i < 0 { 112 i = 0 113 } else if i > pageSize { 114 i = pageSize 115 } 116 117 if j < 0 { 118 j = 0 119 } else if j > pageSize { 120 j = pageSize 121 } 122 123 if i < j { 124 return p.buffer[i:j] 125 } 126 127 return nil 128 } 129 130 func (p *page) Cap() int { return pageSize } 131 132 func (p *page) Len() int { return p.length } 133 134 func (p *page) Size() int64 { return int64(p.length) } 135 136 func (p *page) Truncate(n int) { 137 if n < p.length { 138 p.length = n 139 } 140 } 141 142 func (p *page) ReadAt(b []byte, off int64) (int, error) { 143 if off -= p.offset; off < 0 || off > pageSize { 144 panic("offset out of range") 145 } 146 if off > int64(p.length) { 147 return 0, nil 148 } 149 return copy(b, p.buffer[off:p.length]), nil 150 } 151 152 func (p *page) ReadFrom(r io.Reader) (int64, error) { 153 n, err := io.ReadFull(r, p.buffer[p.length:]) 154 if err == io.EOF || err == io.ErrUnexpectedEOF { 155 err = nil 156 } 157 p.length += n 158 return int64(n), err 159 } 160 161 func (p *page) WriteAt(b []byte, off int64) (int, error) { 162 if off -= p.offset; off < 0 || off > pageSize { 163 panic("offset out of range") 164 } 165 n := copy(p.buffer[off:], b) 166 if end := int(off) + n; end > p.length { 167 p.length = end 168 } 169 return n, nil 170 } 171 172 func (p *page) Write(b []byte) (int, error) { 173 return p.WriteAt(b, p.offset+int64(p.length)) 174 } 175 176 var ( 177 _ io.ReaderAt = (*page)(nil) 178 _ io.ReaderFrom = (*page)(nil) 179 _ io.Writer = (*page)(nil) 180 _ io.WriterAt = (*page)(nil) 181 ) 182 183 type pageBuffer struct { 184 refc refCount 185 pages contiguousPages 186 length int 187 cursor int 188 } 189 190 func newPageBuffer() *pageBuffer { 191 b, _ := pageBufferPool.Get().(*pageBuffer) 192 if b != nil { 193 b.cursor = 0 194 b.refc.ref() 195 } else { 196 b = &pageBuffer{ 197 refc: 1, 198 pages: make(contiguousPages, 0, 16), 199 } 200 } 201 return b 202 } 203 204 func (pb *pageBuffer) refTo(ref *pageRef, begin, end int64) { 205 length := end - begin 206 207 if length > math.MaxUint32 { 208 panic("reference to contiguous buffer pages exceeds the maximum size of 4 GB") 209 } 210 211 ref.pages = append(ref.buffer[:0], pb.pages.slice(begin, end)...) 212 ref.pages.ref() 213 ref.offset = begin 214 ref.length = uint32(length) 215 } 216 217 func (pb *pageBuffer) ref(begin, end int64) *pageRef { 218 ref := new(pageRef) 219 pb.refTo(ref, begin, end) 220 return ref 221 } 222 223 func (pb *pageBuffer) unref() { 224 pb.refc.unref(func() { 225 pb.pages.unref() 226 pb.pages.clear() 227 pb.pages = pb.pages[:0] 228 pb.length = 0 229 pageBufferPool.Put(pb) 230 }) 231 } 232 233 func (pb *pageBuffer) newPage() *page { 234 return newPage(int64(pb.length)) 235 } 236 237 func (pb *pageBuffer) Close() error { 238 return nil 239 } 240 241 func (pb *pageBuffer) Len() int { 242 return pb.length - pb.cursor 243 } 244 245 func (pb *pageBuffer) Size() int64 { 246 return int64(pb.length) 247 } 248 249 func (pb *pageBuffer) Discard(n int) (int, error) { 250 remain := pb.length - pb.cursor 251 if remain < n { 252 n = remain 253 } 254 pb.cursor += n 255 return n, nil 256 } 257 258 func (pb *pageBuffer) Truncate(n int) { 259 if n < pb.length { 260 pb.length = n 261 262 if n < pb.cursor { 263 pb.cursor = n 264 } 265 266 for i := range pb.pages { 267 if p := pb.pages[i]; p.length <= n { 268 n -= p.length 269 } else { 270 if n > 0 { 271 pb.pages[i].Truncate(n) 272 i++ 273 } 274 pb.pages[i:].unref() 275 pb.pages[i:].clear() 276 pb.pages = pb.pages[:i] 277 break 278 } 279 } 280 } 281 } 282 283 func (pb *pageBuffer) Seek(offset int64, whence int) (int64, error) { 284 c, err := seek(int64(pb.cursor), int64(pb.length), offset, whence) 285 if err != nil { 286 return -1, err 287 } 288 pb.cursor = int(c) 289 return c, nil 290 } 291 292 func (pb *pageBuffer) ReadByte() (byte, error) { 293 b := [1]byte{} 294 _, err := pb.Read(b[:]) 295 return b[0], err 296 } 297 298 func (pb *pageBuffer) Read(b []byte) (int, error) { 299 if pb.cursor >= pb.length { 300 return 0, io.EOF 301 } 302 n, err := pb.ReadAt(b, int64(pb.cursor)) 303 pb.cursor += n 304 return n, err 305 } 306 307 func (pb *pageBuffer) ReadAt(b []byte, off int64) (int, error) { 308 return pb.pages.ReadAt(b, off) 309 } 310 311 func (pb *pageBuffer) ReadFrom(r io.Reader) (int64, error) { 312 if len(pb.pages) == 0 { 313 pb.pages = append(pb.pages, pb.newPage()) 314 } 315 316 rn := int64(0) 317 318 for { 319 tail := pb.pages[len(pb.pages)-1] 320 free := tail.Cap() - tail.Len() 321 322 if free == 0 { 323 tail = pb.newPage() 324 free = pageSize 325 pb.pages = append(pb.pages, tail) 326 } 327 328 n, err := tail.ReadFrom(r) 329 pb.length += int(n) 330 rn += n 331 if n < int64(free) { 332 return rn, err 333 } 334 } 335 } 336 337 func (pb *pageBuffer) WriteString(s string) (int, error) { 338 return pb.Write([]byte(s)) 339 } 340 341 func (pb *pageBuffer) Write(b []byte) (int, error) { 342 wn := len(b) 343 if wn == 0 { 344 return 0, nil 345 } 346 347 if len(pb.pages) == 0 { 348 pb.pages = append(pb.pages, pb.newPage()) 349 } 350 351 for len(b) != 0 { 352 tail := pb.pages[len(pb.pages)-1] 353 free := tail.Cap() - tail.Len() 354 355 if len(b) <= free { 356 tail.Write(b) 357 pb.length += len(b) 358 break 359 } 360 361 tail.Write(b[:free]) 362 b = b[free:] 363 364 pb.length += free 365 pb.pages = append(pb.pages, pb.newPage()) 366 } 367 368 return wn, nil 369 } 370 371 func (pb *pageBuffer) WriteAt(b []byte, off int64) (int, error) { 372 n, err := pb.pages.WriteAt(b, off) 373 if err != nil { 374 return n, err 375 } 376 if n < len(b) { 377 pb.Write(b[n:]) 378 } 379 return len(b), nil 380 } 381 382 func (pb *pageBuffer) WriteTo(w io.Writer) (int64, error) { 383 var wn int 384 var err error 385 pb.pages.scan(int64(pb.cursor), int64(pb.length), func(b []byte) bool { 386 var n int 387 n, err = w.Write(b) 388 wn += n 389 return err == nil 390 }) 391 pb.cursor += wn 392 return int64(wn), err 393 } 394 395 var ( 396 _ io.ReaderAt = (*pageBuffer)(nil) 397 _ io.ReaderFrom = (*pageBuffer)(nil) 398 _ io.StringWriter = (*pageBuffer)(nil) 399 _ io.Writer = (*pageBuffer)(nil) 400 _ io.WriterAt = (*pageBuffer)(nil) 401 _ io.WriterTo = (*pageBuffer)(nil) 402 403 pagePool sync.Pool 404 pageBufferPool sync.Pool 405 ) 406 407 type contiguousPages []*page 408 409 func (pages contiguousPages) ref() { 410 for _, p := range pages { 411 p.ref() 412 } 413 } 414 415 func (pages contiguousPages) unref() { 416 for _, p := range pages { 417 p.unref() 418 } 419 } 420 421 func (pages contiguousPages) clear() { 422 for i := range pages { 423 pages[i] = nil 424 } 425 } 426 427 func (pages contiguousPages) ReadAt(b []byte, off int64) (int, error) { 428 rn := 0 429 430 for _, p := range pages.slice(off, off+int64(len(b))) { 431 n, _ := p.ReadAt(b, off) 432 b = b[n:] 433 rn += n 434 off += int64(n) 435 } 436 437 return rn, nil 438 } 439 440 func (pages contiguousPages) WriteAt(b []byte, off int64) (int, error) { 441 wn := 0 442 443 for _, p := range pages.slice(off, off+int64(len(b))) { 444 n, _ := p.WriteAt(b, off) 445 b = b[n:] 446 wn += n 447 off += int64(n) 448 } 449 450 return wn, nil 451 } 452 453 func (pages contiguousPages) slice(begin, end int64) contiguousPages { 454 i := pages.indexOf(begin) 455 j := pages.indexOf(end) 456 if j < len(pages) { 457 j++ 458 } 459 return pages[i:j] 460 } 461 462 func (pages contiguousPages) indexOf(offset int64) int { 463 if len(pages) == 0 { 464 return 0 465 } 466 return int((offset - pages[0].offset) / pageSize) 467 } 468 469 func (pages contiguousPages) scan(begin, end int64, f func([]byte) bool) { 470 for _, p := range pages.slice(begin, end) { 471 if !f(p.slice(begin, end)) { 472 break 473 } 474 } 475 } 476 477 var ( 478 _ io.ReaderAt = contiguousPages{} 479 _ io.WriterAt = contiguousPages{} 480 ) 481 482 type pageRef struct { 483 buffer [2]*page 484 pages contiguousPages 485 offset int64 486 cursor int64 487 length uint32 488 once uint32 489 } 490 491 func (ref *pageRef) unref() { 492 if atomic.CompareAndSwapUint32(&ref.once, 0, 1) { 493 ref.pages.unref() 494 ref.pages.clear() 495 ref.pages = nil 496 ref.offset = 0 497 ref.cursor = 0 498 ref.length = 0 499 } 500 } 501 502 func (ref *pageRef) Len() int { return int(ref.Size() - ref.cursor) } 503 504 func (ref *pageRef) Size() int64 { return int64(ref.length) } 505 506 func (ref *pageRef) Close() error { ref.unref(); return nil } 507 508 func (ref *pageRef) String() string { 509 return fmt.Sprintf("[offset=%d cursor=%d length=%d]", ref.offset, ref.cursor, ref.length) 510 } 511 512 func (ref *pageRef) Seek(offset int64, whence int) (int64, error) { 513 c, err := seek(ref.cursor, int64(ref.length), offset, whence) 514 if err != nil { 515 return -1, err 516 } 517 ref.cursor = c 518 return c, nil 519 } 520 521 func (ref *pageRef) ReadByte() (byte, error) { 522 var c byte 523 var ok bool 524 ref.scan(ref.cursor, func(b []byte) bool { 525 c, ok = b[0], true 526 return false 527 }) 528 if ok { 529 ref.cursor++ 530 } else { 531 return 0, io.EOF 532 } 533 return c, nil 534 } 535 536 func (ref *pageRef) Read(b []byte) (int, error) { 537 if ref.cursor >= int64(ref.length) { 538 return 0, io.EOF 539 } 540 n, err := ref.ReadAt(b, ref.cursor) 541 ref.cursor += int64(n) 542 return n, err 543 } 544 545 func (ref *pageRef) ReadAt(b []byte, off int64) (int, error) { 546 limit := ref.offset + int64(ref.length) 547 off += ref.offset 548 549 if off >= limit { 550 return 0, io.EOF 551 } 552 553 if off+int64(len(b)) > limit { 554 b = b[:limit-off] 555 } 556 557 if len(b) == 0 { 558 return 0, nil 559 } 560 561 n, err := ref.pages.ReadAt(b, off) 562 if n == 0 && err == nil { 563 err = io.EOF 564 } 565 return n, err 566 } 567 568 func (ref *pageRef) WriteTo(w io.Writer) (wn int64, err error) { 569 ref.scan(ref.cursor, func(b []byte) bool { 570 var n int 571 n, err = w.Write(b) 572 wn += int64(n) 573 return err == nil 574 }) 575 ref.cursor += wn 576 return 577 } 578 579 func (ref *pageRef) scan(off int64, f func([]byte) bool) { 580 begin := ref.offset + off 581 end := ref.offset + int64(ref.length) 582 ref.pages.scan(begin, end, f) 583 } 584 585 var ( 586 _ io.Closer = (*pageRef)(nil) 587 _ io.Seeker = (*pageRef)(nil) 588 _ io.Reader = (*pageRef)(nil) 589 _ io.ReaderAt = (*pageRef)(nil) 590 _ io.WriterTo = (*pageRef)(nil) 591 ) 592 593 type pageRefAllocator struct { 594 refs []pageRef 595 head int 596 size int 597 } 598 599 func (a *pageRefAllocator) newPageRef() *pageRef { 600 if a.head == len(a.refs) { 601 a.refs = make([]pageRef, a.size) 602 a.head = 0 603 } 604 ref := &a.refs[a.head] 605 a.head++ 606 return ref 607 } 608 609 func unref(x interface{}) { 610 if r, _ := x.(interface{ unref() }); r != nil { 611 r.unref() 612 } 613 } 614 615 func seek(cursor, limit, offset int64, whence int) (int64, error) { 616 switch whence { 617 case io.SeekStart: 618 // absolute offset 619 case io.SeekCurrent: 620 offset = cursor + offset 621 case io.SeekEnd: 622 offset = limit - offset 623 default: 624 return -1, fmt.Errorf("seek: invalid whence value: %d", whence) 625 } 626 if offset < 0 { 627 offset = 0 628 } 629 if offset > limit { 630 offset = limit 631 } 632 return offset, nil 633 } 634 635 func closeBytes(b Bytes) { 636 if b != nil { 637 b.Close() 638 } 639 } 640 641 func resetBytes(b Bytes) { 642 if r, _ := b.(interface{ Reset() }); r != nil { 643 r.Reset() 644 } 645 }