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

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