Skip to content

Commit 9cd532f

Browse files
committed
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
1 parent e1bb59c commit 9cd532f

File tree

5 files changed

+477
-53
lines changed

5 files changed

+477
-53
lines changed

connection.go

Lines changed: 145 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ package tarantool
55
import (
66
"bufio"
77
"bytes"
8+
"context"
89
"errors"
910
"fmt"
1011
"io"
@@ -125,8 +126,11 @@ type Connection struct {
125126
c net.Conn
126127
mutex sync.Mutex
127128
// Schema contains schema loaded on connection.
128-
Schema *Schema
129+
Schema *Schema
130+
// requestId contains the last request ID for requests with nil context.
129131
requestId uint32
132+
// contextRequestId contains the last request ID for requests with context.
133+
contextRequestId uint32
130134
// Greeting contains first message sent by Tarantool.
131135
Greeting *Greeting
132136

@@ -143,16 +147,57 @@ type Connection struct {
143147

144148
var _ = Connector(&Connection{}) // Check compatibility with connector interface.
145149

150+
type futureList struct {
151+
first *Future
152+
last **Future
153+
}
154+
155+
func (list *futureList) findFuture(reqid uint32, fetch bool) *Future {
156+
root := &list.first
157+
for {
158+
fut := *root
159+
if fut == nil {
160+
return nil
161+
}
162+
if fut.requestId == reqid {
163+
if fetch {
164+
*root = fut.next
165+
if fut.next == nil {
166+
list.last = root
167+
} else {
168+
fut.next = nil
169+
}
170+
}
171+
return fut
172+
}
173+
root = &fut.next
174+
}
175+
}
176+
177+
func (list *futureList) addFuture(fut *Future) {
178+
*list.last = fut
179+
list.last = &fut.next
180+
}
181+
182+
func (list *futureList) clear(err error, conn *Connection) {
183+
fut := list.first
184+
list.first = nil
185+
list.last = &list.first
186+
for fut != nil {
187+
fut.SetError(err)
188+
conn.markDone(fut)
189+
fut, fut.next = fut.next, nil
190+
}
191+
}
192+
146193
type connShard struct {
147-
rmut sync.Mutex
148-
requests [requestsMap]struct {
149-
first *Future
150-
last **Future
151-
}
152-
bufmut sync.Mutex
153-
buf smallWBuf
154-
enc *msgpack.Encoder
155-
_pad [16]uint64 //nolint: unused,structcheck
194+
rmut sync.Mutex
195+
requests [requestsMap]futureList
196+
requestsWithCtx [requestsMap]futureList
197+
bufmut sync.Mutex
198+
buf smallWBuf
199+
enc *msgpack.Encoder
200+
_pad [16]uint64 //nolint: unused,structcheck
156201
}
157202

158203
// Greeting is a message sent by Tarantool on connect.
@@ -262,12 +307,13 @@ type SslOpts struct {
262307
// and will not finish to make attempts on authorization failures.
263308
func Connect(addr string, opts Opts) (conn *Connection, err error) {
264309
conn = &Connection{
265-
addr: addr,
266-
requestId: 0,
267-
Greeting: &Greeting{},
268-
control: make(chan struct{}),
269-
opts: opts,
270-
dec: msgpack.NewDecoder(&smallBuf{}),
310+
addr: addr,
311+
requestId: 0,
312+
contextRequestId: 1,
313+
Greeting: &Greeting{},
314+
control: make(chan struct{}),
315+
opts: opts,
316+
dec: msgpack.NewDecoder(&smallBuf{}),
271317
}
272318
maxprocs := uint32(runtime.GOMAXPROCS(-1))
273319
if conn.opts.Concurrency == 0 || conn.opts.Concurrency > maxprocs*128 {
@@ -286,6 +332,9 @@ func Connect(addr string, opts Opts) (conn *Connection, err error) {
286332
for j := range shard.requests {
287333
shard.requests[j].last = &shard.requests[j].first
288334
}
335+
for j := range shard.requests {
336+
shard.requestsWithCtx[j].last = &shard.requestsWithCtx[j].first
337+
}
289338
}
290339

291340
if opts.RateLimit > 0 {
@@ -387,6 +436,17 @@ func (conn *Connection) Handle() interface{} {
387436
return conn.opts.Handle
388437
}
389438

439+
func (conn *Connection) cancelFuture(fut *Future, err error) error {
440+
if fut == nil {
441+
return fmt.Errorf("passed nil future")
442+
}
443+
if fut = conn.fetchFuture(fut.requestId); fut != nil {
444+
fut.SetError(err)
445+
conn.markDone(fut)
446+
}
447+
return nil
448+
}
449+
390450
func (conn *Connection) dial() (err error) {
391451
var connection net.Conn
392452
network := "tcp"
@@ -582,14 +642,11 @@ func (conn *Connection) closeConnection(neterr error, forever bool) (err error)
582642
conn.shard[i].buf.Reset()
583643
requests := &conn.shard[i].requests
584644
for pos := range requests {
585-
fut := requests[pos].first
586-
requests[pos].first = nil
587-
requests[pos].last = &requests[pos].first
588-
for fut != nil {
589-
fut.SetError(neterr)
590-
conn.markDone(fut)
591-
fut, fut.next = fut.next, nil
592-
}
645+
requests[pos].clear(neterr, conn)
646+
}
647+
requestsWithCtx := &conn.shard[i].requestsWithCtx
648+
for pos := range requestsWithCtx {
649+
requestsWithCtx[pos].clear(neterr, conn)
593650
}
594651
}
595652
return
@@ -721,7 +778,7 @@ func (conn *Connection) reader(r *bufio.Reader, c net.Conn) {
721778
}
722779
}
723780

724-
func (conn *Connection) newFuture() (fut *Future) {
781+
func (conn *Connection) newFuture(ctx context.Context) (fut *Future) {
725782
fut = NewFuture()
726783
if conn.rlimit != nil && conn.opts.RLimitAction == RLimitDrop {
727784
select {
@@ -736,7 +793,7 @@ func (conn *Connection) newFuture() (fut *Future) {
736793
return
737794
}
738795
}
739-
fut.requestId = conn.nextRequestId()
796+
fut.requestId = conn.nextRequestId(ctx != nil)
740797
shardn := fut.requestId & (conn.opts.Concurrency - 1)
741798
shard := &conn.shard[shardn]
742799
shard.rmut.Lock()
@@ -761,11 +818,20 @@ func (conn *Connection) newFuture() (fut *Future) {
761818
return
762819
}
763820
pos := (fut.requestId / conn.opts.Concurrency) & (requestsMap - 1)
764-
pair := &shard.requests[pos]
765-
*pair.last = fut
766-
pair.last = &fut.next
767-
if conn.opts.Timeout > 0 {
768-
fut.timeout = time.Since(epoch) + conn.opts.Timeout
821+
if ctx != nil {
822+
select {
823+
case <-ctx.Done():
824+
fut.SetError(fmt.Errorf("context is done"))
825+
shard.rmut.Unlock()
826+
return
827+
default:
828+
}
829+
shard.requestsWithCtx[pos].addFuture(fut)
830+
} else {
831+
shard.requests[pos].addFuture(fut)
832+
if conn.opts.Timeout > 0 {
833+
fut.timeout = time.Since(epoch) + conn.opts.Timeout
834+
}
769835
}
770836
shard.rmut.Unlock()
771837
if conn.rlimit != nil && conn.opts.RLimitAction == RLimitWait {
@@ -785,12 +851,40 @@ func (conn *Connection) newFuture() (fut *Future) {
785851
return
786852
}
787853

854+
func (conn *Connection) contextWatchdog(fut *Future, ctx context.Context) {
855+
select {
856+
case <-fut.done:
857+
default:
858+
select {
859+
case <-ctx.Done():
860+
conn.cancelFuture(fut, fmt.Errorf("context is done"))
861+
default:
862+
select {
863+
case <-fut.done:
864+
case <-ctx.Done():
865+
conn.cancelFuture(fut, fmt.Errorf("context is done"))
866+
}
867+
}
868+
}
869+
}
870+
788871
func (conn *Connection) send(req Request) *Future {
789-
fut := conn.newFuture()
872+
fut := conn.newFuture(req.Ctx())
790873
if fut.ready == nil {
791874
return fut
792875
}
876+
if req.Ctx() != nil {
877+
select {
878+
case <-req.Ctx().Done():
879+
conn.cancelFuture(fut, fmt.Errorf("context is done"))
880+
return fut
881+
default:
882+
}
883+
}
793884
conn.putFuture(fut, req)
885+
if req.Ctx() != nil {
886+
go conn.contextWatchdog(fut, req.Ctx())
887+
}
794888
return fut
795889
}
796890

@@ -877,25 +971,10 @@ func (conn *Connection) fetchFuture(reqid uint32) (fut *Future) {
877971
func (conn *Connection) getFutureImp(reqid uint32, fetch bool) *Future {
878972
shard := &conn.shard[reqid&(conn.opts.Concurrency-1)]
879973
pos := (reqid / conn.opts.Concurrency) & (requestsMap - 1)
880-
pair := &shard.requests[pos]
881-
root := &pair.first
882-
for {
883-
fut := *root
884-
if fut == nil {
885-
return nil
886-
}
887-
if fut.requestId == reqid {
888-
if fetch {
889-
*root = fut.next
890-
if fut.next == nil {
891-
pair.last = root
892-
} else {
893-
fut.next = nil
894-
}
895-
}
896-
return fut
897-
}
898-
root = &fut.next
974+
if reqid%2 == 0 {
975+
return shard.requests[pos].findFuture(reqid, fetch)
976+
} else {
977+
return shard.requestsWithCtx[pos].findFuture(reqid, fetch)
899978
}
900979
}
901980

@@ -984,8 +1063,12 @@ func (conn *Connection) read(r io.Reader) (response []byte, err error) {
9841063
return
9851064
}
9861065

987-
func (conn *Connection) nextRequestId() (requestId uint32) {
988-
return atomic.AddUint32(&conn.requestId, 1)
1066+
func (conn *Connection) nextRequestId(Context bool) (requestId uint32) {
1067+
if Context {
1068+
return atomic.AddUint32(&conn.contextRequestId, 2)
1069+
} else {
1070+
return atomic.AddUint32(&conn.requestId, 2)
1071+
}
9891072
}
9901073

9911074
// Do performs a request asynchronously on the connection.
@@ -1000,6 +1083,15 @@ func (conn *Connection) Do(req Request) *Future {
10001083
return fut
10011084
}
10021085
}
1086+
if req.Ctx() != nil {
1087+
select {
1088+
case <-req.Ctx().Done():
1089+
fut := NewFuture()
1090+
fut.SetError(fmt.Errorf("context is done"))
1091+
return fut
1092+
default:
1093+
}
1094+
}
10031095
return conn.send(req)
10041096
}
10051097

prepared.go

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package tarantool
22

33
import (
4+
"context"
45
"fmt"
56

67
"gopkg.in/vmihailenco/msgpack.v2"
@@ -58,6 +59,12 @@ func (req *PrepareRequest) Body(res SchemaResolver, enc *msgpack.Encoder) error
5859
return fillPrepare(enc, req.expr)
5960
}
6061

62+
// Context sets a passed context to the request.
63+
func (req *PrepareRequest) Context(ctx context.Context) *PrepareRequest {
64+
req.ctx = ctx
65+
return req
66+
}
67+
6168
// UnprepareRequest helps you to create an unprepare request object for
6269
// execution by a Connection.
6370
type UnprepareRequest struct {
@@ -83,6 +90,12 @@ func (req *UnprepareRequest) Body(res SchemaResolver, enc *msgpack.Encoder) erro
8390
return fillUnprepare(enc, *req.stmt)
8491
}
8592

93+
// Context sets a passed context to the request.
94+
func (req *UnprepareRequest) Context(ctx context.Context) *UnprepareRequest {
95+
req.ctx = ctx
96+
return req
97+
}
98+
8699
// ExecutePreparedRequest helps you to create an execute prepared request
87100
// object for execution by a Connection.
88101
type ExecutePreparedRequest struct {
@@ -117,6 +130,12 @@ func (req *ExecutePreparedRequest) Body(res SchemaResolver, enc *msgpack.Encoder
117130
return fillExecutePrepared(enc, *req.stmt, req.args)
118131
}
119132

133+
// Context sets a passed context to the request.
134+
func (req *ExecutePreparedRequest) Context(ctx context.Context) *ExecutePreparedRequest {
135+
req.ctx = ctx
136+
return req
137+
}
138+
120139
func fillPrepare(enc *msgpack.Encoder, expr string) error {
121140
enc.EncodeMapLen(1)
122141
enc.EncodeUint64(KeySQLText)

0 commit comments

Comments
 (0)