@@ -5,6 +5,7 @@ package tarantool
5
5
import (
6
6
"bufio"
7
7
"bytes"
8
+ "context"
8
9
"errors"
9
10
"fmt"
10
11
"io"
@@ -125,8 +126,11 @@ type Connection struct {
125
126
c net.Conn
126
127
mutex sync.Mutex
127
128
// Schema contains schema loaded on connection.
128
- Schema * Schema
129
+ Schema * Schema
130
+ // requestId contains the last request ID for requests with nil context.
129
131
requestId uint32
132
+ // contextRequestId contains the last request ID for requests with context.
133
+ contextRequestId uint32
130
134
// Greeting contains first message sent by Tarantool.
131
135
Greeting * Greeting
132
136
@@ -143,16 +147,57 @@ type Connection struct {
143
147
144
148
var _ = Connector (& Connection {}) // Check compatibility with connector interface.
145
149
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
+
146
193
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
156
201
}
157
202
158
203
// Greeting is a message sent by Tarantool on connect.
@@ -262,12 +307,13 @@ type SslOpts struct {
262
307
// and will not finish to make attempts on authorization failures.
263
308
func Connect (addr string , opts Opts ) (conn * Connection , err error ) {
264
309
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 {}),
271
317
}
272
318
maxprocs := uint32 (runtime .GOMAXPROCS (- 1 ))
273
319
if conn .opts .Concurrency == 0 || conn .opts .Concurrency > maxprocs * 128 {
@@ -286,6 +332,9 @@ func Connect(addr string, opts Opts) (conn *Connection, err error) {
286
332
for j := range shard .requests {
287
333
shard .requests [j ].last = & shard .requests [j ].first
288
334
}
335
+ for j := range shard .requests {
336
+ shard .requestsWithCtx [j ].last = & shard .requestsWithCtx [j ].first
337
+ }
289
338
}
290
339
291
340
if opts .RateLimit > 0 {
@@ -387,6 +436,13 @@ func (conn *Connection) Handle() interface{} {
387
436
return conn .opts .Handle
388
437
}
389
438
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
+
390
446
func (conn * Connection ) dial () (err error ) {
391
447
var connection net.Conn
392
448
network := "tcp"
@@ -582,14 +638,11 @@ func (conn *Connection) closeConnection(neterr error, forever bool) (err error)
582
638
conn .shard [i ].buf .Reset ()
583
639
requests := & conn .shard [i ].requests
584
640
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 )
593
646
}
594
647
}
595
648
return
@@ -721,7 +774,7 @@ func (conn *Connection) reader(r *bufio.Reader, c net.Conn) {
721
774
}
722
775
}
723
776
724
- func (conn * Connection ) newFuture () (fut * Future ) {
777
+ func (conn * Connection ) newFuture (ctx context. Context ) (fut * Future ) {
725
778
fut = NewFuture ()
726
779
if conn .rlimit != nil && conn .opts .RLimitAction == RLimitDrop {
727
780
select {
@@ -736,7 +789,7 @@ func (conn *Connection) newFuture() (fut *Future) {
736
789
return
737
790
}
738
791
}
739
- fut .requestId = conn .nextRequestId ()
792
+ fut .requestId = conn .nextRequestId (ctx != nil )
740
793
shardn := fut .requestId & (conn .opts .Concurrency - 1 )
741
794
shard := & conn .shard [shardn ]
742
795
shard .rmut .Lock ()
@@ -761,11 +814,20 @@ func (conn *Connection) newFuture() (fut *Future) {
761
814
return
762
815
}
763
816
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
+ }
769
831
}
770
832
shard .rmut .Unlock ()
771
833
if conn .rlimit != nil && conn .opts .RLimitAction == RLimitWait {
@@ -785,12 +847,40 @@ func (conn *Connection) newFuture() (fut *Future) {
785
847
return
786
848
}
787
849
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
+
788
867
func (conn * Connection ) send (req Request ) * Future {
789
- fut := conn .newFuture ()
868
+ fut := conn .newFuture (req . Ctx () )
790
869
if fut .ready == nil {
791
870
return fut
792
871
}
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
+ }
793
880
conn .putFuture (fut , req )
881
+ if req .Ctx () != nil {
882
+ go conn .contextWatchdog (fut , req .Ctx ())
883
+ }
794
884
return fut
795
885
}
796
886
@@ -877,25 +967,10 @@ func (conn *Connection) fetchFuture(reqid uint32) (fut *Future) {
877
967
func (conn * Connection ) getFutureImp (reqid uint32 , fetch bool ) * Future {
878
968
shard := & conn .shard [reqid & (conn .opts .Concurrency - 1 )]
879
969
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 )
899
974
}
900
975
}
901
976
@@ -984,8 +1059,12 @@ func (conn *Connection) read(r io.Reader) (response []byte, err error) {
984
1059
return
985
1060
}
986
1061
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
+ }
989
1068
}
990
1069
991
1070
// Do performs a request asynchronously on the connection.
@@ -1000,6 +1079,15 @@ func (conn *Connection) Do(req Request) *Future {
1000
1079
return fut
1001
1080
}
1002
1081
}
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
+ }
1003
1091
return conn .send (req )
1004
1092
}
1005
1093
0 commit comments