api: proposal to add the context support

This patch adds the support of using context in API. The proposed 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
tarantool · Jul 18, 2022 · 66a9dc2 · 66a9dc2
1 parent e1bb59c
commit 66a9dc2
Show file tree

Hide file tree

Showing 5 changed files with 459 additions and 43 deletions.
diff --git a/connection.go b/connection.go
@@ -5,6 +5,7 @@ package tarantool
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -143,16 +144,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.
@@ -286,6 +328,9 @@ func Connect(addr string, opts Opts) (conn *Connection, err error) {
 		for j := range shard.requests {
 			shard.requests[j].last = &shard.requests[j].first
 		}
+		for j := range shard.requests {
+			shard.requestsWithCtx[j].last = &shard.requestsWithCtx[j].first
+		}
 	}
 
 	if opts.RateLimit > 0 {
@@ -387,6 +432,17 @@ func (conn *Connection) Handle() interface{} {
 	return conn.opts.Handle
 }
 
+func (conn *Connection) cancelFuture(fut *Future, err error) error {
+	if fut == nil {
+		return fmt.Errorf("passed nil future")
+	}
+	if fut = conn.fetchFuture(fut.requestId); fut != nil {
+		fut.SetError(err)
+		conn.markDone(fut)
+	}
+	return nil
+}
+
 func (conn *Connection) dial() (err error) {
 	var connection net.Conn
 	network := "tcp"
@@ -582,14 +638,11 @@ func (conn *Connection) closeConnection(neterr error, forever bool) (err error)
 		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
-			}
+			requests[pos].clear(neterr, conn)
+		}
+		requestsWithCtx := &conn.shard[i].requestsWithCtx
+		for pos := range requestsWithCtx {
+			requestsWithCtx[pos].clear(neterr, conn)
 		}
 	}
 	return
@@ -721,7 +774,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 {
@@ -761,11 +814,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 +847,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.Context())
 	if fut.ready == nil {
 		return fut
 	}
+	if req.Context() != nil {
+		select {
+		case <-req.Context().Done():
+			conn.cancelFuture(fut, fmt.Errorf("context is done"))
+			return fut
+		default:
+		}
+	}
 	conn.putFuture(fut, req)
+	if req.Context() != nil {
+		go conn.contextWatchdog(fut, req.Context())
+	}
 	return fut
 }
 
@@ -877,26 +967,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
+	fut := shard.requests[pos].findFuture(reqid, fetch)
+	if fut == nil {
+		fut = shard.requestsWithCtx[pos].findFuture(reqid, fetch)
 	}
+	return fut
 }
 
 func (conn *Connection) timeouts() {
@@ -1000,6 +1075,15 @@ func (conn *Connection) Do(req Request) *Future {
 			return fut
 		}
 	}
+	if req.Context() != nil {
+		select {
+		case <-req.Context().Done():
+			fut := NewFuture()
+			fut.SetError(fmt.Errorf("context is done"))
+			return fut
+		default:
+		}
+	}
 	return conn.send(req)
 }
 

diff --git a/prepared.go b/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)
 }
 
+// WithContext sets a passed context to the request.
+func (req *PrepareRequest) WithContext(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)
 }
 
+// WithContext sets a passed context to the request.
+func (req *UnprepareRequest) WithContext(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)
 }
 
+// WithContext sets a passed context to the request.
+func (req *ExecutePreparedRequest) WithContext(ctx context.Context) *ExecutePreparedRequest {
+	req.ctx = ctx
+	return req
+}
+
 func fillPrepare(enc *msgpack.Encoder, expr string) error {
 	enc.EncodeMapLen(1)
 	enc.EncodeUint64(KeySQLText)