Skip to content

Commit 98bc358

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 98bc358

File tree

5 files changed

+468
-53
lines changed

5 files changed

+468
-53
lines changed

connection.go

+141-53
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,13 @@ func (conn *Connection) Handle() interface{} {
387436
return conn.opts.Handle
388437
}
389438

439+
func (conn *Connection) cancelFuture(fut *Future, err error) {
440+
if fut = conn.fetchFuture(fut.requestId); fut != nil {
441+
fut.SetError(err)
442+
conn.markDone(fut)
443+
}
444+
}
445+
390446
func (conn *Connection) dial() (err error) {
391447
var connection net.Conn
392448
network := "tcp"
@@ -582,14 +638,11 @@ func (conn *Connection) closeConnection(neterr error, forever bool) (err error)
582638
conn.shard[i].buf.Reset()
583639
requests := &conn.shard[i].requests
584640
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-
}
641+
requests[pos].clear(neterr, conn)
642+
}
643+
requestsWithCtx := &conn.shard[i].requestsWithCtx
644+
for pos := range requestsWithCtx {
645+
requestsWithCtx[pos].clear(neterr, conn)
593646
}
594647
}
595648
return
@@ -721,7 +774,7 @@ func (conn *Connection) reader(r *bufio.Reader, c net.Conn) {
721774
}
722775
}
723776

724-
func (conn *Connection) newFuture() (fut *Future) {
777+
func (conn *Connection) newFuture(ctx context.Context) (fut *Future) {
725778
fut = NewFuture()
726779
if conn.rlimit != nil && conn.opts.RLimitAction == RLimitDrop {
727780
select {
@@ -736,7 +789,7 @@ func (conn *Connection) newFuture() (fut *Future) {
736789
return
737790
}
738791
}
739-
fut.requestId = conn.nextRequestId()
792+
fut.requestId = conn.nextRequestId(ctx != nil)
740793
shardn := fut.requestId & (conn.opts.Concurrency - 1)
741794
shard := &conn.shard[shardn]
742795
shard.rmut.Lock()
@@ -761,11 +814,20 @@ func (conn *Connection) newFuture() (fut *Future) {
761814
return
762815
}
763816
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
817+
if ctx != nil {
818+
select {
819+
case <-ctx.Done():
820+
fut.SetError(fmt.Errorf("context is done"))
821+
shard.rmut.Unlock()
822+
return
823+
default:
824+
}
825+
shard.requestsWithCtx[pos].addFuture(fut)
826+
} else {
827+
shard.requests[pos].addFuture(fut)
828+
if conn.opts.Timeout > 0 {
829+
fut.timeout = time.Since(epoch) + conn.opts.Timeout
830+
}
769831
}
770832
shard.rmut.Unlock()
771833
if conn.rlimit != nil && conn.opts.RLimitAction == RLimitWait {
@@ -785,12 +847,40 @@ func (conn *Connection) newFuture() (fut *Future) {
785847
return
786848
}
787849

850+
func (conn *Connection) contextWatchdog(fut *Future, ctx context.Context) {
851+
select {
852+
case <-fut.done:
853+
default:
854+
select {
855+
case <-ctx.Done():
856+
conn.cancelFuture(fut, fmt.Errorf("context is done"))
857+
default:
858+
select {
859+
case <-fut.done:
860+
case <-ctx.Done():
861+
conn.cancelFuture(fut, fmt.Errorf("context is done"))
862+
}
863+
}
864+
}
865+
}
866+
788867
func (conn *Connection) send(req Request) *Future {
789-
fut := conn.newFuture()
868+
fut := conn.newFuture(req.Ctx())
790869
if fut.ready == nil {
791870
return fut
792871
}
872+
if req.Ctx() != nil {
873+
select {
874+
case <-req.Ctx().Done():
875+
conn.cancelFuture(fut, fmt.Errorf("context is done"))
876+
return fut
877+
default:
878+
}
879+
}
793880
conn.putFuture(fut, req)
881+
if req.Ctx() != nil {
882+
go conn.contextWatchdog(fut, req.Ctx())
883+
}
794884
return fut
795885
}
796886

@@ -877,25 +967,10 @@ func (conn *Connection) fetchFuture(reqid uint32) (fut *Future) {
877967
func (conn *Connection) getFutureImp(reqid uint32, fetch bool) *Future {
878968
shard := &conn.shard[reqid&(conn.opts.Concurrency-1)]
879969
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
970+
if reqid%2 == 0 {
971+
return shard.requests[pos].findFuture(reqid, fetch)
972+
} else {
973+
return shard.requestsWithCtx[pos].findFuture(reqid, fetch)
899974
}
900975
}
901976

@@ -984,8 +1059,12 @@ func (conn *Connection) read(r io.Reader) (response []byte, err error) {
9841059
return
9851060
}
9861061

987-
func (conn *Connection) nextRequestId() (requestId uint32) {
988-
return atomic.AddUint32(&conn.requestId, 1)
1062+
func (conn *Connection) nextRequestId(context bool) (requestId uint32) {
1063+
if context {
1064+
return atomic.AddUint32(&conn.contextRequestId, 2)
1065+
} else {
1066+
return atomic.AddUint32(&conn.requestId, 2)
1067+
}
9891068
}
9901069

9911070
// Do performs a request asynchronously on the connection.
@@ -1000,6 +1079,15 @@ func (conn *Connection) Do(req Request) *Future {
10001079
return fut
10011080
}
10021081
}
1082+
if req.Ctx() != nil {
1083+
select {
1084+
case <-req.Ctx().Done():
1085+
fut := NewFuture()
1086+
fut.SetError(fmt.Errorf("context is done"))
1087+
return fut
1088+
default:
1089+
}
1090+
}
10031091
return conn.send(req)
10041092
}
10051093

prepared.go

+19
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)