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

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