Skip to content

Commit 3169baa

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 3169baa

File tree

5 files changed

+476
-54
lines changed

5 files changed

+476
-54
lines changed

connection.go

+144-54
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,9 @@ type Connection struct {
125126
c net.Conn
126127
mutex sync.Mutex
127128
// Schema contains schema loaded on connection.
128-
Schema *Schema
129-
requestId uint32
129+
Schema *Schema
130+
requestId uint32
131+
contextRequestId uint32
130132
// Greeting contains first message sent by Tarantool.
131133
Greeting *Greeting
132134

@@ -143,16 +145,57 @@ type Connection struct {
143145

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

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

158201
// Greeting is a message sent by Tarantool on connect.
@@ -262,12 +305,13 @@ type SslOpts struct {
262305
// and will not finish to make attempts on authorization failures.
263306
func Connect(addr string, opts Opts) (conn *Connection, err error) {
264307
conn = &Connection{
265-
addr: addr,
266-
requestId: 0,
267-
Greeting: &Greeting{},
268-
control: make(chan struct{}),
269-
opts: opts,
270-
dec: msgpack.NewDecoder(&smallBuf{}),
308+
addr: addr,
309+
requestId: 0,
310+
contextRequestId: 1,
311+
Greeting: &Greeting{},
312+
control: make(chan struct{}),
313+
opts: opts,
314+
dec: msgpack.NewDecoder(&smallBuf{}),
271315
}
272316
maxprocs := uint32(runtime.GOMAXPROCS(-1))
273317
if conn.opts.Concurrency == 0 || conn.opts.Concurrency > maxprocs*128 {
@@ -286,6 +330,9 @@ func Connect(addr string, opts Opts) (conn *Connection, err error) {
286330
for j := range shard.requests {
287331
shard.requests[j].last = &shard.requests[j].first
288332
}
333+
for j := range shard.requests {
334+
shard.requestsWithCtx[j].last = &shard.requestsWithCtx[j].first
335+
}
289336
}
290337

291338
if opts.RateLimit > 0 {
@@ -387,6 +434,17 @@ func (conn *Connection) Handle() interface{} {
387434
return conn.opts.Handle
388435
}
389436

437+
func (conn *Connection) cancelFuture(fut *Future, err error) error {
438+
if fut == nil {
439+
return fmt.Errorf("passed nil future")
440+
}
441+
if fut = conn.fetchFuture(fut.requestId); fut != nil {
442+
fut.SetError(err)
443+
conn.markDone(fut)
444+
}
445+
return nil
446+
}
447+
390448
func (conn *Connection) dial() (err error) {
391449
var connection net.Conn
392450
network := "tcp"
@@ -582,14 +640,11 @@ func (conn *Connection) closeConnection(neterr error, forever bool) (err error)
582640
conn.shard[i].buf.Reset()
583641
requests := &conn.shard[i].requests
584642
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-
}
643+
requests[pos].clear(neterr, conn)
644+
}
645+
requestsWithCtx := &conn.shard[i].requestsWithCtx
646+
for pos := range requestsWithCtx {
647+
requestsWithCtx[pos].clear(neterr, conn)
593648
}
594649
}
595650
return
@@ -721,7 +776,7 @@ func (conn *Connection) reader(r *bufio.Reader, c net.Conn) {
721776
}
722777
}
723778

724-
func (conn *Connection) newFuture() (fut *Future) {
779+
func (conn *Connection) newFuture(ctx context.Context) (fut *Future) {
725780
fut = NewFuture()
726781
if conn.rlimit != nil && conn.opts.RLimitAction == RLimitDrop {
727782
select {
@@ -736,7 +791,7 @@ func (conn *Connection) newFuture() (fut *Future) {
736791
return
737792
}
738793
}
739-
fut.requestId = conn.nextRequestId()
794+
fut.requestId = conn.nextRequestId(ctx != nil)
740795
shardn := fut.requestId & (conn.opts.Concurrency - 1)
741796
shard := &conn.shard[shardn]
742797
shard.rmut.Lock()
@@ -761,11 +816,20 @@ func (conn *Connection) newFuture() (fut *Future) {
761816
return
762817
}
763818
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
819+
if ctx != nil {
820+
select {
821+
case <-ctx.Done():
822+
fut.SetError(fmt.Errorf("context is done"))
823+
shard.rmut.Unlock()
824+
return
825+
default:
826+
}
827+
shard.requestsWithCtx[pos].addFuture(fut)
828+
} else {
829+
shard.requests[pos].addFuture(fut)
830+
if conn.opts.Timeout > 0 {
831+
fut.timeout = time.Since(epoch) + conn.opts.Timeout
832+
}
769833
}
770834
shard.rmut.Unlock()
771835
if conn.rlimit != nil && conn.opts.RLimitAction == RLimitWait {
@@ -785,12 +849,40 @@ func (conn *Connection) newFuture() (fut *Future) {
785849
return
786850
}
787851

852+
func (conn *Connection) contextWatchdog(fut *Future, ctx context.Context) {
853+
select {
854+
case <-fut.done:
855+
default:
856+
select {
857+
case <-ctx.Done():
858+
conn.cancelFuture(fut, fmt.Errorf("context is done"))
859+
default:
860+
select {
861+
case <-fut.done:
862+
case <-ctx.Done():
863+
conn.cancelFuture(fut, fmt.Errorf("context is done"))
864+
}
865+
}
866+
}
867+
}
868+
788869
func (conn *Connection) send(req Request) *Future {
789-
fut := conn.newFuture()
870+
fut := conn.newFuture(req.Ctx())
790871
if fut.ready == nil {
791872
return fut
792873
}
874+
if req.Ctx() != nil {
875+
select {
876+
case <-req.Ctx().Done():
877+
conn.cancelFuture(fut, fmt.Errorf("context is done"))
878+
return fut
879+
default:
880+
}
881+
}
793882
conn.putFuture(fut, req)
883+
if req.Ctx() != nil {
884+
go conn.contextWatchdog(fut, req.Ctx())
885+
}
794886
return fut
795887
}
796888

@@ -877,25 +969,10 @@ func (conn *Connection) fetchFuture(reqid uint32) (fut *Future) {
877969
func (conn *Connection) getFutureImp(reqid uint32, fetch bool) *Future {
878970
shard := &conn.shard[reqid&(conn.opts.Concurrency-1)]
879971
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
972+
if reqid%2 == 0 {
973+
return shard.requests[pos].findFuture(reqid, fetch)
974+
} else {
975+
return shard.requestsWithCtx[pos].findFuture(reqid, fetch)
899976
}
900977
}
901978

@@ -984,8 +1061,12 @@ func (conn *Connection) read(r io.Reader) (response []byte, err error) {
9841061
return
9851062
}
9861063

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

9911072
// Do performs a request asynchronously on the connection.
@@ -1000,6 +1081,15 @@ func (conn *Connection) Do(req Request) *Future {
10001081
return fut
10011082
}
10021083
}
1084+
if req.Ctx() != nil {
1085+
select {
1086+
case <-req.Ctx().Done():
1087+
fut := NewFuture()
1088+
fut.SetError(fmt.Errorf("context is done"))
1089+
return fut
1090+
default:
1091+
}
1092+
}
10031093
return conn.send(req)
10041094
}
10051095

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)