buffer.go 4.8 KB

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