Skip to content

Commit 66a9dc2

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 66a9dc2

File tree

5 files changed

+459
-43
lines changed

5 files changed

+459
-43
lines changed

connection.go

Lines changed: 127 additions & 43 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"
@@ -143,16 +144,57 @@ type Connection struct {
143144

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

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

158200
// Greeting is a message sent by Tarantool on connect.
@@ -286,6 +328,9 @@ func Connect(addr string, opts Opts) (conn *Connection, err error) {
286328
for j := range shard.requests {
287329
shard.requests[j].last = &shard.requests[j].first
288330
}
331+
for j := range shard.requests {
332+
shard.requestsWithCtx[j].last = &shard.requestsWithCtx[j].first
333+
}
289334
}
290335

291336
if opts.RateLimit > 0 {
@@ -387,6 +432,17 @@ func (conn *Connection) Handle() interface{} {
387432
return conn.opts.Handle
388433
}
389434

435+
func (conn *Connection) cancelFuture(fut *Future, err error) error {
436+
if fut == nil {
437+
return fmt.Errorf("passed nil future")
438+
}
439+
if fut = conn.fetchFuture(fut.requestId); fut != nil {
440+
fut.SetError(err)
441+
conn.markDone(fut)
442+
}
443+
return nil
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 {
@@ -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.Context())
790869
if fut.ready == nil {
791870
return fut
792871
}
872+
if req.Context() != nil {
873+
select {
874+
case <-req.Context().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.Context() != nil {
882+
go conn.contextWatchdog(fut, req.Context())
883+
}
794884
return fut
795885
}
796886

@@ -877,26 +967,11 @@ 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+
fut := shard.requests[pos].findFuture(reqid, fetch)
971+
if fut == nil {
972+
fut = shard.requestsWithCtx[pos].findFuture(reqid, fetch)
899973
}
974+
return fut
900975
}
901976

902977
func (conn *Connection) timeouts() {
@@ -1000,6 +1075,15 @@ func (conn *Connection) Do(req Request) *Future {
10001075
return fut
10011076
}
10021077
}
1078+
if req.Context() != nil {
1079+
select {
1080+
case <-req.Context().Done():
1081+
fut := NewFuture()
1082+
fut.SetError(fmt.Errorf("context is done"))
1083+
return fut
1084+
default:
1085+
}
1086+
}
10031087
return conn.send(req)
10041088
}
10051089

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+
// WithContext sets a passed context to the request.
63+
func (req *PrepareRequest) WithContext(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+
// WithContext sets a passed context to the request.
94+
func (req *UnprepareRequest) WithContext(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+
// WithContext sets a passed context to the request.
134+
func (req *ExecutePreparedRequest) WithContext(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)