@@ -14,8 +14,8 @@ import (
14
14
)
15
15
16
16
type connFn func (dst string , opt * ClientOption ) conn
17
- type dialFn func (dst string , opt * ClientOption ) (net.Conn , error )
18
- type wireFn func () wire
17
+ type dialFn func (ctx context. Context , dst string , opt * ClientOption ) (net.Conn , error )
18
+ type wireFn func (ctx context. Context ) wire
19
19
20
20
type singleconnect struct {
21
21
w wire
@@ -38,7 +38,7 @@ type conn interface {
38
38
Close ()
39
39
Dial () error
40
40
Override (conn )
41
- Acquire () wire
41
+ Acquire (ctx context. Context ) wire
42
42
Store (w wire )
43
43
Addr () string
44
44
SetOnCloseHook (func (error ))
@@ -67,12 +67,12 @@ type mux struct {
67
67
68
68
func makeMux (dst string , option * ClientOption , dialFn dialFn ) * mux {
69
69
dead := deadFn ()
70
- connFn := func () (net.Conn , error ) {
71
- return dialFn (dst , option )
70
+ connFn := func (ctx context. Context ) (net.Conn , error ) {
71
+ return dialFn (ctx , dst , option )
72
72
}
73
- wireFn := func (pipeFn pipeFn ) func () wire {
74
- return func () (w wire ) {
75
- w , err := pipeFn (connFn , option )
73
+ wireFn := func (pipeFn pipeFn ) func (context. Context ) wire {
74
+ return func (ctx context. Context ) (w wire ) {
75
+ w , err := pipeFn (ctx , connFn , option )
76
76
if err != nil {
77
77
dead .error .Store (& errs {error : err })
78
78
w = dead
@@ -152,7 +152,7 @@ func (m *mux) Override(cc conn) {
152
152
}
153
153
}
154
154
155
- func (m * mux ) _pipe (i uint16 ) (w wire , err error ) {
155
+ func (m * mux ) _pipe (ctx context. Context , i uint16 ) (w wire , err error ) {
156
156
if w = m .wire [i ].Load ().(wire ); w != m .init {
157
157
return w , nil
158
158
}
@@ -171,7 +171,7 @@ func (m *mux) _pipe(i uint16) (w wire, err error) {
171
171
}
172
172
173
173
if w = m .wire [i ].Load ().(wire ); w == m .init {
174
- if w = m .wireFn (); w != m .dead {
174
+ if w = m .wireFn (ctx ); w != m .dead {
175
175
m .setCloseHookOnWire (i , w )
176
176
m .wire [i ].Store (w )
177
177
} else {
@@ -193,39 +193,39 @@ func (m *mux) _pipe(i uint16) (w wire, err error) {
193
193
return w , err
194
194
}
195
195
196
- func (m * mux ) pipe (i uint16 ) wire {
197
- w , _ := m ._pipe (i )
196
+ func (m * mux ) pipe (ctx context. Context , i uint16 ) wire {
197
+ w , _ := m ._pipe (ctx , i )
198
198
return w // this should never be nil
199
199
}
200
200
201
201
func (m * mux ) Dial () error {
202
- _ , err := m ._pipe (0 )
202
+ _ , err := m ._pipe (context . Background (), 0 )
203
203
return err
204
204
}
205
205
206
206
func (m * mux ) Info () map [string ]RedisMessage {
207
- return m .pipe (0 ).Info ()
207
+ return m .pipe (context . Background (), 0 ).Info ()
208
208
}
209
209
210
210
func (m * mux ) Version () int {
211
- return m .pipe (0 ).Version ()
211
+ return m .pipe (context . Background (), 0 ).Version ()
212
212
}
213
213
214
214
func (m * mux ) AZ () string {
215
- return m .pipe (0 ).AZ ()
215
+ return m .pipe (context . Background (), 0 ).AZ ()
216
216
}
217
217
218
218
func (m * mux ) Error () error {
219
- return m .pipe (0 ).Error ()
219
+ return m .pipe (context . Background (), 0 ).Error ()
220
220
}
221
221
222
222
func (m * mux ) DoStream (ctx context.Context , cmd Completed ) RedisResultStream {
223
- wire := m .spool .Acquire ()
223
+ wire := m .spool .Acquire (ctx )
224
224
return wire .DoStream (ctx , m .spool , cmd )
225
225
}
226
226
227
227
func (m * mux ) DoMultiStream (ctx context.Context , multi ... Completed ) MultiRedisResultStream {
228
- wire := m .spool .Acquire ()
228
+ wire := m .spool .Acquire (ctx )
229
229
return wire .DoMultiStream (ctx , m .spool , multi ... )
230
230
}
231
231
@@ -262,7 +262,7 @@ block:
262
262
}
263
263
264
264
func (m * mux ) blocking (pool * pool , ctx context.Context , cmd Completed ) (resp RedisResult ) {
265
- wire := pool .Acquire ()
265
+ wire := pool .Acquire (ctx )
266
266
resp = wire .Do (ctx , cmd )
267
267
if resp .NonRedisError () != nil { // abort the wire if blocking command return early (ex. context.DeadlineExceeded)
268
268
wire .Close ()
@@ -272,7 +272,7 @@ func (m *mux) blocking(pool *pool, ctx context.Context, cmd Completed) (resp Red
272
272
}
273
273
274
274
func (m * mux ) blockingMulti (pool * pool , ctx context.Context , cmd []Completed ) (resp * redisresults ) {
275
- wire := pool .Acquire ()
275
+ wire := pool .Acquire (ctx )
276
276
resp = wire .DoMulti (ctx , cmd ... )
277
277
for _ , res := range resp .s {
278
278
if res .NonRedisError () != nil { // abort the wire if blocking command return early (ex. context.DeadlineExceeded)
@@ -286,7 +286,7 @@ func (m *mux) blockingMulti(pool *pool, ctx context.Context, cmd []Completed) (r
286
286
287
287
func (m * mux ) pipeline (ctx context.Context , cmd Completed ) (resp RedisResult ) {
288
288
slot := slotfn (len (m .wire ), cmd .Slot (), cmd .NoReply ())
289
- wire := m .pipe (slot )
289
+ wire := m .pipe (ctx , slot )
290
290
if resp = wire .Do (ctx , cmd ); isBroken (resp .NonRedisError (), wire ) {
291
291
m .wire [slot ].CompareAndSwap (wire , m .init )
292
292
}
@@ -295,7 +295,7 @@ func (m *mux) pipeline(ctx context.Context, cmd Completed) (resp RedisResult) {
295
295
296
296
func (m * mux ) pipelineMulti (ctx context.Context , cmd []Completed ) (resp * redisresults ) {
297
297
slot := slotfn (len (m .wire ), cmd [0 ].Slot (), cmd [0 ].NoReply ())
298
- wire := m .pipe (slot )
298
+ wire := m .pipe (ctx , slot )
299
299
resp = wire .DoMulti (ctx , cmd ... )
300
300
for _ , r := range resp .s {
301
301
if isBroken (r .NonRedisError (), wire ) {
@@ -308,7 +308,7 @@ func (m *mux) pipelineMulti(ctx context.Context, cmd []Completed) (resp *redisre
308
308
309
309
func (m * mux ) DoCache (ctx context.Context , cmd Cacheable , ttl time.Duration ) RedisResult {
310
310
slot := cmd .Slot () & uint16 (len (m .wire )- 1 )
311
- wire := m .pipe (slot )
311
+ wire := m .pipe (ctx , slot )
312
312
resp := wire .DoCache (ctx , cmd , ttl )
313
313
if isBroken (resp .NonRedisError (), wire ) {
314
314
m .wire [slot ].CompareAndSwap (wire , m .init )
@@ -366,7 +366,7 @@ func (m *mux) DoMultiCache(ctx context.Context, multi ...CacheableTTL) (results
366
366
}
367
367
368
368
func (m * mux ) doMultiCache (ctx context.Context , slot uint16 , multi []CacheableTTL ) (resps * redisresults ) {
369
- wire := m .pipe (slot )
369
+ wire := m .pipe (ctx , slot )
370
370
resps = wire .DoMultiCache (ctx , multi ... )
371
371
for _ , r := range resps .s {
372
372
if isBroken (r .NonRedisError (), wire ) {
@@ -379,16 +379,16 @@ func (m *mux) doMultiCache(ctx context.Context, slot uint16, multi []CacheableTT
379
379
380
380
func (m * mux ) Receive (ctx context.Context , subscribe Completed , fn func (message PubSubMessage )) error {
381
381
slot := slotfn (len (m .wire ), subscribe .Slot (), subscribe .NoReply ())
382
- wire := m .pipe (slot )
382
+ wire := m .pipe (ctx , slot )
383
383
err := wire .Receive (ctx , subscribe , fn )
384
384
if isBroken (err , wire ) {
385
385
m .wire [slot ].CompareAndSwap (wire , m .init )
386
386
}
387
387
return err
388
388
}
389
389
390
- func (m * mux ) Acquire () wire {
391
- return m .dpool .Acquire ()
390
+ func (m * mux ) Acquire (ctx context. Context ) wire {
391
+ return m .dpool .Acquire (ctx )
392
392
}
393
393
394
394
func (m * mux ) Store (w wire ) {
0 commit comments