Merge branch 'release/v0.6.1'

Author: dancannon

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## v0.6.1 - 13 Feb 2015
+
+- Reduce GC by using buffers when reading and writing
+- Fixed encoding `time.Time` ignoring millseconds
+- Fixed pointers in structs that implement the `Marshaler`/`Unmarshaler` interfaces being ignored
+
 ## v0.6.0 - 1 Feb 2015
 
 There are some major changes to the driver with this release that are not related to the RethinkDB v1.16 release. Please have a read through them:

README.md

@@ -7,7 +7,7 @@
 [Go](http://golang.org/) driver for [RethinkDB](http://www.rethinkdb.com/) 
 
 
-Current version: v0.6.0 (RethinkDB v1.16.0) 
+Current version: v0.6.1 (RethinkDB v1.16) 
 
 **Version 0.6 introduced some small API changes and some significant internal changes, for more information check the [change log](CHANGELOG.md) and please be aware the driver is not yet stable**
 

buffer.go

@@ -0,0 +1,143 @@
+package gorethink
+
+import "io"
+
+const defaultBufSize = 4096
+
+// A buffer which is used for both reading and writing.
+// This is possible since communication on each connection is synchronous.
+// In other words, we can't write and read simultaneously on the same connection.
+// The buffer is similar to bufio.Reader / Writer but zero-copy-ish
+// Also highly optimized for this particular use case.
+type buffer struct {
+	buf    []byte
+	rd     io.Reader
+	idx    int
+	length int
+}
+
+func newBuffer(rd io.Reader) buffer {
+	var b [defaultBufSize]byte
+	return buffer{
+		buf: b[:],
+		rd:  rd,
+	}
+}
+
+// fill reads into the buffer until at least _need_ bytes are in it
+func (b *buffer) fill(need int) error {
+	n := b.length
+
+	// move existing data to the beginning
+	if n > 0 && b.idx > 0 {
+		copy(b.buf[0:n], b.buf[b.idx:])
+	}
+
+	// grow buffer if necessary
+	// TODO: let the buffer shrink again at some point
+	//       Maybe keep the org buf slice and swap back?
+	if need > len(b.buf) {
+		// Round up to the next multiple of the default size
+		newBuf := make([]byte, ((need/defaultBufSize)+1)*defaultBufSize)
+		copy(newBuf, b.buf)
+		b.buf = newBuf
+	}
+
+	b.idx = 0
+
+	for {
+		nn, err := b.rd.Read(b.buf[n:])
+		n += nn
+
+		switch err {
+		case nil:
+			if n < need {
+				continue
+			}
+			b.length = n
+			return nil
+
+		case io.EOF:
+			if n >= need {
+				b.length = n
+				return nil
+			}
+			return io.ErrUnexpectedEOF
+
+		default:
+			return err
+		}
+	}
+}
+
+// returns next N bytes from buffer.
+// The returned slice is only guaranteed to be valid until the next read
+func (b *buffer) readNext(need int) ([]byte, error) {
+	if b.length < need {
+		// refill
+		if err := b.fill(need); err != nil {
+			return nil, err
+		}
+	}
+
+	offset := b.idx
+	b.idx += need
+	b.length -= need
+	return b.buf[offset:b.idx], nil
+}
+
+// returns a buffer with the requested size.
+// If possible, a slice from the existing buffer is returned.
+// Otherwise a bigger buffer is made.
+// Only one buffer (total) can be used at a time.
+func (b *buffer) takeBuffer(length int) []byte {
+	if b.length > 0 {
+		return nil
+	}
+
+	// test (cheap) general case first
+	if length <= defaultBufSize || length <= cap(b.buf) {
+		return b.buf[:length]
+	}
+
+	return make([]byte, length)
+}
+
+// shortcut which can be used if the requested buffer is guaranteed to be
+// smaller than defaultBufSize
+// Only one buffer (total) can be used at a time.
+func (b *buffer) takeSmallBuffer(length int) []byte {
+	if b.length == 0 {
+		return b.buf[:length]
+	}
+	return nil
+}
+
+// takeCompleteBuffer returns the complete existing buffer.
+// This can be used if the necessary buffer size is unknown.
+// Only one buffer (total) can be used at a time.
+func (b *buffer) takeCompleteBuffer() []byte {
+	if b.length == 0 {
+		return b.buf
+	}
+	return nil
+}
+
+var responseCache = make(chan *Response, 16)
+
+func newCachedResponse() *Response {
+	select {
+	case r := <-responseCache:
+		return r
+	default:
+		return new(Response)
+	}
+}
+
+func putResponse(r *Response) {
+	*r = Response{} // zero it
+	select {
+	case responseCache <- r:
+	default:
+	}
+}

connection.go

@@ -1,10 +1,8 @@
 package gorethink
 
 import (
-	"bufio"
 	"encoding/binary"
 	"encoding/json"
-	"fmt"
 	"io"
 	"net"
 	"sync/atomic"
@@ -13,6 +11,10 @@ import (
 	p "github.com/dancannon/gorethink/ql2"
 )
 
+const (
+	respHeaderLen = 12
+)
+
 type Response struct {
 	Token     int64
 	Type      p.Response_ResponseType `json:"t"`
@@ -29,62 +31,53 @@ type Connection struct {
 	token   int64
 	cursors map[int64]*Cursor
 	bad     bool
+
+	headerBuf [respHeaderLen]byte
+	buf       buffer
 }
 
 // Dial closes the previous connection and attempts to connect again.
 func NewConnection(opts *ConnectOpts) (*Connection, error) {
-	conn, err := net.Dial("tcp", opts.Address)
-	if err != nil {
-		return nil, RqlConnectionError{err.Error()}
-	}
+	var err error
 
-	// Send the protocol version to the server as a 4-byte little-endian-encoded integer
-	if err := binary.Write(conn, binary.LittleEndian, p.VersionDummy_V0_3); err != nil {
-		return nil, RqlConnectionError{err.Error()}
+	// New mysqlConn
+	c := &Connection{
+		opts:    opts,
+		cursors: make(map[int64]*Cursor),
 	}
 
-	// Send the length of the auth key to the server as a 4-byte little-endian-encoded integer
-	if err := binary.Write(conn, binary.LittleEndian, uint32(len(opts.AuthKey))); err != nil {
-		return nil, RqlConnectionError{err.Error()}
+	// Connect to Server
+	nd := net.Dialer{Timeout: c.opts.Timeout}
+	c.conn, err = nd.Dial("tcp", c.opts.Address)
+	if err != nil {
+		return nil, err
 	}
 
-	// Send the auth key as an ASCII string
-	// If there is no auth key, skip this step
-	if opts.AuthKey != "" {
-		if _, err := io.WriteString(conn, opts.AuthKey); err != nil {
-			return nil, RqlConnectionError{err.Error()}
+	// Enable TCP Keepalives on TCP connections
+	if tc, ok := c.conn.(*net.TCPConn); ok {
+		if err := tc.SetKeepAlive(true); err != nil {
+			// Don't send COM_QUIT before handshake.
+			c.conn.Close()
+			c.conn = nil
+			return nil, err
 		}
 	}
 
-	// Send the protocol type as a 4-byte little-endian-encoded integer
-	if err := binary.Write(conn, binary.LittleEndian, p.VersionDummy_JSON); err != nil {
-		return nil, RqlConnectionError{err.Error()}
-	}
+	c.buf = newBuffer(c.conn)
 
-	// read server response to authorization key (terminated by NUL)
-	reader := bufio.NewReader(conn)
-	line, err := reader.ReadBytes('\x00')
-	if err != nil {
-		if err == io.EOF {
-			return nil, fmt.Errorf("Unexpected EOF: %s", string(line))
-		}
-		return nil, RqlConnectionError{err.Error()}
-	}
-	// convert to string and remove trailing NUL byte
-	response := string(line[:len(line)-1])
-	if response != "SUCCESS" {
-		// we failed authorization or something else terrible happened
-		return nil, RqlDriverError{fmt.Sprintf("Server dropped connection with message: \"%s\"", response)}
+	// Send handshake request
+	if err = c.writeHandshakeReq(); err != nil {
+		c.Close()
+		return nil, err
 	}
 
-	c := &Connection{
-		opts:    opts,
-		conn:    conn,
-		cursors: make(map[int64]*Cursor),
+	// Read handshake response
+	err = c.readHandshakeSuccess()
+	if err != nil {
+		c.Close()
+		return nil, err
 	}
 
-	c.conn.SetDeadline(time.Time{})
-
 	return c, nil
 }
 
@@ -127,9 +120,8 @@ func (c *Connection) Query(q Query) (*Response, *Cursor, error) {
 		return nil, nil, nil
 	}
 
-	var response *Response
 	for {
-		response, err = c.readResponse()
+		response, err := c.readResponse()
 		if err != nil {
 			return nil, nil, err
 		}
@@ -140,6 +132,8 @@ func (c *Connection) Query(q Query) (*Response, *Cursor, error) {
 		} else if _, ok := c.cursors[response.Token]; ok {
 			// If the token is in the cursor cache then process the response
 			c.processResponse(q, response)
+		} else {
+			putResponse(response)
 		}
 	}
 }
@@ -158,21 +152,8 @@ func (c *Connection) sendQuery(q Query) error {
 		c.conn.SetDeadline(time.Now().Add(c.opts.Timeout))
 	}
 
-	// Send a unique 8-byte token
-	if err = binary.Write(c.conn, binary.LittleEndian, q.Token); err != nil {
-		c.bad = true
-		return RqlConnectionError{err.Error()}
-	}
-
-	// Send the length of the JSON-encoded query as a 4-byte
-	// little-endian-encoded integer.
-	if err = binary.Write(c.conn, binary.LittleEndian, uint32(len(b))); err != nil {
-		c.bad = true
-		return RqlConnectionError{err.Error()}
-	}
-
 	// Send the JSON encoding of the query itself.
-	if err = binary.Write(c.conn, binary.BigEndian, b); err != nil {
+	if err = c.writeQuery(q.Token, b); err != nil {
 		c.bad = true
 		return RqlConnectionError{err.Error()}
 	}
@@ -187,30 +168,24 @@ func (c *Connection) nextToken() int64 {
 }
 
 func (c *Connection) readResponse() (*Response, error) {
-	// Read the 8-byte token of the query the response corresponds to.
-	var responseToken int64
-	if err := binary.Read(c.conn, binary.LittleEndian, &responseToken); err != nil {
-		c.bad = true
-		return nil, RqlConnectionError{err.Error()}
+	// Read response header (token+length)
+	_, err := io.ReadFull(c.conn, c.headerBuf[:respHeaderLen])
+	if err != nil {
+		return nil, err
 	}
 
-	// Read the length of the JSON-encoded response as a 4-byte
-	// little-endian-encoded integer.
-	var messageLength uint32
-	if err := binary.Read(c.conn, binary.LittleEndian, &messageLength); err != nil {
-		c.bad = true
-		return nil, RqlConnectionError{err.Error()}
-	}
+	responseToken := int64(binary.LittleEndian.Uint64(c.headerBuf[:8]))
+	messageLength := binary.LittleEndian.Uint32(c.headerBuf[8:])
 
 	// Read the JSON encoding of the Response itself.
-	b := make([]byte, messageLength)
-	if _, err := io.ReadFull(c.conn, b); err != nil {
+	b := c.buf.takeBuffer(int(messageLength))
+	if _, err := io.ReadFull(c.conn, b[:]); err != nil {
 		c.bad = true
 		return nil, RqlConnectionError{err.Error()}
 	}
 
 	// Decode the response
-	var response = new(Response)
+	var response = newCachedResponse()
 	if err := json.Unmarshal(b, response); err != nil {
 		c.bad = true
 		return nil, RqlDriverError{err.Error()}
@@ -239,6 +214,7 @@ func (c *Connection) processResponse(q Query, response *Response) (*Response, *C
 	case p.Response_WAIT_COMPLETE:
 		return c.processWaitResponse(q, response)
 	default:
+		putResponse(response)
 		return nil, nil, RqlDriverError{"Unexpected response type"}
 	}
 }