buffer.go 3.1 KB

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