api: context support

This patch adds the support of using context in API. The API is based
on using request objects. Added tests that cover almost all cases of
using the context in a query. Added benchamrk tests are equivalent to
other, that use the same query but without any context.

Closes #48

Author: vr009

connection.go

@@ -5,6 +5,7 @@ package tarantool
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -125,8 +126,11 @@ type Connection struct {
 	c     net.Conn
 	mutex sync.Mutex
 	// Schema contains schema loaded on connection.
-	Schema    *Schema
+	Schema *Schema
+	// requestId contains the last request ID for requests with nil context.
 	requestId uint32
+	// contextRequestId contains the last request ID for requests with context.
+	contextRequestId uint32
 	// Greeting contains first message sent by Tarantool.
 	Greeting *Greeting
 
@@ -143,16 +147,57 @@ type Connection struct {
 
 var _ = Connector(&Connection{}) // Check compatibility with connector interface.
 
+type futureList struct {
+	first *Future
+	last  **Future
+}
+
+func (list *futureList) findFuture(reqid uint32, fetch bool) *Future {
+	root := &list.first
+	for {
+		fut := *root
+		if fut == nil {
+			return nil
+		}
+		if fut.requestId == reqid {
+			if fetch {
+				*root = fut.next
+				if fut.next == nil {
+					list.last = root
+				} else {
+					fut.next = nil
+				}
+			}
+			return fut
+		}
+		root = &fut.next
+	}
+}
+
+func (list *futureList) addFuture(fut *Future) {
+	*list.last = fut
+	list.last = &fut.next
+}
+
+func (list *futureList) clear(err error, conn *Connection) {
+	fut := list.first
+	list.first = nil
+	list.last = &list.first
+	for fut != nil {
+		fut.SetError(err)
+		conn.markDone(fut)
+		fut, fut.next = fut.next, nil
+	}
+}
+
 type connShard struct {
-	rmut     sync.Mutex
-	requests [requestsMap]struct {
-		first *Future
-		last  **Future
-	}
-	bufmut sync.Mutex
-	buf    smallWBuf
-	enc    *msgpack.Encoder
-	_pad   [16]uint64 //nolint: unused,structcheck
+	rmut            sync.Mutex
+	requests        [requestsMap]futureList
+	requestsWithCtx [requestsMap]futureList
+	bufmut          sync.Mutex
+	buf             smallWBuf
+	enc             *msgpack.Encoder
+	_pad            [16]uint64 //nolint: unused,structcheck
 }
 
 // Greeting is a message sent by Tarantool on connect.
@@ -262,12 +307,13 @@ type SslOpts struct {
 // and will not finish to make attempts on authorization failures.
 func Connect(addr string, opts Opts) (conn *Connection, err error) {
 	conn = &Connection{
-		addr:      addr,
-		requestId: 0,
-		Greeting:  &Greeting{},
-		control:   make(chan struct{}),
-		opts:      opts,
-		dec:       msgpack.NewDecoder(&smallBuf{}),
+		addr:             addr,
+		requestId:        0,
+		contextRequestId: 1,
+		Greeting:         &Greeting{},
+		control:          make(chan struct{}),
+		opts:             opts,
+		dec:              msgpack.NewDecoder(&smallBuf{}),
 	}
 	maxprocs := uint32(runtime.GOMAXPROCS(-1))
 	if conn.opts.Concurrency == 0 || conn.opts.Concurrency > maxprocs*128 {
@@ -283,8 +329,11 @@ func Connect(addr string, opts Opts) (conn *Connection, err error) {
 	conn.shard = make([]connShard, conn.opts.Concurrency)
 	for i := range conn.shard {
 		shard := &conn.shard[i]
-		for j := range shard.requests {
-			shard.requests[j].last = &shard.requests[j].first
+		requestsLists := []*[requestsMap]futureList{&shard.requests, &shard.requestsWithCtx}
+		for _, requests := range requestsLists {
+			for j := range requests {
+				requests[j].last = &requests[j].first
+			}
 		}
 	}
 
@@ -387,6 +436,13 @@ func (conn *Connection) Handle() interface{} {
 	return conn.opts.Handle
 }
 
+func (conn *Connection) cancelFuture(fut *Future, err error) {
+	if fut = conn.fetchFuture(fut.requestId); fut != nil {
+		fut.SetError(err)
+		conn.markDone(fut)
+	}
+}
+
 func (conn *Connection) dial() (err error) {
 	var connection net.Conn
 	network := "tcp"
@@ -580,15 +636,10 @@ func (conn *Connection) closeConnection(neterr error, forever bool) (err error)
 	}
 	for i := range conn.shard {
 		conn.shard[i].buf.Reset()
-		requests := &conn.shard[i].requests
-		for pos := range requests {
-			fut := requests[pos].first
-			requests[pos].first = nil
-			requests[pos].last = &requests[pos].first
-			for fut != nil {
-				fut.SetError(neterr)
-				conn.markDone(fut)
-				fut, fut.next = fut.next, nil
+		requestsLists := []*[requestsMap]futureList{&conn.shard[i].requests, &conn.shard[i].requestsWithCtx}
+		for _, requests := range requestsLists {
+			for pos := range requests {
+				requests[pos].clear(neterr, conn)
 			}
 		}
 	}
@@ -721,7 +772,7 @@ func (conn *Connection) reader(r *bufio.Reader, c net.Conn) {
 	}
 }
 
-func (conn *Connection) newFuture() (fut *Future) {
+func (conn *Connection) newFuture(ctx context.Context) (fut *Future) {
 	fut = NewFuture()
 	if conn.rlimit != nil && conn.opts.RLimitAction == RLimitDrop {
 		select {
@@ -736,7 +787,7 @@ func (conn *Connection) newFuture() (fut *Future) {
 			return
 		}
 	}
-	fut.requestId = conn.nextRequestId()
+	fut.requestId = conn.nextRequestId(ctx != nil)
 	shardn := fut.requestId & (conn.opts.Concurrency - 1)
 	shard := &conn.shard[shardn]
 	shard.rmut.Lock()
@@ -761,11 +812,20 @@ func (conn *Connection) newFuture() (fut *Future) {
 		return
 	}
 	pos := (fut.requestId / conn.opts.Concurrency) & (requestsMap - 1)
-	pair := &shard.requests[pos]
-	*pair.last = fut
-	pair.last = &fut.next
-	if conn.opts.Timeout > 0 {
-		fut.timeout = time.Since(epoch) + conn.opts.Timeout
+	if ctx != nil {
+		select {
+		case <-ctx.Done():
+			fut.SetError(fmt.Errorf("context is done"))
+			shard.rmut.Unlock()
+			return
+		default:
+		}
+		shard.requestsWithCtx[pos].addFuture(fut)
+	} else {
+		shard.requests[pos].addFuture(fut)
+		if conn.opts.Timeout > 0 {
+			fut.timeout = time.Since(epoch) + conn.opts.Timeout
+		}
 	}
 	shard.rmut.Unlock()
 	if conn.rlimit != nil && conn.opts.RLimitAction == RLimitWait {
@@ -785,12 +845,40 @@ func (conn *Connection) newFuture() (fut *Future) {
 	return
 }
 
+func (conn *Connection) contextWatchdog(fut *Future, ctx context.Context) {
+	select {
+	case <-fut.done:
+	default:
+		select {
+		case <-ctx.Done():
+			conn.cancelFuture(fut, fmt.Errorf("context is done"))
+		default:
+			select {
+			case <-fut.done:
+			case <-ctx.Done():
+				conn.cancelFuture(fut, fmt.Errorf("context is done"))
+			}
+		}
+	}
+}
+
 func (conn *Connection) send(req Request) *Future {
-	fut := conn.newFuture()
+	fut := conn.newFuture(req.Ctx())
 	if fut.ready == nil {
 		return fut
 	}
+	if req.Ctx() != nil {
+		select {
+		case <-req.Ctx().Done():
+			conn.cancelFuture(fut, fmt.Errorf("context is done"))
+			return fut
+		default:
+		}
+	}
 	conn.putFuture(fut, req)
+	if req.Ctx() != nil {
+		go conn.contextWatchdog(fut, req.Ctx())
+	}
 	return fut
 }
 
@@ -877,25 +965,12 @@ func (conn *Connection) fetchFuture(reqid uint32) (fut *Future) {
 func (conn *Connection) getFutureImp(reqid uint32, fetch bool) *Future {
 	shard := &conn.shard[reqid&(conn.opts.Concurrency-1)]
 	pos := (reqid / conn.opts.Concurrency) & (requestsMap - 1)
-	pair := &shard.requests[pos]
-	root := &pair.first
-	for {
-		fut := *root
-		if fut == nil {
-			return nil
-		}
-		if fut.requestId == reqid {
-			if fetch {
-				*root = fut.next
-				if fut.next == nil {
-					pair.last = root
-				} else {
-					fut.next = nil
-				}
-			}
-			return fut
-		}
-		root = &fut.next
+	// futures with even requests id belong to requests with nil context
+	// and can be found only in requests list
+	if reqid%2 == 0 {
+		return shard.requests[pos].findFuture(reqid, fetch)
+	} else {
+		return shard.requestsWithCtx[pos].findFuture(reqid, fetch)
 	}
 }
 
@@ -984,8 +1059,12 @@ func (conn *Connection) read(r io.Reader) (response []byte, err error) {
 	return
 }
 
-func (conn *Connection) nextRequestId() (requestId uint32) {
-	return atomic.AddUint32(&conn.requestId, 1)
+func (conn *Connection) nextRequestId(context bool) (requestId uint32) {
+	if context {
+		return atomic.AddUint32(&conn.contextRequestId, 2)
+	} else {
+		return atomic.AddUint32(&conn.requestId, 2)
+	}
 }
 
 // Do performs a request asynchronously on the connection.
@@ -1000,6 +1079,15 @@ func (conn *Connection) Do(req Request) *Future {
 			return fut
 		}
 	}
+	if req.Ctx() != nil {
+		select {
+		case <-req.Ctx().Done():
+			fut := NewFuture()
+			fut.SetError(fmt.Errorf("context is done"))
+			return fut
+		default:
+		}
+	}
 	return conn.send(req)
 }
 

prepared.go

@@ -1,6 +1,7 @@
 package tarantool
 
 import (
+	"context"
 	"fmt"
 
 	"gopkg.in/vmihailenco/msgpack.v2"
@@ -58,6 +59,12 @@ func (req *PrepareRequest) Body(res SchemaResolver, enc *msgpack.Encoder) error
 	return fillPrepare(enc, req.expr)
 }
 
+// Context sets a passed context to the request.
+func (req *PrepareRequest) Context(ctx context.Context) *PrepareRequest {
+	req.ctx = ctx
+	return req
+}
+
 // UnprepareRequest helps you to create an unprepare request object for
 // execution by a Connection.
 type UnprepareRequest struct {
@@ -83,6 +90,12 @@ func (req *UnprepareRequest) Body(res SchemaResolver, enc *msgpack.Encoder) erro
 	return fillUnprepare(enc, *req.stmt)
 }
 
+// Context sets a passed context to the request.
+func (req *UnprepareRequest) Context(ctx context.Context) *UnprepareRequest {
+	req.ctx = ctx
+	return req
+}
+
 // ExecutePreparedRequest helps you to create an execute prepared request
 // object for execution by a Connection.
 type ExecutePreparedRequest struct {
@@ -117,6 +130,12 @@ func (req *ExecutePreparedRequest) Body(res SchemaResolver, enc *msgpack.Encoder
 	return fillExecutePrepared(enc, *req.stmt, req.args)
 }
 
+// Context sets a passed context to the request.
+func (req *ExecutePreparedRequest) Context(ctx context.Context) *ExecutePreparedRequest {
+	req.ctx = ctx
+	return req
+}
+
 func fillPrepare(enc *msgpack.Encoder, expr string) error {
 	enc.EncodeMapLen(1)
 	enc.EncodeUint64(KeySQLText)