github.com/cellofellow/gopkg@v0.0.0-20140722061823-eec0544a62ad/database/mysql/buffer.go (about)

     1  // Go MySQL Driver - A MySQL-Driver for Go's database/sql package
     2  //
     3  // Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
     4  //
     5  // This Source Code Form is subject to the terms of the Mozilla Public
     6  // License, v. 2.0. If a copy of the MPL was not distributed with this file,
     7  // You can obtain one at http://mozilla.org/MPL/2.0/.
     8  
     9  package mysql
    10  
    11  import "io"
    12  
    13  const defaultBufSize = 4096
    14  
    15  // A buffer which is used for both reading and writing.
    16  // This is possible since communication on each connection is synchronous.
    17  // In other words, we can't write and read simultaneously on the same connection.
    18  // The buffer is similar to bufio.Reader / Writer but zero-copy-ish
    19  // Also highly optimized for this particular use case.
    20  type buffer struct {
    21  	buf    []byte
    22  	rd     io.Reader
    23  	idx    int
    24  	length int
    25  }
    26  
    27  func newBuffer(rd io.Reader) buffer {
    28  	var b [defaultBufSize]byte
    29  	return buffer{
    30  		buf: b[:],
    31  		rd:  rd,
    32  	}
    33  }
    34  
    35  // fill reads into the buffer until at least _need_ bytes are in it
    36  func (b *buffer) fill(need int) error {
    37  	// move existing data to the beginning
    38  	if b.length > 0 && b.idx > 0 {
    39  		copy(b.buf[0:b.length], b.buf[b.idx:])
    40  	}
    41  
    42  	// grow buffer if necessary
    43  	// TODO: let the buffer shrink again at some point
    44  	//       Maybe keep the org buf slice and swap back?
    45  	if need > len(b.buf) {
    46  		// Round up to the next multiple of the default size
    47  		newBuf := make([]byte, ((need/defaultBufSize)+1)*defaultBufSize)
    48  		copy(newBuf, b.buf)
    49  		b.buf = newBuf
    50  	}
    51  
    52  	b.idx = 0
    53  
    54  	for {
    55  		n, err := b.rd.Read(b.buf[b.length:])
    56  		b.length += n
    57  
    58  		if err == nil {
    59  			if b.length < need {
    60  				continue
    61  			}
    62  			return nil
    63  		}
    64  		if b.length >= need && err == io.EOF {
    65  			return nil
    66  		}
    67  		return err
    68  	}
    69  }
    70  
    71  // returns next N bytes from buffer.
    72  // The returned slice is only guaranteed to be valid until the next read
    73  func (b *buffer) readNext(need int) ([]byte, error) {
    74  	if b.length < need {
    75  		// refill
    76  		if err := b.fill(need); err != nil {
    77  			return nil, err
    78  		}
    79  	}
    80  
    81  	offset := b.idx
    82  	b.idx += need
    83  	b.length -= need
    84  	return b.buf[offset:b.idx], nil
    85  }
    86  
    87  // returns a buffer with the requested size.
    88  // If possible, a slice from the existing buffer is returned.
    89  // Otherwise a bigger buffer is made.
    90  // Only one buffer (total) can be used at a time.
    91  func (b *buffer) takeBuffer(length int) []byte {
    92  	if b.length > 0 {
    93  		return nil
    94  	}
    95  
    96  	// test (cheap) general case first
    97  	if length <= defaultBufSize || length <= cap(b.buf) {
    98  		return b.buf[:length]
    99  	}
   100  
   101  	if length < maxPacketSize {
   102  		b.buf = make([]byte, length)
   103  		return b.buf
   104  	}
   105  	return make([]byte, length)
   106  }
   107  
   108  // shortcut which can be used if the requested buffer is guaranteed to be
   109  // smaller than defaultBufSize
   110  // Only one buffer (total) can be used at a time.
   111  func (b *buffer) takeSmallBuffer(length int) []byte {
   112  	if b.length == 0 {
   113  		return b.buf[:length]
   114  	}
   115  	return nil
   116  }
   117  
   118  // takeCompleteBuffer returns the complete existing buffer.
   119  // This can be used if the necessary buffer size is unknown.
   120  // Only one buffer (total) can be used at a time.
   121  func (b *buffer) takeCompleteBuffer() []byte {
   122  	if b.length == 0 {
   123  		return b.buf
   124  	}
   125  	return nil
   126  }