buffer.go 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  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. // 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. nc net.Conn
  23. idx int
  24. length int
  25. timeout time.Duration
  26. }
  27. func newBuffer(nc net.Conn) buffer {
  28. var b [defaultBufSize]byte
  29. return buffer{
  30. buf: b[:],
  31. nc: nc,
  32. }
  33. }
  34. // fill reads into the buffer until at least _need_ bytes are in it
  35. func (b *buffer) fill(need int) error {
  36. n := b.length
  37. // move existing data to the beginning
  38. if n > 0 && b.idx > 0 {
  39. copy(b.buf[0:n], b.buf[b.idx:])
  40. }
  41. // grow buffer if necessary
  42. // TODO: let the buffer shrink again at some point
  43. // Maybe keep the org buf slice and swap back?
  44. if need > len(b.buf) {
  45. // Round up to the next multiple of the default size
  46. newBuf := make([]byte, ((need/defaultBufSize)+1)*defaultBufSize)
  47. copy(newBuf, b.buf)
  48. b.buf = newBuf
  49. }
  50. b.idx = 0
  51. for {
  52. if b.timeout > 0 {
  53. if err := b.nc.SetReadDeadline(time.Now().Add(b.timeout)); err != nil {
  54. return err
  55. }
  56. }
  57. nn, err := b.nc.Read(b.buf[n:])
  58. n += nn
  59. switch err {
  60. case nil:
  61. if n < need {
  62. continue
  63. }
  64. b.length = n
  65. return nil
  66. case io.EOF:
  67. if n >= need {
  68. b.length = n
  69. return nil
  70. }
  71. return io.ErrUnexpectedEOF
  72. default:
  73. return err
  74. }
  75. }
  76. }
  77. // returns next N bytes from buffer.
  78. // The returned slice is only guaranteed to be valid until the next read
  79. func (b *buffer) readNext(need int) ([]byte, error) {
  80. if b.length < need {
  81. // refill
  82. if err := b.fill(need); err != nil {
  83. return nil, err
  84. }
  85. }
  86. offset := b.idx
  87. b.idx += need
  88. b.length -= need
  89. return b.buf[offset:b.idx], nil
  90. }
  91. // returns a buffer with the requested size.
  92. // If possible, a slice from the existing buffer is returned.
  93. // Otherwise a bigger buffer is made.
  94. // Only one buffer (total) can be used at a time.
  95. func (b *buffer) takeBuffer(length int) []byte {
  96. if b.length > 0 {
  97. return nil
  98. }
  99. // test (cheap) general case first
  100. if length <= defaultBufSize || length <= cap(b.buf) {
  101. return b.buf[:length]
  102. }
  103. if length < maxPacketSize {
  104. b.buf = make([]byte, length)
  105. return b.buf
  106. }
  107. return make([]byte, length)
  108. }
  109. // shortcut which can be used if the requested buffer is guaranteed to be
  110. // smaller than defaultBufSize
  111. // Only one buffer (total) can be used at a time.
  112. func (b *buffer) takeSmallBuffer(length int) []byte {
  113. if b.length > 0 {
  114. return nil
  115. }
  116. return b.buf[:length]
  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 nil
  124. }
  125. return b.buf
  126. }