niki/vendor/github.com/go-sql-driver/mysql/buffer.go

// Go MySQL Driver - A MySQL-Driver for Go's database/sql package

//

// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.

//

// This Source Code Form is subject to the terms of the Mozilla Public

// License, v. 2.0. If a copy of the MPL was not distributed with this file,

// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"io"
	"net"
	"time"
)

const defaultBufSize = 4096

const maxCachedBufSize = 256 * 1024

// A buffer which is used for both reading and writing.

// This is possible since communication on each connection is synchronous.

// In other words, we can't write and read simultaneously on the same connection.

// The buffer is similar to bufio.Reader / Writer but zero-copy-ish

// Also highly optimized for this particular use case.

// This buffer is backed by two byte slices in a double-buffering scheme

type buffer struct {
	buf []byte // buf is a byte buffer who's length and capacity are equal.

	nc net.Conn

	idx int

	length int

	timeout time.Duration

	dbuf [2][]byte // dbuf is an array with the two byte slices that back this buffer

	flipcnt uint // flipccnt is the current buffer counter for double-buffering

}

// newBuffer allocates and returns a new buffer.

func newBuffer(nc net.Conn) buffer {

	fg := make([]byte, defaultBufSize)

	return buffer{

		buf: fg,

		nc: nc,

		dbuf: [2][]byte{fg, nil},
	}

}

// flip replaces the active buffer with the background buffer

// this is a delayed flip that simply increases the buffer counter;

// the actual flip will be performed the next time we call `buffer.fill`

func (b *buffer) flip() {

	b.flipcnt += 1

}

// fill reads into the buffer until at least _need_ bytes are in it

func (b *buffer) fill(need int) error {

	n := b.length

	// fill data into its double-buffering target: if we've called

	// flip on this buffer, we'll be copying to the background buffer,

	// and then filling it with network data; otherwise we'll just move

	// the contents of the current buffer to the front before filling it

	dest := b.dbuf[b.flipcnt&1]

	// grow buffer if necessary to fit the whole packet.

	if need > len(dest) {

		// Round up to the next multiple of the default size

		dest = make([]byte, ((need/defaultBufSize)+1)*defaultBufSize)

		// if the allocated buffer is not too large, move it to backing storage

		// to prevent extra allocations on applications that perform large reads

		if len(dest) <= maxCachedBufSize {

			b.dbuf[b.flipcnt&1] = dest

		}

	}

	// if we're filling the fg buffer, move the existing data to the start of it.

	// if we're filling the bg buffer, copy over the data

	if n > 0 {

		copy(dest[:n], b.buf[b.idx:])

	}

	b.buf = dest

	b.idx = 0

	for {

		if b.timeout > 0 {

			if err := b.nc.SetReadDeadline(time.Now().Add(b.timeout)); err != nil {

				return err

			}

		}

		nn, err := b.nc.Read(b.buf[n:])

		n += nn

		switch err {

		case nil:

			if n < need {

				continue

			}

			b.length = n

			return nil

		case io.EOF:

			if n >= need {

				b.length = n

				return nil

			}

			return io.ErrUnexpectedEOF

		default:

			return err

		}

	}

}

// returns next N bytes from buffer.

// The returned slice is only guaranteed to be valid until the next read

func (b *buffer) readNext(need int) ([]byte, error) {

	if b.length < need {

		// refill

		if err := b.fill(need); err != nil {

			return nil, err

		}

	}

	offset := b.idx

	b.idx += need

	b.length -= need

	return b.buf[offset:b.idx], nil

}

// takeBuffer returns a buffer with the requested size.

// If possible, a slice from the existing buffer is returned.

// Otherwise a bigger buffer is made.

// Only one buffer (total) can be used at a time.

func (b *buffer) takeBuffer(length int) ([]byte, error) {

	if b.length > 0 {

		return nil, ErrBusyBuffer

	}

	// test (cheap) general case first

	if length <= cap(b.buf) {

		return b.buf[:length], nil

	}

	if length < maxPacketSize {

		b.buf = make([]byte, length)

		return b.buf, nil

	}

	// buffer is larger than we want to store.

	return make([]byte, length), nil

}

// takeSmallBuffer is shortcut which can be used if length is

// known to be smaller than defaultBufSize.

// Only one buffer (total) can be used at a time.

func (b *buffer) takeSmallBuffer(length int) ([]byte, error) {

	if b.length > 0 {

		return nil, ErrBusyBuffer

	}

	return b.buf[:length], nil

}

// takeCompleteBuffer returns the complete existing buffer.

// This can be used if the necessary buffer size is unknown.

// cap and len of the returned buffer will be equal.

// Only one buffer (total) can be used at a time.

func (b *buffer) takeCompleteBuffer() ([]byte, error) {

	if b.length > 0 {

		return nil, ErrBusyBuffer

	}

	return b.buf, nil

}

// store stores buf, an updated buffer, if its suitable to do so.

func (b *buffer) store(buf []byte) error {

	if b.length > 0 {

		return ErrBusyBuffer

	} else if cap(buf) <= maxPacketSize && cap(buf) > cap(b.buf) {

		b.buf = buf[:cap(buf)]

	}

	return nil

}