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

CHANGELOG.md

@@ -16,6 +16,7 @@ Versioning](http://semver.org/spec/v2.0.0.html) except to the first release.
 - Support decimal type in msgpack (#96)
 - Support datetime type in msgpack (#118)
 - Prepared SQL statements (#117)
+- Context support for request objects (#48)
 
 ### Changed
 

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,43 @@ func (conn *Connection) newFuture() (fut *Future) {
 	return
 }
 
+// This method removes a future from the internal queue if the context
+// is "done" before the response is come. Such select logic is inspired
+// from this thread: https://groups.google.com/g/golang-dev/c/jX4oQEls3uk
+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 +968,11 @@ 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 list with nil context
+	if reqid%2 == 0 {
+		return shard.requests[pos].findFuture(reqid, fetch)
+	} else {
+		return shard.requestsWithCtx[pos].findFuture(reqid, fetch)
 	}
 }
 
@@ -984,8 +1061,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 +1081,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)
 }
 

example_test.go

@@ -1,6 +1,7 @@
 package tarantool_test
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -691,3 +692,33 @@ func ExampleConnection_NewPrepared() {
 		fmt.Printf("Failed to prepare")
 	}
 }
+
+// To pass contexts to request objects, use the Context() method.
+// Pay attention that when using context with request objects,
+// the timeout option for Connection will not affect the lifetime
+// of the request. For those purposes use context.WithTimeout() as
+// the root context.
+func ExamplePingRequest_Context() {
+	conn := example_connect()
+	defer conn.Close()
+
+	timeout := time.Nanosecond
+
+	// this way you may set the common timeout for requests with context
+	rootCtx, cancelRoot := context.WithTimeout(context.Background(), timeout)
+	defer cancelRoot()
+
+	// this context will be canceled with the root after commonTimeout
+	ctx, cancel := context.WithCancel(rootCtx)
+	defer cancel()
+
+	req := tarantool.NewPingRequest().Context(ctx)
+
+	// Ping a Tarantool instance to check connection.
+	resp, err := conn.Do(req).Get()
+	fmt.Println("Ping Resp", resp)
+	fmt.Println("Ping Error", err)
+	// Output:
+	// Ping Resp <nil>
+	// Ping Error context is done
+}