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
9package mysql
10
11import (
12	"io"
13	"net"
14	"time"
15)
16
17const defaultBufSize = 4096
18const maxCachedBufSize = 256 * 1024
19
20// A buffer which is used for both reading and writing.
21// This is possible since communication on each connection is synchronous.
22// In other words, we can't write and read simultaneously on the same connection.
23// The buffer is similar to bufio.Reader / Writer but zero-copy-ish
24// Also highly optimized for this particular use case.
25// This buffer is backed by two byte slices in a double-buffering scheme
26type buffer struct {
27	buf     []byte // buf is a byte buffer who's length and capacity are equal.
28	nc      net.Conn
29	idx     int
30	length  int
31	timeout time.Duration
32	dbuf    [2][]byte // dbuf is an array with the two byte slices that back this buffer
33	flipcnt uint      // flipccnt is the current buffer counter for double-buffering
34}
35
36// newBuffer allocates and returns a new buffer.
37func newBuffer(nc net.Conn) buffer {
38	fg := make([]byte, defaultBufSize)
39	return buffer{
40		buf:  fg,
41		nc:   nc,
42		dbuf: [2][]byte{fg, nil},
43	}
44}
45
46// flip replaces the active buffer with the background buffer
47// this is a delayed flip that simply increases the buffer counter;
48// the actual flip will be performed the next time we call `buffer.fill`
49func (b *buffer) flip() {
50	b.flipcnt += 1
51}
52
53// fill reads into the buffer until at least _need_ bytes are in it
54func (b *buffer) fill(need int) error {
55	n := b.length
56	// fill data into its double-buffering target: if we've called
57	// flip on this buffer, we'll be copying to the background buffer,
58	// and then filling it with network data; otherwise we'll just move
59	// the contents of the current buffer to the front before filling it
60	dest := b.dbuf[b.flipcnt&1]
61
62	// grow buffer if necessary to fit the whole packet.
63	if need > len(dest) {
64		// Round up to the next multiple of the default size
65		dest = make([]byte, ((need/defaultBufSize)+1)*defaultBufSize)
66
67		// if the allocated buffer is not too large, move it to backing storage
68		// to prevent extra allocations on applications that perform large reads
69		if len(dest) <= maxCachedBufSize {
70			b.dbuf[b.flipcnt&1] = dest
71		}
72	}
73
74	// if we're filling the fg buffer, move the existing data to the start of it.
75	// if we're filling the bg buffer, copy over the data
76	if n > 0 {
77		copy(dest[:n], b.buf[b.idx:])
78	}
79
80	b.buf = dest
81	b.idx = 0
82
83	for {
84		if b.timeout > 0 {
85			if err := b.nc.SetReadDeadline(time.Now().Add(b.timeout)); err != nil {
86				return err
87			}
88		}
89
90		nn, err := b.nc.Read(b.buf[n:])
91		n += nn
92
93		switch err {
94		case nil:
95			if n < need {
96				continue
97			}
98			b.length = n
99			return nil
100
101		case io.EOF:
102			if n >= need {
103				b.length = n
104				return nil
105			}
106			return io.ErrUnexpectedEOF
107
108		default:
109			return err
110		}
111	}
112}
113
114// returns next N bytes from buffer.
115// The returned slice is only guaranteed to be valid until the next read
116func (b *buffer) readNext(need int) ([]byte, error) {
117	if b.length < need {
118		// refill
119		if err := b.fill(need); err != nil {
120			return nil, err
121		}
122	}
123
124	offset := b.idx
125	b.idx += need
126	b.length -= need
127	return b.buf[offset:b.idx], nil
128}
129
130// takeBuffer returns a buffer with the requested size.
131// If possible, a slice from the existing buffer is returned.
132// Otherwise a bigger buffer is made.
133// Only one buffer (total) can be used at a time.
134func (b *buffer) takeBuffer(length int) ([]byte, error) {
135	if b.length > 0 {
136		return nil, ErrBusyBuffer
137	}
138
139	// test (cheap) general case first
140	if length <= cap(b.buf) {
141		return b.buf[:length], nil
142	}
143
144	if length < maxPacketSize {
145		b.buf = make([]byte, length)
146		return b.buf, nil
147	}
148
149	// buffer is larger than we want to store.
150	return make([]byte, length), nil
151}
152
153// takeSmallBuffer is shortcut which can be used if length is
154// known to be smaller than defaultBufSize.
155// Only one buffer (total) can be used at a time.
156func (b *buffer) takeSmallBuffer(length int) ([]byte, error) {
157	if b.length > 0 {
158		return nil, ErrBusyBuffer
159	}
160	return b.buf[:length], nil
161}
162
163// takeCompleteBuffer returns the complete existing buffer.
164// This can be used if the necessary buffer size is unknown.
165// cap and len of the returned buffer will be equal.
166// Only one buffer (total) can be used at a time.
167func (b *buffer) takeCompleteBuffer() ([]byte, error) {
168	if b.length > 0 {
169		return nil, ErrBusyBuffer
170	}
171	return b.buf, nil
172}
173
174// store stores buf, an updated buffer, if its suitable to do so.
175func (b *buffer) store(buf []byte) error {
176	if b.length > 0 {
177		return ErrBusyBuffer
178	} else if cap(buf) <= maxPacketSize && cap(buf) > cap(b.buf) {
179		b.buf = buf[:cap(buf)]
180	}
181	return nil
182}
183